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.
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