T HE R ISKS OF I NVARIANT R ISK M INIMIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
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.
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 ﬁrst analysis of classiﬁcation 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 ﬁrst results in the non-linear regime: we demonstrate that IRM can fail catastrophically unless the test data are sufﬁciently similar to the training distribution—this is precisely the issue that it was intended to solve.
Thus, in this setting we ﬁnd that IRM and its alternatives fundamentally do not improve over standard Empirical Risk Minimization.
I NTRODUCTION
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 classiﬁer’s uncertainty at test time (Neal, 2012).
Another technique that has shown promise is data augmentation—this includes both automated data modiﬁcations which help prevent overﬁtting (Shorten & Khoshgoftaar, 2019) and speciﬁc 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.
The justiﬁcation for this is that the causal mechanism of the data generating process of the target variable is unchanging but other effects can vary.
Therefore, learning relationships that are invariant 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 classiﬁer”—this is the classiﬁer which uses and is optimal with respect to only the invariant features in the SEM.
By deﬁnition, such a classiﬁer does not overﬁt to environment-speciﬁc 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.
, but our results can easily be extended to similar recently proposed alternatives (see Appendix E).
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 partially observed covariates or instrumental variables.
Under such a condition, results from causal inference (Maathuis et al., 2009; Peters et al., 2017) allow for formal guarantees of the identiﬁcation 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.
Thus, IRM hopes to “extrapolate” to new test environments, unlike ERM which excels at “interpolating”.
Recent works have pointed to shortcomings of IRM and have suggested modiﬁcations 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 ﬁrst formal analysis of classiﬁcation under the IRM objective under a fairly natural and general model which is similar to 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 conditions under which solving to optimality succeeds or, more often, breaks down in recovering the optimal invariant classiﬁer.
We additionally demonstrate a fundamental 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 classiﬁer; thus it will appear as a more attractive solution, yet its reliance on non-invariant features mean it will completely fail to generalize.
As corollaries, we present similar cases where Risk Extrapolation (REx; Krueger et al. 2020) and similar objectives likewise fail.
Futhermore, we present the ﬁrst results in the non-linear regime : we demonstrate the existence of a classiﬁer with exponentially small sub-optimality which nevertheless heavily relies on non-invariant features on most test inputs, resulting in worse-than-chance performance on distributions that are sufﬁciently dissimilar from the training environments.
These ﬁndings 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 beneﬁt from a more formal treatment.
R ELATED WORK
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 )) , but this is 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.
Further, there is evidence that this may not be sufﬁcient for domain adaptation (Zhao et al., 2019; Johansson et al., 2019).
Other works instead hope to ﬁnd representations that are conditionally domain-invariant, with invariance of p (Φ( x ) | y )
(Gong et al., 2016; Li et al., 2018).
More recent works, including the objectives discussed in this paper, suggest invariance of the featureconditioned 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 p ( y | Φ( x ))
(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.
O UR R ESULTS
We consider an SEM with explicit separation of invariant features z c , whose joint distribution with the label is ﬁxed for all environments, and environmental features z e (“non-invariant”), whose distribution can vary.
We assume that data are drawn from a set of E training environments E = { e 1 , e 2 , . . .
, e E } and that we know from which environment each sample is drawn.
For a given environment e , the data are deﬁned by the following process: ﬁrst, a label y ∈ {± 1 } is drawn according to a ﬁxed probability:
[Equation]
Next, both invariant features and environmental features are drawn according to a Gaussian 1 :
[Equation]
with µ c ∈ R d c , µ e ∈ R d e —typically, for complex, high-dimensional data we would expect E < d c (cid:28) d e .
Finally, the observation x is generated as a function of the latent features:
[Equation]
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 ( z c , z e )) =
[ z c , z e ] T .
We write the joint and marginal distributions as p e ( x, y, z c , z e ) .
When clear from context, we omit the speciﬁc arguments.
Remarks on the model.
This model is natural and ﬂexible; 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 , σ 2 c , 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 , σ 2 c are the same for all environments, 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 speciﬁc to environment i ).
• While we have framed the model as y causing z c , the causation can just as easily be viewed in the other direction.
The log-odds of y are a linear function of z c —this matches logistic regression with an invariant regression vector β c
= 2 µ c /σ 2 c and bias β 0 = log η 1 − η .
We present the model as above to emphasize that the causal relationships between y and the z c , z e are a priori indistinguishable, and because we believe this direction is more intuitive.
Assuming a constant marginal is necessary—in logistic regression, the optimal bias term is log η 1 − η .
Thus the optimal classiﬁer can only be invariant if this value is the same for all environments.
• Modeling class-conditional means as direct opposites is made for clarity: it matches existing models and it greatly simpliﬁes the calculations and results.
None of our proofs require this condition, and it is straightforward to extend our results to the general case of arbitrary means.
We consider the setting where we are given inﬁnite samples from each environment; this allows us to isolate the behavior of the objectives themselves, rather than ﬁnite-sample effects.
Upon observing samples from this model, our objective is thus to learn a feature embedder Φ and regression vector ˆ β to minimize the risk on an unseen environment e :
[Equation]
However is often ; historically, a spurious correlation is one which (a) appears by chance and would disappear given enough samples or (b) is due to an unobserved confounder.
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 the label to cause the features (e.g. as in natural images), the resulting correlation is not spurious.
We therefore avoid using the this work, or
The function (cid:96) can be any loss appropriate to classiﬁcation: in this work we consider the logistic and the 0-1 loss.
Note, however, that we are not hoping to minimize this 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 classiﬁer generalizes to all unseen environments regardless of their parameters.
In other words, the focus is on minimizing risk for the worst-case environment.
We refer to the classiﬁer which will minimize worst-case risk under arbitrary distribution shift as the optimal invariant classiﬁer .
To discuss this formally, we deﬁne precisely what we mean by this term.
Deﬁnition 1.
Under the model described by Equations 1-3, an optimal invariant classiﬁer is a classiﬁer which agrees on all inputs with the composition of a) the featurizer which recovers the invariant features and b) the regression vector which is optimal with respect to those features:
[Equation]
Throughout this work we reference "the" optimal invariant classiﬁer; this simply refers to the classiﬁer deﬁned by Φ ∗ , ˆ β ∗ . Observe that this deﬁnition closely resembles Deﬁnition 3 of Arjovsky et al. ; the only difference is that here an optimal invariant classiﬁer 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 well-deﬁned; technically, a featurizer which outputs the empty set would elicit an invariant predictor, but this would not satisfy our deﬁnition. The above classiﬁer 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 classiﬁer is distinct from the Bayes classiﬁer ; the Bayes classiﬁer for a given environment uses environmental features which are informative of the label but non-invariant; the optimal invariant classiﬁer explicitly ignores these features.
With the model deﬁned, we can informally present our results; we defer the formal statements to ﬁrst give a background on the IRM objective in the next section.
With a slight abuse of notation, we identify a classiﬁer by the tuple Φ , ˆ β which parametrizes it.
First, we show that the usefulness of
IRM exhibits a “thresholding” behavior depending on E and d e :
Theorem 3.1 (Informal, Linear) .
For linear f , consider solving the IRM objective to learn a linear Φ with invariant optimal coefﬁcients ˆ β .
If E > d e , then Φ , ˆ β is precisely the optimal invariant classiﬁer; it uses only invariant features and generalizes to all environments with minimax-optimal risk.
If E ≤ d e , then Φ , ˆ β relies upon non-invariant features.
In fact, when E ≤ d e it is even possible to learn a classiﬁer solely relying on environmental features that achieves lower risk on the training environments than the optimal invariant classiﬁer:
Theorem 3.2 (Informal, Linear) .
For linear f and E ≤ d e there exists a linear classiﬁer Φ , ˆ β which uses only environmental features , yet achieves lower risk than the optimal invariant classiﬁer.
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 non-linear f , there exists a non-linear classiﬁer Φ , ˆ β which is nearly optimal under the penalized objective and furthermore is nearly identical to the optimal invariant classiﬁer on the training distribution.
However, for any test environment with a mean sufﬁciently different from the training means, this classiﬁer 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 classiﬁer has almost trivial performance.
Extensions to other objectives.
Several follow-up works have suggested alternatives to IRM (see
Section 4).
Though these objectives perform better on various baselines, there are few formal guarantees and no results beyond the linear case.
Due to their similarities, it is simple to 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.
B ACKGROUND ON IRM AND ITS A LTERNATIVES
During training, a classiﬁer 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 classiﬁer’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 classiﬁer 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 conditional distributions with the label and therefore a ﬁxed classiﬁer on top of them won’t be optimal in all environments.
To learn this Φ , the IRM objective is the following constrained optimization problem:
[Equation]
This bilevel program is highly non-convex and difﬁcult to solve.
To ﬁnd 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:
[Equation]
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 classiﬁer constitutes a stationary point of each of these objectives:
Proposition 4.1.
Suppose the observed data are generated according to Equations 1-3.
Then recovering the (parametrized) optimal invariant classiﬁer Φ ∗ , ˆ β ∗ is a stationary point for Equation 4.
The stationarity of the optimal invariant classiﬁer 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.
T HE D IFFICULTIES OF IRM IN THE L INEAR R EGIME
In their work proposing IRM, Arjovsky et al.
(2019) present speciﬁc 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 ﬁrst result is similar in spirit but presents a substantially stronger (and simpliﬁed) upper bound in the classiﬁcation 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:
Suppose E > d e .
Consider any linear featurizer Φ which is feasible under the IRM objective (4), with invariant optimal regression vector ˆ β (cid:54) = 0 , and write Φ( f ( z c , z e )) = Az c + Bz e . Then under mild non-degeneracy conditions, it holds that B = 0 .
If E ≤ d e and the environmental means µ e are linearly independent, then there exists a linear Φ —where Φ( f ( z c , z e ))
=
Az c + Bz e with rank ( B ) = d e + 1 − E —which is feasible under the IRM objective.
Further, both the logistic and 0-1 risks of this Φ and its corresponding ˆ β are strictly lower than those of the optimal invariant classiﬁer.
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 classiﬁer is Bayes with respect to the invariant features, by the data-processing inequality the only way a classiﬁer can achieve lower risk is by relying on environmental features.
Thus, Theorem 5.1 directly implies that when E ≤ d e , 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 > d e , any feasible solution will generalize, and the optimal invariant classiﬁer has the minimum (and minimax) risk of all such classiﬁers:
Corollary 5.2.
For both logistic and 0-1 loss, the global minimum of the IRM objective is an optimal invariant classiﬁer if and only if E > d e .
Let us compare our theoretical ﬁndings to those of Arjovsky et al.
Suppose the observations x lie in R d .
Roughly, their theorem says that for a learned Φ of rank r with invariant optimal coefﬁcient ˆ β , 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: ﬁrst, 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 highdimensional (i.e. d (cid:29) d e + d c )—in which case Φ will be low-rank (i.e. r
≤ d e + d c )
—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 = d c + d e , the “ideal” Φ which recovers precisely z c would have rank d c , and therefore their condition for invariance would require
E > d e + d e /d c , 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 classiﬁer 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 =
d e :
Proof Sketch.
Since f has an inverse over its range, we can deﬁne Φ as a linear function directly over the latents [ z c , z e ] .
Speciﬁcally, deﬁne Φ( x ) =
[ z c , p T z e ] .
Here, p is a unit-norm vector such that ∀ e ∈ E , p T µ e = σ 2 e ˜ µ ; ˜ µ is a ﬁxed scalar that depends on the geometry of µ e , σ 2 e —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: p [Equation] (cid:107)
[Equation] For separating two Gaussians, the optimal linear classiﬁer weights are Σ − 1 ( µ 1 − µ 0 ) —here, the optimal classiﬁer weight on p T z e is precisely 2˜ µ , which does not depend on the environment (and neither do the optimal coefﬁcients for z c ).
Though the distribution varies across environments, the optimal classiﬁer 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 classiﬁer, note that this classiﬁer is Bayes with respect to its features and that the optimal invariant classiﬁer uses a strict subset of these features, and therefore it has less information for its predictions.
A purely environmental classiﬁer.
The precise value of ˜ µ in the proof sketch above represents how strongly this non-invariant feature is correlated with the label.
In theory, a classiﬁer 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 sufﬁce.
This result would be less surprising, since achieving low empirical risk might still ensure that we are “close” to the optimal invariant classiﬁer.
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 classiﬁer on all e ∈ E for which ˜ µ is large enough.
Theorem 5.3.
Suppose we observe E ≤ d e 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 classiﬁer on every environment e such that [Equation] and[-
-]
[Equation]
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 (cid:28) d e 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 d e − d c .
T HE F AILURE OF IRM IN THE N ON -L INEAR R EGIME
We’ve demonstrated that OOD generalization is difﬁcult 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 sufﬁciently 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 classiﬁer, which makes reasoning about the constrained solution to Equation 4 difﬁcult.
Instead we turn our attention to Equation 5, the penalized IRM objective.
In this section we demonstrate a foundational ﬂaw 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 deﬁnition about the optimality of a coefﬁcient vector ˆ β :
Deﬁnition 2.
For 0 < γ < 1 , a coefﬁcient vector ˆ β is γ -close to optimal for a label-conditional feature distribution z ∼ N ( y · µ, Σ
[Equation]
if
[Equation]
Since the optimal coefﬁcient vector is precisely 2Σ − 1 µ , being γ -close implies that ˆ β is reasonably aligned with that optimum.
Observe that the deﬁnition 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 sufﬁciently similar to one another.
Assumption 1.
There exists a 0 ≤ γ < 1 such that the ERM-optimal coefﬁcients for the non-invariant features,
[Equation]
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 classiﬁer” is reasonably predictive for each environment individually.
This is a natural expectation: we are employing IRM precisely because we expect the ERM classiﬁer to do well on the training set but fail to generalize.
If the environmental parameters are sufﬁciently 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 simpliﬁed version, assuming that that σ 2 e = 1 ∀ e .
This is purely for clarity of presentation; the full theorem, for which the result still holds, is presented in Appendix D. We make use of two constants in the following proof— the average squared norm of the environmental means, (cid:107) µ (cid:107) 22 := 1 E (cid:80) e ∈E (cid:107) µ e (cid:107) 22 ; and the standard deviation of the response variable of the ERM-optimal classiﬁer, σ ERM := (cid:112) (cid:107) β c
(cid:107) 2 2 σ 2 c
+ (cid:107) β e ; ERM (cid:107) 2 2 σ 2 e .
Theorem 6.1 (Non-linear case, simpliﬁed) .
Suppose we observe E environments E = { e 1 , e 2 , . .
, e E } , where σ 2 e = 1 ∀ e .
Then, for any (cid:15) > 1 , there exists a featurizer Φ (cid:15) which, combined with the ERM-optimal classiﬁer ˆ β
=
[ β c , β e ; ERM , β 0 ] T , satisﬁes the following properties, where we deﬁne p (cid:15),d e := exp {− d e min(( (cid:15) − 1) , ( (cid:15) − 1) 2 ) / 8 } :
The regularization term of Φ (cid:15) , ˆ β as in Equation
is
bounded as
[Equation]
for some constant c (cid:15) that depends only on (cid:15) .
Φ (cid:15) , ˆ β exactly matches the optimal invariant classiﬁer on at least a 1 − p (cid:15),d e 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 sufﬁciently far from the training means:
[Equation]
for some δ > 0 , and deﬁne q := 2 E √ πδ exp {− δ 2 } .
Then the following holds:
Φ (cid:15) , ˆ β is equivalent to the ERM-optimal classiﬁer
on at least a 1 − q fraction of the test distribution.
Under Assumption 1, suppose it holds that µ E +
= − (cid:80)
e ∈E α e µ e for some set of coefﬁcients { α e } e ∈E .
Then so long as
[Equation]
the 0-1 risk of Φ (cid:15) , ˆ β on the new environment is greater than .
975 − q .
Before giving a proof sketch, we give a brief intuition for each of the claims made in this theorem: 1.
The ﬁrst claim says that the classiﬁer we construct will have a gradient squared norm scaling as p 2 (cid:15),d e which is exponentially small in d e .
Thus, in high dimensions, it will appear as a perfectly reasonable solution to the objective (5).
The second claim says that this classiﬁer is identical to the optimal one 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 d e , the “fake” classiﬁer will often be a preferred solution.
In the ﬁnite-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 classiﬁer 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 ( √ d e ) 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 Ω( √ d e
log E ) sufﬁces.
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.
The ﬁnal statement demonstrates a particularly egregious failure case of this classiﬁer: just like ERM, if the correlation between the non-invariant features and the label reverses at test-time, our classiﬁer will have signiﬁcantly worse than chance performance.
With this intuition, we now present a sketch of the proof—the full proof can be found in Appendix D.
Proof Sketch.
The key is a construction which is almost identical to the optimal invariant classiﬁer on the training data, yet behaves like the ERM solution at test time.
We partition the environmental feature space into two sets, B , B c ⊂ R d e .
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, deﬁne Φ( x ) =
[ z c ] , so our construction is equal to the optimal invariant classiﬁer.
Now consider B c = R d e \ B .
We use standard concentration results to upper bound the measure of B c under the training distribution by p (cid:15),d e .
Next, we show how choosing Φ( x ) =
f − 1 ( x ) =
[ z c , z e ] T on this set results in the sub-optimality bound, which is of order p 2 (cid:15),d e .
It is also clear that our constructed classiﬁer is equivalent to the ERM-optimal solution on B c .
Thus, our classiﬁer will have often have lower empirical risk on B c , 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 classiﬁer 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 classiﬁer under such a distribution shift by analyzing the distribution of the response variable.
The proof is completed by observing that our classiﬁer’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 ﬁnd that our result trivially extends to the alternative objectives; we again refer to Appendix E.
C ONCLUSION
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.
This paper demonstrates that Invariant Risk Minimization and subsequent works have signiﬁcant 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?
It would be interesting to investigate ours or similar models further; some possible directions include (a) characterizing in which settings speciﬁcally an “invariance-like” constraint can lead to improved performance over ERM and (b) formulating an objective that encourages invariance in the non-linear setting in a way that can be formally demonstrated and quantiﬁed.
We hope that this work will inspire further theoretical study on the effectiveness of IRM and similar objectives.
A A DDITIONAL N OTATION FOR THE A PPENDIX
To avoid overloading, we use φ, F for the standard Gaussian PDF and CDF respectively.
We write
S c to denote the set complement of a set S .
We write | · | to denote entrywise absolute value.
B P ROOF OF P ROPOSITION 4.
Recall the IRM objective:
[Equation]
subject
[Equation]
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 β coefﬁcients is well known:
[Equation]
Suppose we recover the true invariant features Φ θ ( x )
= (cid:20)
z c (cid:21) and coefﬁcients ˆ β =
(cid:20) β 0 (cid:21) (in other words, we allow for the introduction of new features).
Then the IRM constraint becomes:
[Equation]
Since ˆ β is constant across environments, this constraint is clearly satisﬁed 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 Φ θ :
[Equation]
Then the derivative of the loss with respect to these parameters is
[Equation]
So, the optimal invariant classiﬁer is a stationary point with respect to the feature map parameters as well.
C R ESULTS FROM S ECTIONC. P ROOF OF T HEOREM 5.
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 coefﬁcients α e : µ e = (cid:88)
i α ei µ i .
Then the environments are considered non-degenerate if the following inequality holds for any such set of coefﬁcients:
[Equation]
and furthermore that the following ratio is different for at least two different environments a, b :
[Equation]
The ﬁrst inequality says that none of the environmental means are are an afﬁne combination of the others; in other words, they lie in general linear position , which is the same requirement as Arjovsky et al.
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 ≤ d e 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 = { e 1 , e 2 , . . .
, e E } , each with environmental mean of dimension d e ≥ E , such that all environmental means are linearly independent.
Then there is a unique unit-norm vector p such that
[Equation]
where ˜ µ is the largest scalar which admits such a solution.
Proof.
Let v 1 , v 2 , . .
. , v E be a set of basis vectors for span { µ 1 , µ 2 , . .
. , µ E } .
Each mean can then be expressed as a combination of these basis vectors: u i = (cid:80)
Ej
=
α ij v j .
Since the means are linearly independent, we can combine these coefﬁcients into a single invertible matrix [Equation] 
[Equation]
We can then combine the constraints (11) as
[Equation]
where p α denotes our solution expressed in terms of the basis vectors { v i } E i =1 .
This then has the solution p [Equation]
This deﬁnes the entire space of solutions, which consists of p α plus any element of the remaining ( d e − 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 ˜ µ =
[Equation]
Adding any element of the orthogonal subspace would only increase the norm of p , decreasing ˜ µ .
Thus, the unique maximizing solution is
[Equation]
with
[Equation]
Finally, p α has to be rotated back into the original space by deﬁning p = (cid:80) Ei =1 p αi v i .
Lemma C.2.
Assume f is linear.
Suppose we observe E ≤ d e environments whose means are linearly independent.
Then there exists a linear Φ with rank (Φ) =
d c + d e + 1 − E whose output depends on the environmental features, yet the optimal classiﬁer on top of Φ is invariant.
Proof.
We will begin with the case when E = d e and then show how to modify this construction for when E < d e .
Consider deﬁning
[Equation]
with
[Equation]
Here, p ∈ R d c is deﬁned as the unit-norm vector solution to
[Equation]
such that ˜ µ is maximized—such a vector is guaranteed to exist by Lemma C.1.
Thus we get Φ( x ) =
(cid:20)
z c p T z e (cid:21) , which is of rank d c + 1 as desired.
Deﬁne ˜
z e = p T z e , which means that ˜ z e | y ∼
N ( y · σ 2 e ˜ µ, σ 2 e ) .
For each environment we have
[Equation]
The log-odds of y is linear in these features, so the optimal classiﬁer is
[Equation]
which is the same for all environments.
Now we show how to modify this construction for when E < d e .
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 (cid:48) , and consider redeﬁning M as
[Equation]
The distribution of ˜ z
e in each of the remaining dimensions is normal with mean 0, which means a corresponding coefﬁcient of 0 is optimal for all environments.
So the classiﬁer ˆ β 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 (Φ) =
d c + 1 + ( d e − E ) , as desired.
Note that for E = 1 any Φ will be vacuously feasible.
Lemma C.3.
Suppose we observe E environments E = { e 1 , e 2 , . . .
, e E } whose parameters satisfy the non-degeneracy conditions (9, 10).
Let Φ( x ) =
Az c
+ Bz e 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 > d e , then ˆ β = 0 or B = 0 .
Proof.
Write Φ =
[ A | B ] where A ∈ R d × d c , B ∈ R d × d e and deﬁne
[Equation]
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 coefﬁcient for each environment is 2¯Σ − 1 e ¯ µ 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 (cid:54) = 0 ∀ e ∈ E .
With this fact, we have that ∀ e ∈ E ,
[Equation]
Deﬁne the vector v =
Aµ c − σ 2 c AA T ˆ β .
We will show that for any ˆ β, A , with probability 1 only
B = 0 can satisfy Equation 12 for every environment.
If E > d e , 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 ¯ µ = (cid:80) d e i
=
α i µ i .
However, note that by assumption the means lie in general linear position, and therefore we actually have at least d e sets of coefﬁcients α for which this holds.
Rearranging Equation 12, we get
[Equation]
and rearranging once more yields
[Equation]
By assumption, (1 − (cid:80) α i ) is non-zero.
We can therefore rewrite this as [Equation]
, where ˆ α = ¯ σ 2 − (cid:80) α i σ 2 i 1 − (cid:80) α
i is a scalar.
As the vectors BB T ˆ β
and v are both independent of the environment, this can only hold true if ˆ α is ﬁxed for all environments or if both BB T ˆ β, 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
[Equation]
As the span of the observed µ e is all of R d e , 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:
Suppose E > d e .
Consider any linear featurizer Φ which is feasible under the IRM objective (4), with invariant optimal regression vector ˆ β (cid:54) = 0 , and write Φ( f ( z c , z e )) = Az c + Bz e . Then under mild non-degeneracy conditions, it holds that B = 0 .
If E ≤ d e and the environmental means µ e are linearly independent, then there exists a linear Φ —where Φ( f ( z c , z e ))
=
Az c + Bz e with rank ( B ) = d e + 1 − E —which is feasible under the IRM objective.
Further, both the logistic and 0-1 risks of this Φ and its corresponding ˆ β are strictly lower than those of the optimal invariant classiﬁer.
Proof.
Since Φ , f are linear, we can write Φ( x ) =
Az c + Bz e 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 p e ( z c , y ) is invariant, this classiﬁer has identical risk across all environments.
The existence of such a classiﬁer is proven by Lemma C.2.
It remains to show that the risk of this discriminator is lower than that of the optimal invariant classiﬁer.
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
[Equation]
results in strictly higher risk.
The proof is completed by noting that this deﬁnition is precisely the optimal invariant classiﬁer.
C. E XPERIMENTS FOR T HEOREM 5.
To corroborate our theoretical ﬁndings, we run an experiment on data drawn from our model to see at what point IRM is able to recover a generalizing classiﬁer.
We generated data precisely according to our model in the linear setting, with d c = 3 , d e = 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 classiﬁer which ignores meaningful covariates z e , to the detriment of its training performance but to the beneﬁt of OOD generalization.
We chose equal class marginals ( η = 0 . 5 ).
Figure C.1 shows the result of ﬁve 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 speciﬁc 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 ﬁve runs.
At test time, we evaluated the trained classiﬁers on additional, unseen environments whose parameters were drawn from the same prior.
To simulate distribution shift, we evaluated the classiﬁers on the same data but with the environmental means negated.
Thus the correlations between the environmental features z e and the label y were reversed.
Observe that the results closely track the expected outcome according to Theorem 5.1: up until
E = d e , IRM essentially matches ERM in performance both in-distribution and under distribution shift.
As soon as we cross that threshold of observed environments, the classiﬁer learned via IRM begins to perform drastically better under distribution shift, behaving more like the optimal invariant classiﬁer.
We did however observe that occasionally the invariant solution would be found after only E = d e = 6 environments; we conjecture that this is because at this point the feasible-yet-notinvariant classiﬁer 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 difﬁcult for the optimizer to ﬁnd. Observe that up until E = d e , IRM consistently returns a classiﬁer with performance similar to ERM: good generalization without distribution shift, but catastrophic failure when the correlation is reversed. In contrast, once E > d e , IRM is able to recover a Φ , ˆ β with performance similar to that of the causal classiﬁer.
C. P ROOF OF T HEOREM 5.
Theorem 5.3.
Suppose we observe E ≤ d e 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 classiﬁer on every environment e such that[-σ e ˜ µ > σ − 1 c (cid:107) µ c (cid:107) 2 and 
-]
[Equation]
Proof.
We consider the non-invariant classiﬁer constructed as described in Lemma C.2, but dropping the invariant features and coefﬁcients.
By Lemma F.2, the optimal coefﬁcients for the invariant and non-invariant classiﬁers are
[Equation]
and
[Equation]
, respectively.
Therefore, the 0-1 risk of the optimal invariant classiﬁer is precisely
[Equation]
where F is the Gaussian CDF.
By the same reasoning, the 0-1 risk of the non-invariant classiﬁer is ηF
[Equation]
Deﬁne α = σ − 1 c (cid:107) µ c (cid:107) 2 and γ = σ e ˜ µ .
By monotonicity of the Gaussian CDF, the former risk is higher than the latter if
[Equation]
[Equation]
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 ﬁrst condition).
Then Equation 14 is immediate.
Finally, for Equation 13, observe that
[Equation]
which is the second condition.
C. S IMULATIONS OF M AGNITUDE OF E NVIRONMENTAL F EATURES
As discussed in Section 5, analytically quantifying the solution ˜ µ to the equation in Lemma C.1 is difﬁcult; instead, we present simulations to give a sense of how often these conditions would hold in practice.
For each choice of environmental dimension d e , we generated a “base” correlation b ∼ N (0 , I d e ) as the mean of the prior over environmental means µ e .
Each of these µ e was then drawn from
N ( b, 4 I d e ) —thus, while they all came from the same prior, the noise in the draw of each µ e was signiﬁcantly larger than the bias induced by the prior.
We then solved for the precise value σ e ˜ µ , with the same variance σ 2 e for all environments, chosen as a hyperparameter.
The shaded area represents a 95% conﬁdence interval over 20 runs.
The dotted lines are √ d c .
If we imagine the invariant parameters are drawn from a standard Gaussian prior, then this is precisely E [ σ − 1 c (cid:107) µ c (cid:107) 2 ] .
Thus, the point where σ e ˜ µ crosses these dotted lines is approximately how many environments would need to be observed before the non-invariant classiﬁer has higher risk than the optimal invariant classiﬁer.
We note that this value is quite large, on the order of d e − d c .
[Figure]
Projected Environmental Mean Separation
[Figure]
D T HEOREM 6.1 AND D ISCUSSION D. P ROOF OF T HEOREM 6.
We again begin with helper lemmas.
Our classiﬁer Φ is constructed to recover the environmental features only if they fall within a set B c .
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 classiﬁer.
Lemma D.1.
Suppose we observe environments E = { e 1 , e 2 , . . . } .
Given a set B ⊆ R d e , consider the classiﬁer deﬁned by Equation 19.
Then for any environment e , the penalty term of this classiﬁer in Equation 5 is bounded as
[Equation]
Proof.
We write out the precise form of the gradient for an environment e : ∇
[Equation]
Observe that since z c ⊥⊥ z e | y , the optimal invariant coefﬁcients are unchanged, and therefore the gradient in the invariant dimensions is 0.
We can split the gradient in the environmental dimensions into two integrals:
[Equation]
Since the features are 0 within B , the ﬁrst term reduces to 0.
For the second term, note that ∀ x, y ∈ R , | σ ( x ) − σ ( y ) | ≤ 1 , and therefore
[Equation]
We can marginalize out z c , and noting that we want to bound the squared norm,
[Equation]
This next lemma says that if the environmental mean of the test distribution is sufﬁciently separated from each of the training means, with high probability a sample from this distribution will fall outside of B r , and therefore Φ (cid:15) , ˆ β will be equivalent to the ERM solution.
Lemma D.2.
For a set of E environments E = { e 1 , e 2 , . . .
, e E } and any (cid:15) > 1 , construct B r as in Equation 18 and deﬁne Φ (cid:15) using B r as in Equation 19.
Suppose we now test on a new environment with parameters ( µ E +1 , σ 2 E +1 ) , and assume Equation 15 holds with parameter δ .
Deﬁne k = min
e ∈E σ 2 e [Equation] 2 E +1 .
Then with probability ≥ 1 − 2 E √ kπδ exp {− kδ 2 } over the draw of an observation from this new environment, we have
[Equation]
Proof.
By Equation 15 our new environmental mean is sufﬁciently far away from all the labelconditional means of the training environments.
In particular, for any environment e ∈ E and any label y ∈ {± 1 } , the (cid:96) 2 distance from that mean to µ E +1 is at least ( √ (cid:15) + δ ) σ e √ d e .
Recall that B r is the union of balls ± B e , where B e is the ball of (cid:96) 2 radius (cid:112) (cid:15)σ 2 e d e 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 B e .
This halfspace fully contains B e , and therefore the measure of B e 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 ( √ (cid:15) + δ ) σ e √ d e away from the mean µ E +1 , so accounting for its radius, the distance from the mean to the halfspace is δσ e √ d e .
The variance of the rotated distribution one dimension is σ 2 E +1 d e , so the measure of this halfspace is upper bounded by
[Equation]
using results from Kschischang (2017).
There are 2 E such balls comprising B r , 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 = { e 1 , e 2 , . . .
, e E } .
Then, for any (cid:15) > 1 , there exists a featurizer Φ (cid:15) which, combined with the ERM-optimal classiﬁer ˆ β
=
[ β c , β e ; ERM , β 0 ] T , satisﬁes the following properties, where we deﬁne p (cid:15),d e := exp {− d e min(( (cid:15) − 1) , ( (cid:15) − 1) 2 ) / 8 } :
Deﬁne σ 2max = max e σ 2 e .
Then the regularization term of Φ (cid:15) , ˆ β is bounded as E (cid:88)
e ∈E (cid:107)∇ ˆ β R e (Φ (cid:15)
[Equation]
Φ (cid:15) , ˆ β exactly matches the optimal invariant classiﬁer on at least a 1 − p (cid:15),d e 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 , σ 2 E +1 ) , suppose the environmental mean µ E +1 is sufﬁciently far from the training means:
[Equation]
for some δ > 0 .
Deﬁne the constants:
[Equation]
Then the following holds:
Φ (cid:15) , ˆ β is equivalent to the ERM-optimal classiﬁer on at least a 1 − q fraction of the test distribution.
Under Assumption 1, suppose it holds that
[Equation]
for some set of coefﬁcients { α e } e ∈E .
Then for any c ∈ R , so long as
[Equation]
the 0-1 risk of Φ (cid:15) , ˆ β is lower bounded by F (2 c ) − q .
Proof.
Deﬁne r = (cid:112) (cid:15)σ 2 e d e and construct B r ⊂ R d e as
[Equation]
where B r ( α ) is the ball of (cid:96) 2 -norm radius r centered at α .
Further construct Φ (cid:15) using B r as follows:
[Equation]
and
[Equation]
Without loss of generality, ﬁx an environment e .
By Lemma D.1, the squared gradient norm is upper bounded by
[Equation]
Deﬁne B e := B r ( µ e ) , and observe that B cr ⊆ B ce .
Since | z e | is non-negative,
[Equation]
(this inequality is element-wise).
Plugging this into Equation 20 yields
[Equation]
Deﬁne p
=
P ( z e ∈ B cr ) ≤ P ( z e ∈ B ce ) .
By Lemma F
[Equation]
[Equation]
Combining Lemmas F.4 and F.5 gives
[Equation]
Putting these two bounds together [Equation] we have
[Equation]
and averaging this value across environments gives the result.
Φ (cid:15) , ˆ β is equal to the optimal invariant classiﬁer on B r and the ERM solution on B cr .
The claim then follows from Lemma F.3.
This follows directly from Lemma D.2.
With Equation 16, we have that
[Equation]
where we have applied Assumption 1 in the ﬁrst inequality and Equation 17 in the second.
Consider the full set of features Φ (cid:15) ( x ) =
f − 1 ( x ) , and without loss of generality assume y = 1 .
The label-conditional distribution of the resulting logit is
[Equation]
Therefore, the 0-1 risk is equal to the probability that this logit is negative.
This is precisely
[Equation]
Observe that by the previous part, Φ (cid:15) (cid:54) =
f − 1 on at most a q fraction of observations, so the risk of our classiﬁer Φ (cid:15) , ˆ β can differ from that of f − 1 , ˆ β by at most q .
Therefore our classiﬁer’s risk is lower bounded by F (2 c ) − q .
In particular, choosing c = 1 recovers the statement in the main body.
D. D ISCUSSION OF C ONDITIONS AND A SSUMPTION
To see just how often we can expect the conditions for Theorem D.3 to hold, we can do a Fermi 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 (cid:107) m (cid:107)
(cid:29) Tr(Σ) ).
For simplicity we treat the variances σ 2 c , σ 2 e as constants.
Then the expected separation between any two means from this distribution is
[Equation]
In high dimensions this value will tightly concentrate around the mean, which is in O ( √ d e ) .
On the other hand, even a slight deviation from this separation, to Ω( √ d e log E ) , means δ ∈ Ω( √ log E ) , 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 d e , 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 (cid:80) i α i ∈ Ω(1) , the left-hand side of Equation 17 is approximately d e .
On the other hand, treating γ as a constant, the right-hand side is close to d c + √ d c + d e ∈ O
( d c + √ d e ) .
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 coefﬁcient will be approximately 2 m/σ 2 e with high probability, meaning ˆ β T µ ≈ 2 (cid:107) m (cid:107) 22 /σ 2 e for every environment.
Thus, this vector will be γ -close to optimal with γ ≈ 1 for every environment with high probability.
E E XTENSIONS TO A LTERNATIVE O BJECTIVES E. E XTENSIONS FOR THE L INEAR C ASE
Observe that the constraint of Equation 4 is strictly stronger than that of Bellot & van der Schaar (2020); when the former is satisﬁed, 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.
[Equation]
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 ≤ d e environments with linearly independent means and identical variance σ 2 e .
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 classiﬁer 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,
[Equation]
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 classiﬁer.
As mentioned in Section 5, this additional requirement of constant variance is due to the assumptions underlying the design of the objective—REx expects invariance of the conditional distribution p ( 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 ≤ d e 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 classiﬁer under 0-1 loss on every environment e such that ˜ µ > σ − 1 c (cid:107) µ c (cid:107) 2 + | β 0 | 2 σ − 1 c (cid:107) µ c (cid:107) 2 .
Proof.
We follow the proof of Theorem 5.3, except when solving for p as in Lemma C.
we instead ﬁnd the unit-norm vector such that
[Equation]
Observe that by setting Φ( x ) =
[ p T z e ] and ˆ β = [1] , the 0-1 risk in a given environment is
[Equation]
, which is independent of the environment.
Further, by carrying through the same proof as in Theorem 5.3, we get that this non-invariant classiﬁer has lower 0-1 risk so long as
[Equation]
where α = σ − 1 c (cid:107) µ c (cid:107)
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 modiﬁed equation below. [Figure] Environments observed         ˜ µ
Projected Environmental Mean Separation d c = 40 ,d e = 200 ,σ 2 e =
d c = 100 ,d e = 200 ,σ 2 e =
d c = 50 ,d e = 300 ,σ 2 e =d c = 60 ,d e = 300 ,σ 2 e =√
d e − E E. E XTENSIONS FOR THE N ON -L INEAR C ASE
The failure of these objectives in the non-linear regime is even more straightforward, as we can keep unchanged the constructed classiﬁer from Theorem 6.1.
Observe that parts 2-4 of the theorem do not involve the objective itself, and therefore do not require modiﬁcation.
To see that part 1 still holds, note that since the constructed classiﬁer matches the optimal invariant classiﬁer on 1 − p of the observations, its risk across environments can only vary on the remaining p fraction: as the 0-1 risk on this fraction is bounded between 0 and p , it is immediate that the variance of the environmental risks is upper bounded by p 2 4 , which is in O ( p 2 ) as before.
Applying this argument to the other objectives yields similar results.
F T ECHNICAL L EMMAS
Lemma F.1.
Consider solving the standard logistic regression problem
[Equation]
Assume that none of the latent dimensions are degenerate— ∀ S ⊆
[ k ] , P ( β TS z S (cid:54) = 0) > 0 , and no feature can be written as a linear combination of the other features.
Then for any distribution p ( z ) , any classiﬁer f ( z ) = σ ( β TS z S ) that uses a strict subset of the features S (cid:40)
[ k ] has strictly higher risk with logistic loss than the Bayes classiﬁer f ∗ [Equation]
= σ ( β T z ) .
This additionally holds for 0-1 loss if β T − S z − S has greater magnitude and opposite sign of β TS z S with non-zero probability.
Proof.
The Bayes classiﬁer suffers the minimal expected loss for each observation z.
Therefore, another classiﬁer has positive excess risk if and only if it disagrees with the Bayes classiﬁer on a set of non-zero measure.
Consider the set of values z − S such that β T − S z − S (cid:54) = 0 .
Then on this set we have
[Equation]
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
[Equation]
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
= Az c
+ Bz e , any optimal corresponding coefﬁcient for an environment e is of the form
[Equation]
where G + is a generalized inverse of G .
Proof.
We begin by evaluating a closed form for p e ( y | ˜ z ) .
We have:
[Equation]
Now we need a closed form expression for p ( Az c +
Bz e = ˜
z | y ) .
Noting that z c ⊥⊥ z e | y , this is a convolution of the two independent Gaussian densities, which is the density of their sum.
In other words,
[Equation]
Thus,
[Equation]
Canceling common terms [Equation] we get
[Equation]
Therefore, given a feature vector ˜ z [Equation] the optimal coefﬁcient vector is 2¯Σ + ¯ µ .
Lemma F.3.
For any environment e with parameters µ e , σ 2 e and any (cid:15) > 1 , deﬁne
[Equation]
where B r ( α ) is the ball of (cid:96) 2 -norm radius r centered at α .
Then for an observation drawn from p e , we have
[Equation]
Proof.
Without loss of generality, suppose y = 1 .
We have
[Equation]
Each term in the squared norm of z e is a random variable with distribution χ 21 , which means their sum has mean d e and is sub-exponential with parameters (2 √ d e , 4) .
By standard sub-exponential concentration bounds we have
[Equation]
which immediately implies the claim.
Lemma F.4.
Let z ∼ N ( µ, σ 2 I d ) be a multivariate isotropic Gaussian in d dimensions, and for some r > 0 deﬁne B as the ball of (cid:96) 2 radius r centered at µ .
Then we have
[Equation]
[Equation]
where φ, F are the standard Gaussian PDF and CDF.
Proof.
Observe that
[Equation]
Now, consider the expectation for an individual dimension, and note that |
z i | > r √ d ∀ i = ⇒ (cid:107)
z (cid:107) 2 > r .
So because the dimensions are independent, conditioning on this event can only increase the expectation:
[Equation]
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:
[Equation]
where α = F ( r/ √ d ) .
Combining the above results, squaring with ( a + b ) 2 ≤ 2( a 2 + b 2 ) , and summing over dimensions, we get
[Equation]
as desired
[Equation]
Lemma F.5.
For σ, (cid:15) > 0 , deﬁne r = √ (cid:15)σ .
Then
[Equation]
Proof.
We have
[Equation]
and
[Equation]
(see Kschischang (2017)).
Dividing them gives
[Equation]
Squaring and using ( a + b ) 2 ≤ 2( a 2 + b 2 ) yields the result.