$λ$-JEPA Spectral Anti-Collapse Regularization for Self-Supervised Learning
Berker Demirel, Clémentine Dominé, Valentino Maiorca, Marco Fumero, Marco Mondelli, Francesco Locatello
cs.LG, cs.CV
2026-09-28
ISTA's λ-JEPA adds SACReg on JEPA backbones. ViT-S hits 69.7 ImageNet-1k linear probe and 77.6 mean transfer at 100 epochs, 8.9 points above LeJEPA.
Joint-embedding self-supervised learning pulls augmented views of the same image together, then adds a second mechanism so the representation does not collapse to a constant. Contrastive negatives, teacher-student asymmetry, and explicit regularizers such as VICReg, LeJEPA, and VISReg all do that job. Almost always the loss sits after a projection head. Downstream protocols throw that head away and keep the backbone features.
Those two spaces need not share geometry. On explicitly regularized JE-SSL, ISTA finds that positive pairs stay aligned in both the backbone and the projector, while different images still crowd together in the backbone. RankMe, a singular-value estimate of effective rank, is typically higher after the projector; the backbone stays low rank, especially for VISReg and LeJEPA. Task-agnostic pretraining needs leftover directions for downstream work. Garrido et al. already showed that effective rank tracks transfer in JE-SSL. A full-rank projector does not imply a full-rank backbone.
The design starts from λ-balance in a two-layer linear net: the relative scale of adjacent weights, Δ = W2^T W2 − W1 W1^T. If Δ = λI with λ < 0, the encoder Gram is bounded away from singularity and, with a positive-definite input covariance, the hidden covariance cannot collapse. Unregularized gradient flow conserves that balance. Deep nonlinear nets do not let you pin it with initialization.
A constrained problem does the pinning. For a fixed end-to-end map, minimize symmetric L2 plus a negative log-det on the encoder Gram. The minimizer sits at λbal = −γ/λreg. Under gradient flow the imbalance converges exponentially to that point, independent of the dataset and the initial Δ.
Nonlinear nets do not conserve Δ, so the same spectral penalty is applied to representation covariance. SACReg (Spectral Anti-Collapse Regularization) has three terms: a mean penalty that pulls features to the origin, a trace term that caps total variance, and a negative log-det that penalizes tiny covariance eigenvalues, with ε=1e-4 for stability. Left alone, the unique minimizer is zero-mean isotropic covariance. When ε=0 the log-det is a hard barrier: any vanishing eigenvalue sends the regularizer to infinity.
λ-JEPA is that regularizer on a JEPA-style objective. Invariance is still a pairwise mean-squared distance in projector space, with no negatives. SACReg is computed on image centers: average V global views of the same image first, then spread those centers, so the penalty targets between-image variation rather than augmentation. Both backbone and projector get a term. When the method is dropped onto an existing SSL loss, only the backbone term is added. If feature width exceeds the number of image centers in a batch, the log-det is evaluated on a random orthonormal slice each step (128-d after the projector; 128-d for ViT-S backbones, 256-d for ViT-B). ImageNet-100 uses 4 global views; ImageNet-1k and video use 6. Loss weights come from gradient calibration and were not retuned for long runs.
From-scratch ImageNet-1k at 100 epochs, frozen linear probe and mean linear-probe transfer over eight datasets:
| Method | Backbone | IN-1k linear | Transfer |
| LeJEPA | ViT-S/16 | 62.5 | 68.7 |
| DINO | ViT-S/16 | 70.0 | 75.7 |
| iBOT | ViT-S/16 | 70.9 | 75.9 |
| λ-JEPA | ViT-S/16 | 69.7 | 77.6 |
| LeJEPA | ViT-B/16 | 69.7 | 74.2 |
| VISReg | ViT-B/16 | 70.3 | 76.4 |
| DINO | ViT-B/16 | 73.9 | 78.3 |
| iBOT | ViT-B/16 | 77.0 | 80.7 |
| λ-JEPA | ViT-B/16 | 74.2 | 80.9 |
ViT-S transfer is 8.9 points above LeJEPA. The 69.7 ImageNet probe sits near DINO (70.0) and below iBOT (70.9). ViT-B transfer 80.9 is 4.5 above VISReg and slightly above DINO (78.3), matching iBOT (80.7); the 74.2 probe still trails iBOT's 77.0.
The same hyperparameters at 400 epochs, no retuning: ViT-B reaches 76.0 probe / 82.0 transfer, 2.9 transfer points above 400-epoch VISReg (74.6 / 79.1), still short of 400-epoch DINO (77.2 / 83.1) and iBOT (78.5 / 83.2). ViT-S at 400 epochs hits 79.3 transfer, essentially matching 300-epoch DINO and iBOT. Longer-run baselines are not length-matched.
Video training reuses the LeVJEPA model, optimizer, and data pipeline, swapping only the loss, from scratch on a class-balanced 20% subset of Kinetics-710. ViT-S at 240 epochs: ImageNet 46.8 versus 39.4 LeVJEPA and 38.7 V-JEPA 2. ViT-B at 240 epochs: ImageNet 52.9 (LeVJEPA 50.7, V-JEPA 2 51.6), and Something-Something-v2 13.5 points above LeVJEPA (43.9 vs 30.4). At 1085 epochs, SSv2 is 48.3 versus LeVJEPA 40.4 (+7.9); on the same table, ImageNet goes the other way, 57.1 for λ-JEPA against LeVJEPA's 61.0. Kinetics-400 CLS probe is 45.7 against LeVJEPA's mean-pool 44.6. The same frozen encoder scores 43.4 with mean pooling and 62.0 with an attentive probe. Probe mismatch means that 1.1-point K400 gap is not a clean comparison.
On ImageNet-100, SACReg is added to the backbone of LeJEPA, VICReg, VISReg, SimCLR, DINO, and BYOL, leaving each original loss intact. RankMe rises for all six; kNN improves everywhere; linear probe improves too, with a smaller bump for VICReg. Projector-space statistics move much less. For λ-JEPA, class-center separability improves on 96 of 100 classes; view sensitivity moves both ways. The extra rank is used for between-image structure along class directions, not for tighter positive-pair invariance.
For people on the JEPA / VICReg line this is a concrete patch: do not put the anti-collapse term only on a projector you will discard. SACReg can be added to an existing SSL loss; code is public. Frozen transfer is where the numbers hold up best, and temporal recognition on SSv2 moves more than static classification.
The 100-epoch transfer lead is the cleanest result; longer training without retuning closes some of the gap to DINO and iBOT. The method fits global-view, CLS-token classification and transfer. It is not a recipe for dense prediction.
The paper is explicit: SACReg acts on CLS tokens and global views. Patch features and local crops make per-image view centers noisier, the covariance term unstable, and training became unstable in their runs. Dense prediction is untested.
The non-collapse theorem is for two-layer linear nets. In the nonlinear case the regularizer, by itself, likes isotropic full-rank covariance; the identity W2^T W2 − W1 W1^T = −(γ/λreg)I does not carry over. CKA and empirical NTK drift on ViT are consistent with rich learning, not a proof of it.
Comparisons are uneven. The 400-epoch and 1085-epoch runs were not retuned. Longer SSL checkpoints use different durations. Video K400 mixes CLS and mean-pool probes. Pretraining data is a 20% slice of Kinetics-710. On ImageNet-100 the DINO regularizer weight was tuned separately; the other five methods share gradient calibration.