Self-Supervised Learning from Structural Invariance

AI summary

Overview

Research area: self-supervised representation learning (joint-embedding SSL), with connections to causal representation learning, fine-grained visual understanding, and video world modeling. Technical level: Intermediate. Scope: this paper diagnoses why standard SSL objectives cannot represent the one-to-many, heteroscedastic conditional distributions induced by naturally paired data, and proposes a latent-variable fix, AdaSSL, with two concrete variants.

What This Paper Is About

Joint-embedding SSL learns by making the representations of semantically related pairs similar, but it implicitly assumes that the relationship between the two members of a pair is simple — either nearly invariant (augmented views) or unimodal with one fixed noise scale. When positive pairs come from natural generative processes, such as successive video frames or image–caption pairs, the mapping from one member to the other is one-to-many: a present state can lead to several plausible futures, and caption detail varies with image complexity. The goal of this paper is to let an SSL model represent that conditional uncertainty instead of discarding the features responsible for it.

Key Contributions

1. A framing called "SSL from structural invariance." The authors formulate SSL as learning the conditional distribution p(z⁺ | z) induced by a data generating process (z ~ p(z), z⁺ | z ~ p(z⁺ | z), x = g(z), x⁺ = g(z⁺)), and argue that modelling this structure lets SSL recover more diverse and generalizable features.

2. A theory result that heteroscedasticity is unavoidable (Proposition 2.1). For a C¹ diffeomorphic mixing function g, an almost-everywhere C¹ encoder f mapping into the unit sphere 𝕊^{d_f}, and additive latent noise ε with 𝔼[ε] = 0 and Cov(ε) ≻ 0, the conditional law p(h(z⁺) | h(z)) of the composed map h = f ∘ g is necessarily heteroscedastic: its variance depends on h(z) for p_z-almost every z, regardless of the encoder or embedding dimensionality.

3. AdaSSL, a general objective. A latent variable r parameterizes conditional uncertainty, giving ℒ_AdaSSL = ℒ_SSL((x, r), x⁺) + β ℒ_Reg(r), where ℒ_SSL can be InfoNCE-style contrastive loss or a distillation loss such as BYOL, and ℒ_Reg prevents r from becoming a shortcut that encodes x⁺ directly.

4. Two instantiations. AdaSSL-V (variational) infers r with a variational posterior q_φ(r | x, x⁺) and adds a KL term against a conditional prior p_θ(r | x), yielding a tractable lower bound on I(x; x⁺); AdaSSL-S (sparse) predicts r deterministically with an MLP followed by tanh and penalizes its L0 norm, using a modular editing function built from rank-1 matrices.

Alongside these, the authors introduce H-InfoNCE, a baseline extending AnInfoNCE with a data-dependent weighting matrix Λ_x, to test whether simpler heteroscedastic models suffice.

Main Findings

• Existing SSL methods break under heteroscedastic conditionals. In the numerical linear-regression benchmark (Table 1), plain InfoNCE reaches R² of 0.8553 ± 0.0532 in-distribution but drops to 0.2664 ± 0.0984 under a covariate shift and to −0.1891 ± 0.2545 out-of-distribution (unbounded model space). AnInfoNCE is comparable (0.8447 ± 0.0611 in-distribution, −0.2277 ± 0.3284 OOD). The caption states that values below 0.7 are dimmed.

• A flexible conditional model recovers much of the gap. H-InfoNCE MLP, which predicts Λ_x with an MLP and uses a separate MLP predictor, reports 0.9892 ± 0.0023 in-distribution and 0.9149 ± 0.0348 OOD on the same heteroscedastic setting; H-InfoNCE Affine reports 0.9826 ± 0.0060 and 0.8666 ± 0.0741.

• Unimodal settings still expose OOD fragility. For the simplest conditional, the reported rows show InfoNCE at 0.9912 ± 0.0051 in-distribution yet 0.2683 ± 0.2626 on the hypersphere OOD column, while the H-InfoNCE Affine row reports 0.9154 ± 0.0100 there.

• Natural pairs help on 3D-rendered data. On 3DIdent, AdaSSL identifies latent factors better than all baselines (β-VAE and AdaGVAE are included as classic disentanglement baselines).

• Fine-grained image understanding improves. On CelebA, AdaSSL captures fine-grained features and learns more disentangled representations than baselines including LieSSL, and generalizes to OOD data. On iNat-2021 it is reported to be more robust to noisy data pairings than vanilla SSL.

• Video world modeling retains stochastic detail. On an extended Moving-MNIST dataset with BYOL as the base method, AdaSSL captures stochastic object accelerations that baselines discard without sacrificing class accuracy.

• The specific AdaSSL numbers from most benchmark tables are not included in the excerpt provided. Only the numerical Table 1 values listed above appear; the paper's claim that AdaSSL consistently outperforms baselines on multiview CRL, fine-grained image understanding, and world modeling is stated in the contributions but the corresponding figures are not present in the supplied content.

Methodology in Plain English

The authors start from the observation that mutual information between paired embeddings can be decomposed by the chain rule as I(f(x); f(x⁺)) = I((f(x), r); f(x⁺)) − I(r; f(x⁺) | f(x)). Introducing an auxiliary variable r splits the problem into two parts: an SSL term that uses r to make prediction of f(x⁺) easier, and a regularizer that keeps r from simply memorizing f(x⁺).

For AdaSSL-V, r is inferred by a small network q_φ that sees both members of the pair, a prior p_θ(r | x) is also learned, and the two are tied together with a KL divergence — the familiar variational trick — which produces a tractable lower bound on the mutual information. The similarity function stays a plain dot product, but the first embedding is "edited" by r through a function t, keeping feature extraction simple.

For AdaSSL-S, r is predicted deterministically and pushed toward sparsity with an L0 penalty made differentiable via a Gumbel-Sigmoid estimator, on the assumption that natural transitions change only a few latent factors. Editing is modular: each active component of r applies a small rank-1 affine transformation plus an offset to the embedding.

Evaluation uses synthetic numerical data where the ground-truth latents are known and where the conditional variance is systematically made more complex (unimodal, anisotropic, then heteroscedastic), plus 3DIdent, CelebA, iNat-2021 and extended Moving-MNIST. Encoders are a five-layer MLP for numerical data, ResNet-18 with a two-layer MLP projector for image benchmarks, ResNet-50 with a two-layer MLP projector for iNat-2021, and a five-layer 3D CNN with a three-layer MLP projector for video. Models are trained from scratch, selected using an online affine probe on validation data, and evaluated with affine probes on frozen embeddings or representations; the paper reports mean and standard deviation over three random seeds.

Why This Matters

The work argues that the choice of how positive pairs are constructed should be matched by an SSL objective expressive enough to model the resulting conditional. If correct, this widens SSL beyond handcrafted augmentations toward naturally paired data — video, captions, cross-model embeddings — without losing fine-grained information.

Real-world applications:

• Robotics and autonomous driving: present state to multiple plausible futures, such as a car turning left or right, requires exactly the multimodal conditional the paper targets.
• Video world modeling and simulation: retaining stochastic object accelerations rather than averaging them out supports more faithful predictive models.
• Fine-grained visual domains: agriculture, ecology, and biodiversity monitoring, where datasets like iNat-2021 involve noisy or variable pairings.
• Medical and scientific imaging: settings with heteroscedastic variation across samples, where feature-discarding SSL objectives would remove clinically relevant detail.

Industry relevance: the method is presented as a drop-in regularizer and editing function on top of existing contrastive or distillation-based SSL pipelines, so the practical cost is adding small MLP heads and one hyperparameter β rather than replacing an established training stack. Code is available at https://github.com/SkrighYZ/AdaSSL.

Future Directions

• Scaling the latent variable to large-scale pretraining. The experiments cover numerical data, 3DIdent, CelebA, iNat-2021 and Moving-MNIST; whether the added latent variable pays off at foundation-model scale is not established here.
• Making AdaSSL-S work reliably with distillation-based SSL. The paper reports that AdaSSL-S works well with contrastive learning but requires additional care for distillation-based methods in some settings, without a full resolution in the section summarized.
• Removing the extra augmented view. Inferring an r that is invariant to low-level transformations sometimes requires an additional augmented view (Appendix C.2), which is an extra data requirement to eliminate.
• Extending the theory. The variational lower bound is only theoretically justified for contrastive SSL; the distillation case rests on the intuitive argument that any loss encouraging r to aid prediction will work.

Target Audience

Researchers and graduate students working on self-supervised learning, joint-embedding predictive architectures, and causal representation learning, who want a principled account of why natural pairs produce difficult conditionals; practitioners who train contrastive or BYOL-style models on video, video-plus-action, or image–caption data and want a lightweight objective change that preserves conditional uncertainty.

Authors’ 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.