Self-Supervised Representation Learning as Mutual Information Maximization

TL;DR

This work reframes self-supervised representation learning as mutual information maximization guided by the Donsker–Varadhan bound, yielding two principled optimization paradigms: Self-Distillation MI (SDMI) and Joint MI (JMI). It shows that architectural elements like stop-gradient, predictor networks, and regularizers are not ad hoc choices but natural consequences of the MI objective, with many existing SSRL methods mapping to one of the two paradigms. A Taylor-series derivation connects DV-based objectives to the Barlow Twins loss, unifying mean-variance surrogates with MI maximization. Empirically, canonical SDMI/JMI prototypes exhibit monotonic MI growth and competitive downstream performance, while ablations reveal the necessity of the marginal term and support the interpretation of predictors as marginal surrogates. The framework provides a principled roadmap for designing future SSRL methods tailored to the chosen optimization paradigm and highlights limitations and reproducibility efforts for rigorous evaluation.

Abstract

Self-supervised representation learning (SSRL) has demonstrated remarkable empirical success, yet its underlying principles remain insufficiently understood. While recent works attempt to unify SSRL methods by examining their information-theoretic objectives or summarizing their heuristics for preventing representation collapse, architectural elements like the predictor network, stop-gradient operation, and statistical regularizer are often viewed as empirically motivated additions. In this paper, we adopt a first-principles approach and investigate whether the learning objective of an SSRL algorithm dictates its possible optimization strategies and model design choices. In particular, by starting from a variational mutual information (MI) lower bound, we derive two training paradigms, namely Self-Distillation MI (SDMI) and Joint MI (JMI), each imposing distinct structural constraints and covering a set of existing SSRL algorithms. SDMI inherently requires alternating optimization, making stop-gradient operations theoretically essential. In contrast, JMI admits joint optimization through symmetric architectures without such components. Under the proposed formulation, predictor networks in SDMI and statistical regularizers in JMI emerge as tractable surrogates for the MI objective. We show that many existing SSRL methods are specific instances or approximations of these two paradigms. This paper provides a theoretical explanation behind the choices of different architectural components of existing SSRL methods, beyond heuristic conveniences.

Figures (7)

• Figure 1: Canonical forms of our proposed paradigms: (a) SDMI alternates updates between two encoders using stop-gradients, while (b) JMI jointly updates both views with shared gradients.
• Figure 2: Estimated MI over CIFAR10 training for SDMI-based (top row) and JMI-based (bottom row) methods, using three estimators (cos–DV, InfoNCE and JSD; left to right). Both paradigms exhibit consistent MI growth: SDMI curves feature early fluctuations before trending upward, while JMI estimates rise more uniformly, and to much higher levels.
• Figure 3: Embedding trajectories of the five Gaussian cluster centers. Opacity increases over training, showing how different methods progressively separate the clusters in embedding space.
• Figure 4: Cluster centers on the unit sphere showing how SDMI prototype encoders progressively separate them
• Figure 5: Estimated MI over training using cos–DV, InfoNCE, and JSD bounds for both SDMI methods (top) and JMI methods (bottom). All three estimators show approximately monotonic growth for all methods under both paradigms.