Predictive Embeddings, JEPA, and Energy-Based Models › Collapse Prevention and Target Encoders

Everything converges on one question. What stops the encoder from outputting a constant?

Complete collapse is the extreme case: $f_\theta(o) = \mathbf{v}$ for all $o$, for a fixed vector $\mathbf{v}$. The prediction loss is exactly zero. The representation carries no information at all, and any downstream probe will fail. Complete collapse is easy to detect because the per-dimension variance is zero for every dimension.

Dimensional collapse is subtler. The embedding retains variation but occupies an effectively lower-dimensional subspace; strong coordinate correlations can be one symptom, but correlation alone is not the definition. If a nominally 64-dimensional latent has entropy effective rank 3, it uses roughly three dimensions' worth of spectral capacity, a deficit of 61 relative to uniform use of all 64. That does not mean exactly 61 singular directions are zero. The loss can still fall smoothly while effective dimensionality disappears.

Diagnostics matter here. Cheap moments can run every few epochs on a fixed validation subset; approximate effective-rank checks and trained probes can run at a cadence their cost permits. Useful signals include:

• Per-dimension standard deviation. After fixing the representation scale, a coordinate whose standard deviation is near zero over representative data is nearly constant. Plot every coordinate because a mean can hide coordinate-wise failure.
• Effective rank. When the singular values are not all zero, center the embedding matrix, normalize its singular values as $\tilde\sigma_i = \sigma_i / \sum_j \sigma_j$, and compute $\exp\big(-\sum_i \tilde\sigma_i \log \tilde\sigma_i\big)$, using the continuous-extension convention $0\log 0=0$ for zero components. The normalized spectrum is undefined when every singular value is zero; handle that case separately as complete collapse. Otherwise, effective rank is a smooth measure of spectral dimensionality, not a count of exactly nonzero directions.
• Linear probe performance on a held-out variable the representation should contain, such as the true state or a task label. It tests linear decodability under the chosen split and probe. Poor performance does not prove that the information is absent, and high performance does not establish downstream usefulness or causal relevance.
• Uniformity and alignment, if you are working in the contrastive tradition: how evenly the embeddings spread on the unit sphere, and how close positive pairs are.

The distinction between "non-collapsed" and "useful" is the one that matters. A representation can have full effective rank and still fail at control. Effective rank and variance diagnose degeneracy; a held-out linear probe tests access to one chosen variable. Downstream usefulness ultimately requires task-level evaluation.

If the loss contains terms that explicitly push decoys away, the constant solution is not optimal whenever the model can score positives above sampled negatives. A constant embedding makes all $K+1$ candidate scores equal, so the per-example InfoNCE loss is $\log(K+1)$, the chance-level value. It is not the largest possible loss, because assigning a positive pair a much worse score than its negatives can make the loss arbitrarily larger. The price is the cost of negatives and sensitivity to the sampling distribution. In practice this can motivate large batches, memory banks, or hard-negative mining.

One widely used asymmetric recipe combines a predictor, stop-gradient, and an EMA target. The components interact, but no single three-part recipe is universally necessary: SimSiam, for example, uses a predictor and stop-gradient without an EMA target.

In the simplified centered linear regression above, a sufficiently capable predictor can absorb part of the context-to-target map. Direct unnormalized matching can favor shrinkage in directions where that map is contractive, but omitting a predictor does not universally imply shrinkage.

The stop-gradient breaks the symmetry of the gradient updates while leaving the forward scalar loss unchanged. Without it, the same parameters can receive updates through both encoding branches and move cooperatively toward a constant solution. With it, the online encoder cannot directly update the target representation through that loss.

The EMA target encoder controls the target timescale. A permanently frozen target would reduce training to regression toward fixed, randomly initialized target features, which need not encode useful invariances. Copying online weights into the target after every update would preserve stop-gradient within each step yet remove EMA smoothing. A larger $\tau$ lengthens the nominal smoothing timescale and gives less weight to the current online parameters; a smaller value shortens that timescale and gives them more weight. Individual target-step magnitudes and the pointwise online-target gap also depend on the evolving online trajectory, so neither comparison guarantees monotone realized changes or pointwise closeness. A smaller value may reduce stability or representation quality in recipes that rely on a slow target. Whether this appears as collapse depends on the other mechanisms.

The third family adds explicit statistical pressure. VICReg penalizes per-dimension standard deviations below a threshold and off-diagonal covariance entries. Barlow Twins pushes the cross-correlation matrix between two augmented-view embeddings toward the identity. Whitening methods such as W-MSE normalize embedding covariance using a transformation such as Cholesky, eigendecomposition, or an iterative inverse square root. Sinkhorn normalization instead balances assignment matrices in methods such as SwAV; it is not covariance whitening. These losses are broadly composable, but their effectiveness still depends on batch statistics, normalization, dimension, loss weights, and the surrounding objective.

The hinge form of the variance term is important. Penalizing deviation from a target standard deviation would be symmetric and would fight the encoder's natural scale, whereas the hinge $\max(0, \gamma - \mathrm{std}(z_j))$, where $\gamma>0$ is the target standard-deviation threshold and $z_j$ is the $j$-th embedding coordinate, only fires below that threshold and is silent otherwise. This is what makes the regularizer compatible with a wide range of architectures. A dimension at or above the threshold contributes nothing to the penalty, so the regularizer does not impose an upper target scale on the representation; it only refuses to tolerate dimensions below the chosen minimum.

There is no single best answer, but the tradeoffs are clear enough to guide a first choice.

• If you have a natural contrastive structure, an explicit-negative objective makes its anti-collapse pressure easy to inspect, but requires a sampling scheme and candidate scores.
• If you want to avoid explicit negatives and memory banks, an asymmetric recipe is an option. Its target encoder and predictor still add computation, and successful recipes depend on their surrounding architecture and optimization.
• If you want explicit anti-collapse pressure that composes with many architectures, add variance and covariance regularization when batch statistics and loss weights are reliable.
• Combining mechanisms is common. The full toy variant below uses an EMA target and a variance hinge; that combination is stable in this seeded setup, not a general optimum for small datasets.

Now let's build one and watch the failure mode happen.