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

1. Problem Statement: Partial Prototype Collapse

The paper addresses a critical failure mode in Prototypical Self-Supervised Learning (SSL) frameworks (e.g., DINO, DINOv2, CARP). In these methods, a set of learnable vectors (prototypes) acts as cluster anchors to guide the encoder toward rich, diverse representations.

• The Phenomenon: The authors identify Partial Prototype Collapse, where multiple prototypes converge to nearly identical representations during training.
• The Consequence: This reduces the effective diversity of the embedding space, undermining the core purpose of prototypes (providing diverse targets). It forces practitioners to over-parameterize the number of prototypes or add ad-hoc regularizers (like KoLeo regularization) to mitigate symptoms rather than solving the root cause.
• Current Limitations: Existing solutions often rely on explicit regularization terms that introduce hyperparameter trade-offs (balancing diversity vs. encoder learning) or fail to fully resolve the issue in instance-level objectives.

2. Root Cause Diagnosis: Joint Optimization

Through systematic empirical analysis across various frameworks (DINO, DINOv2, CARP, CAPI, iBOT), the authors diagnose the root cause of the collapse:

• Joint Optimization Shortcut: The standard approach jointly optimizes the encoder ($f_\theta$) and prototypes ($C$) under a shared loss function ($L_f$).
• Mechanism: This joint optimization encourages a form of shortcut learning. Early in training, prototypes drift toward redundant representations that minimize the prediction loss without actually enhancing the diversity or semantic quality of the encoder's features.
• Evidence: The paper demonstrates that methods with partial decoupling (like CAPI) show better diversity than fully joint methods (like DINO), but only full decoupling eliminates the collapse entirely.

3. Methodology: Fully Decoupled Training

The authors propose a fully decoupled training strategy that isolates prototype estimation from encoder learning, breaking the joint optimization loop.

Core Concept

Instead of minimizing a joint loss $\min_{\theta, C} L_f(f_\theta, C)$, the training alternates between two separate objectives:

1. Prototype Update (Unsupervised Estimation): Update prototypes $C_t$ based on the latent features $h_\phi$ from the teacher, independent of the encoder's loss.
   $$C_t = \arg \min_{C} L_C(C_{t-1}, h_\phi)$$
2. Encoder Update: Update the encoder $\theta$ using fixed prototypes $C_t$ to minimize the consistency loss.
   $$\theta_{t+1} = \arg \min_{\theta} L_f(h_\theta, C_t)$$

Implementation: Online Gaussian Mixture Model (GMM)

To solve the prototype estimation problem (Eq. 3) efficiently and robustly, the authors model prototypes as the means of an Online Gaussian Mixture Model (GMM):

• EM-Style Procedure: Prototypes are updated using an Expectation-Maximization (EM) approach on the fly.
• Responsibility-Weighted Forgetting: To handle high-dimensional spaces and large prototype counts (e.g., 65k), they introduce a mechanism where the "forgetting factor" is modulated by the expected responsibility of each component. This prevents components with sparse updates from collapsing to zero.
• Split-Resurrect Regularization: To maintain diversity, if a component becomes too dominant, its mass is split and a "resurrected" component is initialized with random scale-aware parameters.
• No Gradient Flow: Crucially, gradients do not flow from the prototype update back to the encoder, eliminating the incentive for the encoder to "game" the prototypes.

4. Key Contributions

1. Systematic Diagnosis: The first comprehensive analysis showing that partial prototype collapse is a widespread phenomenon across the DINO family and other prototypical SSL methods, not just an artifact of specific architectures.
2. Mechanism Identification: Identifying joint optimization as the primary driver of collapse, explaining it as a shortcut learning phenomenon where prototypes converge to redundant states to minimize loss prematurely.
3. Decoupled Framework: Introducing a principled, fully decoupled training framework using an online GMM that eliminates prototype collapse without explicit regularization or complex hyperparameter tuning.
4. Theoretical Justification: Providing a theoretical perspective showing that decoupled fixed points (where prototypes are centroids of distinct responsibilities) almost surely prevent collapse, whereas joint optimization allows collapsed states to be stationary points.

5. Experimental Results

The authors evaluated their method (CARP + Decoupling) against baselines like DINO, DINOv2, and CARP with KoLeo regularization (KP).

• Prototype Diversity:
  • DINO/DINOv2: Suffered severe collapse (retaining only ~1.5% to 43% of unique prototypes depending on the threshold).
  • CARP + Decoupling: Maintained 100% unique prototypes across all tested thresholds ($\epsilon$), effectively eliminating collapse.
• Training Dynamics:
  • Collapse in joint methods occurs very early (within 10 epochs).
  • The decoupled method maintains full prototype utilization throughout training.
• Downstream Performance:
  • k-NN Accuracy: CARP + Decoupling improved k-NN accuracy on ImageNet-1k by +1.8% over standard CARP.
  • Long-Tailed Robustness: On the iNaturalist-18 dataset (highly imbalanced), the decoupled method improved accuracy across Head, Medium, and Tail classes by +3.0% overall. This confirms that diversity helps model tail classes, not just head classes.
  • Transfer Learning: Performance was comparable to or slightly better than baselines on standard transfer benchmarks (Aircraft, Cars, etc.).
• Efficiency:
  • Memory: The decoupled approach reduces peak GPU memory usage (by up to 1.5GB at large batch sizes) because it avoids storing the computation graph for prototype optimization.
  • Time: Negligible increase in training time (~0.2 hours for 100 epochs).

6. Significance and Conclusion

This work fundamentally shifts the paradigm of prototypical SSL from joint optimization to decoupled estimation.

• Theoretical Insight: It proves that the "shortcut" of collapsing prototypes is a direct result of the optimization landscape created by joint training.
• Practical Impact: The proposed method offers a simple, scalable, and principled way to prevent collapse without the need for ad-hoc regularizers or massive over-parameterization.
• Robustness: By ensuring diverse prototypes, the method significantly improves model robustness to imbalanced data distributions, a critical requirement for real-world applications.

The authors conclude that breaking the joint optimization loop is essential for unlocking the full potential of prototypical self-supervised learning, yielding richer representations and more stable training dynamics.