RankMe

class stable_ssl.callbacks.RankMe(name: str, target: str, queue_length: int, target_shape: int | Iterable[int])

Bases: Callback

RankMe (effective rank) monitor using queue discovery.

RankMe measures the effective rank of feature representations by computing the exponential of the entropy of normalized singular values. This metric helps detect dimensional collapse in self-supervised learning.

Parameters:

• name – Unique name for this callback instance

• target – Key in batch dict containing the feature embeddings to monitor

• queue_length – Required queue length

• target_shape – Shape of the target embeddings (e.g., 768 for 768-dim features)

on_validation_batch_end(trainer: Trainer, pl_module: LightningModule, outputs: dict, batch: dict, batch_idx: int, dataloader_idx: int = 0) → None

Compute RankMe metric on the first validation batch only.

setup(trainer: Trainer, pl_module: LightningModule, stage: str) → None

Find or create the queue callback for target features.