Source-linked AI summary

The Risks of Invariant Risk Minimization

Elan Rosenfeld, Pradeep Ravikumar, Andrej Risteski

arXiv:2010.05761v2cs.LGcs.AIstat.ML

TL;DR

IRM was proposed to learn invariant predictors for out-of-distribution generalization, but formal guarantees for latent-variable settings remain limited. This paper formally analyzes IRM and related objectives under a general classification model, finding frequent linear failures and catastrophic nonlinear failures unless test environments resemble or are covered by training environments.

  • Problem

    Formal guarantees are severely lacking for IRM and related invariant-representation objectives when observations are complex functions of unobserved latent factors.

  • Method

    The paper formally analyzes classification under IRM and related objectives using a natural, general model with invariant and environment-dependent latent features.

  • Results

    IRM frequently fails to recover the optimal invariant predictor: linear solutions can rely on environmental features, while nonlinear solutions can fail catastrophically under sufficiently different test distributions.

  • Takeaways & Limitations

    In this setting, IRM and its alternatives fundamentally do not improve over standard ERM.

  • Takeaways & Limitations

    Nonlinear generalization can be guaranteed only when training distributions approximately cover the full space of environments, a condition that removes the need for ICP.

Abstract

from arXiv · show

Invariant Causal Prediction (Peters et al., 2016) is a technique for out-of-distribution generalization which assumes that some aspects of the data distribution vary across the training set but that the underlying causal mechanisms remain constant. Recently, Arjovsky et al. (2019) proposed Invariant Risk Minimization (IRM), an objective based on this idea for learning deep, invariant features of data which are a complex function of latent variables; many alternatives have subsequently been suggested. However, formal guarantees for all of these works are severely lacking. In this paper, we present the first analysis of classification under the IRM objective--as well as these recently proposed alternatives--under a fairly natural and general model. In the linear case, we show simple conditions under which the optimal solution succeeds or, more often, fails to recover the optimal invariant predictor. We furthermore present the very first results in the non-linear regime: we demonstrate that IRM can fail catastrophically unless the test data are sufficiently similar to the training distribution--this is precisely the issue that it was intended to solve. Thus, in this setting we find that IRM and its alternatives fundamentally do not improve over standard Empirical Risk Minimization.

1 INTRODUCTION

The paper examines whether IRM can recover invariant predictors for out-of-distribution generalization when data contain latent, environment-dependent features. It finds that formal analysis often reveals failures rather than improvement over ERM.

  • Motivation: OOD generalization is difficult because test distributions may differ from training distributions, causing ERM and deep models to fail under minor shifts.Deep networks may rely on statistically informative but non-causal features.
  • Invariant prediction: Invariant Causal Prediction assumes changing environments alter distributions while leaving the underlying causal mechanisms constant.This frames distribution shift through structural causal models and interventions.
  • Invariant prediction: The optimal invariant predictor uses only invariant SEM features and therefore avoids overfitting to environment-specific properties.The paper focuses on IRM and extends its analysis to related objectives.
  • Research gap: Prior invariant-representation objectives for unobserved latent factors have encouraging empirical results but few formal guarantees.This gap is especially important for complex, high-dimensional data such as images.
  • Contribution: The paper formally analyzes classification under IRM and shows that its optimal solution frequently performs no better than ERM.The analysis covers exact linear conditions and failure cases in the nonlinear regime.

2 RELATED WORK

Related work studies domain and conditional invariance through several objectives, but these approaches differ in assumptions and generally lack guarantees matching the paper’s setting.

  • Domain generalization: Recent domain-generalization objectives seek invariance of p(Φ(x) | y), while IRM assumes invariance of E[y | Φ(x)].Follow-up methods impose invariance on higher conditional moments.
  • Domain generalization: These methods resemble covariate shift but restrict shifts to occur between discrete, labeled environments rather than arbitrary train-to-test changes.The distinction concerns how and where changes in p(Φ(x)) are represented.
  • Concurrent theory: Concurrent theory finds that ERM can remain optimal in linear domain generalization, mirroring this paper’s conclusion that studied objectives do not outperform ERM.Related work also reports that extrapolation is exponentially harder than interpolation.

3 MODEL AND INFORMAL RESULTS

The paper models invariant and environmental latent features across observed environments, then analyzes IRM’s linear and nonlinear behavior. Its informal results identify threshold conditions and severe failures under insufficient environmental coverage.

  • Model: The model separates invariant features zc from environmental features ze whose distributions vary across environments.The invariant feature-label relationship is fixed, while environmental parameters may change.
  • Model: The observation x is generated from latent features through an injective function, allowing the latent variables to be recovered in principle.The only stated assumption on the observation function is injectivity.
  • Model: The environmental features ze are deliberately allowed to depend on y, distinguishing label-caused correlations from non-causal spurious features.The paper avoids calling every non-causal correlation spurious.
  • Objective: Infinite samples from each training environment isolate the objectives’ behavior from finite-sample effects.The goal is to minimize risk on an unseen environment while extracting invariant features.
  • Linear results: If E > de, optimal linear IRM recovers the invariant predictor; if E ≤ de, it relies on non-invariant features.The threshold depends on the number of environments and environmental-feature dimension.
  • Linear results: When E ≤ de, an environmental-only linear predictor can achieve lower training risk than the optimal invariant predictor.This makes the non-invariant solution more attractive to the objective despite its intended role.
  • Nonlinear results: In the nonlinear regime, IRM can appear nearly optimal on training data yet behave like ERM on test environments with sufficiently different means.Reversed environmental correlations can reduce its accuracy to almost 0.
  • Implications: IRM and related objectives require training environments to approximately cover possible environments for nonlinear generalization, eliminating their advantage over ERM or DRO.The paper extends the same failure results to suggested IRM alternatives.

4 BACKGROUND ON IRM AND ITS ALTERNATIVES

IRM seeks feature representations whose optimal classifier is invariant across environments, but its alternatives share a potentially misleading stationarity guarantee.

  • The IRM objective: IRM learns environmentally invariant representations by requiring the optimal classifier on top of learned features to be identical across environments.The objective is motivated by excluding features whose label relationships vary across environments.
  • The IRM objective: The original IRM formulation is a constrained, bilevel optimization problem that is highly non-convex and difficult to solve.A Lagrangian approximation penalizes the squared gradients of the inner optimization problems.
  • The IRM objective: With a convex inner problem, feasibility corresponds to a zero penalty, making the constrained and penalized formulations equivalent when λ = ∞.
  • Alternative objectives: Follow-up objectives strengthen invariance from E[y | Φ(x)] to p(y | Φ(x)) and include penalties based on the variance of risks.The cited alternatives include risk-variance penalties and a square-root variance variant.
  • Alternative objectives: The optimal invariant predictor is a stationary point of IRM and related objectives, but stationarity alone is misleading and requires further analysis.

5 THE DIFFICULTIES OF IRM IN THE LINEAR REGIME

In the linear regime, IRM recovers invariant prediction only with sufficiently many environments; with too few, feasible solutions can exploit environmental features and outperform the invariant predictor in training environments.

  • Main linear results: Observing more than de environments forces feasible linear IRM solutions to discard environmental features under mild non-degeneracy conditions.The resulting invariant classifier is optimal for all possible environments.
  • Main linear results: When E ≤ de and environmental means are linearly independent, feasible solutions can retain environmental features while achieving lower logistic and 0-1 risks than the optimal invariant predictor.Such solutions can have rank(B) = de + 1 − E.
  • Generalization consequences: When E ≤ de, the global minimum necessarily uses non-invariant features and therefore does not universally generalize to unseen environments.When E > de, the optimal invariant predictor has minimum and minimax risk among feasible predictors.
  • Main linear results: For logistic and 0-1 loss, the optimal invariant predictor is the global IRM minimum if and only if E > de.
  • Environmental-feature failure: A predictor depending on environmental features can still have an environment-independent optimal classifier because its feature distributions vary while the optimal coefficients remain constant.It can outperform the invariant predictor because it retains more predictive information.
  • Environmental-feature failure: A feasible predictor can use only environmental features and achieve lower 0-1 risk than the optimal invariant predictor when σe˜µ > σ−1.The construction is more likely to matter when E ≪ de and some non-invariant features correlate with labels.

6 THE FAILURE OF IRM IN THE NON-LINEAR REGIME

In the non-linear regime, IRM admits solutions that look invariant on training data but behave like ERM on sufficiently different test environments, so generalization requires covering the environment space.

  • Non-linear failure: Non-linear IRM solutions can appear invariant unless training environments sufficiently cover the space of non-invariant features.This exposes a foundational failure of the penalized IRM objective.
  • Constructed solution: For any ϵ > 1, a constructed featurizer can nearly satisfy the IRM penalty while matching the invariant predictor on at least 1 − pϵ of training data.On the remaining inputs, it matches the ERM-optimal solution, and the penalty scales exponentially small in de.
  • Test-time behavior: The constructed predictor is equivalent to ERM on at least 1 − q of a sufficiently different test distribution and can have 0-1 risk greater than .975 − q.
  • Test-time behavior: A separation of Ω(√de log E) between test and training environmental means can suffice for IRM to fail to use invariant prediction.When test environments are expected to be similar, ERM already guarantees reasonable performance, leaving no improvement for IRM in this regime.
  • Test-time behavior: If non-invariant feature-label correlations reverse at test time, the constructed predictor, like ERM, can perform significantly worse than chance.
  • Proof intuition: The proof constructs a predictor equal to the invariant predictor near training means and equal to the ERM solution outside those regions.Training concentration makes the exceptional region small during training but potentially large under a moderately shifted test distribution.
  • Scope and implication: Generalization can be guaranteed by having training distributions cover the full environment space, but then ERM or DRO may perform just as well.The same failure extends to the alternative objectives.

7 CONCLUSION

The paper finds significant under-explored risks in IRM and related objectives, especially for latent-variable models in nonlinear settings. It calls for further theoretical study of invariant prediction with complex functions of unobserved factors.

  • IRM and related objectives have significant under-explored risks and formulation issues.The conclusion frames these risks as motivating further theoretical work.
  • Formal results remain lacking for latent-variable models, particularly nonlinear settings with fully unobserved covariates.
  • The paper raises the question of how invariant prediction should be formulated for complex nonlinear observations of unobserved latent factors.
  • Out-of-distribution generalization remains an important direction, and Invariant Causal Prediction remains promising.

B PROOF OF PROPOSITION 4.1

The proof shows that the optimal invariant predictor satisfies the IRM stationarity conditions for both classifier and feature-map parameters. Its classifier is constant across environments, making the constraint feasible.

  • The argument begins from the IRM objective and its resulting constraint.
  • The proof obtains the feature-map derivative after parameterizing Φ as Φθ and differentiating logistic-regression negative log-likelihood with respect to β.
  • The IRM constraint is satisfied because the classifier ˆβ is constant across environments.This makes ˆβ minimizing for the training data as a whole.
  • The optimal invariant predictor is also a stationary point with respect to the feature-map parameters.

C.1 PROOF OF THEOREM 5.1

The linear proof characterizes when IRM excludes environmental features and when it instead admits lower-risk, non-invariant solutions. More than de environments force invariant use under mild conditions, whereas E ≤ de permits counterexamples.

  • Assumptions: The non-degeneracy assumptions require environmental means in general linear position and impose a relative-variance condition; they fail only on measure-zero parameter sets.
  • Construction: For E ≤ de, Lemma C.2 constructs a feasible linear feature map of rank dc + de + 1 − E whose optimal classifier remains invariant despite environmental dependence.
  • Construction: The construction works because adding environmental directions can preserve an invariant optimal classifier while enlarging the feature-map span.
  • Exclusion: For E > de, environmental means become linearly dependent, and the proof uses this dependence to show environmental coefficients must vanish.
  • Theorem 5.1: When E > de, any feasible linear predictor with nonzero invariant classifier must have B = 0, so it uses no environmental features.Under mild non-degeneracy conditions, the resulting classifier is optimal for all possible environments.
  • Theorem 5.1: When E ≤ de and environmental means are linearly independent, a feasible predictor can use environmental features and achieve lower logistic and 0-1 risks than the optimal invariant predictor.The construction has rank(B) = de + 1 − E.
  • Risk comparison: The resulting invariant-only predictor has identical risk across environments, while dropping the environmental term from the constructed predictor strictly increases risk.

C.2 EXPERIMENTS FOR THEOREM 5.1

Experiments reproduce the theorem’s threshold behavior: IRM often matches ERM through E = de, then improves under distribution shift once E exceeds de. Simulations also indicate that environmental-only predictors can remain competitive until many environments are observed.

  • Setup: The experiments used dc = 3 and de = 6, with invariant features yielding 71.9% accuracy before evaluating reversed environmental correlations.
  • Experimental stability: IRM optimization across five runs was unstable and often collapsed to ERM unless λ and the optimizer learning rate were carefully tuned.
  • Figure C.1: Up until E = de, IRM essentially matches ERM both in-distribution and under distribution shift; after crossing de, it performs drastically better under shift.After the threshold, IRM behaves more like the optimal invariant predictor.
  • Environmental-feature simulations: Theorem 5.3 states that for E ≤ de, an environmental-only predictor can achieve lower 0-1 risk whenever σe˜µ exceeds σ−1.
  • Environmental-feature simulations: The simulations estimate σe˜µ across environmental dimensions and show the crossover threshold is large, on the order of de − dc.
  • Environmental-feature simulations: In the high-dimensional setting, the simulated conditions for the environmental-only predictor to outperform the invariant predictor are very likely to hold.

D.1 PROOF OF THEOREM 6.1

The proof constructs a featurizer that behaves invariantly near training environmental means but reverts to ERM elsewhere, yielding low IRM penalties while failing under shifted test environments.

  • Construction and penalty control: The constructed featurizer recovers environmental features only inside a designated set, allowing the gradient penalty to be bounded by that set’s measure and geometry.This bound is combined with auxiliary lemmas to control the overall penalty.
  • Construction and penalty control: When the test environmental mean is sufficiently separated from all training means, samples fall outside the constructed region with high probability, where the predictor equals ERM.The separation condition is used to show that the constructed predictor behaves like ERM on most test observations.
  • Theorem properties: The featurizer combined with the ERM-optimal classifier matches the optimal invariant predictor on at least a 1 −pϵ,de fraction of training observations.On the remaining inputs, it matches the ERM-optimal solution.
  • Theorem properties: The same predictor is equivalent to the ERM-optimal predictor on at least a 1 −q fraction of the test distribution.Thus, under the separation condition, its test behavior is predominantly ERM-like.
  • Conditions and implications: For environmental means drawn from a high-dimensional Gaussian prior, separation of order Ω(√de log E) implies δ ∈Ω(√log E) and q ∈O(1/E).The resulting q is described as sufficiently small to ensure worse-than-random test error.
  • Conditions and implications: The conditions can hold when test environments have the same scale as training environments but reversed correlations, yet IRM still does not outperform ERM.The paper identifies this as precisely the setting where IRM is intended to help.

E.1 EXTENSIONS FOR THE LINEAR CASE

The linear-case analysis extends to risk-variance penalties and related objectives, showing that non-invariant predictors can achieve better objectives or lower risk under explicit conditions.

  • Extensions to related objectives: The constraint used by one alternative objective is strictly stronger than the corresponding IRM constraint, so the Section 5 results extend to it.When the stronger constraint holds, the alternative objective’s penalty is necessarily zero.
  • Risk-variance penalties: For E ≤de environments with linearly independent means and identical variance, a predictor using non-invariant features achieves a lower objective than the optimal invariant predictor for every λ ∈[0, ∞].This applies to empirical risk minimization with a penalty on risk variance.
  • Risk-variance penalties: Under constant environmental variance, the risk-variance penalty is zero and the objective is independent of λ, while the resulting risk is lower than that of the optimal invariant predictor.The argument uses a constructed featurizer and the linear-case risk comparison.
  • Risk-variance penalties: REx assumes invariance of Var(y | Φ(x)), which is stronger than IRM’s assumed invariance of E[y | Φ(x)].The additional second-moment requirement arises from the design of the objective.
  • Risk-variance penalties: Even with this stronger invariance requirement, REx admits environmental-only predictors with lower 0-1 loss than the optimal invariant predictor under stated linear conditions.The paper therefore identifies another failure case for REx.
  • Empirical extension: The modified risk-variance condition depends on training-environment geometry similarly to the original condition, and simulations report correspondingly similar results.The comparison concerns the geometry of environmental means and variances.

F TECHNICAL LEMMAS

The technical lemmas establish optimal logistic coefficients, risk comparisons for feature subsets, Gaussian concentration bounds, and the penalty estimates used in the main constructions.

  • Optimal predictors: Under the stated nondegeneracy conditions, any classifier using a strict subset of features has strictly higher logistic risk than the Bayes classifier.The same strict-risk conclusion holds for 0-1 loss when the omitted contribution is nonzero with positive probability.
  • Optimal predictors: For a feature vector formed as a linear function of invariant and environmental features, an optimal environment-specific coefficient has a closed-form expression.The derivation uses conditional independence and Gaussian convolution.
Loading 2010.05761v2…