Invariant Risk Minimization

Invariant Risk Minimization (IRM), introduced by Arjovsky et al. (2019), is a training principle that aims to learn predictors that generalize out-of-distribution by exploiting invariant causal mechanisms rather than spurious statistical correlations. IRM extends empirical risk minimization (ERM) by requiring that the learned representation support the same optimal linear classifier across all training environments.

Motivation: The Problem with ERM

Standard empirical risk minimization minimizes average loss across all training data:

$$\hat{\theta}{\text{ERM}} = \arg\min\theta \frac{1}{n} \sum_{i=1}^n \ell(f_\theta(x_i), y_i)$$

ERM is blind to the source of statistical patterns. Given multiple training environments with different distributions, ERM will exploit any correlation that reduces training loss — including spurious correlations that hold in training environments but fail at test time.

The Cow-Camel Example

A classic illustration: a model trained to classify cows vs. camels learns “camel = sand background.” In training data, camels appear on sand and cows on grass. At test time, a cow on sand is misclassified as a camel. ERM exploited a spurious correlation (background) instead of the invariant feature (animal shape).

The ColorMNIST Benchmark

ColorMNIST is a controlled OOD benchmark: digit color is spuriously correlated with label (e.g., red = 1) in training, but anti-correlated in the test environment. ERM achieves ~10% test accuracy (worse than random) despite high training accuracy. The goal is to learn shape, not color.

IRM: The Invariance Constraint

IRM seeks a feature map $\Phi: \mathcal{X} \to \mathcal{H}$ such that there exists a single classifier $w: \mathcal{H} \to \mathcal{Y}$ that is simultaneously optimal across all environments $e \in \mathcal{E}_{\text{train}}$:

$$\min_{\Phi, w} \sum_{e \in \mathcal{E}} \mathcal{R}^e(w \circ \Phi) \quad \text{subject to} \quad w \in \arg\min_{\bar{w}} \mathcal{R}^e(\bar{w} \circ \Phi) \quad \forall e \in \mathcal{E}$$

where $\mathcal{R}^e$ is the risk in environment $e$.

The constraint says: $\Phi$ must elicit a representation on which the same $w$ is optimal for every environment. This is the invariance condition — if the optimal predictor is different for each environment, the representation must be encoding environment-specific spurious features.

IRMv1: A Practical Relaxation

The bilevel constraint is computationally intractable. Arjovsky et al. proposed a gradient-based relaxation (IRMv1) by fixing $w = 1$ (a scalar dummy classifier on top of $\Phi$) and penalizing the gradient norm of the risk w.r.t. this fixed classifier:

$$\mathcal{L}{\text{IRM}}(\Phi) = \sum{e \in \mathcal{E}} \mathcal{R}^e(\Phi) + \lambda \sum_{e \in \mathcal{E}} \left|\nabla_{w | w=1} \mathcal{R}^e(w \cdot \Phi)\right|^2$$

The gradient penalty $|\nabla_w \mathcal{R}^e|^2$ is zero exactly when $w = 1$ is optimal for environment $e$. Minimizing this penalty across environments encourages $\Phi$ to be a representation where no environment can profitably use a better classifier.

Intuition

If the penalty term is large for environment $e$, it means the gradient of the loss w.r.t. the fixed $w=1$ classifier is large — i.e., we could improve by adjusting $w$ for that environment. This signals that $\Phi$ is encoding environment-specific information that a different $w$ would exploit.

Connection to Causality

IRM has a deep connection to structural causal models (SCMs). In a causal model, the causal parents of $Y$ determine $P(Y | \text{Pa}(Y))$, which is invariant across interventions on non-$Y$ variables. Spurious correlations arise from non-causal pathways (e.g., common causes or selection bias).

IRM operationalizes the intuition: learn a representation that captures causal features by requiring it to support an invariant predictor. The ideal representation $\Phi^*$ outputs exactly the causal parents of $Y$.

Causal Graph View

Consider a simple SCM:

• $S \to X$ (spurious feature $S$ causes $X$)
• $C \to Y$ and $C \to X$ (causal feature $C$ causes both $X$ and $Y$)

ERM learns to use both $C$ and $S$ from $X$. IRM, given environments where $S$ varies, learns to use only $C$ — the invariant mechanism.

Variants and Improvements

Risk Extrapolation (REx)

Krueger et al. (2021) proposed V-REx (Variance-based Risk Extrapolation), which penalizes the variance of risks across environments instead of the gradient norm:

$$\mathcal{L}_{\text{V-REx}} = \sum_e \mathcal{R}^e + \lambda \cdot \text{Var}_e\left(\mathcal{R}^e\right)$$

This encourages training risks to be equal across environments, which is a necessary condition for invariance. V-REx is simpler to implement and can outperform IRMv1 in practice.

IRM Games

Ahuja et al. (2020) reformulated IRM as a game between environment-specific classifiers and a shared representation learner. Each environment player tries to minimize its own risk; the representation player tries to make all environments agree on the optimal classifier. At Nash equilibrium, invariance is achieved.

Fishr

Rame et al. (2022) proposed Fishr, which matches the gradient covariance (Fisher information) across environments:

$$\mathcal{L}{\text{Fishr}} = \sum_e \mathcal{R}^e + \lambda \sum{e, e’} |\nabla^2_\theta \mathcal{R}^e - \nabla^2_\theta \mathcal{R}^{e’}|^2_F$$

The intuition: if the Fisher matrices agree, the loss landscape looks the same for all environments, suggesting invariant features are being used.

EIIL (Environment Inference for Invariant Learning)

Creager et al. (2021) addressed the practical challenge that environment labels are often unavailable. EIIL infers environments by finding partitions of the training data that maximize the IRM penalty — discovering spurious correlations automatically.

Theoretical Analysis and Limitations

When Does IRM Work?

IRM’s theoretical guarantees hold under specific conditions:

• Linear SCMs: For linear models, IRM recovers the causal features when environments provide sufficient variation
• Sufficient diversity: Training environments must differ in ways that make spurious correlations inconsistent across environments
• No hidden common causes: Unmeasured confounders can violate the invariance assumption

Failure Modes

Rosenfeld et al. (2021) and others showed that IRMv1 can fail even in simple settings:

• Small number of environments: With few environments, many representations satisfy the invariance constraint, including some that rely on spurious features
• Linear approximation: The IRMv1 penalty (using fixed $w = 1$) is a first-order approximation that can miss important violations
• Non-identifiability: Without sufficient environment diversity, causal and spurious features cannot be distinguished

The Environments Problem

IRM requires explicit environment labels during training. In practice:

• Collecting environment-labeled data is expensive
• Defining what constitutes an “environment” requires domain knowledge
• Inappropriate environment partitioning can degrade performance relative to ERM

Comparison with Related Methods

Group DRO (Sagawa et al., 2020) is a related and often competitive baseline: minimize the worst-case risk across predefined groups, which encourages the model to not ignore any environment.

Practical Recommendations

Hyperparameter $\lambda$

The penalty weight $\lambda$ controls the trade-off between ERM performance and invariance. Common practice:

1. Start with $\lambda = 0$ (pure ERM) and increase
2. Monitor performance on a held-out OOD validation set
3. Use a warm-up schedule: start with small $\lambda$ and increase after initial ERM convergence

Environment Design

The quality of IRM solutions depends heavily on environments:

• Environments should disagree on spurious features but agree on causal features
• More diverse environments generally yield better invariant representations
• Domain knowledge about which features are spurious helps design informative environments

Implementation Tips

def irm_penalty(logits, y):
    """Compute IRMv1 gradient penalty for a batch."""
    scale = torch.tensor(1.0, requires_grad=True, device=logits.device)
    loss = F.cross_entropy(logits * scale, y)
    grad = torch.autograd.grad(loss, [scale], create_graph=True)[0]
    return grad ** 2

def irm_loss(model, batches_per_env, lambda_irm):
    total_loss = 0
    total_penalty = 0
    for x, y in batches_per_env:
        logits = model(x)
        total_loss += F.cross_entropy(logits, y)
        total_penalty += irm_penalty(logits, y)
    return total_loss + lambda_irm * total_penalty

Impact and Current Status

IRM sparked a large body of research on out-of-distribution generalization and causal machine learning. Key impacts include:

• Establishing environments as a first-class training signal beyond data and labels
• Formalizing the link between causal reasoning and machine learning
• Motivating benchmark development (DomainBed, WILDS)

However, empirical results on real OOD benchmarks have been mixed. On DomainBed (Gulrajani & Lopez-Paz, 2021), IRM and variants often fail to consistently outperform well-tuned ERM with proper data augmentation, suggesting that:

• The causal assumptions underlying IRM may not hold in many real datasets
• Simple techniques like data augmentation that implicitly diversify training distributions can be very effective
• OOD generalization remains an open problem

Summary

Invariant Risk Minimization is a principled approach to learning causal features by exploiting multi-environment training signals. Its key contribution is formalizing the connection between predictive invariance and causal structure, motivating a family of practical algorithms (IRMv1, V-REx, Fishr) and benchmarks. While theoretical and empirical limitations exist, IRM remains a foundational framework for anyone working on distribution shift, causal ML, and robust generalization.