Source-linked AI summary
Invariant Risk Minimization
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, David Lopez-Paz
TL;DR
Machine learning can fit correlations caused by biased data rather than stable, causally relevant structure. IRM learns a representation whose optimal classifier is shared across training environments, and the paper’s theory and experiments connect this invariance to causal prediction and out-of-distribution generalization.
Problem
Training data can contain spurious correlations from selection biases and other peculiarities, making it difficult to identify stable properties that support out-of-distribution prediction.
Method
IRM learns a data representation that balances predictive power with invariance of the optimal classifier across multiple training environments.
Results
IRM’s learned invariances are linked theoretically to causal structure and empirically to more accurate prediction across experimental conditions than ERM and ICP.
Takeaways & Limitations
Invariant correlations can provide a statistically testable route toward discovering causal predictors that generalize across environments.
Takeaways & Limitations
IRM’s generalization theory may require training environments to scale linearly with representation parameters, although two environments often suffice empirically.
Abstract
from arXiv · showhide
We introduce Invariant Risk Minimization (IRM), a learning paradigm to estimate invariant correlations across multiple training distributions. To achieve this goal, IRM learns a data representation such that the optimal classifier, on top of that data representation, matches for all training distributions. Through theory and experiments, we show how the invariances learned by IRM relate to the causal structures governing the data and enable out-of-distribution generalization.
1 Introduction
Machine learning can exploit spurious correlations created by biased training data, undermining out-of-distribution generalization. IRM addresses this by learning representations whose optimal classifier is invariant across distinct training environments.
- Motivation: Training-error minimization can absorb selection biases, confounding factors, and other correlations unrelated to the causal explanation of interest.These inherited data biases contribute to failures when models encounter conditions unlike those represented in training data.
- Motivation: In the cow–camel example, landscape context becomes a shortcut, causing a model trained on pastures and deserts to fail on cows photographed on sandy beaches.The example illustrates how a selection bias links animal identity with background rather than animal shape.
- Motivation: Spurious correlations are properties not expected to remain stable in the future, whereas stable properties are the target of robust prediction.Datasets often do not expose enough distributional variation to identify which properties are stable.
- Invariant environments: Distinct environments preserve information about how data distributions change across sources, locations, times, or conditions, which helps distinguish spurious from stable properties.Pooling or shuffling data can destroy the variation needed to make that distinction.
- Causation and invariance: The paper connects invariant descriptions with causal explanations while noting that prior causal approaches may require meaningful graphs, linear models, or scale poorly with variables.IRM is motivated as a way to seek causal structure from perceptual inputs such as pixels.
- IRM: IRM estimates nonlinear, invariant, causal predictors from multiple training environments to support out-of-distribution generalization.Its principle is to learn a representation such that the optimal classifier on top matches across environments.
2 The many faces of generalization
The paper uses a simple structural-equation example to show why common generalization strategies can select environment-dependent correlations. It then motivates IRM as a method for extracting invariant predictors across multiple environments.
- Problem setup: The goal is to learn a predictor from multiple environment-specific datasets that performs well across a larger set of unseen but related environments.The environments measure the same variables under different conditions, with each dataset sampled from its own distribution.
- Example 1: The structural-equation example defines a setting where environment changes can alter feature equations and target noise while preserving the broader system.Least-squares regressions are evaluated separately across environments to expose coefficient stability or variation.
- Example 1: The regression using X1 is the only listed regression whose coefficients do not depend on environment, making it an invariant correlation.Regressions using X2 or both features have coefficients that vary with the environment and may fail in novel environments.
- Example 1: The invariant predictor ˆY = 1 · X1 + 0 · X2 is the only predictor with finite ROOD and is identified as the causal explanation across environments.The paper links this predictor to how the target responds to interventions on the inputs.
- Existing approaches: ERM and robust learning can assign large weight to X2 when training environments favor its correlation, producing infinite ROOD when that correlation disappears.The failure occurs because optimizing mixtures or worst-case training risk does not enforce the desired invariance.
- Existing approaches: Distribution matching can reject the causal feature when its distribution changes, while residual-distribution matching can fail when noise variance changes across environments.These examples show that enforcing the wrong invariance can harm generalization.
- IRM: IRM is proposed to extract nonlinear invariant predictors across multiple environments, enabling out-of-distribution generalization.The proposal responds to the difficulty of finding invariant predictors even in simple problems.
3 Algorithms for invariant risk minimization
IRM seeks representations whose optimal classifier is shared across environments, balancing predictive performance with invariant relationships. Its practical penalties avoid discontinuities but retain nonconvex optimization challenges and require careful parameterization.
- IRM defines invariance as a representation Φ for which one classifier is simultaneously optimal across all environments.For mean squared error and cross-entropy, this corresponds to stable conditional expectations across environments.
- IRM formulates representation learning as predictive risk minimization subject to environment-wise optimality constraints, then replaces the bi-level problem with a penalized objective.The penalty balances empirical risk against invariance, while the practical version uses a fixed dummy classifier and gradient norm.
- Choosing a penalty D for linear classifiers w: The classifier-distance penalty Ddist is discontinuous at the invariant representation and can approach zero as the representation coefficient diverges, even with severe regularization.These pathologies arise in the linear example as the representation changes its dependence on the spurious variable.
- Choosing a penalty D for linear classifiers w: The alternative penalty Dlin measures violation of the normal equations, is smooth in Φ and w, and reaches an accessible minimum at the invariant representation.Dlin equals zero exactly when the classifier is risk-optimal for the representation, although it is non-convex for general representations.
- Scalar fixed classifiers ˜w are sufficient to monitor invariance: Reparameterization permits restricting the search to representations whose environment-optimal classifiers equal a fixed scalar, yielding the practical IRMv1 formulation.The restriction preserves linear invariant predictors through rank-varying representations, while high-rank decompositions are important for out-of-distribution generalization.
- Extending to general losses and multivariate outputs: IRMv1 generalizes the invariance penalty as the squared gradient of risk at w = 1.0, making the fixed classifier’s optimality measurable for any convex loss.The resulting objective remains nonconvex, with multiple connected solution components and sensitivity to initialization.
4 Invariance, causality and generalization
IRM connects invariance across environments to causal structure and out-of-distribution generalization. Its theory characterizes when training-environment invariance identifies predictors that remain valid across interventions, while also exposing assumptions and open limitations.
- Generalization theory: IRM studies whether invariant predictors across training environments remain invariant across all valid intervention environments.The framework asks when low error and invariance on Etr imply low error and invariance on Eall.
- Causal assumptions: The causal framework assumes environments arise from interventions on a shared structural equation model with an acyclic causal graph.Valid interventions preserve the conditional mean of the target given its parents and keep target noise variance within a finite range.
- Causal interpretation: A predictor is invariant across all valid environments if and only if it achieves optimal robust out-of-distribution risk and uses only the direct causal parents of Y.This establishes a direct link between statistical invariance, causal prediction, and optimal ROOD.
- Theoretical extensions: The theory extends beyond ICP by allowing non-Gaussian data, scrambled latent causal setups, and interventions without requiring a causal graph or specific intervention types.Higher-rank representations can require fewer environments because they discard less information about the learning problem.
- Generalization theory: Under linear-theory conditions, a rank-r representation that elicits an invariant predictor across sufficiently diverse training environments also elicits it across all environments.The relevant training environments must lie in linear general position of degree r.
- Generalization behavior: IRM can extrapolate beyond training-environment interpolations, and its differentiability makes approximately invariant solutions robust to mild model misspecification.In the motivating example, robust learning works for σ ∈[10, 20], whereas IRM works for all σ.
- Limitations: The nonlinear theory remains incomplete, and the general theorem requires the number of training environments to scale linearly with the parameters in Φ.The paper reports that two environments are often sufficient empirically, but identifying when few environments suffice remains open.
- Causation as invariance: IRM treats invariant correlations as a statistically testable proxy for causal structure in settings where observed features, such as pixels, are distant from the underlying causal concepts.The paper emphasizes generalization across novel distributions rather than recovering the entire causal graph.
5 Experiments
IRM is evaluated on synthetic causal-weight recovery and Colored MNIST, where training environments vary spurious correlations. Across both settings, IRM generally recovers invariant predictors and generalizes better than ERM, though invariance can remain approximate.
- Synthetic data: The synthetic experiments vary feature scrambling and whether hidden confounders are observed, testing IRM under fully and partially observed graphs.Scrambled observations use an orthogonal transformation, whereas unscrambled observations use the identity.
- Synthetic data: IRM estimates accurate causal and non-causal weights across all synthetic conditions, usually outperforming ERM by orders of magnitude and ICP by a large margin.The comparison uses causal-weight squared error and the norm of non-causal weights; the y-axes are logarithmic.
- Colored MNIST: In Colored MNIST, color is more predictive than digit during training but reverses its correlation with the label in the test environment.The construction uses two training environments with p_e=0.2 and 0.1, and a test environment with p_e=0.9.
- Colored MNIST: ERM achieves high training accuracy but below-chance test accuracy, whereas IRM sacrifices training performance to rely less on color and generalize better.An oracle that ignores color by construction outperforms IRM only slightly.
- Colored MNIST: IRM is closer to invariance than ERM in Colored MNIST, although finite samples make invariance difficult to estimate in the distribution tails.The figure compares P(y = 1|h, e) across environments for ERM, IRM, and a grayscale oracle.
- Colored MNIST: Conditional domain adaptation could solve this particular task similarly, but it would fail under changes in the digit distribution across environments.This task has identical causal-feature distributions and P(y|e) across environments, conditions that make the alternative viable here.
6 Looking forward: a concluding dialogue
The concluding dialogue frames IRM as a response to ERM’s assumptions and its tendency to absorb spurious correlations under distribution shift. It emphasizes invariant predictors as a route to causal explanations, while noting limits from model capacity and changing conditional distributions.
- ERM and its assumptions: ERM assumes identical training and testing distributions, sufficient data relative to hypothesis-class capacity, and realizability; violating these assumptions can promote spurious correlations.The discussion presents these as three conditions underlying ERM’s usual guarantees.
- Causal and anticausal learning: In supervised annotation problems, P(Y e|X e) is often stable when the annotation process is shared, helping explain why ERM can succeed.The dialogue contrasts this with anticausal settings such as unsupervised learning of latent causal factors, where conditionals may vary.
- Limitations: Linear invariant predictors can be too restrictive: for Y e=(X e)^2 with Gaussian X e, the only linear invariant rule is the null predictor despite X being causal.This example motivates considering richer predictor structures and sufficient capacity.
- Invariant prediction: IRM may recover invariance with only two environments in some problems, whereas ERM can require many environments to achieve invariance in the realizable case.The dialogue presents this as a possibility rather than a universal guarantee.
- Future directions: IRM differs from regularization methods that encode known invariances through architectures or augmentation because it discovers unknown invariances from data.The dialogue also suggests environments could correspond to reinforcement-learning episodes or demographic groups.
A Additional theorems
The supplied passage contains only a fragment identifying a space of symmetric positive semi-definite matrices and a matrix indexed by environment.
- Additional theorems: The appendix fragment references the space of symmetric positive semi-definite matrices and an environment-indexed matrix.No further theorem statement or interpretation is available in the supplied passage.
B.1 Proof of Proposition 2
The supplied passage states KKT-based stationarity conditions for a constrained optimization problem with environment-indexed multipliers.
- Proof of Proposition 2: Under KKT differentiability and qualification conditions, nonnegative environment multipliers exist such that the Lagrangian gradient vanishes at the optimum.The passage denotes these multipliers by λe and the optimum by (f ⋆, M ⋆).
B.2 Proof of Theorem 4
The proof characterizes invariant predictors through the geometry of reachable classifier vectors and the representation’s nullspace. It establishes both directions between simultaneous optimality and orthogonality conditions on environment-specific risk gradients.
- Geometry of reachable classifiers: Reachable classifier vectors GΦ equal the orthogonal complement of the representation’s nullspace, GΦ = Ker(Φ)⊥.The proof shows inclusion by inner-product orthogonality and equality by matching subspace dimensions.
- From invariant optimality to orthogonality: If w simultaneously minimizes every environment risk after Φ, then v = Φ⊤w lies in Ker(Φ)⊥ and satisfies Φ∇Re(v) = 0 for every environment.The stationarity condition follows by differentiating each risk with respect to w.
- Converse construction: Conversely, if v⊤∇Re(v) = 0 for all environments, a representation can be constructed whose nullspace contains the gradients and whose reachable space contains v.This makes v representable as Φ⊤w while forcing the derivative with respect to w to vanish.
- Identifying the predictor: The proof obtains a contradiction when a constructed invariant predictor differs from the target predictor: gradient vectors force a nullspace dimension exceeding the rank assumption.The contradiction uses the span dimension of the environment-specific vectors and the assumed rank of Φ.
B.4 Proof of Theorem 10
The proof uses transversality and dimension counting to show that low-rank environment-gradient configurations are generically absent under suitable dimensional conditions. A measure-zero exception remains for nongeneric covariance choices.
- Geometric setup: The map G sends nonzero x to an m×d matrix, while W is its image and therefore has dimension at most d.The construction omits a single point because G is affine on Rd \ {0}.
- Avoiding low-rank configurations: For generic environment covariance matrices, W does not intersect rank-k matrix manifolds when k < d − r.The argument identifies W and the low-rank manifolds as the relevant geometric objects.
- Transversality: Thom’s transversality theorem implies that W and each low-rank manifold are transversal for generic covariance choices, with nongeneric choices forming a measure-zero set.Surjectivity of the derivative map establishes the required transversality.
- Dimension argument: Dimension counting makes the intersection impossible when k < d − r and m > d r + d − r.Under these conditions, a transversal intersection would have negative dimension.
- Surjectivity: The derivative with respect to each symmetric covariance matrix is constructed explicitly, proving surjectivity of the full derivative map.The construction uses symmetric matrices obtained through an orthogonal transformation and the spectral theorem.
C Failure cases for Domain Adaptation
The section shows that standard domain adaptation can enforce the wrong invariance and fail under label-shift or representation-dependent trade-offs. Conditional domain adaptation can also fail to identify invariant predictors in general.
- Marginal alignment: Domain adaptation matches source and target feature marginals while fitting a source-domain classifier, but this objective can conflict with predictive generalization.The method enforces P(Φ(Xes)) = P(Φ(Xet)) alongside source accuracy.
- Label-shift failure: 20% accuracy is the maximum achieved by the domain-adaptation classifier in a label-shift example, below random guessing, despite source-only training generalizing to the target.The failure follows because equal feature marginals force equal predicted-label marginals across environments.
- Conditional adaptation: Conditional domain adaptation enforces a stronger condition than invariant prediction when class proportions match, yet it cannot identify the invariant predictor for some general problems.Its conditional feature-distribution constraint is therefore not generally sufficient for invariant prediction.
- Generic failure: Almost surely, conditional domain adaptation rejects every nonzero representation eliciting an invariant prediction rule, even though the trivial representation Φ(x) = x succeeds.The admissible distribution pairs for any fixed representation form a measure-zero set.
- Bound limitation: The domain-adaptation bound contains λ⋆, the summed optimal-classifier error over domains, which depends on Φ rather than being a fixed constant.Ignoring this representation-dependent term omits a trade-off in the adaptation objective.
D Minimal implementation of IRM in PyTorch
The implementation computes an IRM penalty from gradients of losses under alternating environment subsets and optimizes it jointly with prediction error. The included example trains a representation and classifier over repeated iterations.
- Penalty computation: The implementation imports PyTorch’s automatic differentiation function to compute gradients for the IRM penalty.The penalty depends on differentiating environment-specific losses with respect to a dummy classifier weight.
- Penalty computation: The penalty multiplies gradients from alternating loss subsets and sums their elementwise products.This gradient agreement term is returned as the scalar penalty.
- Synthetic example: The example generates Gaussian variables across environments, constructs inputs by concatenating x and z, and predicts the sum of y’s coordinates.The environment parameter scales noise in the generated variables.
- Optimization loop: Each training iteration accumulates mean squared error and the gradient penalty across environments after randomly permuting each environment’s examples.The loss uses a representation multiplied by a dummy classifier weight.
- Optimization loop: Optimization backpropagates through 1e-5 * error + penalty and updates the parameters with opt.step().The prediction-error term is strongly downweighted relative to the penalty in the displayed objective.