A Minimal Model of Representation Collapse: Frustration, Stop-Gradient, and Dynamics
Authors: Louie Hong Yao, Yuhao Li, Shengchao Liu
Organizations: Center for Soft Matter and Biological Physics, Virginia Tech, Blacksburg, VA 24061, USA · Department of Computer Science and Engineering, The Chinese University of Hong Kong, Hong Kong, China
Self-supervised representation learning is central to modern machine learning because it extracts structured latent features from unlabeled data and enables robust transfer across tasks and domains. However, it can suffer from representation collapse, a widely observed failure mode in which embeddings lose discriminative structure and distinct inputs become indistinguishable. To understand the mechanisms that drive collapse and the ingredients that prevent it, we introduce a minimal embedding-only model whose gradient-flow dynamics and fixed points can be analyzed in closed form, using a classification-representation setting as a concrete playground where collapse is directly quantified through the contraction of label-embedding geometry. We illustrate that the model does not collapse when the data are perfectly classifiable, while a small fraction of frustrated samples that cannot be classified consistently induces collapse through an additional slow time scale that follows the early performance gain. Within the same framework, we examine collapse prevention by adding a shared projection head and applying stop-gradient at the level of the training dynamics. We analyze the resulting fixed points and develop a dynamical mean-field style self-consistency description, showing that stop-gradient enables non-collapsed solutions and stabilizes finite class separation under frustration. We further verify empirically that the same qualitative dynamics and collapse-prevention effects appear in a linear teacher-student model, indicating that the minimal theory captures features that persist beyond the pure embedding setting.
Figures & tables
Figure 1: Overview of model collapse in representation learning. (a) Standard contrastive-based classification pipeline with fixed one-hot labels, where only data representations are learned. (b) Generative-based representation-learning pipeline used in this work, with learnable data and label embeddings. (c) Embedding model under Eq. ( 5 ) without frustration, where class-wise dynamics remain decoupled. (d) Same model with frustration induced by shared data embeddings, leading to competing alignment constraints across classes. (e) Architecture with projection heads and stop-gradient [Eq. ( 26 )], which modifies the effective gradient flow between branches. (f) Same architecture without stop-gradient, restoring symmetric coupling between embeddings. (g) Teacher–student pipeline where data embeddings are generated by a parametrized encoder. (h) Teacher–student model with a learned encoder, projection head, and stop-gradient, extending the mechanism beyond the embedding-only setting.
Figure 2: Training dynamics of the classification–representation model on MNIST and CIFAR-10. In both experiments, we use an embedding dimension of 10, batch size 128, and stochastic gradient descent without weight decay or other regularization; the learning rate is 0.05 for MNIST and 0.01 for CIFAR-10. Left: training accuracy as a function of epochs for MNIST with LeNet and CIFAR-10 with ResNet-18. Right: corresponding training loss and minimal L2 distance between label embeddings ( MinL2 ), both normalized by their values at epoch zero and shown on a logarithmic scale, highlighting a fast accuracy-improvement regime followed by a slower collapse-dominated regime.
Figure 3: Effect of stop-gradient on representation collapse. Training dynamics on MNIST with LeNet under standard optimization and stop-gradient. Top: training accuracy (left) and loss (right). Bottom: minimal L2 distance between label embeddings before (left) and after (right) projection. Stop-gradient preserves early accuracy improvement while suppressing the decay of inter-label distances, stabilizing non-collapsed representations, and preventing the collapse-dominated regime.
Figure 4: Mean inter-class deviation ⟨∣δvi∣⟩ as a function of training step for unfrustrated runs at learning rate 0.5 . The vertical axis is shown on a logarithmic scale. All curves relax rapidly to small but nonzero plateaus, indicating stable class separation rather than complete collapse.
Figure 5: Training dynamics of the frustrated representation model, showing a clear separation of time scales. Curves show the square root of the loss ℓ , the mean magnitude of sample-level deviations ⟨∣δui∣⟩ , and class-level deviations ⟨∣δvi∣⟩ on a logarithmic scale. Top: n=2 , d=5 , r=0.05 . Bottom: n=10 , d=2 , r=0.10 . All experiments use N=1000 and a learning rate of 0.5 . Fast sample-level relaxation is followed by slow, frustration-controlled inter-class dynamics.
Figure 6: Two-stage relaxation of the training loss in the frustrated representation model. We plot the square root of the mean-squared loss, ℓ , using the fast rescaled time τ1 in (a) and the slow rescaled time τ2 in (b). The early-time decay in (a) reflects rapid sample-to-class alignment, while the long-time relaxation in (b) is controlled by frustration-induced inter-class coupling. Curves correspond to different choices of the number of classes n , embedding dimension d , frustration level r , and effective learning rate γ~ , as indicated in the legend.
Figure 7: Training dynamics of the minimal model with projection head and stop-gradient, shown for different numbers of classes n , embedding dimensions d , and frustration ratios r (legend). Top: the minimum pairwise distance between label embeddings, MinL2 . Bottom: the mean magnitude of label-embedding deviations, ⟨∣δvi∣⟩ . Both panels use a logarithmic y -axis and share the same y -range for direct comparison; the horizontal axis denotes training steps (in units of 102 ).
Figure 8: Training dynamics of the minimal model with projection head but without stop-gradient, shown for different numbers of classes n , embedding dimensions d , and frustration ratios r (legend). Top: the minimum pairwise distance between label embeddings, MinL2 . Bottom: the mean magnitude of label-embedding deviations, ⟨∣δvi∣⟩ . Both quantities rapidly decay toward zero, indicating convergence to a fully collapsed configuration in which label embeddings become indistinguishable. The horizontal axis denotes training steps.
Figure 9: Schematic training trajectories in the minimal model with a single frustrated sample. Shown are the 2D paths of the projected class label embeddings Wvi for n=5 classes (colored curves) together with the frustration embedding s (black dashed curve), with embedding dimension d=2 and N=50 samples per class. The single frustrated sample corresponds to a frustration ratio r=1/50 . Positions are recorded every 1000 optimization steps (open circles), and the final state is marked by a cross.
Figure 10: Eigenvalue spectra of W2 at the end of training (100,000 steps) for the stop-gradient (SG) runs. Each row corresponds to one setting (n,d,r) , and markers denote the eigenvalues of W2 . Across settings, the spectrum concentrates near a small number of well-separated values, with an eigenvalue near 0 and the remaining eigenvalues grouped near 1 and 1−r .
Figure 11: Two-dimensional visualizations of the synthetic teacher-student dataset in d=64 , where different colors denote the two classes. Left: PCA projection. Right: random affine map to R2 .
Figure 12: Training dynamics of the linear student-label embedding model with embedding dimension 5 at learning rate 0.005 for different frustration rates r . (a) ℓ as a function of training steps on a log scale. (b) Rescaled loss ratio ℓ/ℓr=0 versus the effective time τ .
Figure 13: Training dynamics of the model with projection head and stop-gradient in the linear teacher–student model for different frustration rates r . Left: separation between the learned label embeddings, ∥v0−v1∥2 , versus training steps. Right: classification accuracy versus training steps.
Self-supervised learning (SSL) has revolutionized representation learning, with Joint-Embedding Architectures (JEAs) emerging as an effective approach for capturing semantic features. Existing JEAs rely on implicit or explicit batch interaction -- via negative sampling or statistical regularization -- to prevent representation collapse. This reliance becomes problematic in regimes where batch sizes must be small, such as high-dimensional scientific data, where memory constraints and class imbalance make large, well-balanced batches infeasible. We introduce IConE (Instance-Contrasted Embeddings), a framework that decouples collapse prevention from the training batch size. Rather than enforcing diversity through batch statistics, IConE maintains a global set of learnable auxiliary instance embeddings regularized by an explicit diversity objective. This transfers the anti-collapse mechanism from the transient batch to a dataset-level embedding space, allowing stable training even when batch statistics are unreliable, down to batch size 1. Across diverse 2D and 3D biomedical modalities, IConE outperforms strong contrastive and non-contrastive baselines throughout the small-batch regime (from B=1 to B=64) and demonstrates marked robustness to severe class imbalance. Geometric analysis shows that IConE preserves high intrinsic dimensionality in the learned representations, preventing the collapse observed in existing JEAs as batch sizes shrink.
Konstantinos Almpanakis, Anna Kreshuk
European Molecular Biology Laboratory, Heidelberg, Germany
A central problem in self-supervised learning is preventing representation collapse. Most methods avoid it through global mechanisms, such as contrastive expansion, variance constraints, decorrelating dimensions, or enforcing certain output distributions. In this work, we study a different design: short-range repulsion. We introduce Hypersolid, a self-supervised objective that combines view alignment with local collision avoidance. Our method induces a latent geometry of compact, semantically aligned neighborhoods with low anisotropy. This geometry is especially effective for unsupervised clustering and fine-grained separation, although it comes at the cost of weaker transferability.
Esteban Rodríguez-Betancourt, Edgar Casasola-Murillo
Posgrado en Computación e Informática, Universidad de Costa Rica, San José, Costa Rica · Escuela de Ciencias de la Computación e Informática, Universidad de Costa Rica, San José, Costa Rica
We argue that learning visual representations without labels requires a training signal jointly complete across three non-overlapping objectives: semantic invariance across augmented views, patch-level spatial prediction, and representational non-degeneracy. We formalize these as the observation, prediction, and regularization principles and prove (i) that combining observation and prediction without regularization admits the constant encoder as a global minimizer under negative-free alignment; (ii) that the two objectives are gradient-complementary and structurally non-conflicting at the encoder output; and (iii) that the momentum encoder converges to the same fixed point as the online encoder and provides no collapse guarantee at convergence. Contrastive alignment provides only self-limiting collapse resistance, formalized via an explicit gradient-decay argument. Dropping prediction withholds the spatial training signal by construction; dropping observation forfeits cross-view semantic invariance by construction; at the scale we study, no pair substitutes for the third. Every major self-supervised method is a special case of a single unified energy decomposition. We pair every theoretical claim with a controlled experiment, including a patch-retrieval evaluation for the spatial consequence of prediction.