Source-linked AI summary
Representation Learning via Invariant Causal Mechanisms
Jovana Mitrovic, Brian McWilliams, Jacob Walker, Lars Buesing, Charles Blundell
TL;DR
The paper addresses limited theoretical understanding of successful self-supervised representation learning from unlabeled data. It uses a causal framework to enforce invariant proxy prediction across augmentations and generalizes contrastive learning through refinements. RELIC significantly improves robustness and out-of-distribution generalization on ImageNet and outperforms related methods on Atari, reaching above-human performance on 51 of 57 games.
Problem
Self-supervised learning reduces reliance on costly labels or rewards, but the theoretical understanding of its success remains limited.
Method
RELIC uses an explicit invariance regularizer for proxy predictions across augmentations, while causal refinements generalize contrastive learning.
Results
RELIC significantly outperforms competing methods in robustness and out-of-distribution generalization on ImageNet and achieves above-human performance on 51 out of 57 Atari games.
Takeaways & Limitations
Invariant prediction across augmentations provides theoretical support for useful downstream representations and for the role of data augmentations in contrastive learning.
Takeaways & Limitations
The fixed augmentation set cannot generate all possible styles, and choosing augmentations implicitly determines which aspects are treated as style versus content.
Abstract
from arXiv · showhide
Self-supervised learning has emerged as a strategy to reduce the reliance on costly supervised signal by pretraining representations only using unlabeled data. These methods combine heuristic proxy classification tasks with data augmentations and have achieved significant success, but our theoretical understanding of this success remains limited. In this paper we analyze self-supervised representation learning using a causal framework. We show how data augmentations can be more effectively utilized through explicit invariance constraints on the proxy classifiers employed during pretraining. Based on this, we propose a novel self-supervised objective, Representation Learning via Invariant Causal Mechanisms (ReLIC), that enforces invariant prediction of proxy targets across augmentations through an invariance regularizer which yields improved generalization guarantees. Further, using causality we generalize contrastive learning, a particular kind of self-supervised method, and provide an alternative theoretical explanation for the success of these methods. Empirically, ReLIC significantly outperforms competing methods in terms of robustness and out-of-distribution generalization on ImageNet, while also significantly outperforming these methods on Atari achieving above human-level performance on $51$ out of $57$ games.
1 INTRODUCTION
The paper analyzes self-supervised representation learning causally, proposing invariant proxy prediction across augmentations and generalizing contrastive learning through refinements. RELIC improves robustness and out-of-distribution generalization across ImageNet and Atari tasks.
- Motivation: Self-supervised learning seeks to reduce costly reliance on labels or rewards by pretraining representations from abundant unlabeled data.It uses proxy tasks defined on unsupervised data.
- Causal framework: A causal analysis characterizes useful representations as invariant predictors of proxy targets under interventions on features unrelated causally to downstream targets.The framework formalizes this intuition with a causal graph and derives properties of optimal representations.
- RELIC: RELIC enforces invariant proxy prediction across augmentations with an explicit regularizer, yielding improved generalization guarantees under weaker assumptions than prior work.The method treats augmentations as simulated interventions and leverages them more effectively.
- Generalizing contrastive learning: Refinement-based proxy tasks generalize contrastive learning and provide a causal explanation for why invariant prediction can produce representations that generalize to downstream tasks.The paper identifies learning invariant predictors on refinements as sufficient for downstream generalization.
- Empirical evaluation: 4.9% absolute error reduction versus the supervised baseline is reported for robustness, while RELIC achieves above-human performance on 51 out of 57 Atari games.On ImageNet, it also significantly outperforms competing methods for robustness and out-of-distribution generalization.
2 REPRESENTATION LEARNING VIA INVARIANT CAUSAL MECHANISMS
The paper formulates self-supervised representation learning causally, treating content as task-relevant and style as task-irrelevant, then uses augmentations to enforce invariant proxy prediction. RELIC combines proxy-task loss with an explicit invariance regularizer, yielding representations with concentrated within-class variability and stronger linear separability.
- Problem setting: The problem is to pretrain a representation f(X) useful for unknown downstream tasks Y using unlabeled observed data X.
- Causal interpretation: The causal model separates data into content C and style S, with content influencing downstream targets while style does not.Content and style generate X, and the model assumes they are independent.
- Causal interpretation: Interventions on style leave P(Y_t|C) invariant, making content a useful and robust representation across style perturbations.
- Proxy tasks and augmentations: Because downstream targets are unknown, the method constructs a proxy task Y^R and uses content-preserving augmentations as simulated interventions on style.Examples include rotation, grayscaling, translation, and cropping for images.
- RELIC objective: RELIC minimizes proxy-task loss while explicitly regularizing prediction distributions to remain invariant across augmentations.The regularizer uses KL divergence, although other distributional distance measures can replace it.
- Generalization effect: The invariance penalty contracts within-class representation variability, whereas contrastive loss alone does not guarantee generalization.The paper connects this concentration to better downstream linear separability and generalization guarantees.
3 GENERALIZING CONTRASTIVE LEARNING
The paper generalizes contrastive learning by viewing proxy tasks through causal refinements, which are more fine-grained versions of downstream tasks. It proves that invariant prediction on a refinement is sufficient for invariant representations of the corresponding downstream tasks.
- Refinements: A refinement is a more fine-grained version of a task, such as classifying individual cat and dog breeds instead of cats versus dogs.
- Refinements: Instance discrimination, the proxy task commonly used in contrastive learning, is the most fine-grained refinement of the dataset.
- Theoretical connection: If proxy targets Y^R refine all downstream tasks Y, invariant prediction of Y^R under style interventions implies invariant prediction of every task in Y.
- Theoretical connection: Theorem 1 states that learning a representation on a refinement is a sufficient condition for usefulness on the corresponding downstream tasks.For instance discrimination, the theorem’s refinement condition is satisfied for downstream tasks obeying the stated assumptions.
- Implications: This refinement framework provides a causal alternative to mutual-information explanations and applies beyond instance-discrimination contrastive losses.The paper also notes that metadata can guide richer proxy-task constructions.
4 RELATED WORK
Related work includes mutual-information-based contrastive methods, invariant prediction for domain robustness, and learning-theoretic analyses of contrastive objectives. The paper positions RELIC as a causal account that explicitly addresses limitations of prior explanations and concentration assumptions.
- Mutual information: CPC, DIM, and AMDIM use noise contrastive estimation in approaches rooted in mutual-information maximization.
- Mutual information: The relationship between mutual-information bounds and downstream performance is debated, and InfoNCE methods may require custom architectures to avoid non-informative solutions.
- Contrastive methods: SimCLR uses strong augmentations, while BYOL avoids a contrastive objective and uses target-network ideas.
- Learning-theoretic analysis: Prior learning-theoretic analysis identifies strong intra-class concentration assumptions as a fundamental limitation of contrastive objectives.RELIC explicitly enforces intra-class concentration through its invariance regularizer.
- RELIC’s position: The paper explains RELIC’s performance through causal theory and reports more robust representations than SimCLR or BYOL when invariance penalties are used.
- Empirical comparison: Figure 2 compares the distributions of FLDA for RELIC, SimCLR, and AMDIM, with the y-axis clipped for visualization.The caption identifies FLDA as the linear discriminant ratio.
5 EXPERIMENTS
The experiments evaluate RELIC's representation quality, robustness, out-of-distribution generalization, and reinforcement-learning performance. RELIC improves linear separability and robustness, and compares favorably with competing self-supervised methods across ImageNet and Atari.
- Linear evaluation: RELIC's explicit invariance constraint increases linear separability by concentrating within-class variability.FLDA increases when between-class distances grow or within-class variance shrinks.
- Linear evaluation: RELIC outperforms competing approaches on ImageNet linear evaluation when methods use SimCLR augmentations.The comparison includes ResNet-50 and ResNet-50 with a target network.
- Robustness and generalization: RELIC achieves significantly lower mCE than supervised ResNet-50, SimCLR, and BYOL on ImageNet-C.It also has the lowest corruption error on 14 of 15 corruptions compared with SimCLR and BYOL.
- Robustness and generalization: RELIC outperforms SimCLR and BYOL on ImageNet-R, demonstrating stronger out-of-distribution generalization of the learned representation.ImageNet-R evaluates representations trained on clean ImageNet images using different renditions of the same classes.
- Reinforcement Learning: RELIC performs best in 25 out of 57 Atari games when used as an auxiliary representation-learning loss.The comparison includes SimCLR, BYOL, CURL, and feeding augmented observations directly to the agent.
6 CONCLUSION
The paper formulates self-supervised representation learning causally and uses invariant prediction across augmentations to propose RELIC. It also generalizes contrastive methods through refinements and reports stronger robustness, out-of-distribution generalization, and Atari performance.
- 6 CONCLUSION: A causal graph characterizes optimal representations as invariant predictors of proxy targets under interventions on noncausal features.Data augmentations simulate these interventions, motivating an explicit invariance constraint.
- 6 CONCLUSION: RELIC enforces invariant proxy-target prediction across augmentations through an invariance regularizer.The objective is derived from the paper's causal analysis of self-supervised learning.
- 6 CONCLUSION: Contrastive methods are generalized using refinements, with invariant prediction on refinements providing a sufficient condition for downstream-task generalization.This gives an alternative explanation for contrastive-learning success beyond mutual information.
- 6 CONCLUSION: RELIC significantly outperforms related methods in ImageNet robustness and out-of-distribution generalization and achieves superhuman performance on 51 of 57 Atari games.The paper compares RELIC with recent self-supervised methods across prediction and reinforcement-learning tasks.
- 6 CONCLUSION: Future work will investigate more coarse-grained refinements and empirically evaluate different kinds of refinements.These directions concern the refinement-based generalization analysis.
A RELATIONSHIP BETWEEN RELIC AND OTHER METHODS
The RELIC objective can recover several recent state-of-the-art methods through different design choices. The correspondence is summarized in Table 5.
- A RELATIONSHIP BETWEEN RELIC AND OTHER METHODS: RELIC's objective recovers recent state-of-the-art methods through appropriate choices of φ, g, f, and h.The table also distinguishes identity functions and unit-norm constraints on g.
B DISTANCE CONCENTRATION AND GENERALIZATION
The invariance penalty is interpreted as concentrating distances and within-class representations, which supports generalization under stated assumptions. The analysis formalizes this effect through concentration and generalization results.
- Distance concentration: Augmentation overlap can indirectly bring representations of similar images closer by enforcing closeness between each image and its augmentations.This concentrates distances between similar images.
- Objective: For a single negative, the Euclidean distance-regularized objective is equivalent to the standard RELIC objective with an identity critic.The logistic loss is specified for the modified objective.
- Concentration: Under the stated assumptions, the learned representation is sub-Gaussian, providing the concentration property required by the generalization analysis.The assumptions require an L-Lipschitz minimizer with an active constraint and a bounded input variable.
- Generalization: The Euclidean invariance regularizer contracts within-class variance, supporting a generalization result for contrastive learning on unlabelled data.The result is connected to the concentration condition used by Saunshi et al. (2019).
- Generalization: The generalization bound depends on the Rademacher complexity of the function class and the sample size.The paper states that the final result is otherwise identical to the corresponding contrastive-learning result.
C ADDITIONAL RESULTS
The analysis studies how augmentation-induced connectivity contracts representation variance and how task refinements support invariant representations. Random-graph structure yields tighter concentration results, while contraction depends on augmentation overlap and smoothness conditions.
- Variance contraction: ρ < (1−(1−α)L)B ensures the variance bound remains below one under the stated overlap condition.This condition follows from requiring B < 1 − (1 − α)L.
- Limitations: If α is too small or f is insufficiently smooth, α may prevent choosing ρ to induce contraction in σ2.α is determined by the augmentation set rather than directly controlled by the user.
- Random-graph refinement: When the augmentation graph is connected, every representation-space point can be related to every other through a finite sequence of graph edges.Connectivity allows the variance contribution to be decomposed across adjacent and reachable components.
- Random-graph refinement: With high probability, the graph diameter satisfies 3 ≤ D ≤ 4, enabling a refined variance bound based on short paths.The resulting condition requires ρ ≤ B 2Lκ(α+3(1−α)) ≤ B 6Lκ.
- Task refinements: A refinement subdivides an existing task partition, so each coarse equivalence class becomes a union of finer equivalence classes.The paper illustrates this by refining aquatic/non-aquatic and animal/non-animal tasks into four joint classes.
D.2 PROOF OF THEOREM 1
The proof formalizes invariant representations through interventions on style and shows that invariance for a refinement of downstream tasks implies invariance for each original task. The empirical setup applies RELIC with standard image augmentations and linear evaluation protocols.
- Theorem 1: If f(X) is invariant for a refinement Y R under style interventions, it is invariant for every downstream task Yt.The implication holds for all tasks t and all style values si and sj.
- Proof: The proof uses that the mechanism p(Yt|Y R) is independent of style S and that f(X) is invariant for Y R.These assumptions establish equality of interventional task distributions across styles.
- Implementation: RELIC pretraining uses SimCLR-style random crops, flips, color jittering, grayscale, and Gaussian blurring with the original augmentation probabilities.Images are normalized using ImageNet channel means and standard deviations.
- Evaluation: The ImageNet evaluation trains a linear classifier on frozen representations without updating network parameters or batch statistics.ResNet-50 representations are taken from the final average-pooling layer, with dimension 2048.
E.5.2 EVALUATION
The evaluation measures robustness and distributional generalization using ImageNet-C and ImageNet-R alongside standard ImageNet linear evaluation. ImageNet-C varies corruption type and severity to test semantic robustness.
- Evaluation protocol: Linear evaluation trains a classifier on frozen representations, using random crops, resizing, and horizontal flips during training.The protocol follows standard ImageNet evaluation procedures.
- Datasets: ImageNet-C contains 15 corruption types spanning noise, blur, weather, and digital categories, each at five severity levels.This produces 75 distinct corruptions, with increasing severity illustrated for impulse noise.
E.5.3 ROBUSTNESS METRICS AND FURTHER RESULTS
The paper evaluates absolute and relative corruption robustness on ImageNet-C and compares self-supervised representation methods in Atari reinforcement learning. Robustness metrics normalize corruption errors against AlexNet and distinguish absolute from clean-set-relative performance loss.
- Robustness metrics: The mean Corruption Error averages normalized errors across all 15 ImageNet-C corruption types and measures absolute robustness.Individual corruption errors are normalized using AlexNet’s corresponding errors.
- Robustness metrics: The mean relative Corruption Error averages performance loss under corruption relative to each classifier’s clean ImageNet error.This metric measures relative robustness rather than absolute corrupted-image error.
- ImageNet-C results: Table 6 reports mCE and corruption-specific errors for Blur, Weather, and Digital categories on models trained only with clean ImageNet images.Table 7 reports corresponding relative robustness values and includes the RELIC-t architecture variant.
- Reinforcement learning: Atari experiments compare RELIC with other self-supervised auxiliary losses across 57 games using human-normalized scores.The experiments use an R2D2 agent supplemented with a second encoder trained for representation learning.
- Reinforcement learning: RELIC and SimCLR use the encoder’s last layer without a critic embedding, while a momentum-updated target encoder processes the second augmentation.The target encoder is updated with momentum .99.