Unsupervised Representation Learning - an Invariant Risk Minimization Perspective

TL;DR

This work proposes a novel unsupervised framework for IRM, extending the concept of invariance to settings where labels are unavailable, and introduces two methods within this framework: Principal Invariant Component Analysis (PICA), a linear method that extracts invariant directions under Gaussian assumptions, and Variational Invariant Autoencoder (VIAE), a deep generative model that separates environment-invariant and environment-dependent latent factors.

Abstract

We propose a novel unsupervised framework for \emph{Invariant Risk Minimization} (IRM), extending the concept of invariance to settings where labels are unavailable. Traditional IRM methods rely on labeled data to learn representations that are robust to distributional shifts across environments. In contrast, our approach redefines invariance through feature distribution alignment, enabling robust representation learning from unlabeled data. We introduce two methods within this framework: Principal Invariant Component Analysis (PICA), a linear method that extracts invariant directions under Gaussian assumptions, and Variational Invariant Autoencoder (VIAE), a deep generative model that separates environment-invariant and environment-dependent latent factors. Our approach is based on a novel ``unsupervised'' structural causal model and supports environment-conditioned sample-generation and intervention. Empirical evaluations on synthetic dataset, modified versions of MNIST, and CelebA demonstrate the effectiveness of our methods in capturing invariant structure, preserving relevant information, and generalizing across environments without access to labels.

Figures (8)

• Figure 1: IRM Generative Structural Causal Models for supervised (left $3$ figures) and unsupervised (right figure) cases
• Figure 2: Output of the PICA algorithm with $d_r=1$ on the synthetic dataset. The projection captures the invariant component shared across both environments.
• Figure 3: VIAE architecture. A shared invariant encoder produces $Z_\mathrm{inv}$, while environment-specific encoders produce $Z_e$. The decoder reconstructs $X$ from both components.
• Figure 4: Generated samples conditioned on a fixed $Z_\mathrm{inv}$. Top row: samples with different $Z_e$ drawn from $P^1(Z_e)$. Bottom row: samples with different $Z_e$ sampled from $P^2(Z_e)$. Left side: SMNIST dataset, right side: SCMNIST dataset. Invariant features (in our case, the digits) are preserved for all samples, with stable environment for each row.
• Figure 5: Environment Transfer for $e_s\in\mathcal{E}_{\mathrm{train}}$