[Paper Review] The Risks of Invariant Risk Minimization

Elan Rosenfeld, Pradeep Ravikumar+1 more|arXiv (Cornell University)|Oct 12, 2020

Adversarial Robustness in Machine LearningComputer Science43 references73 citations

TL;DR

This paper analyzes Invariant Risk Minimization (IRM) and its nonlinear variants, showing that IRM can fail to outperform ERM under both linear and nonlinear settings, and may rely on non-invariant features, especially with limited environments.

ABSTRACT

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.

Motivation & Objective

• Motivate and formalize the invariant prediction framework as a means to achieve out-of-distribution generalization by leveraging invariant causal mechanisms.
• Provide a rigorous analysis of IRM in a natural SEM-based model to determine when it succeeds or fails.
• Show that IRM can be no better than ERM and, in some cases, can rely on non-invariant features leading to poor generalization.

Proposed method

• Define a structural equation model with invariant features z_c and environmental (non-invariant) features z_e across multiple environments.
• Formalize the optimal invariant predictor that uses only invariant features and is invariant across environments.
• Analyze the IRM objective (bilevel and penalized forms) in the linear setting to establish necessary and sufficient conditions for optimal invariance given the number of environments E and environmental dimension d_e.
• Extend to the nonlinear regime to demonstrate that IRM can degrade to ERM-like behavior unless training environments sufficiently cover the space of non-invariant variations.
• Present corollaries showing that alternatives to IRM inherit similar limitations.

Experimental results

Research questions

• RQ1 Under what conditions does IRM recover the optimal invariant predictor in the linear regime?
• RQ2 Can IRM and its nonlinear variants outperform standard ERM or DRO under distribution shifts?
• RQ3 What happens when the number of training environments is small relative to environmental features (E ≤ d_e)?
• RQ4 How does IRM perform in the nonlinear setting when environments do not cover the space of non-invariant variations?
• RQ5 Do IRM alternatives suffer the same fundamental shortcomings as IRM?

Key findings

• In the linear case, if E > d_e, the IRM solution recovers the invariant predictor that uses only invariant features.
• If E ≤ d_e, there exists a linear predictor using environmental features that achieves lower training risk than the optimal invariant predictor.
• In the nonlinear regime, IRM can fail dramatically unless the training environments approximately cover the space of non-invariant features, making it behave similarly to ERM or DRO.
• There exist feasible predictors using environmental features that achieve lower risk on training environments than the optimal invariant predictor, indicating a failure to generalize via invariance.
• Corollaries show that many proposed IRM alternatives suffer from the same fundamental shortcomings as IRM in both linear and nonlinear settings.