Self-Supervised Representation Learning as Mutual Information Maximization

2025-10-01

AI Summary

This work addresses the lack of theoretical foundations in self-supervised representation learning (SSRL). Starting from variational lower bounds on mutual information (MI), we establish a unified theoretical framework that introduces two canonical paradigms: Self-Distillation MI Maximization (SDMI) and Joint MI Maximization (JMI). For the first time, our framework rigorously justifies—rather than heuristically motivates—the necessity of prediction networks, stop-gradient operations, and statistical regularization terms. We prove that prominent methods including SimCLR, BYOL, and SwAV are either exact instances or approximations of SDMI or JMI. Furthermore, we characterize the intrinsic mechanism underlying representation collapse and derive principled design guidelines for SSRL architectures. Our analysis achieves a theoretically grounded, interpretable unification of SSRL algorithmic paradigms, component functionalities, and optimization trajectories.

Technology Category

• Machine Learning: Unsupervised & Self-Supervised Learning
• Computer Vision: Representation Learning for Vision
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs
• Responsible Web: Machine-in-the-loop, human agency and autonomy

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.

Problem

• Unifying self-supervised learning via mutual information principles
• Explaining architectural choices like stop-gradient theoretically
• Deriving two training paradigms covering existing SSRL methods

Innovation

• Self-Distillation MI requires alternating optimization with stop-gradient
• Joint MI enables symmetric architectures without stop-gradient operations
• Predictor networks and statistical regularizers approximate mutual information objectives