THE RISKS OF INVARIANT RISK MINIMIZATION
Year: 2021

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

heading: 1 INTRODUCTION
text: Prediction algorithms are evaluated by their performance on unseen test data. In classical machine learning, it is common to assume that such data are drawn i.i.d. from the same distribution as the data set on 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 is evidence that they rely on statistically informative but non-causal features in the data (Beery et al., 2018; Geirhos et al., 2018; Ilyas et al., 2019). As a result, such models are prone to errors under surprisingly minor distribution shift (Su et al., 2019; Recht et al., 2019). 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-Distribution (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 is data augmentation—this includes both automated data modifications which help prevent overfitting (Shorten & Khoshgoftaar, 2019) and specific counterfactual augmentations to ensure invariance in the resulting features (Volpi et al., 2018; Kaushik et al., 2020).
A strategy which has recently gained particular traction is Invariant Causal Prediction (ICP; Peters et al. 2016), which views the task of OOD generalization through the lens of causality. This framework assumes that the data are generated according to a Structural Equation Model (SEM; Bollen 2005), which consists of a set of so-called mechanisms or 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 a result, learning mechanisms that are the same across environments ensures recovery of the invariant features which generalize under arbitrary interventions. In this work, we consider objectives that attempt to
learn what we refer to as 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 to environment-specific 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 our results can easily be extended to similar recently proposed alternatives.
Various works on invariant prediction (Muandet et al., 2013; Ghassami et al., 2017; Heinze-Deml et al., 2018; Rojas-Carulla et al., 2018; Subbaswamy et al., 2019; Christiansen et al., 2020) consider regression in both the linear and non-linear 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 a 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 to learn a 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 ERM. 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—under mild conditions, there exists a feasible point that uses only non-invariant 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, we 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 a classifier with exponentially small suboptimality which nevertheless heavily relies on non-invariant features on most test inputs, resulting in worse-than-chance performance on distributions that are sufficiently dissimilar from the training environments. These findings strongly suggest that existing approaches to ICP for high-dimensional latent variable models do not cleanly achieve their stated objective and that future work would benefit from a more formal treatment.

heading: 2 RELATED WORK
text: Works on learning deep invariant representations vary considerably: some search for a domaininvariant representation (Muandet et al., 2013; Ganin et al., 2016), i.e. invariance of the distribution p(Φ(x)), typically used for domain adaptation (Ben-David et al., 2010; Ganin & Lempitsky, 2015; Zhang et 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 domain-invariant, with invariance of p(Φ(x) | 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., 2011; Rosenfeld et al., 2021), 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 feature-conditioned label distribution. In particular, Arjovsky et al. (2019) only assume invariance of E[y | Φ(x)]; follow-up works rely on a stronger assumption of invariance of higher conditional moments (Krueger et al., 2020; Xie et al., 2020; Jin et al., 2020; 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),
which considers the same setting. The main difference is that these more recent works assume that the shifts in p(Φ(x)) occur between discrete, labeled environments, as opposed to more generally from train to test distributions.
Some concurrent lines of work study different settings yet give results which are remarkably similar to ours. Xu et al. (2021) show that an infinitely wide two-layer network extrapolates linear functions when the training data is sufficiently diverse. In the context of domain generalization specifically, Rosenfeld et al. (2021) 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.

heading: 3 MODEL AND INFORMAL RESULTS
text: We consider an SEM with explicit separation of invariant features zc, 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 to 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 of E training environments E = {e1, e2, . . . , eE} and that we know from which environment each sample is drawn. For a given environment e, the data are defined by the following process: first, a label y ∈ {±1} is drawn according to a fixed probability:
y = { 1, w.p. η −1, otherwise. (1)
Next, both invariant features and environmental features are drawn according to a Gaussian:1
zc ∼ N (y · µc, σ2cI), ze ∼ N (y · µe, σ2eI), (2)
with µc ∈ Rdc , µe ∈ Rde—typically, for complex, high-dimensional data we would expect E < dc de. Finally, the observation x is generated as a function of the latent features:
x = f(zc, ze). (3)
The complete data generating process is displayed in Figure 3.1. We assume f is injective, so that it is in principle possible to recover the latent features from the observations, i.e. there exists a function Φ such that Φ(f(zc, ze)) = [zc, ze]T . We remark that this our only assumption on f , even when it 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 f could be applied to any distribution that sufficiently concentrates about its mean (e.g., sub-Gaussian). We write the joint and marginal distributions as pe(x, y, zc, 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 correlations (Schmidt et al., 2018; Sagawa et al., 2020). The fundamental facet of this model is the constancy of the invariant parameters η, µc, σ2c , f across environments—the dependence of µe, σe 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 a prior over the environmental parameters. Observe that µc, σ2c are the same for all environments,
1Note the deliberate choice to have ze 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-opted to refer to any feature that correlates with the label but does not cause it. Thus there is a subtle distinction: if we allow for anti-causality, i.e. the label causing the features, the resulting correlation is not spurious. We therefore avoid using the term “spurious” in this work.
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 (e.g., µi represents the mean specific to environment i). • While we have framed the model as y causing zc, the causation can just as easily be viewed in the other direction. The log-odds of y are a linear function of zc—this matches logistic regression with an invariant regression vector βc = 2µc/σ2c and bias β0 = log η 1−η . We present the model as above
to emphasize that the causal relationships between y and the zc, ze are a priori indistinguishable, and because we believe this direction is more intuitive.
We consider the setting where we are given infinite samples from each environment; this allows us to isolate the behavior of the objectives themselves, rather than finite-sample effects. Upon observing samples from this model, our objective is thus to learn a feature embedder Φ and classifier2 β̂ to minimize the risk on an unseen environment e:
Re(Φ, β̂) := E(x,y)∼pe [ `(σ(β̂TΦ(x)), y) ] .
The function ` 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; Bagnell 2005; Ben-Tal et al. 2009). Rather, we hope to extract and regress on 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 on minimizing risk in the worst-case. We refer to the predictor which will minimize worst-case risk under arbitrary distribution shift as the optimal invariant predictor. To discuss this formally, we define precisely what we mean by this term. Definition 1. Under the model described by Equations 1-3, 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 to those features:
Φ∗(x) := [ I 0 0 0 ] ◦ f−1(x) = [zc], β̂∗ := [ βc β0 ] := [ 2µc/σ 2 c log η1−η ] .
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 a data model, the concept of recovering “all invariant features” is not welldefined for their setting; technically, a featurizer which outputs the empty set would elicit an invariant predictor, but this would not satisfy the above definition. The classifier β̂∗ 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 a slight abuse of notation, we identify a predictor by the tuple Φ, β̂ 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 a linear Φ with invariant optimal classifier β̂. If E > de, then Φ, β̂ is precisely the optimal invariant predictor; it uses only invariant features and generalizes to all environments with minimax-optimal risk. If E ≤ de, then Φ, β̂ relies upon non-invariant features.
In fact, when E ≤ de it is even possible to learn a classifier solely relying on environmental features that achieves lower risk on the training environments than the optimal invariant predictor:
2Following the terminology of Arjovsky et al. (2019), we refer to the regression vector β̂ as a “classifier” and the composition of Φ, β̂ as a “predictor”.
Theorem 3.2 (Informal, Linear). For linear f and E ≤ de there exists a linear predictor Φ, β̂ which uses only environmental features, yet achieves lower risk than the optimal invariant predictor.
Finally, in the non-linear case, we 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 f , there exists a non-linear predictor Φ, β̂ which is nearly optimal under the penalized objective and furthermore is nearly identical to the optimal invariant predictor on the training distribution. However, for any test environment with a mean sufficiently different from the training means, this predictor will be equivalent to 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 IRM—some are described in 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 E contains example corollaries for each of the results presented in this work.

heading: 4 BACKGROUND ON IRM AND ITS ALTERNATIVES
text: During training, a classifier will learn to 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 in future test distributions—worse yet, it may be reversed—harming 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 a feature embedder Φ such that the optimal classifier on top of these features is the same for every environment. The authors argue that such a function will use 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 learn this Φ, the IRM objective is the following constrained optimization problem:
min Φ,β̂
1 |E| ∑ e∈E Re(Φ, β̂) s.t. β̂ ∈ arg min β Re(Φ, β) ∀e ∈ E . (4)
This bilevel program is highly non-convex and difficult to solve. To find an approximate solution, the authors consider a 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 Φ,β̂
1 |E| ∑ e∈E [ Re(Φ, β̂) + λ‖∇β̂Re(Φ, β̂)‖22 ] . (5)
Assuming the inner optimization problem is convex, achieving feasibility is equivalent to the penalty term being equal to 0. Thus, Equations 4 and 5 are equivalent if we set λ =∞.
Alternative objectives. IRM is motivated by the existence of a featurizer Φ such that E[y | Φ(x)] is invariant. Follow-up works have proposed variations on this objective, based instead on the strictly stronger desideratum of the invariance of p(y | Φ(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., 2020; Mahajan et al., 2020; Bellot & van der Schaar, 2020). These objectives are compelling—indeed, 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 1-3. Then the (parametrized) optimal invariant predictor Φ∗, β̂∗ is a stationary point for Equation 4.
The stationarity of the optimal invariant predictor for the other objectives is a trivial corollary. However, in the following sections we will demonstrate that such a result is misleading and that a more careful investigation is necessary.

heading: 5 THE DIFFICULTIES OF IRM IN THE LINEAR REGIME
text: 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 β̂. 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—linear 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 β̂ 6= 0, and write Φ(f(zc, ze)) = Azc +Bze. Then under mild non-degeneracy conditions, it holds that B = 0. Consequently, β̂ is the optimal classifier for all possible environments.
2. If E ≤ de and the environmental means µe are linearly independent, then there exists a linear Φ—where Φ(f(zc, ze)) = Azc +Bze with rank(B) = de + 1− E—which is feasible under the IRM objective. Further, both the logistic and 0-1 risks of this Φ 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.1. 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 on environmental features. Thus, Theorem 5.1 directly implies that when E ≤ 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 E > 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. (2019). Suppose the observations x lie in Rd. Roughly, their theorem says that for a learned Φ of rank r with invariant optimal coefficient β̂, if the training set contains d− r + d/r “non-degenerate” environments, then β̂ will be optimal for all environments. There are several important issues with this result: first, they present no result tying the rank of Φ to 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. d de + dc)—in which case Φ will be comparatively low-rank (i.e. r ≤ de + dc)—their result requires Ω(d) environments, which is extreme. For example, think of images on a low-dimensional manifold embedded in very high-dimensional space. Even when d = dc + de, the “ideal” Φ which recovers precisely zc would have rank dc, and therefore their condition for invariance would require E > de +de/dc, 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 run a 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 [zc, ze]. Specifically, define Φ(x) = [zc, pT ze]. Here, p is a unit-norm vector such that ∀e ∈ E , pTµe = σ2e µ̃; µ̃ is a fixed scalar that depends on the geometry of µe, σ2e—such a vector exists so long as the means are linearly independent. Observe that this Φ also has the desired rank. Since this is a linear function of a multivariate Gaussian, the label-conditional distribution of each environment’s non-invariant latents has a simple closed form: pT ze | y ∼ N (y · pTµe, ‖p‖22σ2e) d = N (y · σ2e µ̃, σ2e).
For separating two Gaussians, the optimal linear classifier is Σ−1(µ1 − µ0)—here, the optimal classifier on pT ze is precisely 2µ̃, which does not depend on the environment (and neither do the optimal coefficients for zc). Though the distribution varies across environments, the optimal classifier is the same! Thus, Φ directly depends on the environmental features, yet the optimal regression vector β̂ 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 a strict subset of these features, and therefore it has less information for its predictions.
A purely environmental predictor. The precise value of µ̃ in the proof sketch above represents how strongly this non-invariant feature is correlated with the label. In theory, a predictor that achieves a lower objective value could do so by a very small margin—incorporating an arbitrarily small amount of information from a 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 on all e ∈ E for which µ̃ is large enough.
Theorem 5.3. Suppose we observe E ≤ de environments, such that all environmental means are linearly independent. Then there exists a feasible Φ, β̂ which uses only environmental features and achieves lower 0-1 risk than the optimal invariant predictor on every environment e such that σeµ̃ > σ −1 c ‖µc‖2 and 2σeµ̃σ−1c ‖µc‖2 ≥ |β0|.
The latter of these two conditions is effectively trivial, requiring only a small separation of the means and balance in class labels. From the construction of µ̃ in the proof of Lemma C.1, we can see that the former condition is more likely to be met when E de and in environments where some non-invariant features are reasonably correlated with the label—both of which can be expected to hold in the high-dimensional setting. Figure C.2 in the appendix plots the results for a few toy examples for various dimensionalities and variances to see how often this condition holds in practice. For all settings, the number of environments observed before the condition ceases to hold is quite high—on the order of de − dc.

heading: 6 THE FAILURE OF IRM IN THE NON-LINEAR REGIME
text: We’ve demonstrated that OOD generalization is difficult in the linear case, but it is achievable given enough training environments. Our results—and those of Arjovsky et al. (2019)—intuitively 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 or DRO.
Non-linear transformations of the latent variables make it hard to characterize the optimal linear 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 we observe enough environments to “cover” the space of non-invariant features, a solution that appears to be invariant can still wildly underperform on a new test distribution. We begin with a definition about the optimality of a coefficient vector β̂:
Definition 2. For 0 < γ < 1, a coefficient vector β̂ is γ-close to optimal for a label-conditional feature distribution z ∼ N (y · µ,Σ) if
β̂Tµ ≥ (1− γ)2µTΣ−1µ.
Since the optimal coefficient vector is precisely 2Σ−1µ, being γ-close implies that β̂ is reasonably aligned with that optimum. Observe that the definition does not account for magnitude—the set of vectors which is γ-close to optimal is therefore a halfspace which is normal to the optimal vector. One of our results in the non-linear case uses the following assumption, which says that the observed environmental means are sufficiently similar to one another.
Assumption 1. There exists a 0 ≤ γ < 1 such that the ERM-optimal classifier for the non-invariant features,
βe;ERM := arg min β̂e
1 |E| ∑ e∈E Ezc,ze,y∼pe [ `(σ(βTc zc + β̂ T e ze + β0), y) ] , (6)
is γ-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 is a natural expectation: we are employing IRM precisely because we expect the ERM classifier to do well on the training set but fail to generalize. If 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 a subset of features, our result still applies by marginalizing out the dimensions for which it does not hold.
We are now ready to give our main result in the non-linear regime. We present a simplified version, assuming that that σ2e = 1 ∀e. This is purely for clarity of presentation; the full theorem is presented in Appendix D. We make use of two constants in the following proof—the average squared norm of the environmental means, ‖µ‖22 := 1E ∑ e∈E ‖µe‖22; and the standard deviation of the response
variable of the ERM-optimal classifier, σERM := √ ‖βc‖22σ2c + ‖βe;ERM‖22σ2e . Theorem 6.1 (Non-linear case, simplified). Suppose we observe E environments E = {e1, e2, . . . , eE}, where σ2e = 1 ∀e. Then, for any > 1, there exists a featurizer Φ which, combined with the ERM-optimal classifier β̂ = [βc, βe;ERM, β0]T , satisfies the following properties, where we define p := exp{−de min( − 1, ( − 1)2)/8}):
1. The regularization term of Φ , β̂ as in Equation 5 is bounded as
1
E ∑ e∈E ‖∇β̂Re(Φ , β̂)‖22 ∈ O ( p2 ( c de + ‖µ‖22 )) ,
for some constant c that depends only on .
2. Φ , β̂ exactly matches the optimal invariant predictor on at least a 1− p fraction of the training set. On the remaining inputs, it matches the ERM-optimal solution.
Further, for any test distribution, suppose its environmental mean µE+1 is sufficiently far from the training means:
∀e ∈ E , min y∈{±1}
‖µE+1 − y · µe‖2 ≥ ( √ + δ) √ de (7)
for some δ > 0, and define q := 2E√ πδ exp{−δ2}. Then the following holds:
3. Φ , β̂ is equivalent to the ERM-optimal predictor on at least a 1−q fraction of the test distribution. 4. Under Assumption 1, suppose it holds that µE+1 = − ∑ e∈E αeµe for some set of coefficients
{αe}e∈E . Then so long as∑ e∈E αe‖µe‖22 ≥ ‖µc‖22/σ2c + |β0|/2 + σERM 1− γ , (8)
the 0-1 risk of Φ , β̂ on the new environment is greater than .975− q.
We give a brief intuition for each of 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 we construct will have a gradient squared norm scaling as p2 which is exponentially small in de. Thus, in high dimensions, it will appear as a perfectly reasonable solution to the objective (5).
2. The second claim says that this predictor is identical to the invariant optimal predictor on 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 is the crux of the theorem; it 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 a Gaussian prior on the training environment means, they will have separation in O( √ de) with high probability. Observe that q will be vanishingly small so
long as δ ≥ polylog(E). 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 Ω( √ de logE) 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 at test-time, our predictor will have significantly worse than chance performance.
Proof 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,Bc. B is the union of balls centered at each µe, each 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 Φ(x) = [zc], so our construction is equal to the optimal invariant predictor. Now consider Bc = Rde \ B. We use standard concentration results to upper bound the measure of Bc under the training distribution by p. Next, we show how choosing Φ(x) = f−1(x) = [zc, ze]
T on this set results in the sub-optimality bound, which is of order p2. It is also clear that our constructed predictor is equivalent to the ERM-optimal solution on Bc. Thus, our predictor will often have lower empirical risk on Bc, countering the regularization penalty. The second part of the proof shows that while B has large measure under the training environments, it will have very small measure under any moderately different test environment. We can see this by considering the separation of means (Equation 7); 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 µE+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 q of the observations (part 3). Finally, we lower bound the 0-1 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 risk can differ from this by at most q.
Theorem 6.1 shows 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 on future distributions. But in such a scenario, there would no longer be a 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 to Appendix E.

heading: 7 CONCLUSION
text: Out-of-distribution generalization is 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 demonstrates 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.

heading: ACKNOWLEDGEMENTS
text: 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 N000141812861.

heading: A ADDITIONAL NOTATION FOR THE APPENDIX
text: To avoid overloading, we use φ, F for the standard Gaussian PDF and CDF respectively. We write Sc to denote the set complement of a set S. We write | · | to denote entrywise absolute value.

heading: B PROOF OF PROPOSITION 4.1
text: Recall the IRM objective:
min Φ,β
E(x,y)∼p(x,y)[− log σ(y · β̂TΦ(x))]
subject to ∂
∂β̂ E(x,y)∼pe [− log σ(y · β̂TΦ(x))] = 0. ∀e ∈ E .
Concretely, we represent Φ as some parametrized function Φθ, over whose parameters θ we then optimize. The derivative of the negative log-likelihood for logistic regression with respect to the β coefficients is well known:
∂
∂β̂
[ − log σ(y · β̂TΦθ(x)) ] = (σ(β̂TΦθ(x))− 1{y = 1})Φθ(x).
Suppose we recover the true invariant features Φθ(x) = [ zc 0 ] and coefficients β̂ = [ β 0 ] (in other
words, we allow for the introduction of new features). Then the IRM constraint becomes:
0 = ∂
∂β̂ E(x,y)∼pe [− log σ(y · β̂TΦθ(x))]
= ∫ Z pe(zc) ∑ y∈{±1} pe(y | zc) ∂ ∂β̂ [ − log σ(y · βT zc) ] dzc
= ∫ Z pe(zc)Φθ(x) [ σ(β̂T zc)(σ(β̂ T zc)− 1) + (1− σ(β̂T zc))σ(β̂T zc) ] dzc.
Since β̂ is constant across environments, this constraint is clearly satisfied for every environment, and is therefore also the minimizing β̂ for the training data as a whole.
Considering now the derivative with respect to the featurizer Φθ:
∂
∂θ
[ − log σ(y · β̂TΦθ(x)) ] = (σ(β̂TΦθ(x))− 1{y = 1}) ∂
∂θ β̂TΦθ(x).
Then the derivative of the loss with respect to these parameters is∫ Z pe(zc) ( ∂ ∂θ β̂TΦθ(x) )[ σ(βT zc)(σ(β T zc)− 1) + (1− σ(βT zc))σ(βT zc) ] dzc = 0.
So, the optimal invariant predictor is a stationary point with respect to the feature map parameters as well.

heading: C RESULTS FROM SECTION 5

heading: C.1 PROOF OF THEOREM 5.1
text: We begin by formally stating the non-degeneracy condition. Consider any environmental mean µe, and suppose it can be written as a linear combination of the others means with coefficients αe:
µe = ∑ i αeiµi.
Then the environments are considered non-degenerate if the following inequality holds for any such set of coefficients: ∑
i
αei 6= 1, (9)
and furthermore that the following ratio is different for at least two different environments a, b:
∃αa, αb. σ 2 a −
∑ i α a i σ 2 i
1−∑i αai 6= σ 2 b −
∑ i α b iσ 2 i
1−∑i αbi . (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 Arjovsky et al. (2019). The other inequality is a similarly lax non-degeneracy requirement regarding the relative scale of the variances. It is clear that the set of environmental parameters that do not satisfy Equations 9 and 10 has measure zero under any absolutely continuous density, and similarly, if E ≤ de then the environmental means will be linearly independent almost surely.
We can now proceed with the proof, beginning with some helper lemmas: Lemma C.1. Suppose we observe E environments E = {e1, e2, . . . , eE}, each with environmental mean of dimension de ≥ E, such that all environmental means are linearly independent. Then there is a unique unit-norm vector p such that
pTµe = σ 2 e µ̃ ∀e ∈ E , (11)
where µ̃ is the largest scalar which admits such a solution.
Proof. Let v1, v2, . . . , vE be a set of basis vectors for span{µ1, µ2, . . . , µE}. Each mean can then be expressed as a combination of these basis vectors: ui = ∑E j=1 αijvj . Since the means are linearly independent, we can combine these coefficients into a single invertible matrix
U =  α11 α21 . . . αE1 α12 α22 . . . αE2
... ...
. . . ...
α1E α2E . . . αEE  . We can then combine the constraints (11) as
UT pα = σ ,  σ21 σ22 ... σ2E  , where pα denotes our solution expressed in terms of the basis vectors {vi}Ei=1. This then has the solution
pα = U −Tσ.
This defines the entire space of solutions, which consists of pα plus any element of the remaining (de − E)-dimensional orthogonal subspace. However, we want p to be unit-norm—observe that the current vector solves Equation 11 with µ̃ = 1, which means that after normalizing we get µ̃ = 1‖pα‖2 . Adding any element of the orthogonal subspace would only increase the norm of p, decreasing µ̃. Thus, the unique maximizing solution is
pα = U−Tσ
‖U−Tσ‖2 , with µ̃ =
1
‖U−Tσ‖2 .
Finally, pα has to be rotated back into the original space by defining p = ∑E i=1 pαivi.
Lemma C.2. Assume f is linear. Suppose we observe E ≤ de environments whose means are linearly independent. Then there exists a linear Φ with rank(Φ) = dc + de + 1− E whose output depends on the environmental features, yet the optimal classifier on top of Φ is invariant.
Proof. We will begin with the case when E = de and then show how to modify this construction for when E < de. Consider defining
Φ = [ I 0 0 M ] ◦ f−1
with
M =  pT
0 ... 0  . Here, p ∈ Rdc is defined as the unit-norm vector solution to
pTµe = σ 2 e µ̃ ∀e
such that µ̃ is maximized—such a vector is guaranteed to exist by Lemma C.1. Thus we get
Φ(x) = [ zc pT ze ] , which is of rank dc + 1 as desired. Define z̃e = pT ze, which means that z̃e | y ∼ N (y · σ2e µ̃, σ2e). For each environment we have
p(y | zc, z̃e) = p(zc, z̃e, y)
p(zc, z̃e)
= σ(y · βTc zc)p(z̃e | y · σe2e µ̃, σ2e)
[σ(y · βTc zc)p(z̃e | y · σe2e µ̃, σ2e) + σ(−y · βTc zc)p(z̃e | −y · σe2e µ̃, σ2e)]
= σ(y · βTc zc) exp(y · z̃eµ̃)
[σ(y · βTc zc) exp(y · z̃eµ̃) + σ(−y · βTc zc) exp(−y · z̃eµ̃)]
= 1
1 + exp(−y · (βTc zc + 2z̃eµ̃)) .
The log-odds of y is linear in these features, so the optimal classifier is
β̂ = [ βc 2µ̃ ] ,
which is the same for all environments.
Now we show how to modify this construction for when E < de. If we remove one of the environmental means, Φ trivially maintains its feasibility. Note that since they are linearly independent, the mean which was removed has a component in a direction orthogonal to the remaining means. Call this component p′, and consider redefining M as
M =  pT p′T
0 ... 0  . The distribution of z̃e in each of the remaining dimensions is normal with mean 0, which means a corresponding coefficient of 0 is optimal for all environments. So the classifier β̂ remains optimal for all environments, yet we’ve added another row to M which increases the dimensionality of its span, and therefore the rank of Φ, by 1. Working backwards, we can repeat this process for each additional mean, such that rank(Φ) = dc + 1 + (de −E), as desired. Note that for E = 1 any Φ will be vacuously feasible.
Lemma C.3. Suppose we observe E environments E = {e1, e2, . . . , eE} whose parameters satisfy the non-degeneracy conditions (9, 10). Let Φ(x) = Azc + Bze be any feature vector which is a linear function of the invariant and environmental features, and suppose the optimal β̂ on top of Φ is invariant. If E > de, then β̂ = 0 or B = 0.
Proof. Write Φ = [A|B] where A ∈ Rd×dc , B ∈ Rd×de and define
µ̄e = Φ [ µc µe ] = Aµc +Bµe,
Σ̄e = Φ
[ σ2cIdc 0
0 σ2eIde
] ΦT = σ2cAA T + σ2eBB T .
Without loss of generality we assume Σ̄ is invertible (if it is not, we can consider just the subspace in which it is—outside of this space, the features have no variance and therefore cannot carry information about the label). By Lemma F.2, the optimal coefficient for each environment is 2Σ̄−1e µ̄e. In order for this vector to be invariant, it must be the same across environments; we write it as a constant β̂. Suppose µ̄e = 0 for some environment e—then the claim is trivially true with β̂ = 0. We therefore proceed under the assumption that µ̄e 6= 0 ∀e ∈ E . With this fact, we have that ∀e ∈ E ,
β̂ = 2(σ2cAA T + σ2eBB T )−1(Aµc +Bµe)
⇐⇒ (σ2cAAT + σ2eBBT )β̂ = 2Aµc + 2Bµe ⇐⇒ σ2eBBT β̂ − 2Bµe = 2Aµc − σ2cAAT β̂. (12)
Define the vector v = 2Aµc − σ2cAAT β̂. We will show that for any β̂, A, with probability 1 only B = 0 can satisfy Equation 12 for every environment. If E > de, then there exists at least one environmental mean which can be written as a linear combination of the others. Without loss of generality, denote the parameters of this environment as (µ̄, σ̄2) and write µ̄ = ∑de i=1 αiµi. However, note that by assumption the means lie in general linear position, and therefore we actually have at least de sets of coefficients α for which this holds. Rearranging Equation 12, we get
σ̄2BBT β̂ − v = 2Bµ̄
= de∑ i=1 αi2Bµi
= de∑ i=1 αi [ σ2iBB T β̂ − v ] ,
and rearranging once more yields( σ̄2 − ∑ αiσ 2 i ) BBT β̂ = ( 1− ∑ αi ) v. By assumption, (1−∑αi) is non-zero. We can therefore rewrite this as α̂BBT β̂ = v,
where α̂ = σ̄ 2−
∑ αiσ 2 i
1− ∑ αi is a scalar. As the vectors BBT β̂ and v are both independent of the
environment, this can only hold true if α̂ is fixed for all environments or if both BBT β̂,v are 0. The former is false by assumption, so the the latter must hold.
As a result, we see that Equation 12 reduces to
Bµe = 0 ∀e ∈ E .
As the span of the observed µe is all of Rde , this is only possible if B = 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 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 β̂ 6= 0, and write Φ(f(zc, ze)) = Azc +Bze. Then under mild non-degeneracy conditions, it holds that B = 0. Consequently, β̂ is the optimal classifier for all possible environments.
2. If E ≤ de and the environmental means µe are linearly independent, then there exists a linear Φ—where Φ(f(zc, ze)) = Azc +Bze with rank(B) = de + 1− E—which is feasible under the IRM objective. Further, both the logistic and 0-1 risks of this Φ and its corresponding optimal β̂ are strictly lower than those of the optimal invariant predictor.
Proof. 1. Since Φ, f are linear, we can write Φ(x) = Azc + Bze for some matrices A,B. Assume the non-degeneracy conditions (9, 10) hold. By Lemma C.3, one of B = 0 or β̂ = 0 holds. Thus, Φ, β̂ uses only invariant features. Since the joint distribution pe(zc, y) is invariant, this predictor has identical risk across all environments.
2. The existence of such a predictor is proven by Lemma C.2. 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 R, and therefore by Lemma F.1, dropping the z̃e term and using
Φ(x) = [zc], β̂ = [ βc β0 ] results in strictly higher risk. The proof is completed by noting that this definition is precisely the optimal invariant predictor.

heading: C.2 EXPERIMENTS FOR THEOREM 5.1
text: To corroborate our theoretical findings, we run an experiment on data drawn from our model to see at what point IRM is able to recover a generalizing predictor. We generated data precisely according to our model in the linear setting, with dc = 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 a 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 (η = 0.5).
Figure C.1 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 rate 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 is 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 ze and the label y were reversed.
Observe that the results closely track the expected outcome according to Theorem 5.1: up until E = de, 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; we conjecture that this is because at this point the feasible-yet-notinvariant predictor with lower objective value presented in Theorem 5.1 is precisely a single point, as opposed to a multi-dimensional subspace, and therefore might be difficult for the optimizer to find.

heading: C.3 PROOF OF THEOREM 5.3
text: Theorem 5.3. Suppose we observe E ≤ de environments, such that all environmental means are linearly independent. Then there exists a feasible Φ, β̂ which uses only environmental features and achieves lower 0-1 risk than the optimal invariant predictor on every environment e such that σeµ̃ > σ −1 c ‖µc‖2 and 2σeµ̃σ−1c ‖µc‖2 ≥ |β0|.
Proof. We consider the 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
β̂caus = [ 2σ−2c µc β0 ] and β̂non−caus = [ 2µ̃ β0 ] ,
respectively. Therefore, the 0-1 risk of the optimal invariant predictor is precisely
ηP(2σ−2c µTc zc + β0 < 0) + (1− η)P(−2σ−2c µTc zc + β0 > 0)
=ηF ( −σ−1c ‖µc‖2 −
β0σc 2‖µc‖2
) + (1− η)F ( −σ−1c ‖µc‖2 +
β0σc 2‖µc‖2
) ,
where F is the Gaussian CDF. By the same reasoning, the 0-1 risk of the non-invariant predictor is
ηF ( −σeµ̃−
β0 2σeµ̃
) + (1− η)F ( −σeµ̃+
β0 2σeµ̃
) .
Define α = σ−1c ‖µc‖2 and γ = σeµ̃. By monotonicity of the Gaussian CDF, the former risk is higher than the latter if
α+ β0 2α ≤ γ + β0 2γ , (13)
α− β0 2α < γ − β0 2γ . (14)
Without loss of generality, we will prove these inequalities for β0 ≥ 0; an identical argument proves it for β0 < 0 but with the ‘≤’ and ‘<’ swapped.
Suppose γ > α (the first condition). Then Equation 14 is immediate. Finally, for Equation 13, observe that
γ + β0 2γ ≥ α+ β0 2α
⇐⇒ γ − α ≥ β0 2α − β0 2γ = (γ − α)β0 2γα
⇐⇒ 2γα ≥ β0, which is the second condition.

heading: C.4 SIMULATIONS OF MAGNITUDE OF ENVIRONMENTAL FEATURES
text: As discussed in Section 5, analytically quantifying the solution µ̃ to the equation in Lemma C.1 is difficult; instead, we present simulations to give a sense of how often these conditions would hold in practice.
For each choice of environmental dimension de, we generated a “base” correlation b ∼ N (0, Ide) as the mean of the prior over environmental means µe. Each of these µe was then drawn from N (b, 4Ide)—thus, while they all came from the same prior, the noise in the draw of each µe was significantly larger than the bias induced by the prior. We then solved for the precise value σeµ̃, with the same variance σ2e for all environments, chosen as a hyperparameter. The shaded area represents a 95% confidence interval over 20 runs. The dotted lines are √ dc. If we imagine the invariant parameters are drawn from a standard Gaussian prior, then this is precisely E[σ−1c ‖µc‖2]. Thus, the point where σeµ̃ crosses these dotted lines is approximately how many environments would need to be 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 − dc.

heading: D THEOREM 6.1 AND DISCUSSION

heading: D.1 PROOF OF THEOREM 6.1
text: We again begin with helper lemmas.
Our featurizer Φ is constructed to recover the environmental features only if they fall within a set Bc. The following lemma shows that since only the environmental features contribute to the gradient penalty, the penalty can be bounded as a 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 E = {e1, e2, . . .}. Given a set B ⊆ Rde , consider the predictor defined by Equation 19. Then for any environment e, the penalty term of this predictor in Equation 5 is bounded as
‖∇β̂Re(Φ, β̂)‖22 ≤ ∥∥∥∥P(ze ∈ Bc)E[|ze| | ze ∈ Bc]∥∥∥∥2
2
.
Proof. We write out the precise form of the gradient for an environment e: ∇β̂Re(Φ, β̂) = ∫ Zc×Ze pe(zc, ze) [ σ(β̂TΦ(f(zc, ze)))− pe(y = 1 | zc, ze) ] Φ(f(zc, ze)) d(zc, ze).
Observe that since zc ⊥⊥ ze | y, the optimal invariant coefficients are unchanged, and therefore the gradient in the invariant dimensions is 0. We can split the gradient in the environmental dimensions into two integrals:∫
Zc×B pe(zc, ze)
[ σ(βTc zc + β0)− pe(y = 1 | zc, ze) ] [0] d(zc, ze)
+ ∫ Zc×Bc pe(zc, ze) [ σ(βTc zc + β T e;ERMze + β0)− σ(βTc zc + βTe ze + β0) ] [ze] d(zc, ze).
Since the features are 0 within B, the first term reduces to 0. For the second term, note that ∀ x, y ∈ R, |σ(x)− σ(y)| ≤ 1, and therefore
|∇β̂Re(Φ, β̂)| ≤ ∫ Zc×Bc pe(zc, ze)[|ze|] d(zc, ze).
We can marginalize out zc, and noting that we want to bound the squared norm, ‖∇β̂Re(Φ, β̂)‖22 ≤ ∥∥∥∥∫ Bc pe(ze)[|ze|] dze ∥∥∥∥2 2
= ∥∥∥∥P(ze ∈ Bc)E[|ze| | ze ∈ Bc]∥∥∥∥2 2 .
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 Br, and therefore Φ , β̂ will be equivalent to the ERM solution. Lemma D.2. For a set of E environments E = {e1, e2, . . . , eE} and any > 1, construct Br as in Equation 18 and define Φ using Br as in Equation 19. Suppose we now test on a new environment with parameters (µE+1, σ2E+1), and assume Equation 15 holds with parameter δ. Define k = mine∈E σ2e σ2E+1 . Then with probability ≥ 1 − 2E√ kπδ
exp{−kδ2} over the draw of an observation from this new environment, we have
Φ (x) = f −1(x) = [ zc ze ] .
Proof. By Equation 15 our new environmental mean is sufficiently far away from all the labelconditional means of the training environments. In particular, for any environment e ∈ E and any label y ∈ {±1}, the `2 distance from that mean to µE+1 is at least ( √ + δ)σe √ de.
Recall that Br is the union of balls±Be, whereBe is the ball of `2 radius √ σ2ede centered at µe. For each environment e, consider constructing the halfspace which is perpendicular to the line connecting
µe and µE+1 and tangent to Be. This halfspace fully contains Be, and therefore the measure of Be is upper bounded by that of the halfspace.
By rotational invariance of the Gaussian distribution, we can rotate this halfspace into one dimension and the measure will not change. The center of the ball is ( √ + δ)σe √ de away from the mean µE+1,
so accounting for its radius, the distance from the mean to the halfspace is δσe √ de. The variance of the rotated distribution one dimension is σ2E+1de, so the measure of this halfspace is upper bounded by
1− Φ  δσe√de√ σ2E+1de  ≤ Φ(−√kδ) ≤ 1√
kπδ exp{−kδ2},
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 E environments E = {e1, e2, . . . , eE}. Then, for any > 1, there exists a featurizer Φ which, combined with the ERM-optimal classifier β̂ = [βc, βe;ERM, β0]
T , satisfies the following properties, where we define p ,de := exp{−de min(( − 1), ( − 1)2)/8}:
1. Define σ2max = maxe σ 2 e . Then the regularization term of Φ , β̂ is bounded as
1
E ∑ e∈E ‖∇β̂Re(Φ , β̂)‖22 ∈ O ( p2 ,de [ deσ 4 max exp{2 σ2max}+ ‖µ‖22 ]) .
2. Φ , β̂ exactly matches the optimal invariant predictor on at least a 1− p ,de fraction of the training set. On the remaining inputs, it matches the ERM-optimal solution.
Further, for any test distribution with environmental parameters (µE+1, σ2E+1), suppose the environmental mean µE+1 is sufficiently far from the training means:
∀e ∈ E , min y∈{±1}
‖µE+1 − y · µe‖2 ≥ ( √ + δ)σe √ de (15)
for some δ > 0. Define the constants:
k = min e∈E σ2e σ2E+1
q = 2E√ kπδ exp{−kδ2}.
Then the following holds:
3. Φ , β̂ is equivalent to the ERM-optimal predictor on at least a 1 − q fraction of the test distribution.
4. Under Assumption 1, suppose it holds that µE+1 = − ∑ e∈E αeµe (16)
for some set of coefficients {αe}e∈E . Then for any c ∈ R, so long as∑ e∈E αe ‖µe‖22 σ2e ≥ ‖µc‖ 2 2/σ 2 c + |β0|/2 + cσERM 1− γ , (17)
the 0-1 risk of Φ , β̂ is lower bounded by F (2c)− q.
Proof. Define r = √ σ2ede and construct Br ⊂ Rde as
Br = [⋃ e∈E Br(µe) ] ∪ [⋃ e∈E Br(−µe) ] , (18)
where Br(α) is the ball of `2-norm radius r centered at α. Further construct Φ using Br as follows:
Φ (x) =  [ zc 0 ] , ze ∈ Br[
zc ze
] , ze ∈ Bcr,
and β̂ = βcβ̂e β0  . (19) Without loss of generality, fix an environment e.
1. By Lemma D.1, the squared gradient norm is upper bounded by ‖∇β̂Re(Φ , β̂)‖22 ≤ ∥∥∥∥P(ze ∈ Bcr)E[|ze| | ze ∈ Bcr]∥∥∥∥2
2
. (20)
Define Be := Br(µe), and observe that Bcr ⊆ Bce . Since |ze| is non-negative, P(ze ∈ Bcr)E[|ze| | ze ∈ Bcr] ≤ P(ze ∈ Bce)E[|ze| | ze ∈ Bce]
(this inequality is element-wise). Plugging this into Equation 20 yields ‖∇β̂Re(Φ , β̂)‖22 ≤ ∥∥∥∥P(ze ∈ Bce)E[|ze| | ze ∈ Bce]∥∥∥∥2
2 = [P(ze ∈ Bce)]2 ∥∥∥∥E[|ze| | ze ∈ Bce]∥∥∥∥2
2
.
Define p = P(ze ∈ Bcr) ≤ P(ze ∈ Bce). By Lemma F.3, p ≤ p ,de = e−de min(( −1),( −1) 2)/8.
Combining Lemmas F.4 and F.5 gives∥∥∥∥E[|ze| | ze ∈ Bce]∥∥∥∥2 2 ≤ 2de [ σ φ(r/ √ de) F (−r/ √ d) ]2 + 2‖µe‖22
≤ deσ2e exp { 2 (σ2e − 1/2) } [ σ2e + 1 ] + 2‖µe‖22.
Putting these two bounds together, we have ‖∇β̂Re(Φ , β̂)‖22 ∈ O ( p2 ,de [ deσ 4 max exp{2 σ2max}+ ‖µe‖22 ]) ,
and averaging this value across environments gives the result.
2. Φ , β̂ is equal to the optimal invariant predictor on Br and the ERM solution on Bcr. The claim then follows from Lemma F.3.
3. This follows directly from Lemma D.2.
4. With Equation 16, we have that βTe;ERMµE+1 = − ∑ e∈E αeβ T e;ERMµe
≤ −2(1− γ) ∑ e∈E αe ‖µe‖22 σ2e
≤ −2(1− γ)‖µc‖ 2 2/σ 2 c + |β0|/2 + cσERM 1− γ = −(2‖µc‖22/σ2c + |β0|+ 2cσERM).
where we have applied Assumption 1 in the first inequality and Equation 17 in the second. Consider the full set of features Φ (x) = f−1(x), and without loss of generality assume y = 1. The label-conditional distribution of the resulting logit is
βTc zc + β T e;ERMze + β0 ∼ N ( βTc µc + β T e;ERMµE+1 + β0, σ 2 ERM ) .
Therefore, the 0-1 risk is equal to the probability that this logit is negative. This is precisely
F ( − βTc µc + β T e;ERMµE+1 + β0
σERM
) ≥ F ( (2‖µc‖22/σ2c + |β0|+ 2cσERM)− 2‖µc‖22/σ2c − |β0|
σERM ) = F
( 2cσERM σERM ) = F (2c).
Observe that by the previous part, Φ 6= f−1 on at most a q fraction of observations, so the risk of our predictor Φ , β̂ can differ from that of f−1, β̂ by at most q. Therefore our predictor’s risk is lower bounded by F (2c)− q. In particular, choosing c = 1 recovers the statement in the main body.

heading: D.2 DISCUSSION OF CONDITIONS AND ASSUMPTION
text: To see just how often we can expect the conditions for Theorem D.3 to hold, we can do a rough approximation based on the expectations of each of the terms. A reasonable prior for the environmental means is a multivariate Gaussian N (m,Σ). We might expect them to be very concentrated (with Tr(Σ) small), or perhaps to have a strong bias (with ‖m‖22 Tr(Σ)). For simplicity we treat the variances σ2c , σ 2 e as constants. Then the expected separation between any two means from this distribution is E[‖µ1 − µ2‖2] = Ex∼N (0,2Σ)[‖x‖2] ≈ √ 2 Tr(Σ). In high dimensions this value will tightly concentrate around the mean, which is in O( √ de). On the
other hand, even a slight deviation from this separation, to Ω( √ de logE), means δ ∈ Ω( √ logE), which implies q ∈ O(1/E); this is plenty small to ensure worse-than-random error on the test distribution.
Now we turn our attention to the second condition (17). 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 is technically not true with a non-centered Gaussian). Then so long as ∑ i αi ∈ Ω(1), the left-hand side of Equation 17 is approximately de. On the other hand, treating γ as a constant, the right-hand side is close to dc + √ dc + de ∈ O(dc + √ de). Thus, Equation 17 is quite likely to hold for any mean µE+1 with the same scale as the training environments but with reversed correlations—again, this is exactly the situation where IRM hopes to outperform ERM, and we have shown that it does not.
We can also do a quick analysis of Assumption 1 under this prior: the ERM-optimal non-invariant coefficient will be approximately 2m/σ2e with high probability, meaning β̂
Tµ ≈ 2‖m‖22/σ2e for every environment. Thus, this vector will be γ-close to optimal with γ ≈ 0 for every environment with high probability.

heading: E EXTENSIONS TO ALTERNATIVE OBJECTIVES

heading: E.1 EXTENSIONS FOR THE LINEAR CASE
text: Observe that the constraint of Equation 4 is strictly stronger than that of Bellot & van der Schaar (2020); when the former is satisfied, the penalty term of the latter is necessarily 0. It is thus trivial to extend all results in the Section 5 to this objective. As another example, consider the risk-variancepenalized objective of Krueger et al. (2020):
min Φ,β̂
1 |E| ∑ e∈E Re(Φ, β̂) + λVare∈E ( Re(Φ, β̂) ) , (21)
It is simple to extend Theorem 5.1 under an additional assumption: Corollary E.1 (Extension to Theorem 5.1). Assume f is linear. Suppose we observe E ≤ de environments with linearly independent means and identical variance σ2e . Consider minimizing empirical risk subject to a penalty on the risk variance (21). Then there exists a Φ, β̂ dependent on the non-invariant features which achieves a lower objective value than the optimal invariant predictor for any choice of regularization parameter λ ∈ [0,∞].
Proof. Consider the featurizer Φ constructed in Lemma C.2. If the environmental variance is constant, then the label-conditional distribution of the environmental features,
ze | y ∼ N (y · µ̃σ2e , σ2e), is also invariant. This implies that the optimal β̂ also has constant risk across the environments, meaning the penalty term is 0, and as a result the objective does not depend on the choice of λ. As in 5.1, invoking Lemma F.1 implies that the overall risk is lower than that of the optimal invariant predictor.
As mentioned in Section 5, this additional requirement of constant variance is due to the assumptions underlying the design of the objective—REx expects additional invariance of the second moment Var(y | Φ(x)), in contrast with the strictly weaker invariance of E[y | Φ(x)] assumed by IRM. This might seem to imply that REx is a 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 to Theorem 5.3). Suppose we observe E ≤ de environments, such that all environmental means are linearly independent. Then there exists a Φ, β̂ which uses only environmental features and, under any choice of λ ∈ [0,∞], achieves a lower objective value than the optimal invariant predictor under 0-1 loss on every environment e such that µ̃ > σ−1c ‖µc‖2 + |β0| 2σ−1c ‖µc‖2 .
Proof. We follow the proof of Theorem 5.3, except when solving for p as in Lemma C.1 we instead find the unit-norm vector such that
pTµe = σeµ̃ ∀e ∈ E . (22) Observe that by setting Φ(x) = [pT ze] and β̂ = [1], the 0-1 risk in a given environment is ηF (−µ̃σe/σe) + (1− η)F (−µ̃σe/σe) = F (−µ̃), which is independent of the environment. Further, by carrying through the same proof as in Theorem 5.3, we get that this non-invariant predictor has lower 0-1 risk so long as
α+ |β0| 2α ≤ µ̃,
where α = σ−1c ‖µc‖2
Though µ̃ 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.

heading: E.2 EXTENSIONS FOR THE NON-LINEAR CASE
text: The failure of these objectives in the non-linear regime is even more straightforward, as we can keep unchanged the constructed predictor from Theorem 6.1. Observe that parts 2-4 of the theorem do not involve the objective itself, and therefore do not require modification.
To see that part 1 still holds, note that since the constructed predictor matches the optimal invariant predictor on 1− p of the observations, its risk across environments can only vary on the remaining p fraction: thus the centered 0-1 risk is bounded between 0 and p. It is immediate that the variance of the environmental risks is upper bounded by p 2
4 ∈ O(p2). Applying this argument to the other objectives yields similar results.
0 50 100 150 200 250 300
Environments observed
0
5
10
15
20
25
30
35
40
µ̃
Projected Environmental Mean Separation
dc = 40, de = 200, σ 2 e = 3 dc = 100, de = 200, σ 2 e = 1 dc = 50, de = 300, σ 2 e = 3 dc = 60, de = 300, σ 2 e = 1√
de − E
Figure E.1: Simulations to evaluate µ̃ for varying ratios of dedc . When σ 2 e = 1 the value closely tracks√
de − E, and the crossover point is approximately de − σ2edc. Due to the similarity of Equation 22 to Equation 11, it makes sense that the results are very similar to those presented in Figure C.2.

heading: F TECHNICAL LEMMAS
text: Lemma F.1. Consider solving the standard logistic regression problem
z ∼ p(z) ∈ Rk,
y = { +1 w.p. σ(βT z), −1 w.p. σ(−βT z).
Assume that none of the latent dimensions are degenerate—∀S ⊆ [k], P(βTS zS 6= 0) > 0, and no feature can be written as a linear combination of the other features. Then for any distribution p(z), any classifier f(z) = σ(βTS zS) that uses a strict subset of the features S ( [k] has strictly higher risk with logistic loss than the Bayes classifier f∗(z) = σ(βT z). This additionally holds for 0-1 loss if βT−Sz−S has greater magnitude and opposite sign of β T S 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 on a set of non-zero measure. Consider the set of values z−S such that βT−Sz−S 6= 0. Then on this set we have
f∗(βT z) = σ(βTS zS + β T −Sz−S) 6= σ(βTS zS) = f(z).
Since these values occur with positive probability, f has strictly higher logistic risk than f∗. By the same argument, there exists a set of positive measure under which
f∗(βT z) = sign(βTS zS + β T −Sz−S) 6= sign(βTS zS) = f(z),
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 z̃ = Azc +Bze, any optimal corresponding coefficient for an environment e is of the form
2(AATσ2c +BB Tσ2e) +(Aµc +Bµe),
where G+ is a generalized inverse of G.
Proof. We begin by evaluating a closed form for pe(y | z̃). We have: pe(y | Azc +Bze = z̃)
= p(Azc +Bze = z̃ | y)p(y)
pe(Azc +Bze = z̃)
= pe(Azc +Bze = z̃ | y)
pe(Azc +Bze = z̃ | y) + pe(Azc +Bze = z̃ | −y)
= 1
1 + p e(Azc+Bze=z̃|−y) pe(Azc+Bze=z̃|y)
.
Now we need a closed form expression for p(Azc +Bze = z̃ | y). Noting that zc ⊥⊥ ze | y, this is a convolution of the two independent Gaussian densities, which is the density of their sum. In other words,
Azc +Bze | y ∼ N (y · ( µ̄︷ ︸︸ ︷ Aµc +Bµe), Σ̄︷ ︸︸ ︷ AATσ2c +BB Tσ2e).
Thus,
pe(Azc +Bze = z̃ | y) = 1( 2π|Σ̄| )k/2 exp{−12(z̃ − y · µ̄)T Σ̄+(z̃ − y · µ̄) } .
Canceling common terms, we get
pe(y = 1 | Azc +Bze = z̃) = 1
1 + p e(Azc+Bze=z̃|−y) pe(Azc+Bze=z̃|y)
= 1 1 + exp { −y · 2z̃T Σ̄+µ̄ } = σ ( y · 2z̃T Σ̄+µ̄ ) .
Therefore, given a feature vector z̃, the optimal coefficient vector is 2Σ̄+µ̄.
Lemma F.3. For any environment e with parameters µe, σ2e and any > 1, define
B := B√ σ2ede (µe),
where Br(α) is the ball of `2-norm radius r centered at α. Then for an observation drawn from pe, we have
P ze∼pe
(ze ∈ Bc) ≤ exp { −de min(( − 1), ( − 1) 2)
8
} .
Proof. Without loss of generality, suppose y = 1. We have
P(ze ∈ B) ≥ P ze∼N (µe,σ2eI)
( ‖ze − µe‖2 ≤ √ σ2ede ) = P ze∼N (0,σ2eI) ( ‖ze‖2 ≤ √ σ2ede
) = P ze∼N (0,I) ( ‖ze‖22 ≤ de ) .
Each term in the squared norm of ze is a random variable with distribution χ21, which means their sum has mean de and is sub-exponential with parameters (2 √ de, 4). By standard sub-exponential concentration bounds we have
P ze∼N (0,I)
( ‖ze‖22 ≥ de ) ≤ exp { −de min(( − 1), ( − 1) 2)
8
} ,
which immediately implies the claim.
Lemma F.4. Let z ∼ N (µ, σ2Id) be a multivariate isotropic Gaussian in d dimensions, and for some r > 0 define B as the ball of `2 radius r centered at µ. Then we have∥∥∥∥E[|z| | z ∈ Bc]∥∥∥∥2
2
≤ 2d [ σ φ(r/ √ d)
F (−r/ √ d)
]2 + 2‖µ‖22,
where φ, F are the standard Gaussian PDF and CDF.
Proof. Observe that
E z∼N (µ,σ2Id) [|z| | z ∈ Bc] = E z∼N (µ,σ2Id) [|z| | ‖z − µ‖2 > r]
= E z∼N (0,σ2Id)
[|z + µ| | ‖z‖2 > r]
≤ E z∼N (0,σ2Id) [|z| | ‖z‖2 > r] + |µ|.
Now, consider the expectation for an individual dimension, and note that |zi| > r√d ∀i =⇒ ‖z‖2 > r. So because the dimensions are independent, conditioning on this event can only increase the expectation:
E z∼N (0,σ2Id) [|zi| | ‖z‖2 > r] ≤ E zi∼N (0,σ2)
[ |zi| | |zi| >
r√ d ] = E zi∼N (0,σ2) [ zi | zi > r√ d ] ,
where the equality is because the distribution is symmetric about 0. This last term is known as the conditional tail expectation of a Gaussian and is available in closed form:
E zi∼N (0,σ2)
[ zi | zi >
r√ d
] = σ φ(F−1(α))
1− α ,
where α = F (r/ √ d). Combining the above results, squaring with (a + b)2 ≤ 2(a2 + b2), and summing over dimensions, we get∥∥∥∥E[|z| | z ∈ Bc]∥∥∥∥2 2 ≤ 2 d∑ i=1 E zi∼N (0,σ2) [ zi | zi > r√ d ]2 + 2‖µ‖22
= 2d [ σ φ(r/ √ d)
F (−r/ √ d)
]2 + 2‖µ‖22,
as desired.
Lemma F.5. For σ, > 0, define r = √ σ. Then[
φ(r)
F (−r)
]2 ≤ 1
2 exp
{ 2 (σ2 − 1/2) } [ σ2 + 1 ] .
Proof. We have
φ(r) = 1√ 2π
exp { −
2 } and
F (−r) ≥ 2 exp{− σ 2}√
π( √ σ + √ σ2 + 2)
(see Kschischang (2017)). Dividing them gives φ(r)
F (−r) ≤ 1 2 √ 2 exp{ (σ2 − 1/2)}
[√ σ + √ σ2 + 2 ] .
Squaring and using (a+ b)2 ≤ 2(a2 + b2) yields the result.

Review summaries:
* 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.

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