A Minimal Model of Representation Collapse: Frustration, Stop-Gradient, and Dynamics

TL;DR

This paper introduces a minimal model to analyze representation collapse in self-supervised learning, showing how frustration and stop-gradient influence the dynamics and stability of embeddings.

Contribution

It provides a theoretical framework with closed-form analysis of collapse mechanisms and demonstrates how stop-gradient prevents collapse, validated by empirical models.

Findings

Collapse occurs when data are not perfectly classifiable.

Stop-gradient stabilizes non-collapsed solutions.

Frustrated samples induce collapse through slow dynamics.

Abstract

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.