The Risks of Invariant Risk Minimization
Authors: Elan Rosenfeld
Pradeep Kumar Ravikumar
Andrej Risteski
Venue: ICLR 2021 Poster (2021)
Keywords: out-of-distribution generalization
causality
representation learning
deep learning

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 is 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.

Published as a conference paper at ICLR 2021
THE RISKS OF INVARIANT RISK MINIMIZATION
Elan Rosenfeld, Pradeep Ravikumar, Andrej Risteski
Machine Learning Department Carnegie Mellon University
elan@cmu. edu, pradeeprêcs emu, edu, aristesk@anddeww cmu edu
ABSTRACT
Invariant Causal Prediction (Peters al., 2016) a technique for out-of-distributtin generalization which assumes that some aspects of the data distribution vary across the training set but that the underlying causal mechanisms remain constant. Re- cently, Arrovssy et al. (2019) proposed Invariant Risk Minimization (IRM), an objective based on this idea for learning deep, invariant features of data which are complex function of latent variables; many alternatives have subsequently been suggested. However, formal guarantees for lacking. In this paper, present the first analysis of classification under the IRM objective-ar as well as these recently proposed alternatives under fairly natural and general model. In the linear case, give 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-tiie 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.
these works are severely
INTRODUCTION
Prediction algorithms are evaluated by their performance on unseen test data. In classical machine learning, it common to assume that such data are drawn i.idd. from the same distribution as the data set which the learning algorithm was trained-in the real world, however, this is often not the case. When this discrepancy occurs, algorithms with strong in-distribution generalization guarantees, such as Empirical Risk Minimization (ERM), can fail catastrophically. In particular, while deep neural networks achieve superhuman performance on many tasks, there evidence that they rely on statistically informative but non-causal features in the data (Beery al.) 2018 Geirhos al. 2018 Ilyas et al. 2019). As a result, such models are prone to errors under surprisingly minor distribution shift (Su et 2019 Recht al., (019).. To address this, researchers have investigated alternative objectives for training predictors which are robust to possibly egregious shifts in the test distribution.
The task of generalizing under such shifts, known as Out-of-Dittiibution (OOD) Generalization, has led to many separate threads of research. One approach is Bayesian deep learning, accounting for a classifier's uncertainty at test time (Neal) 2012). Another technique that has shown promise data augmentation this includes both automated data modifications which help prevent overfitting (Shorten & Khoshgoftaar 2019) and specific counterfactual augmentations ensure invariance in the resulting features (Volpi et al. 2018 Kaushik et al. 2020).
strategy which has recently gained particular traction Invariant Causal Prediction (ICP: Peters ct al.01))) which views the task of OOD generalization through the lens of causality. This framework assumes that the data are generated according to Structural Equation Model (SEM; [Bollen 2005). which consists of set of so-called mechanisms structural equations that specify variables given their parents. ICP assumes moreover that the data can be partitioned into environments, where each environment corresponds to interventions on the SEM (Pearl 2009), but where the mechanism by which the target variable is generated via its direct parents is unaffected. Thus the causal mechanism of the target variable is unchanging but other aspects of the distribution can vary broadly. As result, learning mechanisms that are the same across environments ensures recovery of the invariant features which generalize under arbitrary interventions. In this work, consider objectives that attempt to

---

Published as a conference paper at ICLR 2021
learn what we refer to the "'optimal invariant predictor" this is the classifier which uses and is optimal with respect to only the invariant features in the SEM. By definition, such a classifier does not overfit eppiroi properties of the data distribution, so it will generalize even under major distribution shift at test time. In particular, we focus our analysis on one of the more popular objectives, Invariant Risk Minimization (IRM: Arjovsky et al. (2019)), but results can easily be extended to similar recently proposed alternatives.
Various works on invariant prediction (Muandet al., 2013 Ghassami et al. 2017 Teinze-Demi et al. 2018 Kojas-Carulla et al. 2018 Subbaswamy et al. [019]] Christiansen et al. 2020 consider regression in both the linear and non- setting, but they exclusively focus on learning with fully
or partially observed covariates or some other source of information. Under such a condition, results
from causal inference (Maathuis et al., 2009 Peters et al. 2017) allow for formal guarantees of the identification of the invariant features, or at least strict subset of them. With the rise of deep learning, more recent literature has developed objectives for learning invariant representations when the data are a non-linear function of unobserved latent factors, a common assumption when working with complex, high-dimensional data such as images. Causal discovery and inference with unobserved confounders or latents is a much harder problem (Peters et al., 2017), , so while empirical results seem encouraging, these objectives are presented with few formal guarantees. IRM is one such objective for invariant representation learning. The goal of IRM is learn feature embedder such that the optimal linear predictor on top of these features is the same for every environment- the idea being that only the invariant features will have an optimal predictor that is invariant. Recent works have pointed to shortcomings of IRM and have suggested modifications which they claim prevent these failures. However, these alternatives are compared in broad strokes, with little in the way of theory.
In this work, we present the first formal analysis of classification under the IRM objective under a fairly natural and general model which carefully formalizes the intuition behind the original work. Our results show that despite being inspired by invariant prediction, this objective can frequently be expected to perform no better than FRM. In the linear setting, we present simple, exact conditions under which solving to optimality succeeds or, more often, breaks down in recovering the optimal invariant predictor. We also demonstrate another major failure case case-under mild conditions, there exists feasible point that uses only non-invarian features and achieves lower empirical risk than the optimal invariant predictor; thus it will appear as a more attractive solution, yet its reliance on non-invariant features mean it will fail to generalize. As corollaries,
present similar settings where all recently suggested alternatives to IRM likewise fail. Futhermore, we present the first results in the non-linear regime: we demonstrate the existence of classifier with exponentially small sub- optimality which nevertheless heavily relies on non-invariant features on most test inputs, resulting in forse-than-chane performance on distributions that are sufficiently dissimilar from the training environments. These findings strongly suggest that existing approaches to ICP for high-dimensions latent variable models do not cleanly achieve their stated objective and that future work would benefit from a more formal treatment.
2 RELATED WORK
Works on learning deep invariant representations vary considerably: some search for a domain-
invariant representation (Muandet et al. 2013 Ganin et al. 2016). i.c. invariance of the distribution p({(z)), typically used for domain adaptation (Ben-Davie et al. [2100] Ganin & Lempitsky 2015
Zhang al. 2015 Long et al. 2018), with assumed access to labeled or unlabeled data from the target distribution. Other works instead hope to find representations that are conditionally nomain-invariang with invariance of p( (#))) y) (Gong et al., 2016 Li et al. 2018). However, there is evidence that invariance may not be sufficient for domain adaptation (Zhao et al. 2019 Johansson et al. 2019).
In contrast, this paper focuses instead on domain generalization (Blanchard et al. 201 Rosenfeld
et al. 202I), where access to the test distribution is not assumed.
Recent works on domain generalization, including the objectives discussed in this paper, suggest invariance of the feattre-conditioned label distribution. In particular, Arjovskyy et al. (2019) only assume invariance of Elyy [:: follow-up works rely on a stronger assumption of invariance of higher conditional moments (Krueger et al. 2020; Xie al., 2020 Jin et al. [020 Mahajan et al.) 2020 Bellot & van der Schaar 2020). Though this approach has become popular in the last year, it is somewhat similar to the existing concept of covariate shift (Shimodaira 2000 Bickel et al. 2009).
2

---

Published as a conference paper at ICLR 2021
which considers the same setting. The main difference is that these more recent works assume that the shifts in p( (z)) occur between discrete, labeled environments, as opposed to more generally from train test distributions.
Some concurrent lines of work study different settings yet give results which are remarkably similar to ours. Xu et al. [202T) show that infinitely wide two-layer network extrapolates linear functions when the training data is sufficiently diverse. In the context of domain generalization specifically, Rosenfeld al. 2022)) prove that ERM remains optimal for both interpolation and extrapolation in the linear setting and that the latter is exponentially harder than the former. These results mirror our findings that none of the studied objectives outperform ERM.
MODEL AND INFORMAL RESULTS
We consider an SEM with explicit separation of invariant features Zes whose joint distribution with the label is fixed for all environments, and environmental features ze ("non-invariant", whose distribution can vary. This choice is to ensure that our model properly formalizes the intuition behind invariant prediction techniques such as IRM, whose objective is ensure generalizing predictors by recovering only the invariant features- we put off a detailed description of these objectives until after we have introduced the necessary terminology.
We assume that data are drawn from a set E training environments know from which environment each sample is drawn. For a given environment e, the data defined by the following process: first, a label y E {+1] is drawn according to a fixed probability:
(ei, ,,, éØ} } and that
w.p. Y) otherwise.
y
Next, both invariant features and environmental features are drawn according to Gaussian
N(y Me, oe..)"))
~ Pe: T)
(2)
with Me E Rdd, te E Rdd typically, for complex, high-dimensional data we would expect < de < de Finally, the observation x is generated function of the latent features:
x f ze).
(3)
The complete data generating process is displayed in Figure 3.1 We assume f is injective, so that it principle possible to recover the latent features from the observations, i.c. there exists function & such that & (f 2e)) = [ We remark that this our only assumption on f. even when is non-linear. Further, note that we model class-conditional means as direct opposites merely for clarity, as it greatly simplifies the calculations. None of our proofs require this condition: it is straightforward to extend our results to arbitrary means, and the non-linear setting also allows for arbitrary covariances. In fact, our proof technique for non-linear could be applied to any distribution that sufficiently concentrates about its mean (c.g., sub- Gaussian.. We write the joint and marginal distributions p",, y, Ze, Ze). When clear from context, we omit the specific arguments.
Remarks on the model. This model is natural and flexible; it generalizes several existing models used to analyze learning under the existence of adversarial distribution shift or non-invariant cor- relations (Schmidt et al.) 2018: Sagawa et al., 2020). The fundamental facet of this model is the constancy of the invariant parameters 1), le. oz across environments- the dependence of Pey de on the environment allows for varying distributions, while the true causal process remains unchanged. Here we make a few clarifying remarks:
We do not impose any constraints on the model parameters. In particular, we do not assume prior over the environmental parameters. Observe that Me. o² are the same for all environments,
Note the deliberate choice to have z depend on y. Much work on this problem models spurious features which correlate with the label but are not causal. However, the term "spurious" is often applied incongruously: in recent work, the term has been co- co-opted to refer to any feature that correlates with the label but does not cause it. Thus there is subtle distinction: if we allow for anti-causality, i.c. the label causing the features, the resulting correlation is spurious. We therefore avoid using the term "spurious" in this work.
3

---

Published as a conference paper at ICLR 2021
hence the subscript indicates the invariant relationship. In contrast, with some abuse of notation, the environmental subscript is used to indicate both dependence on the environment and the index of the environment itself (c.g., M4, represents the mean specific environment i).
While we have framed the model as y causing Ze, the causation can just as easily be viewed in the other direction. The log-odds of y are linear function of 2 this matches logistic regression with an invariant regression vector B. 2 and bias Su = log - We present the model as above to emphasize that the causal relationships between y and the and because we believe this direction is more intuitive.
are priori indistinguishable
We consider the setting where we are given infinite samples from each environment; this allows us isolate the behavior of the objectives themselves, rather than finite-sample effects. Upon observing samples from this model, our objective is thus learn a feature embedder & and classifie unseen environment
Ze
3 to minimize the risk on an
y
R : (z.,)))
The function l can be any loss appropriate to classification: in this work we consider the logistic and the 0-1 loss. Note that we are not hoping to minimize risk in expectation over the environments; this is already accomplished via ERM or distributionally robust optimization (DRO; Bagnelll 2005; Ben- Tal et al. 2009). Rather, we hope to extract and regress invariant features while ignoring environmental features, such that our predictor generalizes to all unseen environments regardless of their parameters. In other words, the focus is minimizing risk in the worst-case. We refer to the predictor which will minimize worst-case risk under arbitrary distribution shift the optimal invariant predictor. To discuss this formally, define precisely what we mean by this term.
Sc
Figure 3.1: Bayesian network depicting our model. Shading in- dicates the variable is observed.
Definition 1. Under the model described by Equations LE the optimal invariant predictor is the predictor defined by the composition of a) the featurizer which recovers the invariant features and b) the classifier which is optimal with respect those features:
(a:
I 0
0 o
[ze]
3
3e 30
[2pe/o² π log T-
Observe that this definition closely resembles Definition 3 of [Arjovsky et al. (2019); the only difference is that here the optimal invariant predictor must recover all invariant features. As Arjovsky et al. (2019) do not posit data model, the concept of recovering "all invariant features" is not well- defined for their setting: technically, featurizer which outputs the empty set would elicit an invariant predictor, but this would not satisfy the above definition. The classifier 3 is optimal with respect to the invariant features and so it achieves the minimum possible risk without using environmental features. Observe that the optimal invariant predictor is distinct from the Bayes classifier; the Bayes classifier uses environmental features which are informative of the label but non-invariant; the optimal invariant predictor explicitly ignores these features.
With the model defined, we can informally present our results; we defer the formal statements to first give a background on the IRM objective in the next section. With slight abuse of notation, we identify predictor by the tuple 4. which parametrizes it. First, we show that the usefulness of IRM exhibits a "thresholding" behavior depending on E and de:
Theorem 3.1 (Informal, Linear). For linear f, consider solving the IRM objective to learn linear & with invariant optimal classifier 9. If E de. then & 3 is precisely the optimal invariant predictor; uses only invariant features and generalizes to all environments with minimax-optimal risk. If E < de, then 4 relies upon non-invariant features.
In fact, when F S de it is even possible to learn classifier solely relying on environmental features that achieves lower risk on the training environments than the optimal invariant predictor:
Following the terminology of Arjovsky et al. 20119) we refer to the regression vector as "classifier" and the composition of 0, f as "predictor".
4

---

Published as a conference paper at ICLR 2021
Theorem 3.2 (Informal, Linear). For linear and E S. de there exists a linear predictor & which uses only environmental features, yet achieves lower risk than the optimal invariant predictor.
Finally, in the non-lineat case, show that IRM fails unless the training environments approximately "cover" the space of possible environments, and therefore it behaves similarly to ERM:
Theorem 3.3 (Informal, Non-linear), For arbitrary J, there exists a non-linear predictor 4 3 which nearly optimal under the penalized objective and furthermore nearly identical the optimal invariant predictor on the training distribution. However, for any test environment with mean sufficiently different from the training means, this predictor will be equivalent the ERM solution on nearly all test points. For test distributions where the environmental feature correlations with the label are reversed, this predictor has almost 0 accuracy.
Extensions to other objectives. Many follow-up works have suggested alternatives to are described the next section. Though these objectives perform better on various baselines, there are few formal guarantees and no results beyond the linear case. Due to their collective similarities, we can easily derive corollaries which extend every theorem in this paper to these objectives, demonstrating that they all suffer from the same shortcomings. Appendix EnvendixEcoonteee contains example corollaries for each of the results presented this work.
4 BACKGROUND ON IRM AND ITS ALTERNATIVES
During training, classifier will learn leverage correlations between features and labels in the training data to make its predictions. If a correlation varies with the environment, it may not be present future test distributions -worte worse yet, may be reversed/-armingg the classifier's predictive ability. IRM (Arjovsky et al., 2019) is a recently proposed approach to learning environmentally invariant representations to facilitate invariant prediction.
The IRM objective. IRM posits the existence of feature embedder & such that the optimal classifier on top of these features is the same for every environment. The authors argue that such function will only invariant features, since non-invariant features will have different joint distributions with the label and therefore a fixed classifier on top of them won't be optimal in all environments. To leam this T the IRM objective is the following constrained optimization problem:
min 4.3
Σ (2) 3EE
s.t. ف arg min R ( Ve E.
B
This bilevel program highly non-convex and difficult to solve. To find an approximate solution, the authors consider Langrangian form, whereby the sub-optimality with respect to the constraint is expressed as the squared norm of the gradients of each of the inner optimization problems:
min
(2, + R (2, (113]
(5)
Assuming the inner optimization problem is convex, achieving feasibility is equivalent to the penalty term being equal to O. Thus, Equations 4 and 5 are equivalent we set À oc.
Alternative objectives. IRM is motivated by the existence of a featurizer & such that Eiy (x)) invariant. Follow-up works have proposed variations on this objective, based instead on the strictly stronger desideratum of the invariance of ply (x:)). Krueger et al.) (2020) suggest penalizing the
variance of the risks, while Xie et al. (2020) give the same objective but taking the square root of the variance. Many papers have suggested similar alternatives (Jin et al. [2200] Mahajan et al. 2020)) Bellot & van der Schaar 2020). These objectives are compelling
it is easy to show that the optimal invariant predictor constitutes a stationary point of each of these objectives:
Proposition 4.1. Suppose the observed data are generated according to Equations Then the (parametrized) optimal invariant predictor .. .B* is a stationary point for Equation
The stationarity of the optimal invariant predictor for the other objectives is trivial corollary. However, in the following sections we will demonstrate that such a result misleading and that a more careful investigation is necessary.
5

---

Published as a conference paper at ICLR 2021
5 THE DIFFICULTIES OF IRM IN THE LINEAR REGIME
In their work proposing IRM, Arjovsky et al. (2019) present specific conditions for an upper bound on the number of training environments needed such that a feasible linear featurizer & will have an invariant optimal regression vector B. Our first result is similar in spirit but presents a substantially stronger (and simplified) upper bound in the classification setting, along with a matching lower bound: we demonstrate that observing a large number of environments- in the number of environmental features is necessary for generalization in the linear regime.
Theorem 5.1 (Linear case). Assume f is linear. Suppose we observe E training environments. Then the following hold:
1. Suppose E > de. Consider any linear featurizer which is feasible under the IRM objective 4 with invariant optimal classifier 3 & O, and write [f [()) = Aze + Bzc- Then under mild
non-degeneracy conditions, it holds that B = 0. Consequently, 3 is the optimal classifier for all possible environments.
E S de and the environmental means pe are linearly independent, then there exists a linear d-where (ff ze)) - Aze + Bze with rank(B) = de + E-which feasible under the IRM objective. Further, both the logistic and 0-1 risks of this t and its corresponding optimal are strictly lower than those of the optimal invariant predictor.
Similar to Arjovsky et al. (2019), the set of environments which do not satisfy Theorem 5.1 has measure zero under any absolutely continuous density over environmental parameters. Further details, and the full proof, can be found in Appendix C Since the optimal invariant predictor is Bayes with respect to the invariant features, by the data-processing inequality the only way a predictor can achieve lower risk is by relying environmental features. Thus, Theorem S.I directly implies that when E S de. the global minimum necessarily uses these non- invariant features and therefore will not universally generalize to unseen environments. On the other hand, in the (perhaps unlikely) case that F > de. any feasible solution will generalize, and the optimal invariant predictor has the minimum (and minimax) risk of all such predictors:
Corollary 5.2. For both logistic and 0-1 loss, the optimal invariant predictor is the global minimum of the IRM objective if and only if E > de-
Let us compare our theoretical findings to those of Arjovsky et al. (201)). Suppose the observations x lie in Rd. Roughly, their theorem says that for learned bb of rank r with invariant optimal coefficient 3, if the training set contains d -r- + d/r "monndegenerate" environments, then 9 will be optimal for all environments. There are several important issues with this result: first, they present no result tying the rank of & their actual objective; their theory thus motivates the objective, but does not provide any performance guarantees for its solution. Next, observe when x is high-dimensional (i.e. >> de + de)-in which case & will be comparatively low-rank (i.e. r S de + de) their result requires (2d) environments, which is extreme. For example, think of images on low-dimensional manifold embedded very high-dimensional space. Even when d = d, + de, the "ideal" which recovers precisely &e would have rank de, and therefore their condition for invariance would require E > de + de. a stronger requirement than ours; this inequality also seems unlikely to hold in most real-world settings. Finally, they give no lower bound on the number of required environments prior to this work, there were no existing results for the performance of the IRM objective when their conditions are not met. We also
simple synthetic experiment to verify our theoretical results, drawing samples according to our model and learning a predictor with the IRM objective. Details and results of this experiment can be found in Appendix C.2 We now sketch a constructive proof of part 2 of the theorem for when E de:
Proof Sketch. Since f has an inverse over its range, we can define & as a linear function directly over the latents Jzc, ze]. Specifically, define d(x)) [Te,p" p"zze] ze]. Here, is unit-norm vector such that Ve E E, p" He = pa; j is a fixed scalar that depends on the geometry of μee o² ou-such a vector exists so long the means are linearly independent. Observe that this & also has the desired rank. Since this a linear function of a multivariate Gaussian, the conditionl distribution of each environment's
non-invariant latents has simple closed form: Tf ze y Now p" pe: l - Niw og o?).

---

Published as a conference paper at ICLR 2021
For separating two Gaussians, the optimal linear classifier is Σ- 1 classifier on p" 2e is precisely 2j,, which does not depend on the environment (and neither do the optimal coefficients for 2,). Though the distribution varies across environments, the optimal classifier the same! Thus, & directly depends on the environmental features, yet the optimal regression vector 33 for each environment is constant. To see that it has lower risk than the optimal invariant predictor, note that this classifier is Bayes with respect to its features and that the optimal invariant predictor uses strict subset of these features, and therefore it has less information for its predictions.
Mo)-here, the optimal
A purely environmental predictor. The precise value of j in the proof sketch above represents how strongly this non-invariant feature is correlated with the label. In theory, predictor that achieves a lower objective value could do so by very small margin m -ncorporating an arbitrarily small amount of information from non-invariant feature would suffice. This result would be less surprising, since achieving low empirical risk might still ensure that we are "close" to the optimal invariant predictor. Our next result shows that this is not the case: there exists a feasible solution which uses only the environmental features yet performs better than the optimal invariant predictor all E for which j is large enough.
Theorem 5.3. Suppose we observe F. S de environments, such that all environmental means are linearly independent. Then there exists a feusible D, which uses only environmental features and achieves lower 0-1 risk than the optimal invariant predictor on every environment e such that Get > 0. |elll and 20, Helll 2 |3ol.
The latter of these two conditions effectively trivial, requiring only a small separation of the means and balance in class labels. From the construction of jí in the proof of Lemma C.T we can see that the former condition is more likely to be met when F. < de and in environments where some non-invariant features reasonably correlated with the sibel-both of which can be expected to hold in the high- high-dimensional setting. Figure C.2 in the appendix plots the results for a few toy examples for various dimensionalities and variances to how often this condition holds in practice. For all settings, the number of environments observed before the condition ceases to hold is quite high the order of de de-
THE FAILURE OF IRM IN THE NON-LINEAR REGIME
We' 'e demonstrated that OOD generalization difficult in the linear case. but is achievable given enough training environments. Our results-ang and those of Arjovsky et al. inttiiintitively proceed by observing that each environment reduces a "degree of freedom" of the solution, such that only the invariant features remain feasible if enough environments are seen. In the non-linear case, it's not clear how to capture this idea of restricting the "degrees of freedom" and in fact our results imply that this intuition is simply wrong. Instead, we show that the solution generalizes only to test environments that are sufficiently similar to the training environments. Thus, these objectives present no real improvement over ERM DRO.
Non-linear transformations of the latent variables make it hard to characterize the optimal lincar classifier, which makes reasoning about the constrained solution to Equation 4 difficult. Instead we turn our attention to Equation 5) the penalized IRM objective. In this section we demonstrate a foundational flaw of IRM in the non-linear regime: unless observe enough environments to "cover" the space of non-invariant features, solution that appears to be invariant can still wildly underperform on a new test distribution. We begin with a definition about the optimality of coefficient vector 3:
Definition 2. For < 7 1, a coefficient vector / is y-close to optimal for label-conditional feature distribution ~ Nyy 2) if
8T p 2 (1 7)2
Since the optimal coefficient vector is precisely 2 nu, being 7-close implies that 3 is reasonably aligned with that optimum. Observe that the definition does not account for magnitude the set of vectors which is y-close to optimal is therefore a halfspace which is normal to the optimal vector. One of our results in the non- non-linear case uses the following assumption, which says that the observed environmental means are sufficiently similar one another.
7

---

Published as a conference paper at ICLR 2021
Assumption 1. There exists 0 I y < 1 such that the ERM-oppimal classifier for the non-invariant features,
Be,ERM == arg min
Σ 3EE
E B [clo se + B0).
(6)
3
A.
is 7-close to optimal for every environmental feature distribution in E.
This assumption says the environmental distributions are similar enough such that the optimal "average classifier" is reasonably predictive for each environment individually. This natural expectation: we are employing IRM precisely because we expect the ERM classifier to do well on the training set but fail to generalize. the environmental parameters are sufficiently orthogonal, we might instead expect ERM to ignore the features which are not at least moderately predictive across all environments. Finally, we note that if this assumption only holds for subset of features, our result still applies marginalizing out the dimensions for which it does not hold.
are now ready to give our main result the non-linear regime. present a simplified version, assuming that that o = Ve. This is purely for clarity of presentation; the full theorem is presented in Appendix D We make use two constants in the following proof-the the average squared norm
of the environmental means, I| == _ EnEE ll 1 and the standard deviation of the response variable of the ERM-optimal classifier, OERM = v VI + I|ße:ERMli2
Theorem 6.1 (Non-linear case, simplified). Suppose (e1, e2,...,eEE) where a² = Ve. Then, for any e > 1, there exists featurizer 4. which, combined with the ERM-optimal classifier B = [3e. Be:ERM, BulT. A the following properties, where we define Pe == exp(-d, min(e (e 1)8)/8));
observe E environments E
1. The regularization term of te, 3 as in Equation 5 is bounded as
IV 3)11 (p² (c.de -1l11)))
EEE
for some constant that depends only on e.
Pen exactly matches the optimal invariant predictor on at least Pe fraction of the training set. On the remaining inputs, it matches the ERM-optimal solution.
Further, for any test distribution, suppose its environmental mean μa+1 is sufficiently far from the training means:
Ve E E, min | - Hellz 2 (Vc+ 8 vE(+1]
for some 5 0, and define q
2E Vio
expl
expf-82), 8²). Then the following holds:
&., is equivalent to the ERM-optimal predictor on at least a - q. frucction of the test distribution.
4. Under Assumption 7 suppose it holds that μe++ - Eees Otelte for some set of coefficients {aeeeeE. Then so long as
| 1/01/2 + FERM
Caellell22
1-7
3EE
the 0-1 risk of 2., on the new environment is greater than .975
We give brief intuition for each the claims made in this theorem, followed by a sketch of the proof-the full proof can be found in Appendix D
1. The first claim says that the predictor construct will have gradient squared norm scaling as p which is exponentially small d,. Thus, in high dimensions, it will appear as a perfectly reasonable solution the objective
8

---

Published as a conference paper at ICLR 2021
2. The second claim says that this predictor is identical to the invariant optimal predictor all but an exponentially small fraction of the training data; on the remaining fraction, it matches the ERM-optimal solution, which has lower risk. The correspondence between constrained and penalized optimization implies that for large enough de, the "fake" predictor will often be a preferred solution. In the finite-sample setting, we would need exponentially many samples to even distinguish between the two!
3. The third claim the crux of the theorem; says that this predictor we've constructed will completely fail to use invariant prediction on most environments. Recall, the intent of IRM is to perform well precisely when ERM breaks down: when the test distribution differs greatly from the training distribution. Assuming Gaussian prior on the training environment means, they will have separation in O(vde dd) with high probability. Observe that q will be vanishingly small so long as 5 2 polylog( F.). Part 3 says that IRM fails to use invariant prediction on any environment that is slightly outside the high probability region of the prior; even a separation of (2(vde log F)) suffices. If we expect the new environments to be similar, ERM already guarantees reasonable performance at test-time; thus, IRM fundamentally does not improve over ERM in this regime.
4. The final statement demonstrates a particularly egregious failure case of this predictor: just like ERM, if the correlation between the non-invariant features and the label reverses test-time, our predictor will have significantly worse than chance performance.
Sketch. We give a construction which is almost identical to the optimal invariant predictor on the training data yet behaves like the ERM solution at test time. We partition the environmental feature space into two sets, B, B 1 is the union of balls centered at each Mle. cach with a large enough radius that it contains most of the samples from that environment; thus B represents the vast majority of the training distribution. On this set, define o (x) = [ze]. so our construction is equal to the optimal invariant predictor. Now consider 8 Rd B. We use standard concentration results to upper bound the measure of B under the training distribution by Next, we show how choosing (xx) f (x) = [ee =e]"" on this set results in the sub-optimality bound, which is of order p². is also clear that our constructed predictor is equivalent to the ERM-optimal solution 5 Thus, predictor will often have lower empirical risk on 8", countering the regularization penalty.
The second part of the proof shows that while B has large measure under the training environments, will have very small measure under any moderately different test environment. We can see this by considering the separation of means (Equation [ the measure of each ball in B can be bounded by the measure of the halfspace containing it; if each ball is far enough away from pE+1, then the total measure of these halfspaces must be small. At test time, our predictor will therefore match the ERM solution on all but y of the observations (part 3). Finally, we lower bound the 0- risk of the ERM predictor under such a distribution shift by analyzing the distribution of the response variable. The proof is completed by observing that our predictor's 's risk can differ from this by at most q. n
Theorem 6.1 Ishows that it's possible for the IRM solution to perform poorly on environments which differ even moderately from the training data. We can of course guarantee generalization if the training distributions "cover" (or approximately cover) the full space of environments in order to tie down the performance need for ICP; we could expect ERM or DRO to perform just as well. Once more, we find that our result trivially extends to the alternative objectives; we again refer Appendix
future distributions. But in such scenario, there would no longer be
7 CONCLUSION
Duttof-distiibution generalization an important direction for future research, and Invariant Causal Prediction remains a promising approach. However, formal results for latent variable models are lacking, particularly in the non-linear setting with fully unobserved covariates. This paper demon- strates that Invariant Risk Minimization and subsequent related works have significant under-explored risks and issues with their formulation. This raises the question: what is the correct formulation for invariant prediction when the observations are complex, non-linear functions of unobserved latent factors? We hope that this work will inspire further theoretical study on the effectiveness of IRM and similar objectives for invariant prediction.

---

Published as a conference paper at ICLR 2021
ACKNOWLEDGEMENTS
We thank Adarsh Prasad, Jeremy Cohen, and Zack Lipton for helpful feedback. Special thanks to Adarsh Prasad for noticing our initial formatting error. E.R. and P.R. acknowledge the support of NSF via IIS-1909816, IIS-1955532 and ONR via N000141812866
REFERENCES
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint arXiv:1907.022 2019.
J Andrew Bagnell. Robust supervised learning. In Proceedings of the 20th national conference Artificial intelligence- Volume 2, pp. 714-719, 2005.
Sara Beery, Grant Van Horn, and Pietro Perona. Recognition in terra incognita. In Proceedings of the European Conference on Computer Vision (ECCV), pp. 456-473, 2018.
Alexis Bellot and Mihaela van der Schaar. Generalization and invariances in the presence of unobserved confounding. arXiv preprint arXiv:2007.10633, 2020.
Shai Ben-David, John Blitzer, Koby Crammer, Alex Kulesza, Fernando Percira, and Jennifer Wortman Vaughan. A theory of learning from different domains. Machine learning, 79(1-2)1551-175, 2010.
Aharon Ben-Tal, Laurent El Ghaoui, and Arkadi Nemirovski. Robust optimization, volume 28. Princeton University Press, 2009.
Steffen Bickel, Michael Brückner, and Tobias Scheffer. Discriminative learning under covariate shift.
Journal of Machine Learning Research, 10(9), 2009.
Gilles Blanchard, Gyemin Lee, and Clayton Scott. Generalizing from several related classification tasks to a new unlabeled sample. Advances in neural information processing systems, 24:2178- 2186, 2011.
Kenneth A Bollen. Structural equation models. Encyclopedia of biostatisties, 7, 2005.
Rune Christiansen, Niklas Pfister, Martin Emil Jakobsen, Nicola Gnecco, and Jonas Peters. A causal framework for distribution generalization. arXiv preprint arXiv:2006.07433. 2020.
Yaroslav Ganin and Victor Lempitsky. Unsupervised domain adaptation by backpropagation. In International conference machine learning. pp. 1180-1189, PMLR, 2015.
Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky. Domain-adversarial training of neural networks. The Journal of Machine Learning Research, 17(1):2096-203 2016.
Robert Geirhos, Patricia Rubisch, Claudio Michaelis, Matthias Bethge, Felix A Wichmann, and Wieland Brendel. Imagenet-trained enns are biased towards texture; increasing shape bias improves accuracy and robustness. arXiv preprint arXiv: 181 1.2321 2018.
AmirEmad Ghassami, Saber Salehkaleybar, Negar Kiyavash, and Kun Zhang, Learning causal structures using regression invariance. Advances in Neural Information Processing Systems, pp. 3011-3021, 2017.
Mingming Gong, Kun Zhang, Tongliang Liu, Dacheng Tao, Clark Glymour, and Bernhard Schölkopf. Domain adaptation with conditional transferable components. In International conference on machine learning, pp.
2016.
Christina Heinze-Deml. Jonas Peters, and Nicolai Meinshausen. Invariant causal prediction for nonlincar models. Journal of Causal Inference, 6(2), 2018.
Andrew Ilyas, Shibani Santurkar, Dimitris Tsipras, Logan Engstrom, Brandon Tran, and Aleksander Madry. Adversarial examples not bugs, they features. In H. Wallach, H. Larochelle. A. Beygelzimer, F. dAlché-Bus E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp. 125-136. Curran Associates, Inc., 2019.
10

---

Published as a conference paper at ICLR 2021
Wengong Jin, Regina Barzilay, and Tommi Jaakkola. Domain extrapolation via regret minimization. arXiv preprint arXiv:2006.03908 2020.
Fredrik D Johansson, David Sontag, and Rajesh Ranganath. Support and invertibility in domain- invariant representations. arXiv preprint arXiv: 110..19033334 03348,44, 2019.
Divyansh Kaushik, Eduard Hovy, and Zachary Lipton. Learning the difference that makes a difference with suunteeraccualllyaagmeeeet data, International Conference on Learning Representations, 2020.
David Krueger, Ethan Caballero, Joem--Henrk Jacobsen, Amy Zhang, Jonathan Binas, Remi Le Priol, and Aaron Courville. Out-of-distribution generalization via risk extrapolation (rex). arXiv preprint arXiv:2003. 00688, 2020.
Frank R Kschischang.
The complementary enror function. 2017. URL https:: wwwp///wwwcccoom comm.
pdf
Ya Li, Xinmei Tian, Mingming Gong, Yajing Liu, Tongliang Liu, Kun Zhang, and Dacheng Tao. Deep domain generalization via conditional invariant adversarial networks. In Proceedings the European Conference on Computer Vision (ECCV), pp. 624 009, 2018.
Mingsheng Long, Zhangjie Cao, Jianmin Wang, and Michael Jordan. Conditional adversarial domain adaptation. Advances in Neural Information Processing Systems, pp. 1640-1650, 2018.
Marloes H Maathuis, Markus Kalisch, Peter Bühlmann, et al. Estimating high-dimensional interven- tion effects from observational data. The Annals of Statistics, 37(6A):3133-3) 3164,, 2009.
Divyat Mahajan, Shruti Tople, and Amit Sharma. Domain generalization using causal matching. arXiv preprint arXiv:2006. 20X0:00000.07000, 2020.
Krikamol Muandet, David Balduzzi, and Bernhard Schölkopf. Domain generalization via invariant feature representation. volume 28 of Proceedings of Machine Learning Research, pp. 10-18. Atlanta, Georgia, USA, 17-19 Jun 2013. PMLR. URL htRp/ttttpp//pprceddnggg.. wttp:/eeeednnnn.........m. preessss//
w28/muandet13.htm
Radford M Neal. Bayesian learning for neural networks, volume 118. Springer Science &amp; Business Media, 2012.
Judea Pearl. Causality. Cambridge university press, 2009.
J Peters, D Janzing, and B Schölkopf. Elements of causal inference-foundatiios and learning algorithms. 2017.
Jonas Peters, Peter Bühlmann, and Nicolai Meinshausen. Causal inference by using invariant prediction: identification and confidence intervals. Journal of the Royal Statistical Society, 2016.
Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do imagenet classifiers generalize to imagenet? arXiv preprint arXiv:1902.108111 2019.
Mateo Rojas-Carulla, Carulla, Bernhard Schölkopf, Richard Turner, and Jonas Peters. Invariant models for causal transfer learning. The Journal of Machine Learning Research, 19(1):1309-134 1342, 2018.
Elan Rosenfeld, Pradeep Ravikumar, and Andrej Risteski. An online learning approach interpola- tion and extrapolation in domain generalization. arXiv preprint arXiv:2102.13118, 2021.
Shiori Sagawa, Aditi Raghunathan, Pang Wei Koh, and Percy Liang. An investigation of why overparameterizztioo exacerbates spurious correlations. In Proceedings of the 37th International Conference on Machine Learning, 2020.
Ludwig Schmidt, Shibani Santurkar, Dimitris Tsipras, Kunal Talwar, and Aleksander Madry. Ad- versarially robust generalization requires more data. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems 31, pp- 5014-5026. Curran Associates, Inc., 2018

---

Published as a conference paper at ICLR 2021
Hidetoshi Shimodaira. Improving predictive inference under covariate shift by weighting the log- likelihood function. Journal of statistical planning and inference, 90(2)227-244, 2000.
Connor Shorten and Taghi M Khoshgoftaar. A survey on image data augmentation for deep learning. Journal of Big Data, 6(1):60, 2019.
Jiawei Su, Danilo Vasconcellos Vargas, and Kouichi Sakurai. One pixel attack for fooling deep neural networks. IEEE Transactions on Evolutionary Computation, 23(5):828-88) 841, 2019.
Adarsh Subbaswamy, Peter Schulam, and Suchi Saria. Preventing failures due to dataset shift: Learning predictive models that transport. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 11118-3122 PMLR, 2019.
Riccardo Volpi, Hongseok Namkoong, Ozan Sener, John C Duchi, Vittorio Murino, and Silvio Savarese. Generalizing to unseen domains via adversarial data augmentation. In Advances in neural information processing systems, pp. 5334-5344, 2018.
Chuanlong Xie, Fei Chen, Yue Liu, and Zhenguo Li. Risk variance penalization: From distributional robustness to causality. arXiv preprint arXiv:2006.07544 2020.
Keyulu Xu, Mozhi Zhang, Jingling Li, Simon Shaolei Du, Ken-Ichi Kawarabayashi, and Stefanie Jegelka. How neural networks extrapolate: From feedforward to graph neural networks. In International Conference Learning Representations, 2021. URL httpps///pppeeeeeeeee
cmocIIJC
Kun Zhang, Mingming Gong, and Bernhard Scholkopf. Multi-source domain adaptation: a causal view. In Proceedings of the Twenty-Ninte AAAI Conference on Artificial Intelligence, pp. 3150- 3157, 2015.
Han Zhao, Remi Tachet Des Combes, Kun Zhang, and Geoffrey Gordon. On learning invariant representations for domain adaptation. International Conference on Machine Learning, pp. 7523 1223-7532, 2019.
12

---

Published as a conference paper at ICLR 2021
A ADDITIONAL NOTATION FOR THE APPENDIX
To avoid overloading, we use Ø, F' for the standard Gaussian PDF and CDF respectively. We write
S to denote the set complement of set S. We write to denote entrywise absolute value.
B PROOF OF PROPOSITION
Recall the IRM objective:
min 2,
B B( -og loy 3=x(())
subject to
E (av)) 1- 3 [<))] = 0. Ve E E.
03
Concretely, represent O parametrized function ts over whose parameters A we then optimize. The derivative the negative log-likelihood for logistic regression with respect to the B coefficients is well known:
0
log (y (x))
1 1))))(())
0
Suppose recover the true invariant features (x)))) Ze and coefficients B
3
(in other
words, we allow for the introduction of new features). Then the IRM constraint becomes:
i E ((
lologllooo
(y 3d +o(x))]
03
pe (ze) Σ p' =c) 03
- log [y (TTyy-Taa. xc)c] dāc
ME(+1)
/ p (2c)
(3 zc) 13 z6) 1) -33² 2c)) (3t 2a) dzc.
Since 3 is constant across environments, this constraint is clearly satisfied for every environment, and therefore also the minimizing 3 for the training data whole.
Considering now the derivative with respect to the featurizer py:
00 | logg((y oly 3 [ = (a
2 3 (T(x).)
- 1/y= = 1})
30
Then the derivative of the loss with respect to these parameters is
1 p((cc) ( (o(x)) [oLOT (0 26) + (1-o 23)) 01=0 20) de 0.
So, the optimal invariant predictor is stationary point with respect to the feature map parameters well.
C RESULTS FROM SECTION 5
C.1 PROOF OF THEOREM 5.1
We begin by formally stating the non-degeneracy condition. Consider any environmental mean Pes and suppose it can be written as linear combination of the others means with coefficients a":
Me Σ a pi
13

---

Published as a conference paper at ICLR 2021
Then the environments considered non-degenerate if the following inequality holds for any such set of coefficients:
E #1
(9)
and furthermore that the following ratio different for at least two different environments a, b:
0 E, 0 1- E,
o -E E.
3F
(10)
The first inequality says that none of the environmental means are are an affine combination of the others; in other words, they lie in general linear position, which is the same requirement as (Arjossky similarly lax non-degeneracy requirement regarding scale of the variances. It is clear that the set of environmental parameters that do not satisfy Equations 9 and I0 has measure zero under any absolutely continuous density, and similarly, if E S de then the environmental means will be linearly independent almost surely.
jet al. 2019). The other inequality
relative
can now proceed with the proof, beginning with some helper lemmas:
Lemma C.1. Suppose we observe E environments = = e,,..,, eg), each with environmental mean of dimension de 2 E, such that all environmental means are linearly independent. Then there a unique unit-norm vector p such that
pe = ofti We E E,
(11)
where jı is the largest scalar which admits such a solution.
Proof. Let U1, U2 ..... ug be set of basis vectors for span .g.......): Each mean then be expressed combination of these basis vectors: u, = -Ø= aiyUy. Since the means are linearly independent, combine these coefficients into a single invertible matrix
Ta
021
0g
Q122
022
0g2
U
022
ag
We then combine the constraints
11
TT
Pa = -
where Pax denotes our solution expressed in terms of the basis vectors {w:) This then has the solution
Pa = U-TO
This defines the entire space of solutions, which consists of Pa plus any element of the remaining (de -idddienniion orthogonal subspace. However, we want y to be umit-norm-bbset that the current vector solves Equation| with ğ = which means that after normalizing we get ğú = ||p. laa Adding any element of the orthogonal subspace would only increase the norm of p, decreasing Thus, the unique maximizing solution is
U-To
Pa
1 |U-Ta|
with
o
Finally, Pa has to be rotated back into the original space by defining = E Potu.
F.
Lemma C.2. Assume is linear. Suppose observe E S de environments whose means linearly independent. Then there exists a linear with rank( d, + d, E whose output depends the environmental features, yet the optimal classifier top of t invariant.
14

---

Published as a conference paper at ICLR 2021
Proof. We will begin with the when E de and then show how to modify this construction for when < de. Consider defining
0
f
0
M
with
p²
M
Here, p E Rd is defined as the unit-norm vector solution to o² Vve
p" such that i is maximized-such a vector is guaranteed to exist by Lemma C.1 Thus we get
Ue p²ue
Zc ((x)= p² ze Nyy Jt ).). For each environment we have
which is of rank de + desired. Define & p" Ze, which means that &e y ~
p(ze, Ze.1)
ply Eei 2c)
plzer 2e)
( 2c)p y o. (2)
Jo(y $8 2e)p(E
y oe db)) + a
y 38 20)p(Ee -
aly 30 ze) exp(y E07t)
[o(y (T exp(y 2 () + o 3e 2e) exp( -y 2c(iil
I 1+expp exp(-y y (B z. + 22)
The log-odds of y linear these features, so the optimal classifier is
3g 2
8
which is the same for all environments.
Now show how to modify this construction for when E < de- If we remove one of the environ- mental means, & trivially maintains its feasibility. Note that since they are linearly independent, the mean which removed has component in direction orthogonal to the remaining means. Call this component p', and consider redefining M as
T
The distribution of &. in each of the remaining dimensions is normal with mean 0, which means corresponding coefficient of 0 is optimal for all environments. So the classifier remains optimal for all environments, yet we' ve added another row M which increases the dimensionality of its span, and therefore the rank of 4, by 1. Working backwards, we can repeat this process for each additional mean, such that rank( ()) de + 1+ (de E'), as desired, Note that for E any will be vacuously feasible.
Lemma C.3. Suppose we observe environments & {e1, E2..., eg] whose parameters satisfy the non-degeneracy conditions t T0P. Let ((x)) Aze + Bz, be any feature vector which is linear function of the invariant and environmental featurr,,, and suppose the optimal 3 top of Tb is invariant. If F. > do, then B 0 or B=0 0.
15

---

Published as a conference paper at ICLR 2021
Proof. Write + [A]B]
where A Raxd, B R and define
Be Φ
He
Ape + Bus
He_
o 0 ddId.
Σ. +
Φ
= = + B
Without loss of generality we assume Σ is invertible (if it is not, we can consider just the subspace in which it outsiide of this space, the features have variance and therefore cannot carry information about the label). By Lemma F.2 the optimal coefficient for each environment is 2> Pe. In order for this vector to be invariant, it must be the same across environments; we write it as a constant 3.
Suppose fe for some environment e-then the claim is trivially true with 8 0. We therefore proceed under the assumption that te + 0 Ve E &
With this fact, we have that Ve EE,
8 A0 + (Aple Bue)
AA )B.BBB -2Apc +2B
BB 3 -BB,, = =2Ape -ZAA"
(12)
Define the vector v = 2Apa 0² A. We will show that for any B, A, with probability 1 only B 0 can satisfy Equation 12 for every environment. If E' > de. then there exists at least one linear combination of the othery. Without loss of
environmental mean which can be written generality, denote the parameters of this environment as (Ø, σ²) and write jú = Σi= agH,, However, note that by assumption the means lie in general linear position, and therefore we actually have at least de sets of coefficients a for which this holds. Rearranging Equation 12 we get
o² B
2Bp
de
2,2Bpu
d = Σ
OTBBTT -v
1=1
and rearranging once more yields
Emoo BB B
(1
- v.
By assumption, (1 Σ a;) is non-zero. We therefore rewrite this
aBB"B=V
where å -2=2 environment, this only hold true if à is fixed for all environments if both BBT 3, 0. The former is false by assumption, so the the latter must hold.
-T> is scalar. As the vectors BB 3 and are both independent of the
As result, see that Equation 12 reduces to
Byc 0
Ve E.
As the span of the observed μ. is all of Rdde this only possible if = 0.
We are now ready to prove the main claim. We restate the theorem here for convenience:
Theorem 5.1 (Linear case). Assume f is linear. Suppose we observe training environments. Then the following hold:
16

---

Published as a conference paper at ICLR 2021
1. Suppose E de. Consider any linear featurizer & which is feasible under the IRM objective 4 with invariant optimal classifier 3 + O, and write ( Aze Bže- Then under mild non-degeneracy conditions, it holds that B O. Consequently, 3 the optimal classifier for all possible environments.
If E S de and the environmental means pee are linearly independent, then there exists a linear w-where (f((=c =e)) Aze + Bze with rank(B)= de + E E-wwich is feusible under the IRM objective. Further, both the logistic and 0-1 risks of this & and its corresponding optimal B are strictly lower than those of the optimal invariant predictor.
Proof.
1. Since f are linear, we can write ((x) ))= = Aze + Bzg for some matrices A, B. Assume the non-degeneracy conditions 9 [Op hold. By Lemma C.3 one of B 0 or 3 = 0 holds. Thus, $, 3 uses only invariant features. Since the joint distribution p (2e. y) is invariant, this predictor has identical risk across all environments.
2. The existence of such predictor is proven by Lemma It remains to show that the risk of this discriminator is lower than that of the optimal invariant predictor. Observe that these features are non-degenerate independent random variables with support over all of I, and therefore by Lemma F. dropping the ā. term and using
8e 30
T (x) [=c],
$ =
results in strictly higher risk. The proof is completed by noting that this definition is precisely the optimal invariant predictor.
C.2 EXPERIMENTS FOR THEOREM
To corroborate our theoretical findings, we run an experiment on data drawn from our model to see at what point IRM able to recover a generalizing predictor. We generated data precisely according to our model in the linear setting, with de = 3, de 6. The environmental means were drawn from a multivariate Gaussian prior; we randomly generated the invariant parameters and the parameters of the prior such that using the invariant features gave reasonable accuracy (71.9%) but the environmental features would allow for almost perfect accuracy on in-distribution test data (99.8%). Thus, the goal was to see if IRM could successfully learn predictor which ignores meaningful covariates ze, to the detriment of its training performance but to the benefit of OOD generalization. We chose equal class marginals (1) 0.5).
Figure C.T shows the result of five runs of IRM, each with different environmental parameters but the same invariant parameters (the training data itself was redrawn for each run). We found that optimizing for the IRM objective was quite unstable, frequently collapsing to the ERM solution unless À and the optimizer learning
were carefully tuned. This echoes the results of Krueger et al. (2020) who found that tuning À during training to specific values at precisely the right time essential for good performance. To prevent collapse, we kept the same environmental prior and found a single setting for À and the learning rate which resulted in reasonable performance across all five runs. At test time, we evaluated the trained predictors on additional, unseen environments whose parameters were drawn from the same prior. To simulate distribution shift, we evaluated the predictors on the same data but with the environmental means negated. Thus the correlations between the environmental features =, and the label y were reversed.
Observe that the results closely track the expected outcome according to Theorem 5.1 up until E = dz, IRM essentially matches ERM in performance both in-distribution and under distribution shift. As soon as we cross that threshold of observed environments, the predictor learned via IRM begins to perform drastically better under distribution shift, behaving more like the optimal invariant predictor. We did however observe that occasionally the invariant solution would be found after only E de = 6 environments; invariant predictor with lower objective value presented in Theorem 5. precisely a single point, as opposed to multi-dimensional subspace, and therefore might be difficult for the optimizer to find.
conjecture that this is because at this point the paasible-yet-no
17

---

Published as a conference paper at ICLR 2021
(wo (1
IRM Text (161)) Opti Inariiet Closiiox
(E)
Figure C.I: Performance of predictors learned with IRM (5 different runs) and ERM (dashed lines) on test distributions where the correlation between environmental features and the label is consistent (no shift) or reversed (shift). The dashed green line is the performance of the optimal invariant predictor. Observe that up until E de. IRM consistently returns a predictor with performance similar to ERM: good generalization without distribution shift, but catastrophic failure when the correlation is reversed. In contrast, once E > de, IRM is able to recover a 4 B with performance similar to that of the invariant optimal predictor.
C.3 PROOF OF THEOREM 5.3
Theorem 5.3. Suppose we observe E S de environments, such that all environmental means are linearly independent. Then there exists feasible 4,8 which uses only environmental features and achieves lower 0-1 risk than the optimal invariant predictor on every environment e such that oet > o. [|Hell2 and 20 |||tl|22 " 2 130l-
Proof. We consider the non non-invariant predictor constructed as described in Lemma C.2 but dropping the invariant features and coefficients. By Lemma F.2 the optimal coefficients for the invariant and non-invariant predictors are
20 2pp Bu
2P Bu
Beaus
and
Boon Chon-cuus
respectively. Therefore, the 0-1 risk of the optimal invariant predictor is precisely
1)P(20 F ( ||ue|12 211pel12
2 T zc + 30 < 0) (1 22r0--00-2111 Bo >> 0)
Boga
Bno
+ (1 -)F - ||eell +
211
where the Gaussian CDF. By the same reasoning, the 0-1 risk of the non non-invariant predictor is
So n) (-oait
30
nF --oopt
20.
20
Define - Hella and 7 aef.. By monotonicity of the Gaussian CDF, the former risk higher than the latter if
30
Bo
α +
< 7: +
(13)
2a
2
30
30
<
(14)
2a
27
Without loss of generality, we will prove these inequalities for 8o 2 0; an identical argument proves for Bo < 0 but with the *<" and '<" swapped.
18

---

Published as a conference paper at ICLR 2021
Suppose > (the first condition). Then Equation observe that
14 is immediate. Finally, for Equation 13
80
30
?
+
27
20
60
30
(7 - a:)0
2
2a
2
27.a
27a 2 30.
which is the second condition.
C.4 SIMULATIONS OF MAGNITUDE OF ENVTRONMENTAL FEATURES
As discussed in Section 5 analytically quantifying the solution j to the equation in Lemma C.1 is difficult; instead, present simulations to give a sense of how often these conditions would hold in practice.
For each choice of environmental dimension de, generated "base" correlation b N/ Id,) as the mean of the prior over environmental means He- Each of these jee was then drawn from (b, thus,, while they all came from the same prior, the noise in the draw of each μr was significantly larger than the bias induced by the prior. We then solved for the precise value aeJi with the same variance of for all environments, chosen hyperparameter. The shaded area represents 95% confidence interval over 20 runs.
The dotted lines are dd.. we imagine the invariant parameters are drawn from standard Gaussian prior, then this is precisely EB|o 1 [|c]l2]] Thus, the point where Jeli crosses these dotted lines is approximately how many environments would need to observed before the non-invariant predictor has higher risk than the optimal invariant predictor. We note that this value is quite large, on the order of de de.
Projected Environmental Mean Separation
40
d. d, 200.
de 100, d, 200, 2
d. d, 300,
30
de d. 300, c
vill
20
15
10
50
100
150
200
250
300
Environments observed
Figure C.2: Simulations evaluate ozt for varying ratios of 1. When o the value closely tracks vd - E, and the crossover point is approximately de od These results imply the conditions of Theorem 5.3 are very likely to hold in the high-dimensional setting.
D THEOREM 6. AND DISCUSSION
D. PROOF OF THEOREM 6.1
We again begin with helper lemmas.
19

---

Published as a conference paper at ICLR 2021
Our featurizer & is constructed recover the environmental features only if they fall within set B The following lemma shows that since only the environmental features contribute to the gradient penalty, the penalty can be bounded
function of the measure and geometry of that set. This is used together with Lemmas F.3 and F.4 to bound the overall penalty of our constructed predictor.
Lemma D.1. Suppose we observe environments
(e1,e2...... Given a set B S R consider Then for any environment e, the penalty term of this predictor
the predictor defined by Equation 19
in Equation 5 is bounded as
(2, |lce E BB)BBlZe 2e E B
Proof. We write out the precise form the gradient for an environment e:
R (2, B)
22 (Z ze) |o(8 @f((ze ze))) - p² Zes z6) 2 (f (e)) d(zes ze).
/
Observe that since Zp 1L Ze y. the optimal invariant coefficients are unchanged, and therefore the gradient in the invariant dimensions is O. We can split the gradient in the environmental dimensions into two integrals:
f
2.xB (Zer ze)
1o 2e + Bo) - -p = Zes ze)] [0] d(zc: ze)
+ 20) 10(3%
+ 2XB (0
BT. B + Bb) al3. 2e + Be 2e + [0) [ze] d(zes ze).
Since the features are 0 within B, the first term reduces For the second term, note that V x. RR |o(x) - (v)) I 1, and therefore
/ x.
V ( B)1 S
(z ze llzel d(ze: ze).
We can marginalize out ze, and noting that we want to bound the squared norm,
/ pe (zo)lkeell dze
V ( 3)13) <
= PlEe E B ze E B8]
This next lemma says that if the environmental mean of the test distribution is sufficiently separated from each of the training means, with high probability a sample from this distribution will fall outside of B,, and therefore T., will be equivalent to the ERM solution.
Lemma D.2. For a set of E environments E e2... eE and any e > 1, construct B, as in Equation 18 and define D using B, in Equation 19 Suppose now test new environment with parameters (pE+1, 1).)), and assume Equation 15 holds with parameter exp e kk22 over the draw of an
Define k min.cE ae-i observation from this new environment, we have
o²
Then with probability 2
22
Vkzō
%2c [e
(<())= =
Proof. By Equation 15 our new environmental mean is sufficiently far away from all the label- conditional means of the training environments. In particular, for any environment E and any label y E [=1]. the &, distance from that mean to #s+1 is at least (VE + 8)..vde.
Recall that B, is the union balls BB., where B, is the ball of & radius / 2eo2d, centered at pa. For each environment e, consider constructing the halfspace which is perpendicular the line connecting
20

---

Published as a conference paper at ICLR 2021
Me and ME++ and tangent to Be. This halfspace fully contains Be, and therefore the measure of Be is upper bounded by that the halfspace.
By rotational invariance of the Gaussian distribution, and the measure will not change. The center of the ball is (ve + -0oo away from the mean Mg+1+ so accounting for its radius, the distance from the mean to the halfspace is v The variance of the rotated distribution one dimension is of F+1 de, so the measure of this halfspace is upper bounded
rotate this halfspace into dimension
by
d0
1 - +
(-Vko)
VoB+ide
exp(-k03)
Vkzō
using results from Kschischang (2017) There are 2E such balls comprising Br, which can be combined via union bound.
With these two lemmas, we now state the full version of Theorem 6.1 with the main difference being that it allows for any environmental variance.
Theorem D.3 (Non-linear case, full). Suppose we observe F. environments & {ei, 02....; eF.. Then, for any > 1, there exists featurizer $ which, combined with the ERM-optimal classifier
3 [Be, BEERRM Bo]', satisfies the following properties, where we define Pe.d. expl- min((e-
1), (e /)8)
1. Define d max, o².. Then the regularization term of de is bounded
3)12 EO (p2aa (ede eddomax exp(2eo + |1a112
Σ IV E
2. D., 3 exactly matches the optimal invariant predictor on least a Pe,d, de, fraction of the training set. On the remaining inputs, it matches the ERM-optimal solution.
Further, for any test distribution with environmental parameters (HE+1, of oB41) suppose the environ- mental mean pÆ-1 is sufficiently far from the training means:
Ve E E, min || - y Hellz 2 50oeV vE(+1)
(15)
for some 5 > O. Define the constants:
o²
min 2 cEE E+1
2
Vkuō exp(=ko²)
Then the following holds:
3. D.. 3 is equivalent to the ERM-optimal predictor on at least a fraction of the test distribution.
4. Under Assumption 4 suppose holds that
FE+1
Σ 33E
Ccfe
(16)
for some set of coefficients
(ore JeEE. Then for any c. E R. so long as
| + 1/ol22 + CDERM
(17)
1 7
33E
the 0-1 risk of t 3 is lower bounded by F(2c) q-
21

---

Published as a conference paper at ICLR 2021
Proof. Define v Veo2de and construct B,
R as
B, = U
U
UB
(18)
LeEf
LEEg
where (a) is the ball of lz- norm radius centered at a. Further construct &. using B, as follows:
Ze
Ze
B,
Be
0
Pc(x)=
and
β
Be
(19)
Ze
Re E B
50_
Zec
Without loss of generality, fix an environment e.
1. By Lemma D.I the squared gradient norm is upper bounded by
iP(ze E B[)]]l ze B,]
(20)
Define Bg = B, (ue), and observe that BŞ C B Since [ze non-negative,
P(ze E E &e E B3 < P(ze E B| 2e E B
(this inequality is wisse. Plugging this into Equation 20 yields
Re E B (9. s iP(ze E B
[P(ze E BD)
E|
2e E B
2
Define p (Z BB)) < P(ze By). By Lemma E3
min((e- min((e-D),e-112)))
p Pe.d.
Combining Lemmas F.4 and
gives
Mr\ 1
B Re Re E B
< 2d
+
F(-r/v/l/r) Vd)
2
+ 211 S de exp (2e(0² 1/2)} [co2
Putting these two bounds together, we have
(T, 311 E o ( od max 22c max + 11 )
and averaging this value across environments gives the result.
2. 4. 3 is equal to the optimal invariant predictor on B, and the ERM solution on B,. The claim then follows from Lemma F3
3. This follows directly from Lemma D.2
4. With Equation 16 we have that
S
Σ
ae ST PEEERMP
3EE
||mel12
s 2 7 Σ
al
0²
CEE
S 2(1 7) ||/e11//02 + 1/301/2 + COERM
1-7
=-(21lmell302 1,30 -20FERM)
22

---

Published as a conference paper at ICLR 2021
where have applied Assumption in the first inequality and Equation 17 in the second.
Consider the full set of features L. (z) = f- (x), and without loss of generality assume y 1. The labcl-conditional distribution of the resulting logit is
BL 3 + So - (5' pe + EetRM/E+1 + Po. FRRM)
Therefore, the 0-1 risk is equal to the probability that this logit is negative. This is precisely
5 Me + & + 90
1211ae11//03 + [30l + 200TRM) 2 Mel 1//3///3 -1.0)
ORRM
TERM
200TRM
F
TRRM
= =F(cc).
Observe that by the previous part, P. + f-1 on at most a q fraction of observations, so the risk of our predictor e... 3 can differ from that of f by most q. Therefore our predictor's 's risk is lower bounded by F(2c) In particular, choosing = recovers the statement the main body.
D.2 DISCUSSION OF CONDITIONS AND ASSUMPTION
To see just how often we can expect the conditions for Theorem D.3 to hold, we can do a rough approximation based the expectations of each of the terms. A reasonable prior for the environ- mental means is a multivariate Gaussian [m, Σ). We might expect them to be very concentrated (with Tr(E) small), or perhaps to have strong bias (with || >> Tr(E)). For simplicity we treat the variances o oe as constants. Then the expected separation between any two means from this distribution is
Bljlm n22 l2] R [||=||]] & (r(EE
In high dimensions this value will tightly concentrate around the mean, which is in (V On the other hand, even a slight deviation from this separation, to S( log F),, means E D2 /lo F.), which implies q O(/EE); this is plenty small to ensure worse-than-random error on the test distribution.
Now we our attention to the second condition [7 The expected squared norm of each mean is de, and in high dimensions we expect them to be reasonably orthogonal (as a rough approximation; this technically not true with non-centered Gaussian). Then so long as Σ, 0g ((1), the left-hand side of Equation 17 is approximately de. On the other hand, treating as constant, the right-hand side close de + Vdl + de E O(de + vda). Thus, Equation 17 quite likely to hold for any mean EE+1 with the same scale the training environments but with reversed correlations-agale this is exactly the situation where IRM hopes to outperform ERM, and we have shown that it does not.
can also quick analysis of Assumption T under this prior: the ERM-optimal non-invariant coefficient will be approximately 2m/o? 2 with high probability, meaning 3 μ = 2 2|ml|///2 for every environment. Thus, this vector will be y-close to optimal with 0 for every environment with high probability.
E EXTENSIONS TO ALTERNATIVE OBJECTIVES
E. EXTENSIONS FOR THE LINEAR CASE
Observe that the constraint of Equation 4 is strictly stronger than that of Bellot & van der Schaar (2020); when the former satisfied, the penalty term of the latter necessarily It is thus trivial to extend all results in the Section 5 this objective. As another example, consider the risk-variance- penalized objective Krueger et al. (2020):
min 6.9
R' 3) + AVarees (R)
(21)
18
330
23

---

Published as a conference paper at ICLR 2021
It is simple to extend Theorem 5.1 under an additional assumption:
Corollary E.1 (Extension Theorem [. Assume f is linear. Suppose we observe E de environments with linearly independent means and identical variance o². Consider minimizing empirical risk subject penalty on the risk variance (21). Then there exists $, B dependent on the non-inver non-invariant features which achieves a lower objective value than the optimal invariant predictor for any choice of regularization parameter À E [0, oo).
Proof. Consider the featurizer & constructed in Lemma
C.2 If the environmental variance is constant,
then the label-conditional distribution of the environmental features,
y ~ NNy ;00),
Ze
is also invariant. This implies that the optimal 3 also has constant risk across the environments, meaning the penalty term 0, and as result the objective does not depend on the choice of As F.I implies that the overall risk is lower than that of the optimal invariant
in [. invoking Lemma predictor.
As mentioned in Section 5 this additional requirement of constant variance due to the assumptions underlying the design of the objective-REs expects additional invariance of the second moment Var(y ((r)), in contrast with the strictly weaker invariance of Ely (xx)] assumed by IRM. This might seem to imply that REx is more robust objective, but this does not convey the entire picture. The conditions for the above corollary are just one possible failure case for REx; by extending Theorem 5.3 to this objective, we see that REx is just as prone to bad solutions:
Corollary E.2 (Extension Theorem 5.3). Suppose we observe E de environments, such that all environmental means are linearly independent. Then there exists $ 3 which uses only environmental features and, under any choice of X E [0, oo), achieves a lower objective value than the optimal invariant predictor under 0-1 loss on every environment such that jı > o --=1lppe I|lell2
[80|
20c
Proof. We follow the proof of Theorem find the unit-norm vector such that
5.3 except when solving for p as in Lemma C. we instead
p" Me Jeti Ve E.
(22)
Observe that by setting ((x) = 1p" ze] and B 3= [1], the 0-1 risk a given environment is F(-
n) 7P + (1 P--ooo
which is independent of the environment. Further, by carrying through the same proof in Theo- rem 5.3 we get that this non-invariant predictor has lower 0-1 risk so long as
1.00
a +
S p
2a
where a = o. 1|eel22
Though j here is not exactly the same value because of the slightly different solution 22 it depends upon the geometry of the training environments in the same way-it is the same as taking the square root of each of the variances. We can therefore expect this condition to hold in approximately the same situations, which we empirically verify by replicating Figure C.2 with the modified equation below.
E.2 EXTENSIONS FOR THE NON-LINEAR CASE
The failure of these objectives in the non-linear regime is even more straightforward, as we can keep unchanged the constructed predictor from Theorem Observe that parts 2-4 of the theorem do not involve the objective itself, and therefore do not require modification.
To that part still holds, note that since the constructed predictor matches the optimal invariant predictor on - p of the observations, its risk across environments can only vary on the remaining fraction: thus the centered 0-1 risk is bounded between 0 and p. is immediate that the variance of the environmental risks is upper bounded by & E O(p?). Applying this argument to the other objectives yields similar results.
24

---

Published as a conference paper at ICLR 2021
Projected Environmental Mean Separation
40
de d. 200. o2
d 100, d, 200,
d, 5001, c7
30
de 600, de 2001, of
vd.
25
20
15
50
100
150
200
250
300
Environments observed
Figure E.1: Simulations to evaluate j for varying ratios of d- When o 1 the value closely tracks d.
vd E. and the crossover point is approximately de - Due the similarity of Equation 22 to Equation II] it makes sense that the results are very similar to those presented in Figure
F TECHNICAL LEMMAS
Lemma F.l. Consider solving the standard logistic regression problem
~ p(2) E R
+1 wip. o)
- -1 w.p. o(-
Assume that none of the latent dimensions are degenerate-Nss YS S. [k], P(B + 0) 0, and no feature can be written as a linear combination of the other features. Then for any distribution p(z). any classifier [(=) = a zs) that uses strict subset of the features S " [b] has strictly higher risk with logistic loss than the Bayes classifier (z) = o33"² 2). This additionally holds for 0-1 loss 6T s-- has greater magnitude and opposite sign of B3 zs with non-zero probability.
Proof. The Bayes classifier suffers the minimal expected loss for each observation z. Therefore. another classifier has positive excess risk if and only if it disagrees with the Bayes classifier set of non-zero measure. Consider the set of values z such that 3 3= + O. Then on this set we have
(35 = s2 4 (35 os) f(z).
Since these values with positive probability, / has strictly higher logistic risk than f By the same argument, there exists a set of positive measure under which
(8 sign(3}2s + s_ sign f(2).
and so f also has strictly higher 0-1 risk.
Lemma F.2. For any feature vector which is a linear function of the invariant and environmental features = Az. + Bze, any optimal corresponding coefficient for an environment e is of the form
2(AA 02 +BBTTO' + + Bpe),
where G is generalized inverse of G.
25

---

Published as a conference paper at ICLR 2021
Proof. We begin by evaluating closed form for p" 2). We have:
1e (y Aza + Bze = 2)
pl Aze + B2c= = y)p(y)
pe (Aze + Bze = 8)
pf(Aze + B2c = =)
Df (Aze Bi = 1 11 + yf (Aze + Bāc -v)
1+ -z. Br,
Now need closed form expression for p(Aze Bee y). Noting that &e TL ze this is convolution of the two independent Gaussian densities, which the density of their sum. other words,
Aze Ba 1 ~ NWy (Atta + Buc), AAT o + BBT o.))
Thus,
1
1
y p 2 bズ (8 (a---m)) v-
1¹ (Aze + Bze = =1)) (2T E1)
6/2
exp
2
Canceling common terms, we get
1 = Azr + Big =
I' B
1 p= B2, =llv)
1 c { -y y. 22T 1
o (y 2 Σ* Σ²
Therefore, given feature vector غ, the optimal coefficient vector is 25. jr.
Lemma F33 For any environment with parameters Pe, o² and any >> define
B = B Veaza. (pe).
where By) is the ball of t--norm radius centered at a. Then for an observation drawn from pf, we have
(ze BO) exp { de minn((e-1))( -1((e-1),, 1),
1)2)
Proof. Without loss of generality, suppose y 1. We have
(lize Hellz S veo2de
P22ee E B) 2
IP
(1zell2 S voogde
[0,,00
P =0,1
(1 < ede).
Each term in the squared norm of 2, is random variable with distribution X² which means their sum has mean de and is sub-exponentia with parameters (2v 4). By standard sub-exponential concentration bounds we have
(1 2 ede) s exp {
de min((e - 1). (()))
8
which immediately implies the claim.
26

---

Published as a conference paper at ICLR 2021
Lemma F.4. Let (H, da be multivariate isotropic Gaussian in d dimensions, and for some define B as the ball of Eg radius r centered at μ. Then we have
v(r/vd) ]² F(-r/Vd) +2
= 2 BT]
S 2d
where o, F are the standard Gaussian PDF and CDF
Proof. Observe that
E s~
[1] 22 Bd
E sn
[|]| | /4/22 > r]
E s~.((U,o31a)
||2 + /1 || r]
|||| ||||| > r] + |||-
a~
Now, consider the expectation for an individual dimension, and note that | Va So because the dimensions are independent, conditioning on this event can only increase the expectation:
= || >
2
|
r
[
2
Vd.
r
E
2
where the equality is because the distribution is symmetric about O. This last term is known as the conditional tail expectation of a Gaussian and available in closed form:
(fF-1(())
r
2
2
2y (0,
1
where a = F(r/vd). Combining the above results, squaring with (a + b)2 = 2(a² + B2). and summing over dimensions, we get
2
2
B||=|
= B9
E
24
1
+ 211
2
i=1
(0.02)
Vd
v(r/vo) F(-r/Vd)
2d
+211
desired.
Lemma F.5. For o,e define VEa. Then
)(r)
1 exp (2e(0²
1/2)3 [eo2 + +1].
F(-r)
2
Proof. We have
1
exp {Ee
Φ(r)
V2T
2
and
2expf -c02)
F(-r) Vi(Vto + Vco² + 2)
(see Kschischang (2017)). Dividing them gives
@(r) 1
F(-r) 2/2 exple(o² exple(o"- 1/2)} V Veo= +2
Squaring and using (a + 6)2 < 2(a + b2) yields the result.
27


Reviews:
[{"review_id": "t_vlSgM-0zE", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper, entitled \"The Risks of Invariant Risk Minimization\", provides a theoretical analysis of the IRM learning objective under particular conditions. It claims to provide two theoretical results, namely:\n1) in the linear setting (i.e., the optimal model is linear), they provide a necessary and sufficient condition for IRM to recover the optimal invariant predictor.\n2) in the non-linear setting, they demonstrate that IRM does not perform better than ERM (empirical risk minimization).\n\nStrong points of the paper:\n1) the paper addresses an important question: what are the theoretical guarantees of the IRM objective ?\n2) the paper provides a solid argument against the IRM objective, which is shown to be no better than ERM, in a limited but non-trivial setting\n3) the paper tries to relate IRM to several prior works from the literature\n\nWeak points of the paper:\n1) I believe the paper contains a major flaw: the so-called \"invariant classifier\", which uses only invariant features, is not the optimal invariant predictor sought by IRM. As such the first claim of the authors, necessary and sufficient conditions under which IRM recovers the optimal invariant predictor in the linear setting, is wrong.\n2) The paper suffers from a poor structure, a lack of clarity, and a general lack of rigor in the definition of the key concepts.\n\nDespite commendable efforts from the authors to study IRM theoretically, even in a limited setting, I recommend rejection for the paper.\n\nFirst, while it appears rather theoretical, the paper lacks rigor and clarity in the definition of the problem, the key concepts, and the assumptions. The description of the OOD generalization problem, of IRM and its assumptions, of the assumptions made by the authors in their simplified model, and of the introduced concepts (notably the \"invariant classifier\"), are lacking a formal mathematical description. The main results of the paper are given in Section 3, before the studied object, IRM, is even formally  introduced in Section 4. The authors interchangeably use the adjectives \"causal\" or \"invariant\", without having properly defined the meaning of those concepts in the first place, and with no particular utility. Maybe IRM can be given a causal interpretation, under some assumptions, but clearly this is out of the scope of the paper. The theoretical analysis of IRM conducted in the paper does not require the notion of causality, which only brings confusion to the paper in my opinion.\n\nSecond, and most importantly, I believe that one of the author's claims (which occupies about half the paper) is simply wrong (see my detailed comments below with a counter-example). There seems to be a misunderstanding as to what IRM actually seeks for, which is not the so-called \"invariant classifier\" that uses only invariant features. IRM does not seek to use invariant features only, but to learn those invariant features, possibly by using non-invariant ones. As such, results such as Theorem 5.1, Corollary 5.2, and Theorem 5.3, have very little value I think. The second claim of the authors remains interesting, however I believe the paper requires a substantial revision and in its current state can not be accepted at ICLR.\n\nI would appreciate if the authors could formally describe what they mean by \"invariant classifier\", and point me to the evidence where the IRM authors argue that IRM will recover such a classifier.\n\n**Post-discussions**: the authors have clarified the meaning of \"invariant classifier\", which is now tied to their specific toy problem, and I now believe such a classifier is indeed the minimax OOD classifier supposedly seeked by IRM. While I still think the paper does not shine on the side of clarity, my main concern about the incorrectness of the presented theorems has been answered and I believe the paper will be of interest to the community. I therefore raise my recommendation towards acceptance.\n\nBelow are my detailed comments:\n\np.1 §3: each environment corresponds to interventions on the SEM (Pearl, 2009) on just the non-causal mechanisms -> What does non-causal mechanism mean here ? If the SEM is given a causal interpretation, then all mechanisms are causal. Do you mean the mechanisms which are not causaly impacting the variable of interest Y ?\n\np.1 §3: invariant features in the SEM -> Do you mean the causal features ? Or the invariant mechanisms ? Invariant features does not make sense. If interpreted as constant features, those are uninformative. If interpreted as having the same distribution across environments, this does not make sense either, since it does not guarantee such a feature will be causal.\n\np.2 §5:  invariance of the feature-conditioned label distribution -> I like the effort the authors put in grouping the existing works in a unified framework and notation. There is however a bit of inconsistency and ambiguity in the terminologies employed throughout the paper. Invariant distribution? Domain-invariant representation ? Invariant features ? A proper definition would help, and enforcing consistency throughout the paper, by using the same names when those refer to the same concepts, would benefit readability.\n\np.2 §5: The main difference [...] train to test distributions. -> I like that the authors make this comparison, but I also would have liked this discussion to be further extended. There is a clear relationship indeed. If IRM looks for a representation $\\phi(x)$ such that $p(\\phi(x))$ can change between domains, but $p(y|\\phi(x))$ remains the same, then this is exactly the assumption of covariate shift: $p(x)$ can change, but $p(y|x)$ remains the same. However, finding such a representation $\\phi$ does not entirely solve the problem. First, the representation has to be predictive enough for $y$, ideally so that $y \\perp x | \\phi(x)$. Otherwise, even a constant $\\phi$ might result in $p(y|\\phi(x))$ being invariant if $p(y)$ is. But then marginal predictions from $p(y)$ would probably be sub-optimal. Second, while in theory (infinite model capacity, infinite data) covariate shift is not an issue (Bayes-optimal predictors do not need $p(x)$, only $p(y|x)$ over the support of $x$), in practice it is another story, and additional methods such as sample reweighting can be beneficial. See paragraphs 3,4 from the introduction in V. Tran, and A. Aussem. A Practical Approach to Reduce the Learning Bias Under Covariate Shift. ECML/PKDD 2015.\n\np.2 §6: whose joint distribution with the label is fixed for all environments -> this is a strong assumption which IRM, as far as I understand, does not make\n\np.3 §3: invariant (causal) relationship -> I am a bit uncomfortable with the authors mixing the invariant and causal adjectives in this context. In your SCM, if given a causal interpretation, $y$ is a cause of $z_c$, and $p(y|z_c) \\neq p(y|do(z_c))$. I would suggest for the authors to stick to the concept of invariance (which they still have not defined properly...), and drop the causal interpretation, which brings nothing but confusion to the paper. Causality really is a different business, which is not required to study the capabilities of the IRM objective.\n\np.3 §4: the causation can just as easily be viewed in the other direction -> Exactly. So why even talking about causality ? Do you want a model of $p(y|x)$, which generalize across different environments with different distributions for p(x) ? Or do you want a causal model of $p(y|do(x))$ ? Those are different objectives, and using them interchangeable really is confusing, as they do not require the same assumptions. Here, as you say, we do not care whether $x$ causes $y$ or $y$ causes $x$.\n\np.3 §5: a constant marginal -> Marginal of what ? Please be precise.\n\np.3 §5: is necessary: Necessary to what ? Which assumptions of IRM would be violated as a result of not having $y \\perp e$ your model ? The IRM assumptions were never clearly stated. Would that make finding $\\psi$ such that $p(y|\\psi(x))$ is invariant impossible ?\n\np.3 §8: whose joint or conditional distribution -> The conditional is implied by the joint. You can remove \"or conditional\" here.\n\np.3 §8: Definition 1 -> The paper would gain a lot in clarity if this \"invariant classifier\" was given a formal, mathematical definition. For example, what do you mean by \"invariant features\" ? The same holds for you assumptions, which could be defined more clearly: $z_c,y \\perp e$, $e$ uniform. According to your definition, I can think of the following:\n\n$$\nh^\\star = \\arg \\min_h \\mathbb{E}_{(x,y,z_c,z_e,e)}[l(h(x), y)], \\text{subject to } h(f(z_c,z_e)) = h(f(z_c,z'_e)), \\forall z_c,z_e,z'_e \\text{,}\n$$\n\nThe \"invariant classifier\" from this definition is not the one that IRM seeks. It suffices to consider a scenario where $f=I$ $Z_c=\\emptyset$, and two environments such that $\\mu_{e_1}, \\sigma_{e_1}=(100, 1)$ and $\\mu_{e_2}, \\sigma_{e_2}=(1000, 1)$. Then $z_e$ is not an invariant feature, but clearly the representation $f(z_e)=sign(z_e)$ is invariant, and is very informative of $y$. I think this is a major confusion in the article, which becomes clear only when ones writes down this definition formally.\n\np.4 §1: Observe that [...] the invariant classifier does not. -> This is I believe the major flaw of the article. The objective in OOD generalization is to find the Bayes classifier, not the \"invariant classifier\". Likewise, the objective of IRM is not to learn the \"invariant classifier\". As such, showing that IRM does not recover this \"invariant classifier\" has little value.\n\np.4 §3: Since those are informal theorems, I would rather name them differently, e.g.,  Result 1, Result 2, Result 3.\n\np.4 Section 4: The definition of IRM should be presented before the Section 3: results.\n\np.4 Section 4 §2: The authors argue that such a function will use only invariant features -> I do not find this argument in the original IRM paper.\n\np.5 Theorem 5: E training environments -> and how many environments in the test distribution ?\n\np.5 Appendic C.1: In Appendix Figure C.1, I believe the y axis should read $\\eta$, not Accuracy ? Also, I fail to understand why ERM is a flat line in this Figure. If more environments are provided, even with unlimited data, then the ERM classifier should change, and the test performance as well. Maybe I am missing something here, but the experimental setup (such as how model parameters are chosen for each new environment) is not clearly described.\n\np.5 Corollary 5.2: I am not sure what this implies. Maybe IRM does not recover the \"invariant classifier\" in that case, but what we really want for OOD generalization is the Bayes-optimal classifier. Again, it is not clearly stated what IRM wants to achieve...\n\np.6 Theorem 5.3: The initial assumptions these results rely on should be stated formally. Do these results hold for any $p(x,y,e)=p(e)p(y)p(x|e,y)$ distribution ? For any $p(x,y,e)$ distribution satisfying the IRM assumptions ? Only for the specific model the authors consider ?\n\np.6 Theorem 5.3: feasible -> Feasible for IRM ?\n\np.6 Figure C.2: Some labels are missing from Figure C.2. What do the doted horizontal lines represent ?\n\np.8 Conclusion: the number of training environments needed to ensure good generalization under arbitrary distribution shift is linear in the number of non-invariant features -> I believe this is wrong. The Bayes classifier is the one ensuring good OOD generalization. The \"invariant classifier\", which you prove requires a number of environments equal to the number of non-invariant features, is not the Bayes classifier.", "extra_scores": {}, "extra_text": {}}, {"review_id": "XEYCgdQvHig", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Main comments:\n\nThis paper studies a theoretical aspect of IRM and how will it fail. Main contribution is pointing out that IRM is ineffective when the number of environments $E$ is smaller than the dimension of environmental feature $d_e$. A simple but universal model assumption is built, where environmental feature $z_e$ and causal feature $z_c$ is sampled from Gaussian conditional on label $y$. The analysis is two-fold: linear regime and non-linear regime. In the former part, given the feature extractor $\\Phi$ is linear, a constructed solution to IRM is built to demonstrate the result. For the latter part, the other show the failure of IRM via several results in Thm 6.1 / D.3. The whole analysis is clear and easy to follow, thus the reviewer believe this submission deserves to be accepted. \n\nMain comments:\n\nThe simulated experiments in Appendix C demonstrate the theoretical results. However, it's a toy run and  thus one drawback to the reviewer is that, for Colored MNIST experiments in IRM, it's obvious that $d_e > E$ while IRM still works. It would be great if the authors could give some explanation on this point.\n\nAnalysis in the paper is more a constructed one, whereas lacks of how optimization algorithm can lead to such a solution, which is another drawback of this paper.\n\nAnother minor comment:\n\nOne seminal reference seems missing: Causal inference using invariant prediction: identification and confidence intervals. Jonas Peters, Peter Bühlmann, Nicolai Meinshausen.", "extra_scores": {}, "extra_text": {}}, {"review_id": "ThTgSL_Nqme", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Pros:\n* The work gives extended theoretical analysis on the effect of invariant risk minimization scheme, which is an increasingly popular framework for robust prediction. The work considerably extends the results in the original IRM paper. The results seem reasonable, and clarify implausible beliefs on the framework. The intuitions and proving techniques could also inspire new methods to improve the current framework.\n\nCons (minor issues):\n* I know the main contribution of the work is on the theoretical side, but a simple experiment or simulation for the nonlinear case to demonstrate the empirical consequence would be appreciated.\n* In the considered data generating model, it is assumed that $z_c \\perp z_e | y$ (conditional independence). Is it required in the proofs? How natural/general can it be? For one example, it is mentioned that the causation can be viewed in the other direction, i.e. $z \\to y$. In that case, observing $y$ would render components of $z$ correlated. Also, a graphical illustration of the data generating model would be helpful.\n* OOD generalization to an arbitrary environment in the nonlinear case (even with the invariant optimal predictor on top of representations) may not be possible anyway, and some recent works on OOD generalization does not expect the learned model to work well in environments out of the convex hull/combination of training environments. So I am wondering how arbitrary can the test environment be in Theorem 6.1, e.g., can it be out of the convex hull/combination of training environment? The result would seem more interesting to me if the test environment can be.\n\n=== EDIT: post-rebuttal ===\n\nThanks for the additional explanations.", "extra_scores": {}, "extra_text": {}}, {"review_id": "I9U_Pj4o45W", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "- pros:\n    - The paper examined the claims of IRM thoughtfully and critically.\n    - the exposition is clear --- I particularly enjoyed the informal results in section 3.1. I applaud the authors' effort to make the results intuitive!\n    - The appendix included illustrative and convincing empirical studies.\n    - The paper brought up a failure case of IRM, which shows that there exists a nearly optimal classifier, similar to the invariant predictor in the training environment, but performs equal to the ERM solution in the testing set. \n- cons\n    - There is only one DGP used in throughout the paper.  In particular, the DGP  has the same complexity in the y->z_c and z_c -> y relationship. This seems counter-intuitive to me. I would be interested in learning if this phenomenon generalizes to other DGPs\n\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors have made significant efforts to thoroughly address all the concerns. Due to the amount of discussions, I had to go through the paper myself and agree with the authors on many of the points. In my opinion, this is a solid theoretical work on the pitfalls of IRM. 
Decision: Accept (Poster)