$\lambda$-JEPA: Spectral Anti-Collapse Regularization for Self-Supervised Learning

Intro

Representational collapse remains one of the major issues in the self-supervised learning task. Different inputs can collapse onto the same latent point, limiting representational rank of the model. Recent solutions, such as LeJEPA (with SIGReg)¹, VICReg², and VISReg³ try to solve this issue by adding explicite representational regularization.

In regular JEPA (or even DINO) style trainings, the outputs of the encoder get processed by an additional model before putting them into a predictive objective.

$x$Context Encoder $f_\theta$$h_x$Predictor $g_\phi$$\hat{z}_y$$y$Target Encoder $f_{\bar\theta}$$z_y$$\|\hat{z}_y - z_y\|_{2}$

More like DINO, than JEPA

Instead of a regular JEPA like design, the pretraining of $\lambda$-JEPA is more like DINO design; multiple views of the inputs, their augmentations, and a projection of both embeddings. Instead of predicting target embeddings from context ones, they prefer an alignment loss between their projected latent embeddings.

$x_{l}$Encoder $f_\theta$$h_{x_{l}}$Projector $g_\psi$$z_{x_{l}}$$x_{g}$Encoder $f_{\bar\theta}$$h_{x_{g}}$Projector $g_\psi$$z_{x_{g}}$$\|z_{x_{g}} - z_{x_{l}}\|_{2}$

$\lambda$-JEPA tries to solve the collapse issue by directly regularizing representations with their scale, variance, and dimensional collapse prevention. Specificaly, in their proposed regularization loss, they have 3 penalties: mean, variance, and spectral.

$$\begin{aligned} \mathcal{R}_{\text{SAC}}(\mu, C_h) &= \frac{\lambda_{\text{mean}}}{2}\|\mu\|_2^2 + \frac{\lambda_{\text{cov}}}{2}\operatorname{tr}(C_h) - \frac{\gamma}{2}\log\det(C_h + \epsilon I_{N_h}), \\ \mathcal{L}_{\text{SAC}} &= \mathcal{L}_{\text{task}} + \mathcal{R}_{\text{SAC}}(\mu, C_h) \end{aligned}$$

$\lambda$-Balance

In a simple linear toy example, when the input is projected with a double weighted metrices; $W_{1}$ and $W_{2}$, the spectral rank of these weight matrices can describe collapse or non-collapse scenarios. Specifically, for a simple projection of $\hat{y} = W_{2}W_{1}x$, the gram matrices difference $\Delta = W^{T}_{2}W_{2} - W_{1}W^{T}_{1}$ describes the spectral gap between two weight matrices. When $\Delta = \lambda I$, a network is called $\lambda$-balanced.

When this $\lambda < 0$, a non-collapse of $W_{1}$ projection is guaranteed, because $W_{1}W^{T}_{1} \succeq - \lambda I$, which also means all eigenvalues of $W_{1}W^{T}_{1}$ are positive, hence, all singular values of $W_{1}$ are positive. Having all positive singular value means there is no spectral collapse of $W_{1}$ encoding projections.

Linear SACReg

In a linear projection scenario, the loss is defined in the following way

$$\begin{aligned} \mathcal{R}_{\text{LSAC}}(W_{1}, W_{2}) &= \frac{\lambda_{\text{reg}}}{2}(\|W_{1}\|_{F}^{2} + \|W_{2}\|_{F}^{2}) - \frac{\gamma}{2}\log\det(W_{1}W^{T}_{1}), \\ \mathcal{L}_{\text{LSAC}} &= \mathcal{L}_{\text{task}} + \mathcal{R}_{\text{LSAC}}(W_{1}, W_{2}) \end{aligned}$$

However, looking at the objective, one can notice that anti-collapse regularization comes from the last component of $\mathcal{R}_{\text{LSAC}}$ loss, it tries to maximize $\log\det(W_{1}W^{T}_{1})$ spectral term, which means all singular values of $W_{1}$ will be pushed above (positive).

Non-Linear SAC

In a non-linear case ($\mathcal{R}_{\text{SAC}}$), as model weights can not be directly regularized like the linear case, they are replaced by the model’s predictions covariance matrix. Meaning, if the covariance is regularized in the same way, model avoids collapsing scenarios. However, there is a major question:

Is prediction level regularization enough, to propagate anti-collapsing capabilities into 
network's weights? Meaning, making spectral collapse of the model weights less likely.

To answer that question, we would need the eigenvalue spectrum of frozen encoder representations ($C_{h}$) across increasingly different test distributions.

paper: https://arxiv.org/abs/2609.35288

Footnotes

1. LeJEPA: Provable and Scalable Self-Supervised Learning Without the Heuristics. R. Balestriero, et al. ↩

2. VICReg: Variance-Invariance-Covariance Regularization for Self-Supervised Learning. A. Bardes, et al. ↩

3. VISReg: Variance-Invariance-Sketching Regularization for JEPA training. H. Wu, et al. ↩