[Paper Review] Representation Learning via Invariant Causal Mechanisms

Jovana Mitrovic, Brian McWilliams+3 more|arXiv (Cornell University)|Oct 15, 2020

Domain Adaptation and Few-Shot LearningComputer Science39 references41 citations

TL;DR

ReLIC introduces an invariant causal mechanism regularizer for self-supervised representation learning, enforcing proxy-target predictions to be invariant across augmentations, with theoretical support and strong empirical gains in robustness, OOD generalization, and Atari performance.

ABSTRACT

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.

Motivation & Objective

• Formalize self-supervised representation learning with a causal perspective (content vs. style) to improve generalization.
• Propose an invariant prediction objective (ReLIC) that uses data augmentations as style interventions.
• Generalize contrastive learning via causal refinements to explain and improve downstream transfer.
• Provide theoretical guarantees that invariant representations under refinements generalize to downstream tasks.
• Empirically demonstrate robustness, OOD generalization on ImageNet, and superhuman Atari performance.

Proposed method

• Model data with content C and style S, where C influences downstream targets Y t but S does not; assume C ⟂ S and X is generated from C and S.
• Use data augmentations as interventions on S to simulate style changes and learn invariant predictors of proxy targets Y^R from f(X).
• Define ReLIC objective that minimizes a contrastive proxy loss while enforcing p^{do(a)}(Y^R|f(X)) to be invariant across augmentations a, via a KL-based regularizer with tolerance ρ.
• Use instance discrimination as proxy task Y^R with pairs of augmentations to compute similarities; relate encoder f, critic g, and possibly a target network h.
• Show that the invariance constraint improves intra-class concentration and generalization, and provide a theoretical justification (Theorem 1) that invariant refinements suffice for downstream tasks.
• Generalize contrastive learning by introducing refinements, proving that invariant representations for refinements under style interventions extend to downstream tasks.

Experimental results

Research questions

• RQ1 Can explicit invariance under data augmentations improve the usefulness and generalization of self-supervised representations?
• RQ2 Do representations learned via ReLIC generalize to downstream tasks and exhibit robustness and out-of-distribution generalization beyond prior contrastive methods?
• RQ3 How do refinements provide a causal explanation for the success of contrastive learning and relate to mutual information-based views?

Key findings

• ReLIC achieves strong linear evaluation performance on ImageNet, outperforming several self-supervised baselines under comparable augmentations.
• ReLIC yields significantly better robustness and out-of-distribution generalization on ImageNet-C and ImageNet-R compared to supervised and other unsupervised methods.
• In Atari reinforcement learning, ReLIC-based representations lead to superior performance, achieving above-human results on 51 of 57 games.
• The invariance regularizer encourages tighter intra-class representation and greater linear separability (as shown by F_LDA analyses).
• Theorem 1 establishes that learning on refinements with invariant predictions under augmentations suffices to generalize to downstream tasks, offering an alternative explanation to mutual information for contrastive success.