Year: 2021

Abstract:
The Invariant Risk Minimization (IRM) principle was first proposed by Arjovsky et al. (2019) to address the domain generalization problem by leveraging data heterogeneity from differing experimental conditions. Specifically, IRM seeks to find a data representation under which an optimal classifier remains invariant across all domains. Despite the conceptual appeal of IRM, the effectiveness of the originally proposed invariance penalty has recently been brought into question through stylized experiments and counterexamples. In this work, we investigate the relationship between the data representation, invariance penalty, and risk. In doing so, we propose a novel invariance penalty, and utilize it to design an adaptive rule for tuning the coefficient of the penalty proposed by Arjovsky et al. (2019). Moreover, we provide practical insights on how to avoid the potential failure of IRM considered in the nascent counterexamples. Finally, we conduct numerical experiments on both synthetic and real-world data sets with the objective of building invariant predictors. In our non-synthetic experiments, we sought to build a predictor of human health status using a collection of data sets from various studies which investigate the relationship between human gut microbiome and a particular disease. We substantiate the effectiveness of our proposed approach on these data sets and thus further facilitate the adoption of the IRM principle in other real-world applications.

heading: 1 INTRODUCTION
text: Under the learning paradigm of Empirical Risk Minimization (ERM) (Vapnik, 1992), data is assumed to consist of independent and identically distributed (iid) samples from an underlying generating distribution. As the data generating distribution is often unknown in practice, ERM seeks predictors with minimal average training error (i.e., empirical risk) over the training set. However, shuffling and treating data as iid risks possibly losing important information about the underlying conditions of the data generating process. Despite becoming a ubiquitous paradigm in machine learning, a growing body of literature (Arjovsky et al., 2019; Teney et al., 2020) has revealed that ERM and the the common practice of shuffling data inadvertently results in capturing all correlations found in the training data, whether spurious or causal, and produces models that fail to generalize to test data. The potential variation of experimental conditions that can exist at training time and during deployment in real-world applications, manifests in discrepancies between training and testing distributions. This, in turn, highlights the need for machine learning algorithms to generalize out-of-distribution (OoD).
Shuffling and treating data as iid risks possibly losing important information about the underlying conditions of the data generating process. As will be demonstrated in this work, partitioning training data into environments, e.g., based on the conditions under which data is generated, can exploit these differences to enhance generalization. One promising approach based on this observation is that of Arjovsky et al. (2019), in which the principle of Invariant Risk Minimization (IRM) is introduced. The objective of IRM is to find a predictor that is invariant across all training environments (see Definition 1 and Equation 1). Because of the conceptually appealing nature of IRM and its potential to address the OoD-generalization problem, there is a stream of literature scrutinizing various facets of the original framework, e.g., extensions to other settings including online learning (Javed et al., 2020) and treatment effect estimation (Shi et al., 2020), fairness (Adragna et al., 2020), introducing game-theoretic interpretations (Ahuja et al., 2020), and raising concerns on the drawbacks and limitations of current IRM implementations (Rosenfeld et al., 2021; Kamath et al., 2021). For an indepth overview of the broader generalization literature, we refer the interested reader to (Arjovsky, 2020) and the references therein, and for an empirical evaluation of the performance of a number of the state-of-the-art methods on various test cases, we refer the reader to (Gulrajani & Lopez-Paz, 2020).
In this paper, we introduce two practical implementations of the IRM principal to increase its appeal and applicability in real-world settings. First, we propose an invariance penalty that is directly related to risk. More precisely, we show that the risk in each environment under an arbitrary classifier equals to the risk under the optimal classifier for that environment plus the newly proposed invariance penalty between the said classifier and the optimal one. Second, we build on these initial findings to provide practitioners who currently use the original IRMv1 implementation with an adaptive rule by which to choose the penalty coefficient appropriately. In doing so, we characterize the difference between our proposed invariance penalty and the one proposed by Arjovsky et al. (2019) in terms of the eigenvalues of the Gram matrix of the data representation. This eigenstructure plays a significant role in the failure of invariance penalties including the one proposed by Arjovsky et al. (2019).
In addition to providing practitioners with two valid approaches for implementing IRM, this work serves to illustrate the importance of the eigenstructure of the Gram matrix of the data representation for IRM. In particular, we revisit the counterexample of Rosenfeld et al. (2021) where the invariance penalty of Arjovsky et al. (2019) can be made arbitrarily small for a non-invariant representation. We show that the Gram matrix is ill-conditioned in such cases. We then provide a practical solution to alleviate such behavior, in particular for the case where the representation is parameterized by a neural network. Moreover, we show that the proposed framework finds an invariant predictor for the setting in which the data is generated according to a linear Structural Equation Model (SEM) when provided a sufficient number of training environments under a mild non-degeneracy condition, which is similar in nature to the ones considered in (Arjovsky et al., 2019; Rosenfeld et al., 2021). Finally, we evaluate our method on various test cases including InvarianceUnitTests (Aubin et al., 2020) and HealthyGutTests. InvarianceUnitTests is a test bed with three synthetically generated data sets capturing different structures of spurious correlations. HealthyGutTests is a curated collection of biomedical data sets based on a prior meta-analysis of various microbiome studies (Gupta et al., 2020) in which the relationships between human gut microbiome composition and various disease phenotypes can be investigated.
The remainder of the paper is organized as follows. In Section 2, we formally define the notion of invariant prediction, the invariant risk minimization principle, and its relaxation proposed by Arjovsky et al. (2019). In Sections 3 and 4, we introduce our more practical implementation and the rationale for its design. In Section 5, we evaluate the efficacy of our proposed model and compare it with other variations of IRM over a series of experiments. We conclude the paper in Section 6. All mathematical proofs are presented in the Appendix.

heading: 2 BACKGROUND: INVARIANT PREDICTION
text: In this paper, we consider data (Xe, Y e) collected from multiple training environments e ∈ Etr where the distribution of (Xi, Y i) and (Xj , Y j) may be different for i ̸= j with i, j ∈ Etr. We denote by Re the risk under environment e. That is, for predictor f : X → Y , and loss function ℓ : Y × Y → R, the risk under environment e is defined as Re(f) = EXe,Y e [ℓ(f(Xe), Y e)].

heading: 2.1 INVARIANT RISK MINIMIZATION
text: Arjovsky et al. (2019) define the notion of invariant predictors under a multi-environment setting as follows.
Definition 1 (Invariant Predictor). A data representation φ : X → H is said to elicit an invariant predictor w ◦ φ across environments E if there exists a classifier w : H → Y , which is optimal for all environments, i.e., w ∈ argminw̃:H→Y Re (w̃ ◦ φ) for all e ∈ E .
To find such invariant predictors, Arjovsky et al. (2019) introduce the notion of the Invariant Risk Minimization (IRM) principle:
min φ:X→H w:H→Y
∑ e∈Etr Re (w ◦ φ)
subject to w ∈ argmin w̃:H→Y
Re (w̃ ◦ φ) , ∀e ∈ Etr. (1)
As this bi-leveled optimization problem is rather intractable, Arjovsky et al. (2019) propose a practical implementation of IRM by relaxing the invariance constraint (which itself requires solving an optimization problem) to an invariance penalty. We review its derivation in what follows.

heading: 2.2 IRMV1: A RELAXATION OF IRM
text: In order to provide an implementation of IRM, Arjovsky et al. (2019) restrict the classifier w to linear functions, i.e.,
min φ:X→H w∈Rdφ
∑ e∈Etr Re ( w⊤φ ) subject to w ∈ argmin
w̃∈Rdφ Re ( w̃⊤φ ) , ∀e ∈ Etr.
(2)
To motivate their proposed penalty, Arjovsky et al. (2019) first consider the squared loss, i.e., ℓ(f(x), y) = ∥f(x)− y∥2 where ∥ · ∥ denotes the Euclidean norm. Define matrix Ie(φ) as
Ie(φ) := EXe [ φ(Xe)φ(Xe)⊤ ] . (3)
Assuming that Ie(φ) is full rank for a fixed φ, its respective optimal classifier is unique, i.e., argminw̃∈RdφRe ( w̃⊤φ ) = w⋆e(φ) where w⋆e(φ) := Ie(φ)−1EXe,Y e [φ(Xe)Y e] . (4) Hence, in this setting, the invariant constraint of IRM in equation 2 can be simplified to w = w⋆e(φ) for all e ∈ Etr. A natural relaxation of the constraint w−w⋆e(φ) = 0 to a penalty is ∥w−w⋆e(φ)∥2.1 However, Arjovsky et al. (2019) show that this penalty may not capture invariance by constructing an example for which ∥w − w⋆e(φ)∥2 is not well-behaved (see Section 4.3 for more details). They argue that undoing the matrix inversion Ie(φ)−1, which appears in the computation of w⋆e(φ) could improve the behavior of the invariance penalty. That is, considering ∥Ie(φ)(w − w⋆e(φ))∥2 as the invariance penalty. Moreover, for the squared loss, one can show that
∥Ie(φ)(w − w⋆e(φ))∥ 2 = (1/4) ∥∥∇wRe(w⊤φ)∥∥2 . (5) Hence, Arjovsky et al. (2019) propose the following invariance penalty
ρIRMv1e (φ,w) := ∥∥∇wRe(w⊤φ)∥∥2 . (6)
Using the penalty equation 6, the relaxation of IRM is given by
min φ, w ∑ e∈Etr Re(w ⊤φ) + λρIRMv1e (φ,w), (7)
where λ ≥ 0 is the penalty coefficient. Notice that for a given w and φ, the predictor w ◦ φ can be expressed using different classifiers and data representations, i.e., w◦φ = w̃◦φ̃where w̃ = w◦ψ−1 and φ̃ = ψ ◦ φ for some invertible mapping ψ : H → H. Hence, in principle, it is possible to fix w without loss of generality. By relying on this observation, Arjovsky et al. (2019) fix the classifier as a scalar w = 1, and, thus, search for an invariant data representation of the form φ ∈ R1×dx . Their relaxation of IRM, which they refer to by IRMv1 is given by
min φ ∑ e∈Etr Re (φ) + λρ IRMv1 e (φ, 1.0). (IRMv1)
Although equation 5 only holds for squared loss, Arjovsky et al. (2019, Theorem 4) justify the choice of ∥∇w|w=1.0Re(w⊤φ)∥2 as an invariance penalty for other loss functions, e.g., cross-entropy loss. More precisely, let Φ be the matrix that parameterizes the data representation. They show that for all convex differentiable loss functions, (w⊤Φ)⊤∇wR(w⊤Φ) = 0 if and only if w is optimal for all environments.
Although the penalty ρIRMv1e (φ, 1.0) seems to be an appropriate invariance penalty, which is easily implementable, its effectiveness to capture invariance has been brought into question. In particular, Rosenfeld et al. (2021) construct a non-invariant data representation such that the penalty ρIRMv1e (φ, 1.0) is arbitrarily small. In what follows, we propose an alternative invariance penalty that is directly comparable to risk. Hence, if the invariance penalty of a classifier is small, then so is the difference between its risk and the risk under the optimal classifier.
1In the presence of known exogenous (environment dependent) variables, one can utilize anchor regression Rothenhäusler et al. (2021) that is conceptually related to the invariance penalization. Instead of an invariance penalty, anchor regression relies on the projection onto the span of the said variables.

heading: 3 IRMV2: AN ALTERNATIVE PENALTY
text: We revisit the structure of the risk in order to propose an alternative penalty. In particular, in the following Lemma, we provide the sub-optimality gap of risk under an arbitrary classifier in comparison to the optimal classifier.
Lemma 1. Consider squared loss function. Let w ∈ Rdφ and w⋆e(φ) as defined in equation 4. Then,
Re ( w⊤φ ) = Re ( w⋆e(φ) ⊤φ ) + ∥∥∥Ie(φ)1/2 (w − w⋆e(φ))∥∥∥2 . (8)
Based on Lemma 1, we propose an invariance penalty that is directly comparable to risk.
ρIRMv2e (φ,w) := ∥∥∥Ie(φ)1/2 (w − w⋆e(φ))∥∥∥2 . (9)
The relaxation of IRM using the penalty equation 9 is then given by
min φ, w ∑ e∈Etr Re(w ⊤φ) + λρIRMv2e (φ,w). (10)
We further simplify the relaxation equation 10 by finding its optimal classifier for a fixed data representation defined as
w⋆(φ) := argminw ∑ e∈Etr Re ( w⊤φ ) + λρIRMv2e (φ,w). (11)
In the following Lemma, we leverage on the structure of the squared loss to find w⋆(φ).
Lemma 2. Consider squared loss function and fix φ. Let w⋆e(φ) and w⋆(φ) as defined in equation 4 and equation 11, respectively. Then,
w⋆(φ) = (∑ e∈Etr Ie(φ) )−1(∑ e∈Etr Ie(φ)w⋆e(φ) ) . (12)
Moreover, it holds that
argminw ∑ e∈Etr Re ( w⊤φ ) = w⋆(φ). (13)
Based on Lemmas 1 and 2, we propose the following relaxation of IRM, which we refer to by IRMv2.
min φ ∑ e∈Etr Re ( w⋆(φ)⊤φ ) + λρIRMv2e (φ,w ⋆(φ)). (IRMv2)
We provide the pseudo-code for IRMv2 in Algorithm 1 in Appendix A.
There are several comments in order. It is worth noting that Equation equation 13 in Lemma 1 reveals that for a fixed data representation, the optimal classifiers of ERM and IRMv2 are equal. However, this equality would not imply the equality of the ERM and IRMv2 predictors (i.e., the composition of the classifier and the data representation). The latter is due to the presence of the invariance penalty in the optimization of φ in IRMv2 in comparison to ERM.
There are a number of factors distinguishing IRMv2 from IRMv1. First, IRMv2 relies on the optimal classifier w⋆(φ) while w = 1.0 in IRMv1. Second, the loss function in IRMv2 is squared loss while IRMv1 allows for utilization of other loss functions. Although this additional flexibility of IRMv1 may seem appealing, the counterexample of Rosenfeld et al. (2021) shows the failure of the penalty of IRMv1 to capture invariance for logistic loss. More importantly, for squared loss, Ie(φ) is incorporated differently in the invariance penalty of IRMv1 and IRMv2. We formalize this latter observation in the Section 4.3.1.

heading: 3.1 IRMV1A: AN ADAPTIVE PENALTY COEFFICIENT
text: We first bound the invariance penalty of IRMv1 in terms of the penalty of IRMv2 and the eigenvalues of Ie(φ). Then, based on this comparison, we propose an adaptive approach in choosing the penalty coefficient for IRMv1, which we refer to as IRMv1-Adaptive (IRMv1A). Lemma 3. Let ρIRMv1e (φ,w) and ρIRMv2e (φ,w) be the invariance penalties of the IRMv1 and IRMv2 defined in Equations equation 6 and equation 9, respectively. Then,
λmin (Ie(φ)) ρIRMv2e (φ,w) ≤ ρIRMv1e (φ,w) ≤ λmax (Ie(φ)) ρIRMv2e (φ,w). (14)
The proof of Lemma 3 directly follows from the definition of the invariance penalties ρIRMv1e (φ,w) and ρIRMv2e (φ,w), and the fact that for a symmetric matrix A ∈ Rd×d and a vector u ∈ Rd, it holds that λmin(A)∥u∥2 ≤ u⊤Au ≤ λmax(A)∥u∥2. As IRMv1 is used in practice in conjunction with losses other than squared loss for classification, in order to justify our adaptive penalty coefficient, we bound the difference of risk for cross-entropy loss under two different classifiers in terms of the invariance penalty of IRMv2. Lemma 4. Consider binary classification with cross-entropy loss. Then, for w1, w2 ∈ Rdφ it holds that
|Re(w⊤1 φ)−Re(w⊤2 φ)| ≤ ∥∥∥Ie(φ)1/2(w1 − w2)∥∥∥ .
Using Lemmas 3 and 4, we suggest the following rule for the penalty coefficient of IRMv1.
λe := 1
λ0 + λmin(Ie(φ)) , (15)
where λ0 ≥ 0 is a user-specified parameter. The role of λ0 is to avoid numerical instability when λmin(Ie(φ)) is small. Note that this is an adaptive rule, as φ may change throughout training.

heading: 4 THEORETICAL RESULTS AND ANALYSIS
text: In this section, we investigate the effectiveness of IRM and its practical implementations in capturing invariance by focusing on a setting in which the data is generated according to a Structural Equation Model (SEM) (Pearl, 2009). SEM refers to a set of equations specifying the relationship between variables, which is a common structural assumption in causal inference. In particular, under an SEM, variables are related through a causal graph such that each variable is only a function of its parents and an exogenous random variable. Similar to Arjovsky et al. (2019) and Rosenfeld et al. (2021), we establish conditions under which IRM recovers an invariant predictor for data generated according to SEMs with linear functions. We then investigate the role of the eigenstructure of Ie(φ), in particular, in relation to counterexamples of Arjovsky et al. (2019) and Rosenfeld et al. (2021), and provide practical solution to alleviate the possibility of such failure modes.

heading: 4.1 SEM FOR CLASSIFICATION
text: For each environment e, data (Xe, Y e) is generated according to the following SEM:
Xe = S [ Zc Ze ] , Y e = { 1, with prob. η, −1, with prob. 1− η, , (16)
where η ∈ [0, 1], and S ∈ Rd×(de+dc) is a left invertible matrix, i.e., there exists S† such that S†S = I . In this model, Zc captures the causal variables that are invariant across environments, and Ze captures the spurious environment dependent variables.
The variables Zc and Ze are generated as follows
Zc = µcY +Wc, Wc ∼ N (0, σ2cI), (17) Ze = µeY +We, We ∼ N (0, σ2eI). (18)
Here, µc ∈ Rdc , µe ∈ Rde , and N (µ,Σ) denotes multi-variate Gaussian distribution with mean µ and covariance Σ. We further assume that Wc, We, and Y e are independent for all environments.

heading: 4.2 INVARIANT REPRESENTATION UNDER IRM
text: For the setting introduced in 4.1, the invariant data representation is linear. In particular, for any d ≥ dc, φ(Xe) = [ Idc 0 0 0 ] S†Xe = Zc is an invariant data representation. Given that X is also linear
in Y , the class of linear representations may be sufficiently reach to elicit an invariant predictor. Naturally, the possibility of finding an invariant predictor depends on the number and the diversity of training environments. We now introduce non-degeneracy conditions on training environments under which IRM is guaranteed to find an invariant predictor, provided sufficient number of training environments.
Let |Etr| > de. As span({µe}e∈Etr) ≤ de, for each e ∈ Etr there exists a set of coefficients αei for i ∈ Etr\e such that
µe = ∑
i∈Etr\e
αeiµi. (19)
We say that Etr is a non-degenerate set of environments if for all e ∈ Etr it holds that∑ i∈Etr\e αei ̸= 1, (20)
rank (Γe) = de, (21)
where Γe is defined as
Γe := 1 1− ∑
i∈Etr\e α e i σ2eI + µeµ⊤e − ∑ i∈Etr\e (σ2i I + µiµ ⊤ i )α e i  . The conditions of equation 20 and equation 21 specify that the span of covariance matrices of Ze’s is Rde . This is a natural requirement to eliminate the degrees of freedom on the dependency of the data representation on the environment dependent features. We note that the non-degeneracy conditions considered in (Rosenfeld et al., 2021) are similar to the ones introduced here, with the difference that instead of depending on covariance matrices of Ze as in equation 21, their condition relies only on the variances σ2e . This difference in the non-degeneracy requirements is due to the fact that they consider logistic loss and we consider squared loss. Theorem 1. Assume that |Etr| > de where (Xe, Y e) generated according to equation 16. Consider a linear data representation ΦX = AZc +BZe and a classifier w(Φ) on top of Φ that is invariant, i.e., w(Φ) = w⋆e(Φ) for all e ∈ Etr. If non-degeneracy conditions Eqs. (19-21) holds, then either w(Φ) = 0 or B = 0.
Theorem 1 characterizes the connection between invariant predictors and the diversity of training environments in the linear setting. Its conclusions are directly applicable to IRM and its relaxations including IRMv1, v1A, and v2.
4.3 THE ROLE OF Ie(φ)
Recall that the main difference between the invariance penalty of IRMv1 (with squared loss) and IRMv2 is the way that Ie(φ) is incorporated. Although at a first glance this may not seem significant, in what follows we demonstrate the importance of Ie(φ).

heading: 4.3.1 THE POTENTIAL FAILURE OF INVARIANCE PENALIZATION
text: Arjovsky et al. (2019) consider a linear Structural Equation Model (SEM) of the following form.
X1 ∼ N (0, σ2), Y = X1 + Z1 with Z1 ∼ N (0, σ2), X2 = Y + Z2 with Z2 ∼ (0, 1),
and X = [X1, X2]⊤. The variance σ2 is the only parameter that changes across environments, and we assume that Z1, Z2, and X1 are independent in and across all environments. Hence, in this setting, X1 models the invariant features and X2 models the environmental features as the correlation between X2 and Y varies across environments.
Consider the data representation φc(x) parameterized by a variable c ∈ R as
φc(X) = [ X1 cX2 ] , (22)
where c = 0 attains the invariant data representation, i.e., φ0(X) = [X1, 0]⊤. Based on this model, Arjovsky et al. (2019) argue that ∥winv − w⋆e(φc)∥2 is a poor choice for the invariance penalty as it is discontinuous at the invariant representation with c = 0, and vanishes as c → ∞. Interestingly, Ie(φc) is ill-conditioned for both small and large c’s. More precisely, in Appendix D.1 we show that
lim c→0 κ(Ie(φc)) = lim c→+∞ κ(Ie(φc)) = +∞,
where κ(·) denotes the condition number. That is, for a normal matrix A, its condition number is κ(A) := |λmax(A)|/|λmin(A)| where λmax and λmin denote its maximum and minimum eigenvalues, respectively.
We now examine the counterexample introduced by Rosenfeld et al. (2021), which is based on the SEM setting introduced in Section 4.1 and equation 16. They consider the data representation φϵ(X e) defined as
φϵ(X e) := [ Zc 0 ] + [ 0 Ze ] 1{Ze /∈Zϵ}, (23)
where {Ze /∈ Zϵ} is an event with P(Ze ∈ Zϵ) ≤ pe,ϵ where pe,ϵ := exp(−de min{ϵ − 1, (ϵ − 1)2}/8). The parameter ϵ ≥ 1 control the degrees to which the environmental variable Ze appears in the data representation by controlling the set Zϵ (see equation 40 in the Appendix). In particular, φϵ(X
e) converges to the invariant representation as ϵ → ∞, i.e., limϵ→∞ φϵ(Xe) = φinv(Xe) := [Zc, 0] ⊤.
Rosenfeld et al. (2021) show that IRMv1’s invariance penalty ∥∇wRe(w⊤φ)∥2 = O(p2e,ϵ), and can be made arbitrarily small, hence it is poor discrepancy as an invariance penalty. This, together with the fact that the training risk under φϵ is close to the risk under the invariant representation, they argue that φϵ is a plausible representation under IRMv1. However, this observation is only true for the case where λ is a constant. More specifically, φϵ is not a plausible solution of IRMv1 with adaptive penalty coefficient, e.g., λ(φ) = 1/λmin(Ie(φ))2.

heading: 4.3.2 PRACTICAL IMPLICATIONS
text: In both of the examples discussed in Section 4.3.1, I(φ) = E[φ(X)φ(X)⊤] is singular for the invariant representation. This is due to the fact that for both of the examples, the span of data representations φ(X) is the same as the span of inputs X while the span of the invariant features is only a small subset of that space. As a consequence, in presence of environment dependent variables, one should avoid utilizing data representations with codomain that is the same as the input space. In particular, when φ is parameterized by a neural network, the common approach of setting the size of the last layer equal to K − 1 for K-class classification, and 1 for regression problem avoids such phenomenon. We further illustrate such behavior on the Colored MNIST experiment, a standard benchmark for evaluating the performance of domain generalization methods.
Figure 1 illustrate that the test accuracy of IRMv1 drastically drops as we increase the dimension of the representation. However, the test accuracy of IRMv2 gradually increases first before dropping in a considerably overflexible representation. It is worth noting that the condition number of I(φ) increases with dφ and is ill-conditioned for large dφ, e.g., κ(I(φ)) > 1012 when dφ = 2048 (see Figure 2 in the Appendix). Note that in this experiment, images are down-sampled and contain 14 × 14 pixels, i.e., dX = 196.

heading: 5 EXPERIMENTS
text: In this section, we empirically evaluate the efficacy of our proposed implementations of IRM, namely, IRMv2 and IRMv1A with IRMv1 and ERM. We demonstrate the appeal of our approach on InvarianceUnitTests (Aubin et al., 2020), a synthetic test bed for evaluation of OoD methods with multiple experiments and data sets, each of which embedding a different causal relationship. We then evaluate each method on our curated data set HealthyGutTests to illustrate how the principle of invariant risk minimization can be used to minimize disparities and discrepancies in biomedical research. For completeness, we also perform numerical experiments on DomainBed, an extensive framework released by (Gulrajani & Lopez-Paz, 2020) to test domain generalization algorithms for image classification tasks on various benchmark data sets. Similar to Gulrajani & Lopez-Paz (2020), we also observe that none of the domain generalization algorithms significantly outperform ERM. We refer the reader to Appendix E for the detailed results and further discussion.

heading: 5.1 INVARIANCEUNITTESTS
text: In this section, we evaluate the efficacy of our proposed approaches for invariance discovery on the InvarianceUnitTests recently proposed by Aubin et al. (2020). These unit-tests entail three classes of low-dimensional linear problems, each capturing a different structure for inducing spurious correlations. The data set for each problem falls within the multi-environment setting described in Section 2 with ne = 104. For all problems, the input xe ∈ Rd is constructed as xe = (xeinv, xespu) where xeinv ∈ Rdinv and xespu ∈ Rdspu denote the invariant and the spurious features, respectively. To make the problems more realistic, Aubin et al. (2020) repeat each experiment and scramble the inputs by multiplying xe by a rotation matrix. In each problem, the spurious correlations that exist in the training environments are discarded in the test environment by random shuffling. As a basis for comparison, similar to Aubin et al. (2020), we implement an Oracle ERM where the spurious correlations are shuffled in the training data sets as well, and hence, ERM can readily identify them.
Table 1 illustrates the training and test errors of the various methods on linear unit test experiments. One can observe that IRMv2 is the only method that discovered a nearly invariant predictor on all experiments. However, on classification tasks (Example 2), all other methods except IRMv1A attained near perfect training accuracy by relying on the spurious correlations while their test accuracy was comparable to random guessing as the spurious features were randomly shuffled in the test set. The stability of IRMv2 in discovering invariant predictors makes it a suitable fit for real-world applications as we will demonstrate in the following section.

heading: 5.2 HEALTHYGUTTESTS
text: In this section, we further evaluate the efficacy of our proposed approaches on a set of experiments that highlight the importance of generalizability in real-world applications. An increasing number of microbiome studies have associated various alterations in human gut microbiota composition with both chronic and acute diseases, providing strong evidence that the gut microbiome is an essential factor in the maintenance of human health. HealthyGutTests is a curated set of microbiomes extracted from an extensively phenotyped and standardized meta-analysis cohort (Gupta et al., 2020) of 1770 individuals across 10 independent studies (Supplementary Table 7). Study selection was
limited to case-control studies in adult populations wherein Whole-Genome Sequencing (WGS) metagenome data of human stool (gut microbiome) and corresponding subject meta-data were publicly available. To reduce class imbalance, studies consisting of only control samples (healthy) or those with fewer than 100 total samples were excluded.
This data set reflects the characteristics of the domain generalization problem setting, as the data for each study is collected by different researchers, across various timepoints, locations and demographic populations. With this gut microbiome data set, we seek to illustrate a use-case in which the IRM principle can be used to minimize the effect of spurious correlations confounding the discovery of real explanatory variables and thus build a fair and trustworthy ML model that is appropriate for real-world deployment in a biomedical setting. We consider the task of human health status classification: predicting the binary label (healthy/unhealthy) associated with a sample. The data set falls within the multi-environment setting described in Section 2, with ne = 10 and inputs xe consisting of 992 features, where each feature corresponds to a unique microbial species. When data from multiple studies are collectively analyzed, systematic variation across different studies can give rise to biases known as “batch-effects”. These batch effects may obscure biologically relevant and under-represented sub-population effects in the data that are difficult to identify. When viewed in a clinical setting, these sub-population effects may represent differential etiologies, and thus require different treatment strategies. An ERM classifier trained on all such data (i.e. by shuffling and treating data as iid) may learn these batch effects as predictors of health status, and consequently ignore invariant features that are indicative of under-represented subpopulation effects. We seek to build an invariant predictor that is robust in its predictions when trained and tested on different experimental splits across the data set’s ten studies. Table 2 illustrates the test accuracy of the various methods on this binary classification task. One can observe that IRMv2 and IRMv1A methods demonstrate superior performance when compared to ERM and perform as good if not better than IRMv1 across 9 out of 10 test splits.

heading: 6 CONCLUSION
text: In this paper, we have presented IRMv2, an alternative implementation of the IRM principle that aims to enable out-of-distribution generalization by finding environment invariant predictors. We establish theoretical results on the effectiveness of our approach in the linear setting. In doing so, we bring forward the importance of the eigenstructure of the Gramian matrix of the data representation. In particular, we show that for the existing counterexample on the potential failure of IRMv1, the aforementioned matrix is ill-conditioned for the invariant representation. This highlights the significance of the span of the data representations in relation to the span of the underlying true invariant features of the data. That is, if the data representation allows for more degrees of freedom than needed to capture invariance, the Gramian matrix of the invariant representation would be ill-conditioned. This observation provides intuition on the underlying reasons why current implementations of IRM may fail. Our work demonstrates the effectiveness of IRMv2 and IRMv1A on both synthetic and real-world data sets. While this investigation of the IRM principle has lead to the development of a more pragmatic solution, we leave for future work, establishing theoretical guarantees beyond linear regimes.

heading: REPRODUCABILITY STATEMENT
text: We have provided detailed mathematical proofs of all of our theoretical claims in the appendices. The data processing, neural network architecture, and the model selection for Colored MNIST, InvarianceUnitTests, and DomainBed experiments are exactly replicated from their respective author repositories. For HealthyGutTests experiments, we utilized the data sets curated and standardized by domain experts Gupta et al. (2020). We discuss our additional data processing step in Section 5.2. We also provide details of the neural network architecture and the model selection in Appendix F.

heading: 3: for t = 1, 2, . . . , T do
text: 4: for e ∈ Etr do 5: compute the LSE w⋆e(φθt) according to Eq. equation 4
6: compute the optimal classifier w⋆(φθt) according to Eq. equation 12 7: Lt(φθt)← ∑ e∈Etr Re(w ⋆(φθt) ⊤φθt) + λρ IRMv2 e (φθt , w ⋆(φθt))
8: θt+1 ← θt − ηt∇θtLt(φθt) 9: Output predictor w⋆(φθT )⊤φθT .

heading: B CONSIDERATIONS FOR PRACTICAL IMPLEMENTATION
text: As the data generating distributions for each environment are unknown, we provide further details on practical implementations of IRMv2. As training is often done by gradient descent on mini-batches, we focus on this setting. For simplicity of notation, assume that the sample size of all environments are equal, each of which is divided into b mini-batches of size m. Then, ŵ⋆e the empirical estimate of the optimal classifier w⋆e is given by
ŵ⋆e(φ) :=
( b∑
k=1
Φek ⊤Φek )−1 b∑ k=1 Φek ⊤Y ek , (24)
where Φek ∈ Rm×dφ is a matrix each of its row is φ(xei )⊤ of the k-th mini-batch of environment e, and with some abuse of notation Y ek ∈ Rm is defined as the vector containing yie for the same mini-batch.
Bias: Note that whether the estimator equation 24 is unbiased or not depends on the underlying data generating assumption. In particular, it may be natural to assume that for the underlying invariant representation, it holds that E[Y e|φinv(Xe)] = w⊤invφinv(Xe) with the error Y e−E[Y e|φinv(Xe)] being independent of Xe and Y e, e.g., yie = w ⊤ invφinv(x e i )+ ηi where {η}ni=1 is a sequence of zeromean i.i.d. random variables. However, this may not hold for an arbitrary representation φ. One can modify the estimator equation 24 to derive an unbiased estimator by using different mini-batches for each term, e.g., by randomly dividing the mini-batches into two parts. More specifically, let B be the set of randomly chosen mini-batches of size ⌊b/2⌋, and [b]\B be the rest of the mini-batches. Let B be a set of such Bs. Then,
ŵ⋆e(φ) = 1 |B| ∑ B∈B  (∑ k∈B Φek ⊤Φek )−1 ∑ k∈[b]\B Φek ⊤Y ek  . (25) The estimator equation 25 is unbiased if xei are i.i.d. and y e i is independent of all other x e j conditioned on xei .
Numerical Stability: Although ∑
k∈B Φ e k ⊤Φek is positive semi-definite almost surely, it is singular
for m|B| < dφ and may even be singular for m|B| > dφ. Hence, to improve numerical stability, we consider the following empirical estimate of Ie(φ)
Îe,λ,B(φ) := λI + ∑ k∈B Φek ⊤Φek,
where λ ≥ 0. That is, to improve the numerical stability, we consider the ℓ2 regularized least squares estimate2
ŵ⋆e,B(φ) = Îe,λ,B(φ)−1 ∑
k∈[b]\B
Φek ⊤Y ek . (26)
Then,
ŵ⋆B(φ) = (∑ e∈Etr Îe,λ,B(φ) )−1 ∑ e∈Etr Îe,λ,B(φ)ŵ⋆e,B(φ).
We now provide an empirical estimate of the invariant penalty ρIRMv2e . We further partition B into B1, B2, and B3. Then,
ρ̂IRMv2e = 1 |B| ∑ B∈B {( ŵ⋆e,B1(φ)− ŵ ⋆ B1(φ) )⊤ Îe,λ,B2(φ) (ŵ⋆e,B3(φ)− ŵ⋆B3(φ))} . Note that ρ̂IRMv2e is an unbiased estimator of ρ IRMv2 e for λ = 0.

heading: C MATHEMATICAL PROOFS

heading: C.1 PROOF OF LEMMA 1
text: First, the risk under environment e with w = w⋆e(φ) is given by Re ( w⋆e(φ) ⊤φ ) = EY e [ ∥Y e∥2 ] −EXe,Y e [φ(Xe)Y e]⊤ Ie(φ)−1EXe,Y e [φ(Xe)Y e] .
Then, Re ( w⊤φ ) −Re ( w⋆e(φ) ⊤φ )
= w⊤Ie(φ)w − 2w⊤EXe,Y e [φ(Xe)Y e] +EXe,Y e [φ(Xe)Y e]⊤ Ie(φ)−1EXe,Y e [φ(Xe)Y e] = ∥∥∥Ie(φ)1/2 (w − w⋆e(φ))∥∥∥2 .

heading: C.2 PROOF OF LEMMA 2
text: Recall the definition of w⋆(φ) w⋆(φ) = argminw ∑ e∈Etr Re ( w⊤φ ) + λρIRMv2e (φ,w).
That is,
w⋆(φ) = argminw ∑ e∈Etr E [∥∥w⊤φ(Xe)− Y e∥∥2]+ λ ∥∥∥Ie(φ)1/2 (w − w⋆e(φ))∥∥∥2
= argminw ∑ e∈Etr { (1 + λ)w⊤Ie(φ)w − 2w⊤ (E [φ(Xe)Y e] + λIe(φ)w⋆e(φ))
+ w⋆e(φ) ⊤Ie(φ)w⋆e(φ) +E
[ ∥Y e∥2 ]} .
2The ℓ2 regularized least squares estimate of w⋆e is defined as
ŵ⋆e(φ) = argminw b∑ k=1 m∑ i=1 (yei,k − w⊤φ(xei,k))2 + λ∥w∥22.
Note that the objective function is a convex quadratic function of w. Hence, using the first-order optimality condition we have that
(1 + λ) (∑ e∈Etr Ie(φ) ) w⋆(φ)− ∑ e∈Etr (E [φ(Xe)Y e] + λIe(φ)w⋆e(φ)) = 0.
Then,
w⋆(φ) = 1
1 + λ (∑ e∈Etr Ie(φ) )−1(∑ e∈Etr (E [φ(Xe)Y e] + λIe(φ)w⋆e(φ)) ) .
Recall that w⋆e(φ) = Ie(φ)−1 [φ(Xe)Y e]. Then, [φ(Xe)Y e] = Ie(φ)w⋆e(φ). Thus,
w⋆(φ) = (∑ e∈Etr Ie(φ) )−1(∑ e∈Etr Ie(φ)w⋆e(φ) ) .
Finally, using a similar argument, we get argminw ∑ e∈Etr Re ( w⊤φ ) = argminw ∑ e∈Etr E [∥∥w⊤φ(Xe)− Y e∥∥2]
= argminw ∑ e∈Etr w⊤Ie(φ)w − 2w⊤E [φ(Xe)Y e] +E [ ∥Y e∥2 ] =
(∑ e∈Etr Ie(φ) )−1(∑ e∈Etr Ie(φ)w⋆e(φ) ) .

heading: C.3 PROOF OF LEMMA 3
text: For a symmetric matrix A ∈ Rd×d and a vector u ∈ Rd, it holds that λmin(A)∥u∥2 ≤ u⊤Au ≤ λmax(A)∥u∥2. Let u = Ie(φ)1/2 (w − w⋆e(φ)) and A = Ie(φ). Then,
∥Ie(φ) (w − w⋆e(φ))∥ 2 ≤ λmax(Ie(φ)) ∥∥∥Ie(φ)1/2 (w − w⋆e(φ))∥∥∥2 ∥Ie(φ) (w − w⋆e(φ))∥ 2 ≥ λmin(Ie(φ)) ∥∥∥Ie(φ)1/2 (w − w⋆e(φ))∥∥∥2 .

heading: C.4 PROOF OF LEMMA 4
text: Consider binary classification where Y ∈ {0, 1} with binary cross-entropy loss, i.e., ℓ(f(X), Y ) = −Y f(X) + log(1 + exp(f(X))).
Our objective is to bound the difference |R(w⊤1 φ) − R(w⊤2 φ)| for two arbitrary classifiers w1 and w2.
Let g(w) := log(1+exp(w⊤φ(X))). Then, using the multivariate mean-value theorem, there exists c ∈ (0, 1) such that
g(w2) = g(w1) +∇g((1− c)w1 + cw2)⊤(w2 − w1). We have that
∇g(w) = exp(w ⊤φ(X))
1 + exp(w⊤φ(X)) φ(X) = σ(w⊤φ(X))φ(X),
where σ(·) denotes the sigmoid function. Let w̄ := (1− c)w1 + cw2. Then,
ℓ(w⊤2 φ(X), Y )− ℓ(w⊤1 φ(X), Y ) = −Y w⊤2 φ(X) + log(1 + exp(w⊤2 φ(X))) + Y w⊤1 φ(X)− log(1 + exp(w⊤1 φ(X))) = −Y (w2 − w1)⊤φ(X) + σ(w̄⊤φ(X))φ(X)⊤(w2 − w1) = − ( Y − σ(w̄⊤φ(X)) ) (w2 − w1)⊤φ(X),
where the second equality follows from the application of mean-value theorem. Then, |R(w⊤1 φ)−R(w⊤2 φ)|2 = ∣∣E [(Y − σ(w̄⊤φ(X))) (w2 − w1)⊤φ(X)]∣∣2
≤ E [( Y − σ(w̄⊤φ(X)) )2] E [( (w2 − w1)⊤φ(X) )2] ≤ (w2 − w1)⊤E [ φ(X)φ(X)⊤ ] (w2 − w1),
where the first equality follows from Cauchy-Schwarz and the second one follows from the fact that Y ∈ {0, 1} and σ : R→ [0, 1], hence |Y − σ(w̄⊤φ(X)| ≤ 1 almost surely.

heading: C.5 PROOF OF THEOREM 1
text: In order to prove Theorem 1, we first find the optimal classifier for a given data representation for the linear setting in the following Lemma.
Lemma 5. Let w⋆e(Φ) = argminw Re(w⊤Φ) where Re(w⊤Φ) is defined as
Re(w ⊤Φ) = E [ tr (( w⊤ΦXe − [ 1{Y=+1} 1{Y=−1} ])( w⊤ΦXe − [ 1{Y=+1} 1{Y=−1} ])⊤)] .
Then,
w⋆e(Φ) = [ηβ ⋆ e (Φ) (1− η)β⋆e (Φ)] ,
where
β⋆e (Φ) = 1
1 + µ̄⊤e Σ̄ −1 e µ̄e
Σ̄−1e µ̄e. (27)

heading: Here, Σ̄e and µ̄e are defined as
text: Σ̄e := ΦS [ σ2cI 0 0 σ2eI ] S⊤Φ⊤,
µ̄e := ΦS [ µc µe ] .
Proof of Lemma 5: We have that w⋆e(Φ) = Ie(Φ)−1EXe,Y e [ ΦXeỸ e⊤ ] .
First, we have that
EXe,Y e [ ΦXeỸ e⊤ ] = ΦSEXe,Y e [[ Zc Ze ] [ 1{Y = 1} 1{Y = −1} ]⊤]
= ΦS [ ηµc (1− η)µc ηµe (1− η)µe ] = [ηµ̄e (1− η)µ̄e] .
Thus,
w⋆e(Φ) = [ηβ ⋆ e (Φ) (1− η)β⋆e (Φ)] , (28)
where
β⋆e (Φ) = Ie(Φ)−1µ̄e. (29)
We now compute Ie(Φ)−1 in terms of Σ̄−1e and µ̄e. We have that
Ie(Φ) = EXe [ ΦXeXe⊤Φ⊤ ] = ΦSE [[ Zc Ze ] [ Zc Ze ]⊤] S⊤Φ⊤.
From the definition of Zc and Ze, it follows that
E [[ Zc Ze ] [ Zc Ze ]⊤] = [ σ2cI 0 0 σ2eI ] + [ µc µe ] [ µc µe ]⊤ .
Thus,
Ie(Φ) = Σ̄e + µ̄eµ̄⊤e Using Sherman-Morrison formula, we have that
Ie(Φ)−1 = Σ̄−1e − Σ̄−1e µ̄eµ̄ ⊤ e Σ̄ −1 e
1 + µ̄⊤e Σ̄ −1 e µ̄e
. (30)
Finally, using Equation equation 29 it follows that
β⋆e (Φ) = ( Σ̄−1e − Σ̄−1e µ̄eµ̄ ⊤ e Σ̄ −1 e
1 + µ̄⊤e Σ̄ −1 e µ̄e
) µ̄e =
1
1 + µ̄⊤e Σ̄ −1 e µ̄e
Σ̄−1e µ̄e.
■
The proof of Theorem 1 closely follows from the proof of Rosenfeld et al., 2021, Lemma C.3.. In what follows, we include a proof to keep the manuscript self-contained.
Proof of Theorem 1: First, notice that the decomposition φ(Xe) = AZc +BZe (or ΦS = [A B]) is without loss of generality under the assumption of left-invertibilibity of S. Then,
Σ̄e = σ 2 cAA ⊤ + σ2eBB ⊤, (31)
µ̄e = Aµc +Bµe. (32)
Recall from Lemma 5 that β⋆e (Φ) = Σ̄ −1 e µ̄e/(1 + µ̄ ⊤ e Σ̄ −1 e µ̄e). If β ⋆ e (Φ) is invariant, then β ⋆ = β⋆e (Φ) for all e ∈ Etr. Then, by reorganizing terms, we get
Σ̄eβ ⋆ = ( 1− µ̄⊤e β⋆ ) µ̄e.
Thus, using Equation equation 31 and equation 32, we have that
(σ2cAA ⊤ + σ2eBB ⊤)β⋆ = ( 1− (Aµc +Bµe)⊤β⋆ ) (Aµc +Bµe).
Let v := −σ2AA⊤β⋆ + (1− µ⊤c Aβ⋆)Aµc. Then, B ( σ2eI + µeµ ⊤ e ) B⊤β⋆ − v = ( 1− µ⊤c A⊤β⋆ ) Bµe + µ ⊤ e B ⊤β⋆Aµc. (33)
Similar to proof of Lemma C.3. in (Rosenfeld et al., 2021), we show that for all fixed β⋆ and A Eq. equation 33 for all environments only holds (with probability 1) if B = 0. If |Etr| > de. Then, by the degeneracy assumption of the training sets, there exists at least one environment for which Eq. equation 19 holds. Let µ̃ and σ̃2 be the mean of Ze and variance of We for this environment. Then, we have that µ̃ = ∑de i=1 αiµi. By applying this linear combination to Eq. equation 33 for this environment, we get
B ( σ̃2I + µ̃µ̃⊤ ) B⊤β⋆ − v = ( 1− µ⊤c A⊤β⋆ ) B de∑ i=1 αiµi + ( de∑ i=1 αiµi )⊤ B⊤β⋆Aµc
= de∑ i=1 αi (( 1− µ⊤c A⊤β⋆ ) Bµi + µ ⊤ i B ⊤β⋆Aµc )
= de∑ i=1 αi ( B ( σ2i I + µiµ ⊤ i ) B⊤β⋆ − v ) , (34)
where in the last identity, we applied Eq. equation 33 for all i = 1, . . . , de. By rearranging the terms in Eq. equation 34, we get
B ( σ̃2I + µ̃µ̃⊤ −
de∑ i=1 αi ( σ2i I + µiµ ⊤ i
)) B⊤β⋆ = ( 1−
de∑ i=1 αi
) v. (35)
From the non-degeneracy condition equation 20, Eq. equation 35 is equivalent to BΓαB
⊤β⋆ = v, (36) where Γα is defined as
Γα := σ̃2I + µ̃µ̃⊤ −
∑de i=1 αi ( σ2i I + µiµ ⊤ i ) 1− ∑de i=1 αi .
Note that B, β⋆, and v are environment independent and Γα is an environment dependent matrix for which it holds that rank(Γα) = de from the nondegeneracy condition equation 21. Thus, Eq. equation 36 holds if and only v = BΓαB⊤β⋆ = 0. Then, Eq. equation 33 reduces to(
1− µ⊤c A⊤β⋆ ) Bµe + β ⋆⊤BµeAµc = 0
for all e ∈ Etr. Thus, Bµe = 0 for all e ∈ Etr, which holds if and only if B = 0 given that the span of µes is Rde .
D THE ROLE OF THE EIGENSTRUCTURE OF Ie(φ)
In this section, we elaborate on the discussions on the eigenstructure of Ie(φ), and in particular, its condition number in the examples of (Arjovsky et al., 2019) and (Rosenfeld et al., 2021).

heading: D.1 EXAMPLE 1 OF (ARJOVSKY ET AL., 2019)
text: Arjovsky et al. (2019) consider data that is generated according to the following SEM X1 ∼ N (0, σ2), Y = X1 + Z1,
X2 = Y + Z2, where Z1 ∼ N (0, σ2), Z2 ∼ (0, 1), and X1 are independent, and X = [ X1 X2 ] . Consider the following data representation
φc(X) = [ X1 cX2 ] . (37)
Then,
I(φc) = E [[ X1 cX2 ] [ X1 cX2 ]⊤] = [ σ2 cσ2 cσ2 c2(2σ2 + 1) ] .
We now find the eigenvalues of I(φc). That is, the solutions to det(I(φc)− λI) = 0. λ2 − λ ( σ2 + c2(2σ2 + 1) ) + c2σ2(2σ2 + 1)− c2σ4 = 0.
Then,
λ = 1
2
( σ2 + c2(2σ2 + 1)± √ (σ2 + c2(2σ2 + 1)) 2 − 4c2σ2(σ2 + 1) )
= 1
2
( σ2 + c2(2σ2 + 1)± √ σ4 + c4(2σ2 + 1)2 − 2c2σ2 ) .
Hence,
κ(I(φc)) = λmax(I(φc)) λmin(I(φc))
= σ2 + c2(2σ2 + 1) +
√ σ4 + c4(2σ2 + 1)2 − 2c2σ2
σ2 + c2(2σ2 + 1)− √ σ4 + c4(2σ2 + 1)2 − 2c2σ2
=
( σ2 + c2(2σ2 + 1) + √ σ4 + c4(2σ2 + 1)2 − 2c2σ2 )2 (σ2 + c2(2σ2 + 1)) 2 − (σ4 + c4(2σ2 + 1)2 − 2c2σ2)
= 1
2(σ2 + 1)
( 1
c σ2 + c(2σ2 + 1) +
√ 1
c2 σ4 + c2(2σ2 + 1)2 − 2σ2
)2 .
Finally,
lim c→∞ κ(I(φc)) = lim c→0 κ(I(φc)) =∞.
It is worth noting that Arjovsky et al. (2019) discuss that ∥w − w⋆e(φ)∥2 is poor discrepancy both for the invariant data representation, i.e., c = 0, and for a data representation that heavily rely on the spurious features, i.e., large c.

heading: D.2 EXAMPLE OF (ROSENFELD ET AL., 2021)
text: We show that in the setting that Rosenfeld et al. (2021) demonstrate that the invariance penalty is arbitrarily small, the risk is also inevitably arbitrarily small and is of the same order as the invariance penalty.
φϵ(X e) = [ Zc 0 ] + [ 0 Ze ] 1{Ze /∈Zϵ}.
Here, Zϵ is defined as
Zϵ := ⋃ e∈E (Br(µe) ∪ Br(−µe)) , (38)
where r := √ ϵσ2ede and Br(µ) denotes the ℓ2 ball of radius r centered at µ. Note that the invariant representation is given by
φinv(X e) = [ Zc 0 ] .
We compare the risk of the optimal predictor given the invariant and the Rosenfeld et al. (2021) representations.
D.2.1 INVARIANT REPRESENTATION
Ie(φinv) = E[φinv(Xe)φinv(Xe)⊤] = E [ ZcZ ⊤ c 0
0 0
] = [ µcµ ⊤ c + σ 2 cI 0
0 0 ] Given that Ie(φinv) is singular, we consider ℓ2 regularized LSE with regularization parameter γ > 0, i.e.,
w∗inv = (Ie(φinv) + γI)−1E[φinv(Xe)Ye].
Then,
w∗inv =
[ µcµ ⊤ c + (σ 2 c + γ)I 0
0 γI
]−1 E [ ZcY 0 ] = [ (µcµ ⊤ c + (σ 2 c + γ)I) −1µc 0 ] Recall from the proof of Lemma 1 that
Re(w ∗ inv ⊤φinv) = E[∥Y ∥2]−E[φinv(X)Y ]⊤Ie(φinv)E[φinv(X)Y ].
Then,
Re(w ∗ inv ⊤φinv) = 1− µ⊤c (µcµ⊤c + (σ2c + γ)I)−1µc. (39)

heading: D.2.2 ROSENFELD ET AL. (2021) REPRESENTATION
text: We have that
Ie(φϵ) = E[φϵ(Xe)φϵ(Xe)⊤]
= E
[ ZcZ ⊤ c 0
0 0
] +E [[ 0 ZcZ ⊤ e
ZeZ ⊤ c ZeZ ⊤ e
] 1{Ze /∈Zϵ} ] .
Recall that the first term is equal to Ie(φinv). For the second term, we have that
E
[[ 0 ZcZ ⊤ e
ZeZ ⊤ c ZeZ ⊤ e
] 1{Ze /∈Zϵ} ] = P(Ze /∈ Zϵ)E [[ 0 ZcZ ⊤ e
ZeZ ⊤ c ZeZ ⊤ e ] ∣∣∣Ze /∈ Zϵ] . Now define matrix A and parameter p as
A := E
[[ 0 ZcZ ⊤ e
ZeZ ⊤ c ZeZ ⊤ e ] ∣∣∣Ze /∈ Zϵ] , p := P(Ze /∈ Zϵ).
Then,
Ie(φϵ) = Ie(φinv) + pA. Also,
w∗ϵ = (Ie(φϵ) + γI)−1E[φϵ(Xe)Ye] = (Ie(φinv) + γI + pA)−1E [
ZcY ZeY 1{Ze /∈Zϵ}
] .
Let a = E[ZeY |Ze /∈ Zϵ]. Then, w∗ϵ = (Ie(φinv) + γI + pA)−1 [ µc pb ] .
For an invertible matrix S, we have that
(S +∆)−1 = S−1 − S−1∆S−1 +O(∥∆∥2). Thus,
w∗ϵ = ((Ie(φinv) + γI)−1 − p(Ie(φinv) + γI)−1A(Ie(φinv) + γI)−1 +O(p2)) [ µc pb ] = w∗inv + p ([ 0 b ] − (Ie(φinv) + γI)−1Aw∗inv ) +O(p2).
Then,
Re(w ∗ ϵ ⊤φϵ) = 1− [ µc pb ]⊤ w∗ϵ
= 1− µ⊤c (µcµ⊤c + (σ2c + γ)I)−1µc +O(p2). That is,
Re(w ∗ ϵ ⊤φϵ)−Re(w∗inv ⊤φinv) = O(p 2).
Recall that
φϵ(X e) = [ Zc 0 ] 1{Ze∈Zϵ} + [ Zc Ze ] 1{Ze /∈Zϵ}.
Here, Zϵ is defined as
Zϵ := ⋃ e∈E (Br(µe) ∪ Br(−µe)) , (40)
where r := √ ϵσ2ede and Br(µ) denotes the ℓ2 ball of radius r centered at µ. Then,
Ie(φϵ) = E [ φϵ(X e)φϵ(X e)⊤ ] = Ic + Ie.
where Ic and Ie are defined as
Ic :=
[ E [ ZcZ ⊤ c |Ze ∈ Zϵ ] 0
0 0
] P (Ze ∈ Zϵ) ,
Ie := E [[ Zc Ze ] [ Zc Ze ]⊤ ∣∣∣Ze ∈ Zϵ]P (Ze /∈ Zϵ) .
Here, we establish a lower bound on the condition number of Ie(φϵ) in terms of the probability of event 1{Ze /∈Zϵ}. Using Weyl’s inequality, we have that
λmax(Ie(φϵ)) ≥ λmax(Ic) + λmin(Ie), λmin(Ie(φϵ)) ≤ λmin(Ic) + λmax(Ie).
As Ie is positive semidefinite, λmin(Ie) ≥ 0. Moreover, λmin(Ic) = 0. Then,
λmax(Ie(φϵ)) ≥ λmax(Ic), λmin(Ie(φϵ)) ≤ λmax(Ie).
For the first term, we have
λmax(Ic) ≥ 1
de + dc tr(Ic)
= 1 de + dc E [ ∥Zc∥2|Ze ∈ Zϵ ] P (Ze ∈ Zϵ) = 1
de + dc E [ ∥µc∥2Y 2 + ∥Wc∥2 + 2µ⊤c WcY |Ze ∈ Zϵ ] P (Ze ∈ Zϵ)
= 1
de + dc (∥µc∥2 + dcσ2c )P (Ze ∈ Zϵ) , (41)
where the last identity follows from the fact that Y 2 = 1 almost surely, and the fact that Wc is independent of Y and We, and hence is independent of the event 1{Ze /∈Zϵ}.
For the second term, we have
λmax(Ie) ≤ tr(Ic) = E [ ∥Zc∥2 + ∥Ze∥2|Ze /∈ Zϵ ] P (Ze /∈ Zϵ)
= ( ∥µc∥2 + dcσ2c +E [ ∥Ze∥2|Ze /∈ Zϵ ]) P (Ze /∈ Zϵ) ,
where the last identity follows similarly as Equation equation 41. Then, E [ ∥Ze∥2|Ze /∈ Zϵ ] = E [ ∥µe∥2Y 2 + ∥We∥2 + 2µ⊤e WeY |Ze /∈ Zϵ ] = ∥µe∥2 +E [ ∥We∥2 + 2µ⊤e WeY |Ze /∈ Zϵ ] .
We have that
E [ ∥We∥2|Ze /∈ Zϵ ] ≤ deσ2e .
Moreover, using Cauchy-Schwarz inequality, we get E [ µ⊤e WeY |Ze /∈ Zϵ ] ≤ ∥µe∥E [∥We∥|Y ||Ze /∈ Zϵ] ≤ ∥µe∥ √ deσ2e .
Hence,
λmax(Ie) ≤ ( ∥µc∥2 + dcσ2c + (∥µe∥+ √ deσ2e) 2 ) P (Ze /∈ Zϵ) .
Thus,
κ(Ie(φϵ)) ≥ λmax(Ic)
λmax(Ie)
≥ (∥µc∥ 2 + dcσ 2 c )P (Ze ∈ Zϵ) /(de + dc)( ∥µc∥2 + dcσ2c + (∥µe∥+ √ deσ2e) 2 ) P (Ze /∈ Zϵ)
= ∥µc∥2 + dcσ2c (de + dc) ( ∥µc∥2 + dcσ2c + (∥µe∥+ √ deσ2e) 2 ) ( 1 P (Ze /∈ Zϵ) − 1 ) .
Note that (Rosenfeld et al., 2021, Lemma F.3.) show that
P (Ze /∈ Zϵ) ≤ pϵ,e,
where
pϵ,e := exp ( −1 8 min{ϵ− 1, (ϵ− 1)2} ) (42)
Then,
κ(Ie(φϵ)) ≥ ∥µc∥2 + dcσ2c (de + dc) ( ∥µc∥2 + dcσ2c + (∥µe∥+ √ deσ2e) 2 ) ( 1 pϵ,e − 1 ) .
Rosenfeld et al. (2021) show that the invariance penalty of (Arjovsky et al., 2019) is no greater than O(p2ϵ,e), which can be made arbitrarily small by choosing appropriately large ϵ. However, for such choices of ϵ, matrix Ie(φϵ) is ill-conditioned. In particular,
lim ϵ→∞
κ(Ie(φϵ)) =∞.

heading: D.3 EXTENDED COLORED MNIST EXPERIMENTS
text: The model and all its respective parameters are replicated from Arjovsky et al. (2019)’s GitHub repository3. The penalty coefficient for IRMv2 is set at λ = 10 and for λ0 = 10−5 for IRMv1A.
Figure 2 depicts demonstrate the significance of the dimension of the representation in relation to I(φ) being ill-conditioned, and, in turn, the test accuracy of IRMv1, v1A, and v2. In particular, by allowing the representation to be overly flexible

heading: E FULL LINEARUNITTESTS EXPERIMENTS
text: In this section, we provide more detailed tables for the numerical experiments on LinearUnitTests data sets. In addition to ERM, IRMv1, IRMv1A, and IRMv2, we include two other invariance discovery methods, namely Inter-environmental Gradient Alignment (IGA) (Koyama & Yamaguchi, 2020) and AND-Mask (Parascandolo et al., 2020), considered in (Aubin et al., 2020). The IGA method seeks to elicit invariant predictors by an invariance penalty in terms of the variance of the risk under different environments. The AND-Mask method, at each step of the training process, updates the model using the direction where gradient (of the loss) signs agree across environments.
All experiments are run using default parameters of the experiments of Aubin et al. (2020) available on their GitHub repository.4
We note that in the majority of these experiment (13 out of 18) one of our proposed methods IRMv1A or IRMv2 outperforms all other methods.

heading: F DETAILS OF THE HEALTHYGUTTESTS EXPERIMENTS
text: We summarize the details of each data set of the HealthyGutTests in Table 7.
3https://github.com/facebookresearch/InvariantRiskMinimization 4https://github.com/facebookresearch/InvarianceUnitTests

heading: F.1 EXPERIMENT DETAILS

heading: F.1.1 NEURAL NETWORK ARCHITECTURE AND HYPERPARAMETER TUNING
text: The neural network considered consists of 5 layers with input layer dX × dh, hidden layers of size dh × dh/2, dh/2 × dh/2, dh/2 × dh/4, respectively, and the output layer dh/4 × dφ. We use the ReLU action function and the dropout parameter of each layer is set as p. We treat dh, dφ, p, the learning rate, and the invariance penalty coefficient as hyperparameters, which we optimize using the validation accuracy. More specifically, we randomly generate 200 hyperparameters. For each set of parameters, we independently repeat 5 times training the model over 500 epochs. We then select the model with best average validation accuracy (over 5 experiments), and report the average test accuracy in Table 2.
For each data set in Table 2 that is considered as the test set, we treat each of the remaining data sets as an environment. Before doing so, we split each of training data sets into (0.8, 0.2) split of training and validation sets, respectively. We then concatenate all the validation subsets into a unified validation set for model selection.

heading: G FULL DOMAINBED EXPERIMENTS

heading: G.1 EXPERIMENT SUMMARY
text: DomainBed is an extensive framework released by Gulrajani & Lopez-Paz (2020) to test domain generalization algorithms for image classification tasks on various benchmark data sets. In a series of experiments, Gulrajani & Lopez-Paz (2020) show that enabled by data augmentation various stateof-the-art generalization methods perform similar to each other and ERM on several benchmark data sets.
Although the integration of additional data sets and algorithms to DomainBed is straightforward, we note that performing an extensive set of experiments requires significant computational resources as also pointed out by Krueger et al. (2020). For this reason, we limit the scope of our experiments to the comparison of ERM, IRMv1, IRMv1A, and IRMv2.
Similar to Gulrajani & Lopez-Paz (2020), we observe that no method significantly outperforms others on any of the benchmark data sets (see Table 8). For a complete set of results on DomainBed with various model selection methods, we refer the reader to Appendix G. As these data sets are image based and equipped with data augmentation, they may not provide comprehensive insight on the strengths and weaknesses of domain generalization algorithms on other modes of data, e.g., gathered in real-world applications.

heading: Algorithm A C P S Avg

heading: Algorithm A C P S Avg

heading: G.2 FULL DOMAINBED RESULTS

heading: Algorithm ColoredMNIST RotatedMNIST PACS VLCS Avg

heading: Algorithm C L S V Avg

heading: Algorithm A C P S Avg

heading: Algorithm C L S V Avg

heading: Algorithm ColoredMNIST RotatedMNIST PACS VLCS Avg

heading: Algorithm C L S V Avg

heading: Algorithm ColoredMNIST RotatedMNIST PACS VLCS Avg

Review summaries:
In this paper, the authors propose an alternative Invariant Risk Minimization (IRM) penalty for the particular case of a linear output layer trained with a mean-squared error loss. In this particular case, the original formulation can be understood as penalizing the squared Euclidean norm of the output layer's gradient, that is,  $$\rho_{e}^{\mathrm{IRMv1}}(\Phi, w) := \vert\vert \Phi_{e}^{T}(\Phi_{e} w - y_{e}) \vert\vert_{2}^{2} = \vert\vert \Phi_{e}^{T}\Phi_{e}(w -w_{e}^{*})\vert\vert_{2}^{2},$$  where $\Phi_{e} \in \mathbb{R}^{n_{e} \times d}$  is a matrix with the (learnt) data representation for environment $e$, $y_{e}$ the corresponding targets and $w_{e}^{*} = (\Phi_{e}^{T}\Phi_{e})^{-1}\Phi_{e}^{T}y_{e}$ the optimal, unregularized linear regressor. Witth this in mind, this paper proposes to use instead:  $$\rho_{e}^{\mathrm{IRMv2}}(\Phi, w) := \vert\vert (\Phi_{e}^{T}\Phi_{e})^{1/2}(w -w_{e}^{*})\vert\vert_{2}^{2},$$  which interestingly can be trivially shown to lead to an optimal classifier $w^{*}$ identical to that of the ERM formulation. Crucially, however, the authors propose to explicitly backpropagate through the dependence of this classifier on the data representations $\Phi_{e}$ when minimizing a cost function that incorporates $\rho_{e}^{\mathrm{IRMv2}}(\Phi, w)$.  The authors prove that, under similar non-degeneracy conditions as those assumed when analyzing other IRM penalties, $\rho_{e}^{\mathrm{IRMv2}}(\Phi, w)$ will either find an invariant data representation (or a trivial classifier) for a linear SEM as in Rosenfeld et al. 2021. Furthermore, the authors argue that there might be a connection between ill-conditioning of the Gram matrix $\Phi_{e}^{T}\Phi_{e}$ and failure modes of IRM penalties and propose a heuristic rule to adapt the penalty weight in the original IRM formulation.

This paper studies the problem of generalizing out of distribution (OOD), i.e., when data is not iid. Motivated by the error decomposition of the least-squares scheme, this paper proposes a novel invariance penalty that is directly related to risk. Based on the newly proposed penalty, this paper introduces another implementation of the IRM, named IRMv2. By characterizing the difference between the proposed penalty and the one in [1], this paper provides an adaptive rule for the choice of penalty parameter in the original IRMv1. Theoretical results guarantee that the proposed scheme can capture invariance when  in the linear case. The conditional number of Gram matrix of the data representation is analyzed.

This paper investigates the failure case of Invariant Risk Minimization and proposes a new invariance penalty for Out-of-Distribution Generalization. The authors provide some insights on how to avoid the potential failure of IRM based on the failure cases. And empirical results also validate the effectiveness of the proposed method.