Why VISReg Excels: Solving Representation Collapse in JEPA Models

Research Background

Self‑supervised learning (SSL) for JEPA world models, advocated by Yann LeCun since 2017, suffers from representation collapse, where different inputs map to identical or few vectors, yielding non‑discriminative features. Existing methods rely on heuristic tricks (EMA, teacher‑student, stop‑gradient, layer freezing) that make training fragile and hard to tune.

VICReg introduced variance, invariance, and covariance regularization, but its covariance term captures only second‑order statistics and cannot distinguish distributions with identical mean and variance but different shapes. SIGReg improved this by using sketching based on the Cramér–Wold theorem to align the entire embedding distribution to a standard Gaussian, yet it still suffers from gradient vanishing during collapse and coupling between scale and shape.

Core Method: Decoupling Regularization into Scale and Shape

VISReg combines the strengths of VICReg and SIGReg. It retains VICReg’s variance term to control scale and replaces the covariance term with a sketching objective based on the Sliced Wasserstein Distance (SWD) to control shape . The three‑part regularization objective is:

Scale Regularization : Constrains per‑dimension variance to prevent amplitude collapse. When the model collapses, the gradient of this term approaches a constant, ensuring stable recovery.

Shape Regularization : Normalizes embeddings (stop‑gradient on the standard deviation) and aligns the distribution shape to an isotropic Gaussian via SWD. The Cramér–Wold theorem (Lemma 3.1) guarantees that matching all one‑dimensional projections is equivalent to matching the full high‑dimensional distribution.

Centering Loss : Pulls the batch mean toward the origin (‖μ‖₂²).

The combined loss is weighted by a single hyper‑parameter λ to balance prediction and regularization.

Implementation (≈15 lines of PyTorch)

def visreg(z, K=64):
    # 1. Centering loss
    mu = z.mean(dim=0)
    L_center = mu.pow(2).mean()
    # 2. Scale loss
    z_cent = z - mu
    std = z_cent.std(dim=0, unbiased=False)
    L_scale = (1.0 - std).pow(2).mean()
    # 3. Shape loss: sliced Wasserstein distance
    z_norm = z_cent / (std.detach())
    W = torch.randn(D, K)
    W /= W.norm(p=2, dim=0)
    p_sorted = torch.sort(z_norm @ W, dim=0).values
    u = torch.arange(1, N+1) / (N+1)
    target = Normal(0, 1).icdf(u)
    L_shape = (p_sorted - target).pow(2).mean()
    return L_scale + L_shape + L_center

Computational Complexity and Scalability

The regularization cost is O(NDK) (N: batch size, D: feature dimension, K: number of slices), linear in all factors. By contrast, VICReg’s covariance term is O(ND²), quadratic in D. VISReg’s slice computation can be distributed across multiple GPUs; using 8 GPUs with 128 slices each (total 1024) reduces the accuracy gap from ~2.4% to 0.22% compared with a single‑GPU 1024‑slice configuration.

Experimental Results

In‑Domain Linear Probing : Without any heuristic tricks, VISReg achieves 75.7% (ViT‑B/16) and 77.0% (ViT‑L/14) top‑1 accuracy, surpassing MAE and LeJEPA. On the texture dataset DTD, VISReg outperforms all methods, indicating superior cross‑domain generalization.

Out‑of‑Distribution (OOD) Generalization : Evaluated on six OOD datasets (medical, astronomy, remote sensing, texture), VISReg attains the highest average OOD accuracy across all backbones, even beating larger models that use heuristics.

Data Efficiency : Pre‑training VISReg (ViT‑L/14) on ImageNet‑22K (~14M images) yields 72.94% average OOD accuracy, matching DINOv2 trained on ten times more data (LVD‑142M, 142M images) while using only ~1/10 of the data.

Fine‑Tuning Transfer : After fine‑tuning, VISReg surpasses DINO and supervised pre‑training on CIFAR‑10, CIFAR‑100, Flowers, ImageNet‑1K, and Galaxy10, demonstrating more uniform and less redundant representations.

Dense Prediction & Generation Guidance : On ADE20K linear semantic segmentation (ViT‑B/16), VISReg reaches mIoU 30.16, competitive without heuristics. In generation guidance (SiT‑B/2, iREPA, 100k steps), VISReg‑guided images achieve better gFID, Precision, and Recall than DINO.

Robustness on Low‑Quality Data : On long‑tail (ImageNet‑LT) and low‑rank (Galaxy10) datasets, VISReg consistently prevents collapse and learns meaningful features, whereas DINO fails without extensive tuning.

Conclusion

VISReg demonstrates that decoupling representation regularization into independent scale and shape components yields a self‑supervised learning method that is more stable, computationally efficient, and generalizes better than existing approaches. It solves the long‑standing representation collapse problem in JEPA world models, achieving leading results on classification, segmentation, and generation tasks while requiring only about one‑tenth of the data needed by competing methods.