PaperlayerView exact revision on arXiv

Source-linked AI summary

Does Invariant Risk Minimization Capture Invariance?

Pritish Kamath, Akilesh Tangella, Danica J. Sutherland, Nathan Srebro

arXiv:2101.01134v2stat.MLcs.AIcs.LG

TL;DR

This paper examines whether IRM captures natural invariances and whether those invariances support out-of-distribution generalization. It analyzes simple environments, practical linear and empirical variants, and the relationship between training and unseen environments, finding failures of linear IRM, sub-optimal invariant predictor selection, and sampling fragility.

  • Problem

    The paper addresses unresolved questions about whether IRMv1 approximates full IRM, whether invariant predictors generalize to unseen environments, and how IRM behaves with finite samples.

  • Method

    The paper studies IRM, its linear and practical variants, and invariant-predictor identification across simple and structured families of environments.

  • Results

    Linear IRM can miss natural invariances and generalize worse than ERM, while IRM can select an invariant predictor with worse out-of-distribution performance and training-sample invariance can fail to reflect population invariance.

  • Takeaways & Limitations

    Capturing invariance is insufficient for robust out-of-distribution generalization: the function class, predictor-selection criterion, environmental structure, and sampling regime all matter.

  • Takeaways & Limitations

    The stronger joint-distribution invariance that would avoid the loss mismatch is too strict for some problems, including settings where outcome marginals differ across environments.

Abstract

from arXiv · show

We show that the Invariant Risk Minimization (IRM) formulation of Arjovsky et al. (2019) can fail to capture "natural" invariances, at least when used in its practical "linear" form, and even on very simple problems which directly follow the motivating examples for IRM. This can lead to worse generalization on new environments, even when compared to unconstrained ERM. The issue stems from a significant gap between the linear variant (as in their concrete method IRMv1) and the full non-linear IRM formulation. Additionally, even when capturing the "right" invariances, we show that it is possible for IRM to learn a sub-optimal predictor, due to the loss function not being invariant across environments. The issues arise even when measuring invariance on the population distributions, but are exacerbated by the fact that IRM is extremely fragile to sampling.

1 INTRODUCTION

IRM is motivated by the need to ignore spurious correlations that cause failures under distribution shift, but this paper shows that its practical variants can miss invariances and worsen out-of-distribution generalization.

  • Motivation: Spurious background correlations can make image classifiers fail when test inputs differ from the training environments.A cow-versus-camel classifier may rely on grassy or desert backgrounds and then misrecognize cows on beaches.
  • Motivation: IRM seeks representations that discard spurious signals while retaining features that remain predictive across environments.The framework is motivated by identifying stable correlations across data collected from different users, regions, or years.
  • Open questions: The paper asks whether IRMv1 approximates full IRM, whether invariant predictors generalize, when training environments identify broader invariances, and how finite samples affect the framework.These questions cover both the formulation and its practical algorithmic behavior.
  • Contributions: On a simple Colored-MNIST-like problem, linear IRM can fail to find a truly invariant predictor even with population loss and infinitely many training environments.The resulting predictor can generalize worse on out-of-distribution environments than unrestricted ERM, and the issue persists for IRMv1.
  • Contributions: Even when IRM identifies the right invariance, minimizing training-environment loss can select an invariant predictor with worse out-of-distribution generalization.The paper also studies when invariances from a small training range can generalize and shows that this is generally impossible without additional environmental structure.
  • Contributions: Empirical-sample distributions can violate invariance even when the corresponding population distributions are invariant.This exposes an additional sampling-related issue for the IRM paradigm.

2 INVARIANT RISK MINIMIZATION

IRM learns representations whose predictor is simultaneously optimal across environments, while practical restrictions and relaxations create gaps between the ideal formulation, IRMS, IRMv1, and ERM.

  • IRM formulation: The framework seeks a representation that removes spurious features while preserving predictive information shared across environments.In the cow-versus-camel example, the representation might discard background color.
  • IRM formulation: IRM defines invariance by requiring one predictor on a representation to be optimal in every environment.The representation and predictor are initially unrestricted measurable functions, with the same predictor simultaneously minimizing each environment’s population loss.
  • Function classes: Restricted function classes define which predictors count as invariant, so the choice of representation and predictor classes is part of the practical formulation.The paper distinguishes general IRM from variants such as linear invariant predictors.
  • Linear variants: Linear invariant predictors use Z = R^d and linear predictors, with the scalar case providing a simple characterization of the corresponding invariant predictor class.The paper relates this linear formulation to the scalar invariant-predictor setting.
  • Algorithms: IRM minimizes training-environment risk over invariant predictors, while ERM simply minimizes loss over the training environments without an invariance constraint.IRMS remains bi-level, and IRMv1 softens its hard constraint for practical optimization.
  • Algorithms: IRMv1 replaces population loss with empirical loss and optimizes a parameterized representation, such as a deep neural network, using gradient-based local search.These changes address the lack of direct population-loss access and the impracticality of optimizing over unrestricted representations.
  • Gaps and related work: IRMS can differ from full IRM even with population loss and unrestricted representations, while prior work also shows near-optimal IRMv1 solutions can perform no better than ERM far from training distributions.Empirical domain-generalization results cited by the paper likewise report strong practical performance from ERM.

3 COLORED-MNIST AND TWO-BIT ENVIRONMENTS

The two-bit analysis shows that IRM’s linear relaxation can admit predictors that are not truly invariant, producing worse out-of-distribution performance than ERM. Practical IRMv1 reproduces this failure on Colored-MNIST, while the motivating setting also illustrates when IRM succeeds.

  • Motivating example: IRM succeeds in the motivating Colored-MNIST-style setting when it selects the predictor based on X1, whereas ERM exploits the stronger but spurious X2 signal.When X2’s correlation flips, ERM has square loss 0.985, compared with 0.375 for fIRM and 0.5 for the null predictor.
  • Two-bit environments: IRM’s linear formulation can select a predictor that is worse out of distribution than both IRM and unrestricted ERM.For Etr={(0.1,0.2),(0.1,0.25)}, the selected predictor relies heavily on the non-invariant feature X2 despite X1 being the stronger signal.
  • Two-bit environments: Adding more training environments does not necessarily fix the failure: fIRMS remains preferred when all added environments have βe<0.28.For these two-bit environments, any two distinct environments suffice to recover the full invariant sets, yet IRMS can still choose the wrong feasible solution.
  • Two-bit environments: The feasible set of IRMS solutions can be strictly broader than the set of truly invariant predictors, so satisfying the gradient condition need not capture invariance.For E0.1, IS(Etr)⊋I(Etr), and the extra solutions can achieve lower training loss while generalizing worse to environments such as (0.1,0.9).
  • Practical experiments: IRMv1 interpolates from ERM toward IRMS as λ increases, with reliance on X2 increasing toward the large-penalty limit.In neural-network experiments, the failure is essentially reproduced on Colored-MNIST; in the problematic setting, ERM and IRMv1 perform about the same in test loss.

4 CAN IRM FAIL TO CHOOSE THE RIGHT PREDICTOR?

IRM can identify the invariant predictors yet still select one with worse out-of-distribution performance, because invariant predictors may have different losses across environments.

  • IRM can prefer an invariant predictor that performs worse on unseen environments, even when training environments identify all invariant predictors.The paper constructs environments where f2 has lower loss on every training environment, while f1 has lower worst-case loss across all environments.
  • In the constructed family, f1(x)=0.3x1 and f2(x)=0.3x2 are both invariant, but environments with θe<0 favor f2 during training while f1 minimizes worst-case loss.For almost all training sets containing at least two distinct environments, the invariant set is identified, yet predictor choice remains misaligned with out-of-distribution risk.
  • The loss of an invariant predictor need not remain invariant across environments, so training-environment optimization can select the wrong predictor.IRM preserves conditional-label invariance but does not require the predictor’s loss to be equal across environments.
  • The appropriate invariance notion depends on the full environment set and cannot guarantee optimal out-of-distribution loss merely by capturing the right invariances.Stronger joint-distribution invariance would make losses invariant, but can be too restrictive and may force a trivial constant predictor.

5 WHEN DOES INVARIANCE GENERALIZE?

Whether invariance learned from training environments generalizes depends on how environments are parameterized: continuity alone is insufficient, whereas analyticity can provide guarantees under stated conditions.

  • Generalization of invariance requires representative training environments; without such structure, adding an irrelevant environment can invalidate the extension from Etr to E.The inclusion I(EΠ)⊆I(Etr) always holds, but the reverse inclusion needs assumptions on the environment family and training subset.
  • Continuity alone does not ensure that invariances found on training environments generalize, even when the training set contains uncountably many environments.A continuous environment map can change behavior outside the training parameter subset.
  • For analytic environment parameterizations, invariances can generalize from training environments to the full parameterized environment family under finite-space and connected-domain conditions.Proposition 8 assumes analytic Π, finite X and Y, a connected open parameter domain, and Etr=Π(Θtr).
  • The paper’s earlier examples satisfy analyticity, so their training-environment invariances extend to all environments for almost all training sets with at least two distinct environments.The result is stronger for unrestricted IRM than for its restricted variant, whose analysis requires uncountably many training environments.

6 IRM WITH FINITE SAMPLES

Finite-sample estimation makes IRM especially fragile: exact invariance constraints can become infeasible, while IRMv1 may approach a trivial predictor as its penalty grows.

  • With empirical distributions, exact IRM and IRMS constraints are likely to admit no invariant predictors because sampling noise perturbs correlations across environments.The population-defined invariant set can become empty when applied directly to finite samples.
  • IRMS learns the trivial zero predictor on a perturbed proxy for finite-sample training environments, illustrating failure of exact invariance under sampling.The proxy environments are {(0.245,0.105),(0.255,0.195),(0.251,0.302)}.
  • For a fixed empirical distribution, IRMv1 is likely to approach IRMS as λ→∞ and therefore may fail to find a good invariant predictor.Taking the sample size to infinity first at fixed λ instead approaches the population IRMv1 objective, so the order and rate of limits matter.
  • Making IRM robust to finite samples may require relaxing exact optimality constraints, but a practical ε-IRM implementation remains an open challenge.The paper proposes an ε-relaxation of the predictor optimality requirement without providing a practical implementation.

7 DISCUSSION

The discussion highlights substantial gaps between IRM's full formulation and practical variants, including failures of linear methods, ambiguity among invariant predictors, and severe sampling brittleness.

  • IRM and IRMv1 can differ substantially even on simple environments, motivating practical algorithms for nonlinear IRM classes.The discussion presents this gap as central to realizing IRM's promise.
  • Even when predictors satisfy a stronger conditional invariance notion, IRM may select an invariant predictor with worse out-of-distribution robustness.Joint-distribution invariance would suffice for worst-case robustness, but the paper describes it as overly stringent.
  • With significant structure on the environment family, invariances learned from training environments can generalize so that I(E) = I(Etr), but broader settings remain unresolved.The paper specifically notes remaining questions for IS and more realistic assumptions on environments.
  • IRM and IRMv1 are surprisingly brittle on samples rather than population distributions, indicating that further analysis or new algorithms are needed for practical use.The discussion contrasts population-based formulations with empirical-sample behavior.
  • For logistic loss, the framework must accommodate cases where an exact minimizer does not exist, such as separable problems requiring unbounded predictors.The appendix modifies the definition using simultaneous ε-optimality and also discusses closure-based workarounds.

A.2 Proofs of Observation 2 and Lemma 4

The appendix proves that invariance corresponds to shared environment-wise optimality and establishes relationships among unrestricted, restricted, and linear invariant predictor classes.

  • The appendix frames the invariance criterion as simultaneous membership of the same predictor in each environment's risk minimizer set.This follows directly from the definition of invariant representation and the environment-wise conditional optimization problem.
  • Under square loss, invariant representations require equal conditional expectations across environments on shared representation values, which yields a simultaneously optimal predictor.Conversely, matching conditional expectations defines a predictor optimal in every environment.
  • The proof constructs a scalar representation from any invariant predictor, showing that invariant predictors under a function class closed under scalar multiplication also belong to the scalar invariant class.This establishes the inclusion IW(E) ⊆ IS(E), with the linear case following from scalar-multiplication closure.
  • A linear representation can be embedded into Rd along a fixed unit vector, preserving the predictor while placing it in the linear invariant class.The construction defines ϕ′(x) = ϕ(x) · v and composes it with a linear map.

B More details on Two-Bit Environments (from Section 3)

For the two-bit environment family, two distinct training environments suffice to identify both IS and I, although the restricted and unrestricted invariant classes differ.

  • For every α ∈ (0, 1), any two distinct environments determine the full invariant classes: IS(Etr) = IS(Eα) and I(Etr) = I(Eα).Thus the failure studied earlier is not caused by insufficiently representative training environments.
  • The restricted class is characterized by gradients affine in βe, so vanishing at two distinct βe values forces vanishing across the entire environment family.The decomposition ∇w|w=1 Le(w · ϕ) = F(ϕ) + βeG(ϕ) supplies the key algebraic step.
  • The unrestricted invariant representations over the training environments are precisely those whose conditional label expectations do not vary with βe on shared representation fibers.Enumerating all nonempty subsets identifies the resulting invariant partitions.
  • Among the identified partitions, constant and x1-based representations yield the zero predictor or an x1 predictor, while the x1x2 representation again yields the zero predictor.These predictors are invariant over Etr and, in the cases described, over Eα as well.
  • All invariant representations found over Etr remain invariant over Eα, establishing I(Etr) ⊆ I(Eα) before the proposition proves equality.The corresponding predictors are simultaneously optimal across Eα.

B.1 Case of square loss

Under square and logistic loss, the two-bit construction produces multiple restricted invariant predictors, and IRMS can select one that performs worse than IRM on shifted environments.

  • B.1 Case of square loss: The square-loss phenomenon extends beyond α = 0.1: analogous behavior occurs for α < 0.1464 or α > 0.8536.These cutoff values arise from the closed-form analysis of the additional linear solutions.
  • B.2 Case of logistic loss: For logistic loss at α = 0.05, IRMS again selects f1 under βe < 0.25, while IRM with two environments learns fIRM.The logistic-loss construction reproduces the same restricted-versus-unrestricted distinction seen under square loss.
  • B.2 Case of logistic loss: The logistic-loss phenomenon is not specific to α = 0.05; similar behavior occurs for every α < 0.077.Figure 7 reports losses of the odd predictors across the environment family.

C More Colored-MNIST experiments

Additional Colored-MNIST experiments vary the loss, architecture, and which feature is invariant, showing that the expected IRM failure is less visible under ℓlog and depends on architectural details.

  • Architectures: The original architecture is fully connected and uses ReLU layers, while the split architecture gives color and grayscale inputs separate branches before combining them.The original network maps 2 · 14 · 14 inputs through hidden layers to a scalar prediction.
  • Experimental setup: The experiments reproduce Colored-MNIST with ℓsq and ℓlog under both the original and split architectures.The split architecture separately processes grayscale digits and one-hot color before combining them.
  • Experimental setup: The comparisons include ERM models forced to use only color or only grayscale digits, isolating variation caused by the learning process.Color-only receives a two-dimensional one-hot color indicator, while digit-only receives the flattened grayscale image.
  • Results: Under ℓlog, the expected IRM failure mode is generally less visible, but it is more evident when color rather than digit identity is invariant.These results are shown for both the split architecture and the architecture used by Arjovsky et al. (2019).

D More details on failure of IRM (Section 4)

The failure analysis shows that, in the two-bit setting, invariant predictors depend on at most one feature, yet IRM can select the feature that generalizes worse across environments.

  • Invariant predictors: For almost all training-environment sets containing at least two environments, the invariant predictors equal those for the full environment family and depend on at most one input bit.The possible invariant predictors are functions of x1 alone, x2 alone, or neither.
  • Invariant predictors: There are 37 non-empty subsets with environment-independent conditional expectations, but only 6 yield non-zero expectations.These cases are enumerated symbolically and summarized in Table 2.
  • Loss comparison: For squared loss, the best predictors restricted to x1 or x2 are f1(x)=0.3x1 and f2(x)=0.3x2, so IRM selects one of these two.The same selection argument applies under logistic loss.
  • Failure mode: When training environments have θe<0, f2 has lower loss throughout training even though f1 minimizes worst-case loss over all environments.In this case, the worst-case loss of f2 equals that of the zero predictor at 0.5.

E More details on Generalization of Invariance (from Section 5)

The generalization analysis gives conditions under which invariance identified from training environments extends to a broader environment family, while showing that continuous settings require stronger coverage assumptions.

  • Failure outside training range: A training-restricted invariant predictor can fail to remain invariant globally when its conditional expectation changes in environments outside the training range.The example uses f1(X)=0.8X1, which is invariant on Etr but not on E.
  • Finite environments: For finite spaces under analytic environment parameterizations, almost all training sets with at least two environments recover the invariant predictors of the full environment family.The proof uses analyticity and the fact that non-identically-zero analytic functions have zero sets of measure zero.
  • Finite environments: The same recovery holds for invariant representations when the training-parameter set has non-zero Lebesgue measure.This establishes equality for both I(Etr) and the representation-level set IS(Etr) in the stated finite-space setting.
  • Infinite spaces: The infinite-space proof assumes bounded outputs and bounded representations, enabling analytic arguments involving conditional expectations and gradient penalties.The boundedness assumption permits the relevant interchange used in the proof.
  • Infinite spaces: For infinite measurable spaces, the corresponding result requires a training-parameter set of non-zero Lebesgue measure rather than merely almost any set containing two environments.The paper leaves open whether the stronger condition can be weakened for finite training sets.
Loading 2101.01134v2…