Self-Supervised Learning from Structural Invariance

Abstract: Joint-embedding self-supervised learning (SSL), the key paradigm for unsupervised representation learning from visual data, learns from invariances between semantically-related data pairs. We study the one-to-many mapping problem in SSL, where each datum may be mapped to multiple valid targets. This arises when data pairs come from naturally occurring generative processes, e.g., successive video frames. We show that existing methods struggle to flexibly capture this conditional uncertainty. As a remedy, we introduce a latent variable to account for this uncertainty and derive a variational lower bound on the mutual information between paired embeddings. Our derivation yields a simple regularization term for standard SSL objectives. The resulting method, which we call AdaSSL, applies to both contrastive and distillation-based SSL objectives, and we empirically show its versatility in causal representation learning, fine-grained image understanding, and world modeling on videos.

• The paper introduces AdaSSL, which augments joint-embedding frameworks with latent variable modeling to capture stochastic, multimodal conditional invariance.
• It proves that mapping to curved embedding manifolds induces unavoidable heteroscedasticity, challenging standard SSL similarity functions.
• Empirical results show AdaSSL outperforms traditional SSL methods in tasks like world modeling, disentanglement, and robust image understanding.

Self-Supervised Learning from Structural Invariance: A Technical Examination

Motivation and Problem Statement

Joint-embedding self-supervised learning (SSL) methods have become standard for unsupervised representation learning in visual domains, frequently through contrastive or distillation-based protocols. The conventional recipe relies on crafting positive pairs—instances presumed to encode equivalent semantics via augmentations like cropping or color jitter. However, such synthetic augmentations fail to replicate structured changes in real-world data-generating processes, leading to inadequate modeling of semantic invariance and excessive information loss. This paper identifies fundamental limitations in prevailing SSL methods: their inability to flexibly model one-to-many mappings, multimodal or heteroscedastic conditional distributions present in naturally paired data such as successive video frames or related image-caption pairs.

The work poses the technical question of how to enable SSL architectures to correctly learn the structural invariance underlying naturally paired data, where the conditional $p(x^+|x)$ can be stochastic and multimodal, reflecting real-world generative factors. This is formalized within a causal representation learning (CRL) lens, where the model must invert the underlying data-generating process to recover latent factors, even when positive pairs differ according to structured transformations mediated by latent variables.

Theoretical Contribution: Heteroscedasticity in Structural SSL

A significant theoretical advance presented in the manuscript is the proof that, under realistic conditions, there exists unavoidable heteroscedasticity in the conditional law between paired embeddings, irrespective of the encoding function or dimensionality. Specifically, mapping flat latent spaces to curved embedding manifolds (e.g., a hypersphere) necessarily induces input-dependent conditional variances, invalidating the implicit homoscedasticity assumptions in current SSL similarity functions and predictors. This observation, grounded in differential geometry, highlights the technical necessity of modeling latent factors that govern structured transitions between positive pairs, especially in natural data settings.

AdaSSL: Adaptive SSL via Latent Variable Modeling

The proposed solution is AdaSSL, an augmentation of joint-embedding SSL frameworks via explicit latent variable modeling. AdaSSL introduces a latent variable $r$—either variational or sparsity-regularized—that conditions the representation of one input based on another, capturing aspects of $x^+$ not predictable solely from $x$. This induces new machinery for SSL objectives: maximizing a variational lower bound on mutual information between paired embeddings, while regularizing the degree of information $r$ may encode.

Two principal variants are developed:

• AdaSSL-V: Employs variational inference using $q_\phi(r|x, x^+)$ with a regularizer based on KL divergence to $p_\theta(r|x)$, yielding a tractable lower bound on mutual information and enabling more expressive modeling of $p(x^+|x, r)$.
• AdaSSL-S: Implements a deterministic mask $m(f(x), f(x^+))$ with sparse regularization (via $L_0$ penalty), motivated by the empirical observation that natural transitions often correspond to sparse latent changes.

Both variants maintain compatibility with InfoNCE and distillation-based objectives, and crucially allow the similarity function to remain simple while expressing arbitrarily complex conditional distributions.

Figure 1: Adaptive SSL (AdaSSL) generalizes standard contrastive/distillation-based architectures by adding latent variable modeling, enabling flexible adaptation to structured conditional uncertainty.

Figure 2: Conditioning on latent variables transforms complex or multimodal conditional noise into tractable forms, e.g., removing irrelevant modes in temporal prediction tasks.

Empirical Evaluation

Comprehensive experiments validate the theoretical intuition and practical utility of AdaSSL:

Numerical Experiments

AdaSSL variants substantially outperform InfoNCE and AnInfoNCE, especially under complex multimodal and heteroscedastic conditional laws. Only methods that allow flexible modeling of conditional uncertainty sustain high regression fidelity on out-of-distribution tasks, as demonstrated in both unbounded and hyperspherical embedding spaces.

Causal Representation Learning (CRL) and Disentanglement

On the 3DIdent benchmark, AdaSSL recovers data-generating factors more accurately than both classic disentanglement methods ($\beta$-VAE, AdaGVAE) and SSL baselines. AdaSSL consistently exhibits superiority in both latent recovery ($R^2$) and DCI disentanglement scores, particularly when simpler editing functions are used, indicating more effective and efficient utilization of the latent variable.

Figure 3: AdaSSL-V enables conditional, controllable retrieval of images along interpretable latent dimensions that correspond to structured generative factors.

Natural Image Understanding

On the CelebA dataset, AdaSSL-V and AdaSSL-S outperform standard SSL protocols across both weak and strong augmentation regimes, particularly in probing fine-grained facial attributes and in generalizing to out-of-distribution settings. The performance gap is most pronounced when models are trained on naturally paired images as opposed to augmentation-based pairs. Importantly, AdaSSL achieves these improvements without explicit access to ground-truth transformation labels.

Robustness and Scalability

AdaSSL-V demonstrates increased robustness to noisy pairings in large-scale image experiments (iNat-1M), with more graceful degradation as structured pairing corruption increases. This highlights the practical utility in real-world settings where label noise and imperfect pairings are inevitable.

Figure 4: AdaSSL-V maintains higher classification accuracy than InfoNCE as the proportion of corrupted pairings increases, evidencing robustness to label and pairing structure noise.

World Modeling and Video Prediction

In world modeling tasks on stochastic Moving-MNIST, AdaSSL variants outperform baselines in both digit recognition and velocity decoding, especially under multimodal future uncertainty. Detailed analyses show that AdaSSL enables accurate prediction of both invariant and variant factors, with improved diversity in sampled future trajectories.

Figure 5: AdaSSL representations facilitate accurate decoding of both invariant (digit class) and variant (velocity) factors in stochastic video tasks.

Architectural Insights and Analysis

The empirical results align with the theoretical prediction: standard SSL models relying on fixed similarity functions cannot model conditional multimodality or input-dependent noise, leading to information collapse or misalignment under covariate shift. AdaSSL, by incorporating latent variable regularization and editing functions, enables flexible adaptation to conditional structure, preserves more content factors, and supports better downstream generalization.

Ablation studies further attest to the necessity and impact of the proposed mechanisms, such as the utility of additional surrogate views ($x^{++}$) to avoid shortcut solutions and achieve more semantic invariance, and the roles of regularization hyperparameters in latent sparsity and utilization.

Implications and Future Directions

This research has several important implications:

• Generalization and Robustness: The ability to explicitly model conditional uncertainty structures enables SSL models to generalize more robustly to real-world, structured latent changes and distribution shifts.
• Causal Discovery: AdaSSL bridges SSL and CRL paradigms, showing promise for scalable, weakly-supervised causal discovery in high-dimensional data.
• World Modeling: Adaptive SSL frameworks support more diverse and realistic predictive modeling in video domains, suggesting utility for autonomous planning, implicit action space discovery, and controllable generation.
• Scalable Extension: The approach is compatible with modern architectures (e.g., ResNet, ViT) and data modalities (images, videos, captions); future work can extend AdaSSL to web-scale or multimodal settings.
• Theoretical Grounding: The proof of necessary heteroscedasticity provides a rigorous foundation for further investigation into SSL objective design and representation identifiability.

Conclusion

This paper exposes a critical limitation of existing SSL paradigms in modeling the conditional uncertainty of naturally paired data and establishes, both theoretically and empirically, the necessity of latent variable modeling for structural invariance. The proposed AdaSSL method extends joint-embedding architectures with latent variable mechanisms, achieves strong performance gains across benchmarks (including CRL and world modeling), and demonstrates robustness under noise and distribution shift. This contributes a framework for more general, flexible, and causal self-supervised representation learning, with prospective impact in scalable unsupervised understanding, robust AI planning, and controllable generative modeling.