Why Prototypes Collapse: Diagnosing and Preventing Partial Collapse in Prototypical Self-Supervised Learning

TL;DR

Partial prototype collapse undermines prototypical SSL by forcing prototypes toward redundant representations. The authors diagnose the root cause as the joint optimization of encoders and prototypes under a shared loss and propose a fully decoupled training framework that models prototypes as a Gaussian mixture updated via online EM. This decoupling eliminates collapse, yields diverse prototypes, and strengthens downstream transfer, including robustness to long-tailed data. Across extensive experiments, the decoupled approach consistently improves prototype diversity and representation quality, supporting decoupling as a principled, scalable path forward for prototypical SSL.

Abstract

Prototypical self-supervised learning methods consistently suffer from partial prototype collapse, where multiple prototypes converge to nearly identical representations. This undermines their central purpose -- providing diverse and informative targets to guide encoders toward rich representations -- and has led practitioners to over-parameterize prototype sets or add ad-hoc regularizers, which mitigate symptoms rather than address the root cause. We empirically trace the collapse to the joint optimization of encoders and prototypes, which encourages a type of shortcut learning: early in training prototypes drift toward redundant representations that minimize loss without necessarily enhancing representation diversity. To break the joint optimization, we introduce a fully decoupled training strategy that learns prototypes and encoders under separate objectives. Concretely, we model prototypes as a Gaussian mixture updated with an online EM-style procedure, independent of the encoder's loss. This simple yet principled decoupling eliminates prototype collapse without explicit regularization and yields consistently diverse prototypes and stronger downstream performance.

Figures (4)

• Figure 1: Uniformity of prototype representations. Prototypes are projected to $\mathbb{R}^2$ using principal component analysis (PCA). The resulting prototype distributions are visualized using (top) Gaussian kernel density estimation (KDE), and (bottom) the angular distributions are estimated with a von Mises-Fisher KDE with $\kappa=20$.
• Figure 2: (a) Traditional joint-embedding architectures with a prototypical formulation: encoders and prototypes are optimized jointly under the same loss, which can lead to shortcut learning and prototype collapse---prototypes converge to similar representations, reducing the effective representation space. (b) Our proposed solution: decouples the gradient flow to the prototypes and updates them with a separate objective, mitigating shortcut learning and preserving prototype diversity.
• Figure 3: (a) Unique prototypes versus $\epsilon$ (Definition \ref{['def:partial_collapse']}). At $\epsilon=0$ all initialized prototypes are counted as unique, whereas increasing $\epsilon$ enforces stricter criteria (e.g., at $\epsilon=0.5$ only prototypes separated by at least $60^{\circ}$ are considered unique). (b) k-NN performance on IN1k for CARP showcasing an improvement with decoupling.
• Figure 4: (a) Number of unique prototypes over epochs. Prototypical methods without explicit diversity mechanisms (e.g., DINO and CARL) exhibit substantial prototype collapse early in training. (b) Linear evaluation accuracy on ImageNet-1k. Higher prototype diversity generally correlates with stronger final performance, though some methods (e.g., DINO + KP) imposes a cost of slower early convergence.

Theorems & Definitions (1)

• Definition 2.1: Partial prototype collapse