Self-Supervised Learning with Data Augmentations Provably Isolates Content from Style

TL;DR

The work treats SSL with data augmentations as a latent-variable problem with a content block that should be invariant across views and a style block that may change, and proves block-identifiability results for the invariant content partition under mild assumptions, in both generative and discriminative SSL frameworks. It introduces a causal, high-dimensional dataset (Causal3DIdent) to study practical augmentation effects and validates the theory through controlled numerical simulations and high-dimensional image experiments, demonstrating when and how augmentations isolate content. The findings provide a theoretical grounding for the empirical success of contrastive SSL methods like InfoNCE and offer guidance on designing augmentations and regularizers to recover content representations in the presence of dependent latent factors. These results have practical implications for robust, causally aware SSL in vision and beyond, including safety-critical domains where invariant content extraction is essential.

Abstract

Self-supervised representation learning has shown remarkable success in a number of domains. A common practice is to perform data augmentation via hand-crafted transformations intended to leave the semantics of the data invariant. We seek to understand the empirical success of this approach from a theoretical perspective. We formulate the augmentation process as a latent variable model by postulating a partition of the latent representation into a content component, which is assumed invariant to augmentation, and a style component, which is allowed to change. Unlike prior work on disentanglement and independent component analysis, we allow for both nontrivial statistical and causal dependencies in the latent space. We study the identifiability of the latent representation based on pairs of views of the observations and prove sufficient conditions that allow us to identify the invariant content partition up to an invertible mapping in both generative and discriminative settings. We find numerical simulations with dependent latent variables are consistent with our theory. Lastly, we introduce Causal3DIdent, a dataset of high-dimensional, visually complex images with rich causal dependencies, which we use to study the effect of data augmentations performed in practice.

Figures (12)

• Figure 1: Overview of our problem formulation. We partition the latent variable $\mathbf{z}$ into content $\mathbf{c}$ and style $\mathbf{s}$, and allow for statistical and causal dependence of style on content. We assume that only style changes betweenthe original view$\mathbf{x}$ and the augmented view$\Tilde{\mathbf{x}}$, i.e., they are obtained by applying the same deterministic function$\mathbf{f}$ to $\mathbf{z}=(\mathbf{c},\mathbf{s})$ and $\Tilde{\mathbf{z}}=(\mathbf{c},\Tilde{\mathbf{s}})$.
• Figure 2: (Left) Causal graph for the Causal3DIdent dataset. (Right) Two samples from each object class.
• Figure 3: 40 random samples from the marginal distribution of the Teapot object class.
• Figure 4: 40 random samples from the marginal distribution of the Hare object class.
• Figure 5: 40 random samples from the marginal distribution of the Dragon object class.

Theorems & Definitions (9)

• Definition 4.1: Block-identifiability
• Theorem 4.2: Identifying content with a generative model
• Theorem 4.3: Identifying content with an invertible encoder
• Theorem 4.4: Identifying content with discriminative learning and a non-invertible encoder
• Theorem A.1: Identifying content with a generative model
• Theorem A.1: Identifying content with an invertible encoder
• Theorem A.1: Identifying content with discriminative learning and a non-invertible encoder
• Claim A.1
• Proposition A.2: Proposition 5 of zimmermann2021contrastive