Unifying the SSL Zoo: Why Latent Distribution Matching is the Secret Sauce of Representation Learning
Understanding Self-Supervised Learning via Latent Distribution Matching
This paper introduces Latent Distribution Matching (LDM) as a unifying theoretical framework for Self-Supervised Learning (SSL), achieving SOTA results by aligning learned representations with assumed latent models. The authors demonstrate that contrastive, non-contrastive, and predictive SSL methods (like SimCLR, VICReg, and BYOL) can all be derived as specific instances of maximizing the log-probability of data under a latent model while enforcing representational uniformity through entropy maximization.
TL;DR
Why do seemingly different Self-Supervised Learning (SSL) methods—like the contrastive SimCLR, the regularization-heavy VICReg, and the predictive BYOL—all work so well? This paper argues they are all doing the same thing: Latent Distribution Matching (LDM). By framing SSL as a task of matching data to a latent statistical model, the authors unify the field, provide proofs for the "identifiability" of latent factors, and even derive a new, uncertainty-aware Kalman-SSL model.
The "Mutual Information" Mirage
For years, the community has leaned on Mutual Information (MI) maximization to explain SSL. The idea was simple: maximize the information shared between two views of the same image. However, the authors point out a glaring theoretical flaw: MI is invariant under any invertible transformation. You could scramble your latent space into total nonsense, and as long as the scramble is reversible, the MI remains identical.
The real hero, they argue, isn't MI; it's Latent Entropy. By maximizing the entropy of the representations while aligning them to a prior, LDM forces the encoder to become invertible on the data manifold, effectively "unfolding" complex data into a structured latent space.
Methodology: The Unified LDM Framework
The core proposal is a deceptively simple objective function:
The authors show that by tweaking the Latent Model and the Entropy Estimator, you can derive almost every major SSL algorithm:
- VICReg: LDM with a Gaussian conditional model and parametric entropy.
- SimCLR: LDM on a hypersphere (vMF distribution) with Kernel Density Estimation (KDE).
- Predictive SSL (JEPA/BYOL): LDM for temporal data where the entropy is implicitly handled by a "stop-gradient" operator.
Figure 1: LDM unifies alignment and uniformity by matching the transformed data distribution R(z, z') to a latent model Pθ.
One of the most exciting extensions is the Kalman-based SSL. By assuming the latent space follows linear dynamics, the authors could use a Kalman filter to predict future states. Unlike traditional SSL, this allows the model to output uncertainty, telling us not just what it thinks is happening, but how sure it is.
Provable Identification: Can We Recover Reality?
A major contribution of this paper is Theorem 1, which proves that under a Gaussian predictor, LDM-trained models recover the true underlying latent variables up to a simple affine transformation (scaling/rotation).
The intuition is that forcing prediction errors into a Gaussian shape "linearizes" the relationship between the raw data (like pixels) and the true latent factors (like the position of an object).
Figure 2: Predictive LDM successfully recovers the true latent space of a non-linear dynamical system up to an affine transformation.
Experiments: Performance and Insights
The authors tested their framework on image datasets (CIFAR, ImageNet) and biological data (rat hippocampal spikes).
- Image Recognition: They found that maximizing MI (the extra term in many SSL papers) actually contributed almost nothing to the final accuracy. It was the choice of entropy estimator and latent space (Plane vs. Sphere) that dictated performance.
- Neural Spiking Data: The Kalman-SSL model successfully mapped rat brain activity to a circular track, accurately predicting the rat's position and providing confidence intervals that tracked with the "noise" in the neural data.
Table 1: Comparison of different entropy estimators and MI maximization across datasets.
Critical Insight: Discarding vs. Reparameterizing
Standard SSL wisdom says we should "discard irrelevant info." The LDM framework suggests the opposite: SSL doesn't perform lossy compression. Instead, it exerts fine-grained control over how the data manifold is reparameterized. It doesn't throw away the pixels; it just reshapes the geometry so that high-level concepts are "untangled" and linearly accessible to a predictor.
Conclusion & Future Work
This paper provides the strongest theoretical glue we have for the fragmented world of SSL. It tells us that our choice of "Prior" is our most powerful inductive bias. The bottleneck now? Entropy Estimation. As we move to higher-dimensional latent spaces, we need more efficient ways to estimate and maximize representation entropy without the high computational cost of K-Nearest Neighbors or the biased assumptions of Log-Determinant methods.