Information theoretic underpinning of self-supervised learning by clustering

Abstract: Self-supervised learning (SSL) is recognized as an essential tool for building foundation models for Artificial Intelligence applications. The advances in SSL have been made thanks to vigorous arguments about the principles of SSL and through extensive empirical research. The aim of this paper is to contribute to the development of the underpinning theory of SSL, focusing on the deep clustering approach. By analogy to supervised learning, we formulate SSL as K-L divergence optimization. The mode collapse is prevented by imposing an optimisation constraint on the teacher distribution. This leads to normalization using inverse cluster priors. We show that using Jensen inequality this normalization simplifies to the popular batch centering procedure. Distillation and centering are common {heuristics-based} practices in SSL, {but our work underpins them theoretically.} The theoretical model developed not only supports specific existing successful SSL methods, but also suggests directions for future investigations.

• The paper introduces a formal SSL framework by minimizing the KL divergence between teacher and student clustering distributions.
• It derives closed-form solutions linking entropy maximization with centering operations to prevent mode collapse and achieve balanced representations.
• The study unifies heuristic normalization techniques used in methods like DINO and SwAV with rigorous statistical regularization for robust SSL.

Information Theoretic Foundations of SSL via Clustering

Overview

This work provides an information-theoretic formalization of self-supervised learning (SSL) by clustering, focusing particularly on the mechanisms underlying successful deep clustering methods in SSL frameworks. By recasting SSL as a distillation learning paradigm and optimizing the Kullback-Leibler (K-L) divergence between teacher and student cluster assignment distributions, the authors derive principled closed-form solutions for avoiding mode collapse. The analysis connects heuristic centering and normalization practices in state-of-the-art SSL to formal mathematical justification, elucidating their role in stable cluster assignment and balanced representation learning.

Formalization of SSL as Distillation via K-L Divergence

The authors construct an analogy to supervised distillation learning by introducing parallel teacher-student networks, each assigning data points to clusters with probability distributions $Q(y|\mathbf{z})$ (teacher) and $P(y|\mathbf{z})$ (student). SSL is then formulated as minimizing $J_{KL}(Q||P)$ over both distributions, with alternating-direction optimization updating both network parameters and cluster weights.

This approach highlights that unlike supervised settings where the teacher distribution is fixed and one-hot, SSL requires the teacher signal to be estimated and iteratively regularized to guide balanced learning. The K-L divergence optimization, thus, encapsulates not just cross-entropy alignment but entropy maximization for the teacher distribution—a critical mechanism for preventing mode collapse and ensuring uniform coverage over latent space clusters.

Regularization and Mode Collapse Mitigation

A central issue in SSL clustering is mode collapse, where the model assigns nearly all points to a subset of clusters, trivializing the representation and undermining generalization. The paper introduces entropy-based regularization enforcing near-uniform cluster assignment priors, mathematically formalized as maximizing the entropy of cluster assignments.

The regularized objective includes terms for cross-entropy alignment, negative teacher entropy, and normalization constraints for the teacher signal. The key insight is that the closed-form solution for the teacher update involves scaling by the inverse cluster priors: $Q(y|\mathbf{z}) = \frac{P(y|\mathbf{z})}{Q(y)}$ where $Q(y)$ is the empirical prior over cluster assignments. This scaling counteracts imbalance: underpopulated clusters are boosted, while overpopulated ones are dampened, effectively distributing representations across the latent space.

Connection to Centering and Heuristic Normalization

A major contribution of the work is establishing a rigorous bridge between scaling by inverse cluster priors and empirical centering operations commonly used in batch normalization approaches. Applying Jensen's inequality, the normalization can be approximated by centering the embeddings—subtracting the mean of the data representations: $$Q(y|\mathbf{z}) \propto \exp \left( \frac{(\mathbf{z} - \bar{\mathbf{z}})^\top \mathbf{w}_y}{\tau_T} \right)$$ where centering the cluster weight vectors further stabilizes the normalization.

This formal link provides a theoretical rationale for batch centering and normalization, grounding previously heuristic techniques used in methods like DINO, DeepCluster, and SwAV. The analysis demystifies why these normalization strategies prevent collapse and enhance learning dynamics by ensuring balanced cluster assignment and robust teacher-student signal propagation.

Comparative Analysis with Existing SSL Clustering Frameworks

The information-theoretic approach is compared to deep clustering architectures and regularization mechanisms used in major SSL methods:

• DeepCluster/SwAV: Online clustering with regularization via cluster assignment weights, empirically shown to prevent mode collapse, is here shown to be theoretically equivalent to entropy regularization and normalization by cluster priors.
• DINO and Batch Centering: The analysis proves that centering operates as an approximation to scaling by inverse priors, thus theoretically justifying its empirical success.
• Contrastive vs. Non-Contrastive Clustering: The formal K-L framework subsumes both branches, highlighting that entropy maximization and balanced teacher signal are central to all successful SSL clustering.

The authors emphasize that while practical approximations may introduce noise (especially in small batches or extreme temperature settings), the theoretical formulation provides bounds and stability guarantees for SSL clustering.

Numerical and Empirical Claims

While the paper is theoretical, it cites strong empirical findings from prior methods such as DeepCluster, DINO, and Amrani et al., indicating that normalization by cluster priors does not lead to mode collapse and consistently yields optimal or near-optimal performance in large-scale SSL pre-training [Amrani2022]. The authors underline that centering offers computational efficiency and practical robustness, though normalization by the inverse cluster prior can in some settings outperform centering as measured experimentally.

Implications and Future Directions

The information-theoretic formalization of SSL clustering has several far-reaching implications:

• Unified Theory: Bridging practical heuristics and formal statistical regularization brings clarity for designing new SSL algorithms, ensuring balanced representation learning and preventing mode collapse.
• Parameter Selection: The framework suggests new strategies for parameter tuning (e.g., cluster number $K$, temperature $\tau$) and optimizing update frequency/momentum in teacher-student networks.
• Extensions: Future research can explore the relationship between cluster weight centering, dimensionality effects, concept fragmentation, and merging, as well as hybrid approaches integrating transport optimization.
• Foundation Models: The theoretical guarantees for balanced representation reinforce the role of SSL clustering as a cornerstone for scalable, transferable foundation models for computer vision and other domains.

Conclusion

This paper delivers an authoritative information-theoretic underpinning for self-supervised learning via clustering, rigorously explaining the necessity of entropy maximization and normalization in preventing collapse and achieving balanced, transferable representations. The formal connection between normalization by cluster prior and centering provides a theoretical rationale for common SSL practices, with direct implications for future research in representation learning. These results clarify the learning dynamics in large-scale SSL clustering and promise to guide principled algorithm design and parameter tuning in next-generation foundation models (2605.11870).