On Invariance Penalties for Risk Minimization
Authors: Kia Khezeli
Arno Blaas
Frank Soboczenski
Nicholas Chia
John Kalantari
Venue: ICLR 2022 Submitted (2022)
Keywords: domain generalization
out-of-distribution generalization
invariant risk minimization
invariant representation learning

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). More- over, 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.

Under review as a conference paper at ICLR 2022
ON INVARIANCE PENALTIES FOR RISK MINIMIZATION
ABSTRACT
The Invariant Risk Minimization (IRM) principle was first proposed by Arjovsky et al. (2119) to address the domain generalization problem by leveraging data het- erogeneity 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 origi- nally proposed invariance penalty has recently been brought into question through stylized experiments and counterexamples. tionship 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). More- over, we provide practical insights on how to avoid the potential failure of IRM considered in the nascent counterexamples. Finally, conduct numerical exper- iments on both synthetic and real-world data sets with the objective of building invariant predictors. In our non-synthetic experiments, we sought build a pre- dictor of human health status using collection of data sets from various studies which investigate the relationship between human gut microbiome and a particu- lar 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.
this work, we investigate the rela-
INTRODUCTION
Under the learning paradigm of Empirical Risk Minimization (ERM) (Vapnik 1992), data is sumed to consist of independent and identically distributed (iid) samples from an underlying gen- erating 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, 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 test- ing distributions. This, in turn, highlights the need for machine learning algorithms to generalize out-of di-distribtttinn (OoD).
Shuffling and treating data 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 envirouments, 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 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 and Equation []. Because of the conceptually appealing nature of IRM and its potential problem, there is a stream of literature scrutinizing various facets
and to address the DoD-generalizatte 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). introduc- ing 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 in- depth overview the broader generalization literature, we refer the interested reader to (Arjovsky) 2020) and the references therein, and for empirical evaluation of the performance of number of the state-of-the-a methods on various test cases, we refer the reader to (Gulrajani & Lopez-Pat 2020

---

Under review as a conference paper at ICLR 2022
In this paper, introduce two practical implementations of the IRM principal to increase its ap- peal and applicability in real-world settings. First, we propose an invariance penalty that is directly show that the risk in each environment under an arbitrary clas- related to risk. More precisely, sifier 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 ini- tial findings to provide practitioners who currently use the original IRMv1 I implementation with an adaptive rule by which choose the penalty coefficient appropriately. In doing so, we character- ize 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 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, penalty of Arjovsky et al. 2019) can be made arbitrarily small for non-invariant representation. We show that the Gram matrix is ill-conditioned in such cases. We then provide practical solution to alleviate such behavior, 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 mild non-degeneracy condition, which is similar in nature to the ones considered in (Arjovsky et al. [2119] Rosenfeld et al. 2021). Finally, we evaluate our method on various test cases including Invariance Tests (Aubin al., 2020) and Tealsss Invariancel UnitTesss is a test bed with three synthetically generated data sets capturing different structures of spurious correlations. 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.
revisit the counterexample Rosenfild et al. (2021) where the invariance
SealthyyGutTese is curated collection
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 we introduce our more practical implementation and the rationale for its design. In Section 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.
2 BACKGROUND: INVARIANT PREDICTION
In this paper, we consider data (X", Y*) collected from multiple training environments E Eir where the distribution of (X', Y") and (X', Y*) may be different for + j with i,j E & We denote by Re the risk under environment That is, for predictor f X V, and loss function :: y x y + I., the risk under environment is defined as Re( = -xx (e(((())))
2.1
INVARIANT RISK MINIMIZATION
Arjovsky et al. (2019) define the notion of invariant predictors under a multi-environmen setting as follows.
Definition 1 (Invariant Predictor). A data representation & X + H. is said to clicit an invariant
predictor &p across environments & if there exists classifier w H + V, which is optimal for all environments, i.e., w E argmin,,.--> Rs (w o up) for all e E E.
To find such invariant predictors, Arjovsky et al. (009 introduce the notion of the Invariant Risk Minimization (IRM) principle:
min - w:
Σ Re (w o))
eEE..
(1)
subject to
w argmin Rs (wo Ve Er. a:K-y
2

---

Under review as a conference paper at ICLR 2022
As this bi-leveled optimization problem is rather intractable, AArovsky et al. (2019) propose Seeessaan prac- tical 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.
2.2 IRMvI: A RELAXATION OF IRM
In order to provide an implementation of IRM, Arjovsky et al. (2019) restrict the classifier w to linear functions, i.e.,
min pa WER
Σ
Re (w ())
eEE.r
(2)
subject to w E argmin Re (w (wTT)) 4), Ve E Exr
#E
To motivate their proposed penalty, Arjovsky et al. (2019) first consider the squared loss, i.e., {(f(x),)) = [l) y/12 where || denotes the Euclidean norm. Define matrix Ip(p)
[=(+) == Ex
(3)
Assuming that Ie (pp) is full rank for a fixed y,, its respective optimal classifier is unique, i.c., argmin Re (0" p) = " (5p where
(x) Ie( - yx [x(X))
Hence, in this setting, the invariant constraint of IRM in equation 2 can be simplified to w == w (v0
for all Etr. A natural relaxation the constraint w = (xp) = to penalty || - w [+)
However, Arjovsky et al. (2119) show that this penalty may not capture invariance by constructing an example for which w w (+)] is not well-behaved (see Section 4.3 for more details). They
argue that undoing the matrix inversion (p) >' which appears in the computation of w (अ) could
improve the behavior of the invariance penalty. That is, considering
w ( as the
invariance penalty. Moreover, for the squared loss, one can show that
(I)) (P))) (1/4) | R. ((wwT()))]) =))]"
(5)
Hence, Arjovsky et al.
2019 propose the following invariance penalty
Pe IRMMv1 (vp, w) : || R. (w))) the relaxation of IRM is given by
(6)
Using the penalty equation 6
min G. or
Σ Re(w" (p) + !RMM
(4p, w),
where À 2 is the penalty coefficient. Notice that for given w and (p, the predictor yp can be expressed using different classifiers and data representations, i.c., wo = wo where w = woo- and ; yp for some invertible mapping vi H + H. Hence, in principle, is possible fix w without loss of generality. By relying on this observation, Arjovssy et al. (2019) fix the classifier as scalar w 1. and, thus, search for an invariant data representation of the form yp E R Iixdd Their relaxation of IRM, which they refer to by IRMvl is given by
min P
ERe Re (v) !RMv1 (vp, 1.0).
(IRMv1)
REEL
Although equation (nnuaaiioSonnll, 5lonly holds for squared loss, Arjovsky et al. 2019 Theorem 4) justify the choice of || w 1.m (w 12 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 .)) V R(w Φ) = 0 if and only if w is optimal for all environments.
Although the penalty Pe IRMy1 (yp 1.0) seems to be an appropriate invariance penalty, which is cas- ily implementable, its effectiieness to capture invariance has been brought into question. In par- ticular, [Rosenfeld et al.] (202T) construct a non-invariant data representation such that the penalty Pe IRMvi (yp, 1.0) is arbitrarily small. In what follows, we propose 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.
In the presence of known exogenous (environment dependent) variables, one can utilize anchor regression Rothenhäusler et al. (202T) that is conceptually related to the invariance penalization. Instead of an invariance penalty, anchor regression relies on the projection the span of the said variahles.
3

---

Under review as a conference paper at ICLR 2022
3 IRMv2: AN ALTERNATIVE PENALTY
We revisit the structure of the risk in order propose an alternative penalty. In particular, in the fol- lowing Lemma, we provide the sub-optimality gap of risk under an arbitrary classifier in comparison the optimal classifier.
Lemma 1. Consider squared loss function. Let w E Rdd and (4p) as defined in equation Then,
Re (w) x) = =Re
[v)) (w w² (v))
(8)
Based on Lemmal we propose an invariance penalty that is directly comparable to risk.
IRMv2 (0, w) = I=) (w = (
(9)
The relaxation of IRM using the penalty equation is then given by
min
Σ
Re T () Xpe. LRM (5p, w)
(10)
BE
We further simplify the relaxation equation 10 by finding its optimal classifier for a fixed data rep- resentation defined as
W (5) == argmin,
Σ
Re (w T () + Ape IRMv2 (v, w).
(11)
eEEi
In the following Lemma, we leverage on the structure of the squared loss to find w" (p)).
Lemma 2. Consider squared loss function and fix yp Let w (4p) and w" (sp) defined in equation 4 and equation 71 respectively. Then,
w" w²(p)= = Σ Ic(p))
EI((p)
Σ I(()))u((1))
(12)
CeEF
leEf
Moreover, it holds that
argmin w
Σ
Re (w)x)= w=' (1p)))
(13)
e9E.
Based on Lemmas I and IRMv2.
2
we propose the following relaxation of IRM, which we refer to by
min Σ Re (w*(pp) (4p) (p) + Xpr 1RMv (e,, w (()).
(IRMv2)
eEEt
We provide the pseudo-code for IRMv2 in Algorithm T in Appendix A
There are several comments in order. It is worth noting that Equation equation 13 Lemma 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 oo IRMv2 in comparison to ERM.
There are a number of factors distinguishing IRMv2 from IRMv1. First, IRMv2 relies on the optimal classifier (pp) while w = 1.0 IRMv1. Second, the loss function in IRMv2 is squared loss while IRMvl 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, I, (sp) incorporated differently in the invariance penalty of IRMv1 and IRMv2. We formalize this latter observation in the Section 43.1
4

---

Under review as a conference paper at ICLR 2022
3.1 IRMvIA: AN ADAPTIVE PENALTY COEFFICIENT
We first bound the invariance penalty of IRMvl in terms of the penalty of IRMv2 and the eigenvalues of (xp). Then, based on this comparison, propose adaptive approach in choosing the penalty coefficient for IRMv1, which we refer to as Addativee (RMMy1A).
Lemma Let Pe. pRRMul ((p, w) and pe pIRMv2 (4,,0 w) be the invariance penalties of the IRMv1 and IRMv2 defined in Equations equation and equation respectively. Then,
IRMv1 (ve, w) S Amax (Ia(z)) Pe (pp, (p,w). Amin (Ie(P)) Pe (R (p, Pe
(14)
The proof of Lemma 3 directly follows from the definition of the invariance penalties p.RMvl (p, w) and pe @RMMv2(pp (ip, w) and the fact that for a symmetric matrix A Raxd and vector R holds that Amin (A)l|u|l² S u Au S Amax (A]
As IRMv1 used 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 terms of the invariance penalty of IRMv2.
Lemma 4. Consider binary classification with cross-entropy loss. Then, for W1, W2 RR it holds that
[R (p) (us (0) S I=(l))
w22)
Using Lemmas and we suggest the following rule for the penalty coefficient of IRMv1.
Ae
(15)
Ao + Amin (p))
where Ao 2 O is a user-specified parameter. The role of do is to avoid numerical instability when Amin (Te (e)) is small. Note that this is an adaptive rule, as &o may change throughout training.
THEORETICAL RESULTS AND ANALYSIS
In this section, 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 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 function of its parents and an exogenous random variable. Similar to Arjovsky et al. (2019) and Rosenfeld et al. (202T), 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 I, (vp), in particular, in relation to countercxamples of AArjossy et al. (2019) and Rosenfeld et al. 202 and provide practical solution to alleviate the possibility of such failure modes.
4.1
SEM FOR CLASSIFICATION
For each environment data (Xe, ye) is generated according to the following SEM:
2, Ze
{L,
1 with prob. 7,
X = S
y -== 11, with prob. 1
(16)
where E 10, and S E (dx(tc)de) is left invertible matrix, i.e., there exists St such that I. In this model, Ze captures the causal variables that are invariant across environments, and 2, captures the spurious environment dependent variables.
The variables Zz and Ze are generated as follows
Ze = P.Y + We,
We ((0,(0002T).
(17)
Z. = = + W
We Wa~N(0,022)
(18)
Here, Me E Rde Pa. E Rd, and N(,, Σ) denotes valti-variate Gaussian distribution with mean and covariance E. We further assume that W., W., and ye are independent for all environments.
5

---

Under review as a conference paper at ICLR 2022
4.2 INVARIANT REPRESENTATION UNDER IRM
For the setting introduced in 4.1 the invariant data representation is linear. In particular, for any d 2 Td. des ((X*) 0 0 stx Ze an invariant data representation. Given that 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 training environments under which IRM is guaranteed to find an invariant predictor, provided sufficient number of training environments.
Let |Eu| > de. As span( (ume(ueeeee.. S de, for each e E Eir there exists a set of coefficients of for i Etr le such that
He E of
(19)
iEEroi
We say that &sr is non-degenerate set of environments if for all E &,r it holds that
Σ of + 1, iEBLLe
(20)
rank (Te) = de:
(21)
where I defined as
To:- 1- 1-E Eiesure
of = Melle
Σ VEBuLe
(0,1 + pine Joi
The conditions of equation 20 and equation [1 specify that the span of covariance matrices of Ze is Rde, This is 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. 202I) 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 o 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 |Err| > de where (X Y generated according equation 16 Consider a linear data representation x = AZ. + BZ, and a classifier (b) on top of + that is invariant, i.e, w {p) w. (b) for all Exr. If non-degeneracy conditions Eqs. 27 holds, then either w(4 (4) = 0 or 0.
Theorem D 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, AIA, and v2.
4.3 THE ROLE OF I, (vp)
Recall that the main difference between the invariance penalty of IRMv1 (with squared loss) and IRMv2 is the way that (4)) ) is incorporated. Although at a first glance this may not seem significant, in what follows we demonstrate the importance of I. (vp).
4.3.1 THE POTENTIAL FAILURE OF INVARIANCE PENALIZATION
Arjovsky et al.
2019 consider a linear Structural Equation Model (SEM) of the following form.
X, N(0,03), Y = X, + Z1 with Z₁ (N, a)); X₂ = + Zz with Zg (0, 1),
and X = X2] The variance x² is the only parameter that changes across environments, and we assume that Z1, Ze, and X2 are independent in and across all environments. Hence, in this setting, X, models the invariant features and X2 models the environmental features as the correlation between X₂ and Y varies across environments.

---

Under review as a conference paper at ICLR 2022
Consider the data representation (c(x) parameterized by variable R
X1 CX2
(x(X)
(22)
where 0 attains the invariant data representation, i.e., (p = [X,, 0],, Based on this model,
Arjovsky et al. (2019) argue that ||W w (4==)]] poor choice for the invariance penalty as
discontinuous the invariant representation with 0, and vanishes as co. Interestingly,
Ie (1p.) is ill-conditioned for both small and large c's. More precisely, in Appendix D.I we show that
lim (Ze((p))) c-
= lim <(In()P)) = +00, c-+oc
where x<()) denotes the condition number. That is, for a normal matrix A, its condition number is <(A) := Amax (A)l/IAmin(()) where Amax and Amin denote its maximum and minimum eigenval- ues, respectively.
We now examine the counterexample introduced by Rosenfeld et al. (202) which is based on the SEM setting introduced Section 4.1 and equation 16 They consider the data representation
(Xx)) defined as
(x()) = Z
+ Ze
1. (2.42.)
(23)
where [Ze $ 2.) is an event with P(Ze E 2.) S Pe.c where Pe,e exp(-de minfe fc- 1))) The parameter 2 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() converges to the invariant representation as e + xc, i.., lim,-+o (x( [2.00]
(X*)
Rosenfeld et al. (021) show that IRMv1's invariance penalty V Re -] = Op and can be made arbitrarily small, hence it is poor discrepancy the fact that the training risk under &e is close to the risk under the invariant representation, they argue that se is a plausible representation under IRMv1. However, this observation is only true for the case where À is a constant. More specifically, ş2 is not plausible solution of IRMvl with adaptive penalty coefficient, e.g., X(ạ) 1/Amin mii (
invariance penalty. This, together with
4.3.2 PRACTICAL IMPLICATIONS
In both of the examples discussed in Section 4.3.1 I(p) = (ip(())((XT] is singular for the invariant representation. This is due to the fact that for both of the examples, the span of data representations xxX) is the same as the span of inputs X while the span of the invariant features only 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 &p is parameterized by neural network, the common approach of setting the size of the last layer equal to K 1 for K K-class classification, and 1 for regression problem avoids such phenomenon. We further illustrate such behavior on the Colored MNIST ex- periment, standard benchmark for evaluating the performance of domain generalization methods.
Figure illustrate that the test accuracy of IRMvl drastically drops as we increase the di- mension of the representation. However, the test accuracy of IRMv2 gradually increases first before dropping in considerably overflexible representation. It is worth noting that the con- dition number of I[pp) increases with di and ill-conditioned for large dg. e.g., x(I(())) > 10 when d = 2048 (see Figure 2 in the Ap- pendix). Note that in this experiment, images are down-sampled and contain 14 14 pixels, i.e. dx 196.
0.7
0.6-
Test Arruray
0.5-
0.4-
IRMvJA IRHU2
0.1
dp
Figure Effect of representation dimension on test accuracy for IRMv1, IRMvIA, and IRMv2.

---

Under review as a conference paper at ICLR 2022
5 EXPERIMENTS
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 InvariancUnnillen (Aubin et al., (020), a synthetic test bed for evaluation of OoD methods with multiple experiments and data sets, each of which embedding different causal relationship. We then evaluate each method on our curated data set HealthyGu/Tests to illustrate how the principle of invariant risk minimization can be used minimize disparities and discrepancies in biomedical research. For completeness, we also perform numerical experiments on DomainBed, an extensive framework released by (Gulrajani & Lopez-Pate 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.
5.1 INVARIANCRUNITTESS
In this section, we evaluate the efficacy of our proposed approaches for invariance discovery on the InvarianceUniflest 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 cor- relations. The data set for each problem falls within the multi-environment setting described in Section 2 with n2 = 10 For all problems, the input x' E IRd is constructed as x= = (4f spu)
where Tun E dd and zpn E pd 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 xº by 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), correlations are shuffled in the training data sets as well, and hence, ERM can readily identify them.
implement Oracle ERM where the spurious
Table I Illustrates the training and test errors of the various methods linear unit test experiments. One can observe that IRMv2 is the only method that discovered nearly invariant predictor on all experiments. However, on classification tasks (Example 2). all other methods except IRMv IA attained near perfect training accuracy by relying on the spurious correlations while their test accu- racy was comparable random guessing as the spurious features were randomly shuffled in the test set. The stability of IRMv2 in discovering invariant predictors makes it suitable fit for real-world applications as we will demonstrate the following section.
ERM
IRMv1
IRMvIA
IRMv2
Oracle
Training
Test
Training
Test
Training
Test
Training
Test
Training
Test
Examplel
7.35
14,40
9.47
10.72
12.82
12.87
11.14
11.26
10.49
10.32
Example
7.32
14.37
9.36
10.76
12.76
13.24
11.89
12.25
10.43
10.53
Example2
0.03
0.42
0.03
0.45
0.11
0.31
0.26
0.34
0.00
0.00
Example2s
0.03
0.45
0.03
0.45
0.04
0.27
0.22
0.38
0.00
0.00
Example3
0.00
0.37
0.00
0.36
0.49
0.26
0.31
0.40
0.01
0.01
Example3s 0.00
0.37
0.00
0.37
0.49
0.34
0.30
0.41
0.01
0.01
Table 1: Training and Test Errors on InvarianceUni/Tests for all algorithms and examples with
(diwv, dapa: denv) = (5, 3).). The errors for Examples and 1s are in MSE and all others clas- sification error. The empirical mean and the standard deviation are computed using 10 independent experiments. An 's" indicates the scrambled variation of its corresponding problem setting.
5.2 HEALTHYGUTTESTS
In this section, we further evaluate the efficacy of our proposed approaches on 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 curated set of mierobiomes tracted from an extensively phenotyped and standardized meta-am mnaa-snalysis cohort (Gupta et al. (0200) of 1770 individuals across 10 independent studies (Supplementary Table 7). Study selection was
8

---

Under review as a conference paper at ICLR 2022
limited to case-control studies in adult populations wherein Genomee Sequencing (WGS) metagenome data of human stool (gut microbiome) and corresponding subject meta-data were pub- licly available. To reduce class imbalance, studies consisting of only control samples (healthy) or those with fewer than 100 total samples were excluded.
This data reflects the characteristics of the domain generalization problem setting, the data for each study is collected by different researchers, across various timepoints, locations and demo- graphic populations. With this gut microbiome data set, seek illustrate use-case in which the IRM principle can be used to minimize the effect of spurious correlations confounding the discov- ery of real explanatory variables and thus build fair and trustworthy ML model that is appropriate for real-world deployment in biomedical setting. We consider the task of human health status classification: predicting the binary label (healthy/unhealthy) associated with sample. The data set falls within the multi-enviromment setting described in Section 2, with ne = 10 and inputs x 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 popplooultiion effects in the data that are difficult to identify. When viewed in 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 treat- ing data as rid) may learn these hatch effects as predictors of health status, and consequently ignore invariant features that are indicative of under--epresented 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 len studies. Table 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 ERM and perform as good if not better than IRMv1 across 9 out of 10 test splits.
ERM
IRMv1
IRMvIA
IRMv2
Feng (2015)
0.785
0.856
0.848
0,856
(2017)
0.604
0.604
0.604
0.606
(2017)
0.733
0.726
0.724
0.734
Karisson (2013)
0.813
0.861
0.849
0,866
Le Chatelier (2013)
0.642
0.655
0.653
0.653
Nielsen (2014)
0.732
0.744
0.739
0.744
Qin (2012)
0.763
0.785
0.779
0.795
Vogtmann (2016)
0.692
0.696
0,700
0.700
Zeller (2014)
0.775
0.755
0.768
0.787
Zhang (2015)
0.666
0.700
0.701
0.701
Average
0.721 = 0.064
0.738 + 0.077
0.736 0.074
0.744 0.079
Table 2: Test Accuracy on HealthyGurTests health status classification task. Each test accuracy is averaged over 5 independent experiments.
CONCLUSION
In this paper, we have presented IRMv2, alternative implementation of the IRM principle that aims to enable duittf-diittibuution 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 representa- tion. 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 free- dom than needed to capture invariance, the Gramian matrix the invariant representation would be ill-conditioned. This observation provides intuition on the underlying reasons why current im- plementations 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.

---

Under review as a conference paper at ICLR 2022
REPRODUCABILITY STATEMENT
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, InvarianceUniiTeest, and DomainBed experiments are exactly replicated from their respective author repositories. For HealthyGurTests 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 We also provide details of the neural network architecture and the model selection in Appendix
5.2
REFERENCES
Robert Adragna, Elliot Creager, David Madras, and Richard Zemel. Fairness and robustness in invariant learning: A case study in toxicity classification. arXiv preprint arXiv:2011.06485 2020.
Kartik Ahuja, Karthikeyan Shanmugam, Kush Varshney, and Amit Dhurandhar. Invariant risk min- imization games. In Proceedings of the 37th International Conference Machine Learning, pp. 145-155, 2020.
Martin Arjovsky. Out of Distribution Generalization in Machine Learning, PhD thesis, New York University, 2020.
Martin Arjovsky, Léon Bottou, Ishaan Guirajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint 19Xiv:1907.02899 2019.
Benjamin Aubin, Martin Arjovsky, Leon Bottou, and David Lopez-Paz. Linear unit tests for in- variance discovery. In Causal Discovery and cousality-Inspiree Machine Learning Workshop at NeurIPS, 2020.
Ishaan Gulrajani and David Lopez-Paz. In search of lost domain generalization. arXiv preprint arXiv:2007.01434, 2020.
Vinod K Gupta, Minsuk Kim, Utpal Bakshi, Kevin Y Cunningham, John M Davis, Konstantinos N Lazaridis, Heidi Nelson, Nicholas Chia, and Jaeyun Sung. A predictive index for health status using species-level gut microbiome profiling. Nature communications, 11(1):1-16, 2020.
Khurram Javed, Martha White, and Yoshua Bengio. Learning causal models online. arXiv preprint arXiv:2006.07465 07461, 2020.
Pritish Kamath, Akilesh Tangella, Danica Sutherland, and Nathan Srebro. Does invariant risk minimization capture invariance? arXiv preprint arXiv:2/01.01134 2021.
Masanori Koyama and Shoichiro Yamaguchi. -uttof-dist generalization with maximal in- variant predictor. arXiv preprint urXiv:2008.0188 2020.
David Krueger, Ethan Caballero, Priol, and Aaron Courville. out-of-distribution generalization via risk extrapolation (rex). arXiv preprint arXiv:2003. 00688, 2020.
Jacobsen, Amy Zhang, Jonathan Binas, Remi Le
Giambattista Parascandolo. Alexander Neitz, Antonio Orvieto, Luigi Gresele, and Bernhard Schölkopf. Learning explanations that are hard to vary. arXiv preprint wwXi:2209900039, 2020.
Judea Pearl. Causality. Cambridge university press, 2009.
Elan Rosenfeld, Pradeep Kumar Ravikumar, and Andrej Risteski. The risks of invariant risk min-
imization. In International Conference Learning Representations, 2021. URL https:
//openreview.
Dominik Rothenhäusler, Nicolai Meinshausen, Peter Bühlmann, and Jonas Peters. Anchor regres- sion: Heterogeneous data meet causality. Journal of the Royal Statistical Society: Series (Sta- fistical Methodology). 83(2):215-246, 2021.
10

---

Under review as a conference paper at ICLR 2022
Claudia Shi, Victor Veitch, and David Blei. Invariant representation learning for treatment effect estimation. arXiv preprint arXiv:2011.12379, 2020.
Damien Teney, Ehsan Abbasnejad, and Anton van den Hengel. Unshuffling data for improved generalization. arXiv preprint arXiv:2002. 2020.
Vladimir Vapnik. Principles of risk minimization for learning theory. In Advances neural infor- mation processing systems, pp. 831-838, 1992.

---

Under review as a conference paper at ICLR 2022
A
IRMv2 ALGORITHM PSEUDO-CODE
Algorithm IRMv2
I: Input: Data set: De for e E Etr. Loss function: Squared loss, Parameters: penalty coefficient
À 2 O, data representation parameters 0 E RRd, learning rate 7h, training horizon T.
2: Initialize 01 randomly
3: for 1.2...., do
4: for e E Eir do
5: compute the LSS wi (Po.) according to Eq. equation 4
6: compute the optimal classifier w= (ppe,) according to Eq. equation 12
7: Le (((o)) t Eeeen Re(w" (sPo) )oo)) + Ape ARRMM22 (ypo, w* (spo))
8: 044 0g 1/t (PPOL(
9: Output predictor w² ((pr) 1p07
B CONSIDERATIONS FOR PRACTICAL IMPLEMENTATION
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 focus this setting. For simplicity of notation, assume that the sample size of all environments are equal, each of which is divided into mini-batches of size m. Then, W the empirical estimate of the optimal classifier w is given
8
(P) =
Df
YE,
(24)
where &6 E Rmxdd is a matrix each of its row is T of the k-tth mini-batch of environment e, and with some abuse of notation Yg E Th.m is defined as the vector containing y² for the same mini-batch.
Note that whether the estimator equation 24 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 |Pav (Xt) = wh Piny PPivv(X' (X*) with the error y E[Y PPny (X*) being independent of X and Ye, e.g.- Vi "inv'Pin' (x$) + 7)2 where (n) is sequence of zero- mean i.i.d. random variables. However, this may not hold for an arbitrary representation şp. 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 batc-b of size [b/2], and [b];B B be the rest of the mini-bachhes. Let B be set of such Bs. Then,
Bias:
(v)
Σ
& T Y
(25)
ME
KEIN
The estimator equation 25 is unbiased if af are i.i.d. and of is independent of all other of conditioned on zf.
Numerical Stability: Although Σten & of is positive semi-definite almost surely, it is singular for m| B < d, and may even be singular for ml B dg. Hence, to improve numerical stability, we consider the following empirical estimate of Ie (xp)
[ex.(() = M
Σ
4
k
KE
12

---

Under review as a conference paper at ICLR 2022
where A > 0. That is, to improve the numerical stability, we consider the E regularized least squares estimate
(p))(p) - (,))
Σ
Df Y
(26)
kE
Then,
-1
Σ (
(b(+)=
Σ
(e).B(p)
eEEtr
We now provide an empirical estimate the invariant penalty Pe IRMv2, We further partition B into B1, B2, and By. Then,
RRM22
CB ( (x) (TT, - -B (1))
Σ
( (w
B
BEB
Note that 2RRV22 is an unbiased estimator of IRMv2 for X = 0.
C MATHEMATICAL PROOFS
C.I PROOF LEMMA T
First, the risk under environment with (x) is given by
Re (w)( x) Ey ||| E (=(x)) Ie(x)- Ex -x(,)-'xx,x>- (x(
Then,
Rs (w) () (wt(p)
w Io 2w Ex., (x(x + Ex (ø(X))V9 Te(1))'Exx (
= | (w w
C.2 PROOF OF LEMMA 2
Recall the definition of (pp)
w=* (x²) w
argmin,..
Σ eEEr
Re (wTx) x)) Xpe !RM22 (12, (v)
That is,
(vp) = argmin
Σ eEEr
E
I w ((()) + > I< (w =
argmin w
Σ eE
(1 A)w Io(p)) - 2w (E (x(x)))) + XI (v))
In II [
+
"The la regularized least squares estimate of w, is defined as
क (p) = argmin 1 ΣΣ (of, T 22² + A1/w0
13

---

Under review as a conference paper at ICLR 2022
Note that the objective function is a convex quadratic function of w. Hence, using the first-order optimality condition we have that
(1+
Σ Ie(p))
10 (0()
-
Σ CEEur
E (x(X)))++
( (Zx())) = 0.
Then,
w" (1)
1
Σ Ie(oo)
Σ E (x(x))) + AIe (1)) (10))
I+
NeEEer
CeEEtr
Recall that w (vp) = Ic(r [ Then, (F(()) Ie (z))) Thus.
w² (p()
Σ EeE
In(o)
Σ In()))=(())) leEEe
Finally, using a similar argument, we get
argmin Σ
Re (w = argmin
E
| T pp(x) -
eEEu
REEie
argmin
Σ w Ie [e( - 2w 'o''Ele [ - [[
eEEir
Σ II(T)
Σ ((0)) leEEr
LeEET
C.3 PROOF OF LEMMA
For symmetric matrix A E Raxd and a vector E RR. it holds that Amin Amin (A)]lu|1 = TT Au s Amaa (A)]lu||? Let 1/2 (w (x)) and A (x)). Then,
[+x( (w w (+)))1 s Amax Am [/2 (w -w
[It( (w -
Amin (I.(()) 1/( (w- wz(()) (1))) (w
C.4 PROOF OF LEMMA
Consider binary classification where Y E {0, 1} with binary cross- loss, i.e.,
{lf(X)YY) = -Xx(() + log(1 + ex)(f
Our objective is bound the difference [R(w] (p) R(w] (p] for two arbitrary classifiers wi and W2.
Let g log(1 exp(1+ ((x))) Then, using the multivariate mean-value theorem, there exists ((X)))
E (0, 1) such that
g(w2) g(w1) Vg((1 - c)wn + cwz) (wr w1).
We have that
exp(w T (x)))
Vg(w)) 1 1 cxp(w" xxX)) (X) aw ((X))s(()),
where o(-) denotes the sigmoid function. Let w := (1 - c)wi + cw2. Then,
[mw] ((X), ((w] Y)
= YYwj ((X) + log(1 + exp(w) v(x))) + Yuj x(X) log(1 + exp(wF ((X)))
= =w - w1) T ((() o(w (x))) (wy - -w))
=- 1Y-oaww' (x()) (w2 w1) (X)),
14

---

Under review as a conference paper at ICLR 2022
where the second equality follows from the application of mean value theorem. Then, [R[w] (0) R(wg x)| = E [(Y ( (xX))) (w2 -w1) xxx]]]
E (Y ( T p(x)) E ((w2 wi) (x)))
(w w1) E (x ( (w2 w1),
where the first equality follows from Cauchy-Schwarz and the second one follows from the fact that Y E (0.1) and a I2 + [0, 1]. hence |Y - a( (xX) S almost surely.
C.5 PROOF OF THEOREM 1
order to prove Theorem we first find the optimal classifier for given data representation for the linear setting in the following Lemma.
Lemma 5. Let w ( argmin,, Re(w" T 4) where (w +) is defined as
1 [r==+1)
)
w Txx-
1 (x=+1))
(w ②) E ir w xx
(Y=-1)
(Y=-13)
Then,
w = (1))(
(1 - 1))
where
(2(
Ee tee
(27)
1 1 Mo DE D
Here, De and the are defined as
To?I
Σ. = 45
ST
0
Te = &5
Me
[
Proof of Lemma 5
We have that
(=) -xx
Yx x
First, we have that
Ex @xoreT
Zco
1(Y 13
1 y
[ee
1( -1]
-5
Tppe
(1 me
Mike
(1 17) Pr
|n/
(1 [e]
Thus,
( (4) (T) (1 (1-1))3(0)) (
(28)
where
(:(2) = Io (2) Be.
(29)
We now compute I. (+) -1 in terms of Σー and Pe. We have that
Z
Z.
[ Ex x &
&SE
S
Ze
Z
15

---

Under review as a conference paper at ICLR 2022
From the definition of Ze and Ze, it follows that
Ze
Ze
Ue
Ue Me
E
Ze
Ze
0
+
Me
Thus,
Ie [ = Eg + M
Using Sherman- Morrison formula, we have that
Σ- Meffe Eze
== 1 + Pee 2e the
(30)
Finally, using Equation equation 29 it follows that
Σ. Re P.
(1 = Σ,
Tee
Ee Mee
1+
Pe De Me
1 Re
Ee Me
The proof of Theorem I closely follows from the proof of Rosenfeld et al. 2021 what follows, we include proof to keep the manuscript self-contained.
Lemma C.3.. In
Proof of Theorem 7
First, notice that the decomposition ((X") = AZ. + BZ. (or 45 = [A B)) without loss of generality under the assumption of iene-invertiiilibbty of Then,
Σ. =AAA + BB
(31)
Tee = Ape + Blee
(32)
Recall from Lemma 5 that (+) = The + Me 1..5a1i p.). If (3( is invariant, then 3** = 3 (T) for all &tr. Then, by reorganizing terms, we get
Egg (1- fee 8*) Pe.
Thus, using Equation equation 31 and equation 32 we have that
(o.AAT BBBBB ) (1 (Ape + BYc) (y) (Apo + Bue)-
Let UtU:: o2 AA 3r (1 1 AB*) App- Then,
n. B ( Melie) B- (1- A 3*) By + p. B 8* Ape-
(33)
Similar to proof of Lemma C.3. (Rosenfeld et al., 202T), we show that for all fixed 3 and A Eq. equation 33 for all environments only holds (with probability 1) if B = O. If [Err > de. Then, by the degeneracy assumption of the training sets, there exists at least one environment for which Eq. equation 19 holds. Let j and Ø2 be the mean of Z, and variance of W. for this environment. Then, we have that f environment, we get
Ede E aspi By applying this linear combination to Eq. equation 33 for this
T
de
B ( ( BT - == (1 H.T A 3+) 0itli +
Σ
Ciilis
BT T Apo
de
((1 p.- A 3*) Bu, + n1 BT TB* Ape)
0i (B (o + + B ----) w)
(34)
where in the last identity, we applied Eq. equation 33 for all i i= 1,..., de. By rearranging the terms Eq. equation 34 we get
d_ 5 I
(o mmNT) ) B BTB²= 1- 021
B
+
up T
(35)
16

---

Under review as a conference paper at ICLR 2022
From the non-degeneracy condition equation
20 Eq. equation 35 is equivalent
BI B =
(36)
where I is defined as
Ta : 5-1 juj
Ene a (o) + p.p.))
- a;
Note that B, _*,, and v are environment independent and Ia is an environment dependent matrix for which it holds that rank(T..) de from the nondegeneracy condition equation 21 Thus, Eq. equation 36 holds if and only BT.B' 8 O. Then, Eq. equation 33 reduces to
(1
Bp + 3 T BpeAp = 0
for all e Eu. Thus, Byee = for all E Etr, which holds and only if B given that the span of Pes is R
D THE ROLE OF THE EIGENSTRUCTURE OF (y)
In this section, we elaborate on the discussions on the eigenstructure of (Z)),, and in particular, its condition number in the examples of (Arjovsky al. 2019) and (Rosenfeld et al. 2021.
D.1
EXAMPLE 1 OF (ARJOVSKY ET AL., 2019)
Arjovsky et al. 2019 consider data that is generated according to the following SEM
Xi A/ o²) = X1 + Z1.
X₂ Y Zz.
where Z1
N( Z2 (0, 1), and X2 independent, and X
Consider the
X2
following data representation
X1 Px(X)
(37)
CX2
Then,
X1
X1
0²
ca2
I(px) E
CC2
cX2
co2
c202 +1)
We now find the eigenvalues of I((cc). That is, the solutions to det(T[(p.) - \T) = 0.
- + 22022 + 1)) + 020022002 + +1) 0²0º = 0.
Then,
+1)) 0² + (202 + +1) / +2200 4c20 +1)
A
(20 +1) +
Hence,
Ama (I(})))
<< <<(((pp))
Amin (I(p)))
0² + 2 2202 + (
0² + 2 (202 +1) +3
2
+ + +
(0 + 2 + 1)) (o -)0 2ci20)
2
1
1 + c(20² +1 + 2 cc2(202 1)2 20
2(02+11 +1)
c
17

---

Under review as a conference paper at ICLR 2022
Finally,
lim x(I(po)) lim (I(())))
is worth noting that Arjovsky et al. (2019) discuss that - for the invariant data representation, I.e., 0, and for data representation that heavily rely on the spurious features, i.c., large c.
- w (+)) is poor discrepancy both
D.2 EXAMPLE OF (ROSENFELD ET AL. 2021
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 of the same order the invariance penalty.
Ze
0
(x
Z 1 (Ze$2.).
0
Here, 2, is defined as
Z. = UB UB,(-e))),
(38)
AEE
where r := Veo2de and (p) denotes the lg ball of radius centered at μ.. Note that the invariant representation given by
Za
Pine (X*)
0
We compare the risk of the optimal predictor given the invariant and the Rosenfeld al. (2021) representations.
D.2.1 INVARIANT REPRESENTATION
Ie (p(in))) = Elypime ( Pinn (X Z.Z.T 0 =
Mape + oZ. 0
0]
Given that (yPinv) is singular, consider lg regularized LSE with regularization parameter ? 0, i.e.,
w W = (Ie (Vemm) + (X*)
Then,
-02 + 7
Z,Y
(HENT + 'e
in
7I
E
0
Recall from the proof of Lemma 1 that
(w 'Pimm) || Elpino (X){ Ig ( (X)Y).
Then,
Re Re( (Pin)) = ne (Man + (0? + +))) pe-
(39)
D.2.2
ROSENFELD ET AL.
(2021) REPRESENTATION
We have that
I. [p.) Ex
(X+))
7.7T
ZoZZ
=
E
1 1Zo022)
0
Z.Z.
ZoZ.
18

---

Under review as a conference paper at ICLR 2022
Recall that the first term is equal to Ie (fimm)). For the second term, we have that
ZeZ
Z.Z T c
Z
ZeZe
[2.
P(ze & 2
ZeZ
Ze $ 2
Now define matrix and parameter as
0
Z.Z_
A = E
Ze $ 2
Ze
= P(Ze $ Z.).
Then,
Ic( Ic((pine) PA.
Also,
" (Ie( + "111
Z.Y
(Ie (Ieln)) + + PA) E
ZeY1 [Z.
Let a El & 2.). Then.
-1 [g] pb
- (Io((inn)) + +
For an invertible matrix S, we have that
(S + -S-
+ +(1AA13))
Thus,
w = ((Im() + - p(IT (PPim))) + +71) Pe
0
Winv +D
(Iol(Pin)) + + AAWin +
Then,
Re(w)
=1 ==)==
[e pb) w²
= 1 Mi. (Mofle (o + - O0(p))
That is,
(pe) (ww
Re (u T "iin))
= 0(p²).
Recall that
[d px(x))
Zc
(
1 112,22,1 +
Ze] 1(2,22))
Here, Z, is defined as
Z. = (B, U UBB(-)e)),
(40)
aEE
where de and B, (p) denotes the l ball of radius centered at μ. Then,
Ix() E (X") (x9)) = Te + Te.
where Le and L. are defined as
E Z.Z
Ze E EZ]
1
P (Ze 2,),
0
Zo
2
I. E
Ze Z.
P (Ze 2.)
Ze
Ze
19

---

Under review as a conference paper at ICLR 2022
Here, establish a lower bound the condition number of Ie (fec) in terms of the probability of event Il (2. Using Weyl's 'ss inequality, we have that
Xmax (Ie (p.)) 2 Amax (lc) + Amin (le),
Amin (Ie (p.)) I Amin (Le) + Amax (Le).
As Le is positive semidefinite, Amin (1e) 2 0. Moreover, Amin (lo) = 0. Then,
Amax (Ie (fee) 2 Amax (Le), Amia (Ie (fe)) S Amax (Te).
For the first term, we have
Amax (le) 2 tr((o)
de da
E [116,11
Ze E 2.] P (Ze E Z.)
de + de
E [1lue]| 2| + 1| + 2/e W W E 2. (Ze E 2.)
de de
de de (llmoll
+ d.o22) P (Z E 2,)
(41)
where the last identity follows from the fact that y² = almost surely, and the fact that We is independent of Y and We, and hence is independent of the event (Z. For the second term, we have
Amax (Ie) S tr((c)
= E 12 1 & 2.] P (Ze $ 2c)
(1|ael2 2 + d. E [ Ze 2.)) P (Ze & 2c),
where the last identity follows similarly as Equation equation 41 Then,
E
|1 12 & z.] = E [| + ||& 1l + 2 W Y IZ $ 2.]
|| + E W| + 2 W.Y Uee $ z.).
We have that
E |1 $ z S de
Moreover, using Cauchy-Schwaer Schwarz inequality, we get
E [ee Wo $ z.] s l| E W. 2e $ Zd] S lllell V.oz
Hence,
Amax (Te) S (Iuual2 12 + d.o (llmell + vd.o23e) P (Zc $ 2x).
Thus,
Amax (Ic)
<(Ic(p))) 2
Amax (Ie)
[ (Ze E 2.) /do + do)
| + d.o2 + (lluell + vd.03)2 P ( $ Z.)
|| + doo2
( (p 2.) -1)
(de + de)
[Ieell2 + d.o2 + (llHell + vooo?22)
Note that (Rosenfeld et al. 2021 Lemma F.3.) show that
P (Ze & Z.) S Pe,e
20

---

Under review as a conference paper at ICLR 2022
where
Pee i= exp (-3 minff 1, 1)23)
exp
(42)
Then,
l + doo? (dg + do) (1uu112 do + (llsell / /d.o3)))
(1e -1)-1) Pe,e
16 (((c(1)) N
Rosenfeld et al. (2021) show that the invariance penalty of (Arjovsky et al. 2019) is no greater than O(PE,)) which can be made arbitrarily small by choosing appropriately large However, for such choices of matrix Ie (fec) is illlconditioned. In particular,
lim x(Io( = oc. 6-00
D.3 EXTENDED COLORED MNIST EXPERIMENTS
The model and all its respective parameters are replicated from Arjovsky et al. (2019)'s GitHub repository 3
The penalty coefficient for IRMv2 is set at À 10 and for λo 10-3 for IRMvIA.
Figure depicts demonstrate the significance of the dimension of the representation in relation to I(p) being ill-conditioned, and, in turn, the test accuracy of IRMv1, vlA. and v2. In particular, by allowing the representation be overly flexible
E FULL LINEARUNITTESTS EXPERIMENTS
In this section, we provide more detailed tables for the numerical experiments on LinearUnitTests data sets. In addition to ERM, IRMv1, IRMv1A, and IRMv2, include other invariance discovery methods, namely Inter-environmental Gradient Alignment (IGA) (Koyama & Yamaguchi] 2020) and AND-Mask (Parascandolo 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, 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
We note that in the majority of these experiment (13 out of 18) one of our proposed methods IRMy1A or IRMv2 outperforms all other methods.
F DETAILS OF THE
EXPERIMENTS
We summarize the details of each data set of the HurTthy in Table
Sample
Glogmphical Repion
Feng (2015)
145
Austrin
Colorecral Cancer, Advanoed Adenoma, Ohesiry, Overweight, Healthy
He (2017)
98
China
Crohns Disease, Obesity. Overweight, Underweight, Healthy
Je (2017)
281
China
Atheresclerotic Cardiovascular Disease, Obesity, Overweight. Underweight, Healthy
Karisson (2013)
Europe
Type Disbetes, Impuired Gluease Tolerance. Ohesity, Overweight, Underweight, Healthy
Le Chatelier (2013)
112
Denmark
Ohesity, Ovenweight, Underweight. Healthy
Nielsen (2014)
226
Europe
Crohns Disease, Ulceraive Colitis. Obesity, Overweight, Underweight, Healthy
(2012)
China
Overoeight. Underweight, Healthy Coloresral Cancer, Obesity. Overweight. Underweight. Healmh
Type Dinheres,
Vogemann (2016)
Zeller (2014)
196
Europe
Colorectal Cancer. Obesity, Overweight. Underweight, Healthy
Zhang (2015)
183
China
Rheumatoid Artritis, Ohesity, Overweight, Underwwiihh, Healihy
Table 7: Details of the HealthyGurTests data sets.
httte:////ihhbb..cc//ccosoteeeee wtwp:e/rwweereeeeeeiriiratteeeee skMi i zat
21

---

Under review as a conference paper at ICLR 2022
ANDMask
ERM
IGA
IRMv1A
IRMv1
IRMv2
Oracle
Example1.ED
0.14 0.03
2.01
0.27
4.68 + 0.72
0.07 = 0.01
0.13 0.01
0.05
0.00
0.05 0.00
Examplel.EI
11.56 +: 0.20
15.37 0.52
13.1
1.19
12.15 - 0.81
11.25
Example1.12 1.2
20.38 + 0.27
25.82 + 0.89
31.13 2.30
24.74 + 2.20
20.43
21.56 + 1.11
19.66 + 0.38
Examplets.E0
0.12 0.05
1.96 + 0.29
3,88 +:
0.07 = 0.01
0.13
0.06 0.01
0.05 0.00
ExampleIs.E1
13.68 1.00
15,47
18:21 0.55
14.35 1.62
11.72
13.30 1.87
11.03
Example
24.13 1.40
25.69 1.06
28.92 0.53
25.32 2.61
20.44
0.32
23.38 = 3.09
20.53 + 0.28
Exumple2. EO
0
0.01
0.3
0.30 0.09
0.43 0.00
0.09
0.00 - +.00
Example2. H1
0.01
0.00
0.50
0.05
Example2.E3
=
0.00
0.00
0.31 0.09
0.42 0.00
0.23 0.02
0.00 0.00
Example2sED E0
0.00
0.00
0.00
0.17 0.25
0.43
0.06
0.00 0.00
Example2s.E1
0.50 0.00
0.00
0.00
0.50
0.05
Example2s.E2
0.42 0.00
0.00
0.00
0.26 0.17
0.42 0.00
0.29
0,00 + 0.00
Example3.E ED
0.34 0.23
0.21
0.22
0.25 0.02
0.26
0.40
0.01 0.00
Example3.EI E.
0.36 0.21
0.14
0.17
0.02
0.37
0.30
0.00
Exumple3.E3
0.34 0.23
0.20
0.21
0.26 0.02
0.20
0.00
0,01 0.00
Example3s.H0 E0
0.28 0.17
0.35
+.2 0.21
0.35 0.21
0.34 - 0.11
0.36 0.20
0.41
0.01 0.00
Example3x.E1
0.47 0.06
0.41 + 0.14
0.40 0.15
0.34 0.11
0.38 0.19
0.41 0.01
0,01 0.00
Examplc3s.E2
0.38 0.17
0.36 0.20
0.36 0.20
0.34 0.11
0.36
0.41 0.01
0.01 0.00
Table 3: Test errors for all algorithms and examples with (divv, dapu: denv) = (5. 5, The errors for Examples and Is MSE and all others are classification error. The empirical mean and the standard deviation are computed using 10 independent experiments. An 's' indicates the scrambled variation of its corresponding problem setting.
ANDMask
ERM
IGA
IRMvIA
IRMv1
IRMv2
Oracle
Examplel
10.69 0.17
14.40 0.56
18.48 +: 1.46
12,87 +: 1111
10.72
0.13
11.26
10.32 1 0.13
Examplels
12.64 0.82
17.00 0.73
0.42
13.24 141
10.76 0.14
12.25
1.66
10.53 1 0.18
Example2
0.42 1442000 0.03
0,42
0.45
0.31 = 0.10
0.45 0.00
0.34 = 0.05
0.00 = 0.00
Example2s
0.45 = 0.00
0.5
0.00
0,27 1 0.18
0.45 0.00
0.38 0.06
0.00 - 0.00
Example3
0.34 0.22
0.18
0.26
0.26 0.02
0.36 0.19
0,40 0.01
0.00
Example3s
0.38 0.13
0.18
0.19
0.34 0.11
0.37
0.41
0.01 0.00
Table 4: Test errors averaged for each experiment.
ANDMask
ERM
IGA
IRMv1A
IRMv1
IRMv2
Oracle
Examplel.E
0.13 : 0.03
1.98 + 0.29
4.67
0.74
0.07 = 0.01
0.12
0.05 1 0.00
0.05 :: 0.00
Examplel.E1
10.38 : 0.43
7.80 0.51
0.75
1382 1.26
10.16
12.01 - 0.68
11.23
Example1.F2
18.45 + 0.63
12.27 + 1.02
14.13 1.39
24.59
2.26
18.12
21.36 + 0.92
20.19 0.13
Example1s.E0
0.12 0
1.91 2 0.27
0.12
0.07 = 0.01
0.12
0.06::0.01
0.05 0.00
Examplels.1
12.14
7.83 0.55
9.26
1385 1.37
1..1
12.87 = 1.81
11.27 0.02
Exumple1s.22
21.33
12.22 1.17
13:23
0.91 0.05
24.37 243
17.80 0.24
22.73 3.27
19.96 + 0.17
Example2, EO
0.05 2 0.01
0.04 0.00
0.00
0.12 = 0.06
0.05 - 0.00
0.36
0.00 + 0.00
Example2.EI
0.03 0.00
0.03 0.00
0.03
0.11
0.03
0.29
0.00 0.00
Example2.
0.01 0.00
0.00
0.00
0.10 0.07
0.00
0.05
0.00 0.00
Exumple2s.E0
0.05 0.00
0.000
0.05
0.05 = 0.03
0.05 - 0.00
0.31
0.00 + 0.00
Example2s.E1
0.03 0.00
0.00
0.03 0.00
0.04 0.03
0.03 0.00
0.12
0.00 0.00
Example2s
0.00
0.00
0.01 0.00
0.09
Example3.E EO
0.00 0.01
0.00
0.00 0.00
0.40
0.00
0.35
0.01 0.00
Example3.1I E1
0.00 + 0.00
0.00
0.00 0.00
0.49 0.00
0.00::0000
0.29
0.01 4 0.00
Example3.E3
Example3s.E0
0.18 0.09
0.00
0.00 0.00
0.49 = 0.00
0.00 0.00
0.00
0.01 + 0.00
Exumple3s.El
0.00 = 0.00
0.00
0.00 + 0.00
0.49 = 0.00
0.00 1 0.00
0.29 0.00
0.01 + 0.00
Example3s.F2
0.00
0.30
Table 5: Training errors for all algorithms and examples with (div, dapu denv) 5, The
errors for Examples and Is are MSE and all others are classification error. The empirical mean and the standard deviation are computed using 10 independent experiments. An 's' indicates the scrambled variation of its corresponding problem setting.
22

---

Under review as a conference paper at ICLR 2022
ANDMask
ERM
IGA
IRMvIA
IRMv1
IRMv2
Oracle
Examplel
9.65 0.36
7.35
0.61
9.62 + 0.96
12.82 1.17
9.47 + 0.12
11.14 0.53
10.49 0.11
Examplels
11.20 0.19
7.32 0.66
8.77 0.52
12.76 1.27
9.36
11.99 1.70
10.43
Example2
0.03 0.00
0.03 0.00
0.03=0000
0.11
0.03
0.26 = 0.08
0.00 = 0.00
Example2s
0.03 1: 0.00
0.03 0.00
0.03 1 0.00
0.04 0.03
0.0310000
0.22 1 0.14
0.00 1 0.00
Example3
0.00 0.00
0.00 0.00
0.00 0.00
0.00
0.49 0.00
0.00
0.31 = 0.04
0.01 = 0.00
Example3s
0.07 1: 0.04
0.00 0.00
0.00 1 0.00
0.49 1 0.00
0.00 1: 0.00
0.30 1 0.03
0.01 1 0.00
Table 6: Training emoor averaged for each experiment.
F.1 EXPERIMENT DETAILS
F.1.1 NEURAL NETWORK ARCHITECTURE AND HYPERPARAMETEA TUNING
The neural network considered consists of 5 layers with input layer dx x dR. hidden layers of size d2 x dn/2,, d,/2 x da/2,, 2, d, x /4, respectively, and the output layer da.4 dg. We use the ReLU action function and the dropout parameter of each layer is set p. We treat da, dg. p, the learning rate, and the invariance penalty coefficient as hyperparameters, which we optimize using the validation accuracy. More specifically, 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 T
For each data set in Table 2 that is considered the test 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 unified validation set for model selection.
treat each of the remaining data
G FULL DOMAINBED EXPERIMENTS
G.1 EXPERIMENT SUMMARY
DomainBed is an extensive framework released by Gulrajani & Lopez-Paz (020)) to test domain generalization algorithms for image classification tasks on various benchmark data sets. In series of experiments, Gulrajani & Lopez-Par Paz (2020) show that enabled by data augmentation various state- of-the-art art generalization methods perform similar each other and ERM on several benchmark data sets.
Although the integration of additional data sets and algorithms to DomainBed straightforward, we note that performing an extensive set of experiments requires significant computational resources as also pointed out by Krueger al. (2020). For this reason, we limit the scope of our experiments to the comparison of ERM. IRMv1, IRMv IA, and IRMv2.
Similar to Guirajani Lapez-Paz (2020), we observe that no method significantly outperforms others any of the benchmark data sets (see Table For complete set of results 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.
23

---

Under review as a conference paper at ICLR 2022
Algorithm
ColoredMNIST
RotatedMNIST
PACS
VL.CS
Avg
ERM
51.7 + 0.1
96.7 = 0.0
0.1
78.8
0.4
77.0
IRMv1
51.8 + 0.2
95.2 + 0.4
78.6
76.0 + 0.5
75.4
IRMv1A
50.9 + 0.1
64.7 20.1
80.9
77.3
0.2
68.4
IRMv2
50.8 + 0.4
97.1 = 0.0
82.6 0.9
76.5 + 0.4
76.8
Table 8: The test accuracy of ERM and different implementations of IRM on benchmark data-sets. Model selection of the DomainBed chosen trining-domann validation set.
Algorithm
+90%
+80 %
-90 %
Avg
ERM
72.8 +: 0.1
72.6 + 0.2
9.8 +: 0.0
51.7
IRMv1
72.5 0.3
72.9 0.1
9.9 - 0.1
0.1
51.8
IRMv1A
70.7 0.3
72.3
- 0.5
9.7
50.9
IRMv2
69.8 + 0.8
72.9 - 0.3
9.8 0.1
50.8
Table 9: Benchmark: ColoredMNIST, Model selection: training-domain validation set.
G.2 FULL DOMAINBED RESULTS
Algorithm
A
P
S
Avg
ERM
84.5 + 1.6
77.1 0.8
96.9 + 0.3
65.8 = 1.9
81.1
IRMv1
77.0 + 3.0
76.7 + 1.1
96.4 0.4
64.4 0.3
78.6
IRMv1A
82.6 + 0.5
77.7 0.7
96.6 0.4
66.7 = 0.5
80.9
IRMv2
86.0 + 0.9
76.6
0.7
96.9 + 0.0
70.8 2.0
82.6
Table 11: Benchmark: PACS, Model selection: training-domani validation set.
Algorithm
15
30
45
60
75
Avg
ERM
90.3 1.8
+
98.2
0.1
98.2
0.1
1
93.5 0.4
96.0
IRMv1
89.6 2.1
96.0 0.5
97.9 0.1
97.2 0.0
97.0 0.2
90.9 - 0.5
94.8
IRMv1A
75.9 + 5.9
64.2 - 22,4
59.4 19
26.1
59.8 + 26.3
59.0 + 26.7
55.9 + 22.5
62.4
IRMv2
94.1 1 0.0
98.1 1 0.3
98.5 1 0.1
98.2
1 0.1
98.3 1 0.0
94.4 1 0.3
97.0
Table 15: Benchmark: RotatedMNIST, Model selection: leave-one-domain-out cross-validation.
Algorithm
A
C
P
S
Avg
ERM
79.7 0.6
73.0 3.8
97.1 0.9
62.1 1.5
78.0
IRMv1
67.4 6.7
72.3
+ 4.3
5.7
64.1 4.5
72.9
IRMvLA
78.8 3.2
78.9 0.9
96.0
42.3
16.3
74.0
IRMv2
86.3 + 0.3
76.8 + 0.5
97.0 + 0.4
69.7 + 2.6
82.5
Table 16: Benchmark: PACS, Model selection: leave-one-domain-ouuuoroo cwt.
24

---

Under review as a conference paper at ICLR 2022
Algorithm
30
45
Avg
ERM
93.1 + 0.1
97.8 0.0
98.4 0.0
98.3 - 0.1
08.2
94.3 0.1
96.7
IRMv1
89.6 + 2.1
96.8 0.1
97.9 0.1
97.8 + 0.1
97.5
91.6
95.2
IRMvIA
75.9 1: 5.9
71.1
17.7
60.8
+ 25.1
25.8
60.2 25.8
59.8 20.5
64.7
IRMv2
94.1 = 0.0
98.2
0.0
98.5 = 0.1
98.4
- 0.1
98.3 0.0
95.1 + 0.2
97.1
Table 10: Benchmark: RotatedMNIST, Model selection: training-domain validation set.
Algorithm
C
S
V
Avg
ERM
97.4 + 0.1
65.0 + 0.9
74.3 + 1.1
78.7 + 0.1
78.8
IRM
96.3 = 0.6
61.7
70.1
76.0 = 1.8
76.0
IRMA
96.9 0.8
64.8
7.
77.0 = 0.4
77.3
IRMv2
96.6 - 1.1
6..4
73.5
70.6 = 2.4
76.5
Table 12: Benchmark: VLCS, Model selection: training-domain validation set.
Algorithm
15
30
45
60
75
Avg
ERM
92.5 + 0.6
97.8 + 0,0
97.9 + 0.1
97.9 + 0.2
0.2
98.3 + 0.1
94.3 + 0.1
96.4
IRMv1
89.6 - 2.1
96.8 + 0.1
98.0 + 0.2
97.5
0.3 0.3
97.5 1 0.1
91.6 0.0
95.2
IRMvIA
77.9 = 7.4
71.1 = 17.7
59.4 26.1
59.8
26.3
59.0 26.7
63.9
IRMv2
94.7 0.4
98.0
0.2
98.5 0.1
98.3 0.0
98.3 0.0
95.0 0.2
97.2
Table 20: Benchmark: RotatedMNIST, Model selection: test-domain validation set (oracle).
Algorithm
A
C
P
S
Avg
ERM
83.7 + 0.5
82.1 = 0.2
97.5 + 0.2
69.1
83.1
IRMv1
66.7 + 4.3
6..5
8.3
6.07
72.5
IRMv1A
83.5 0.2
75.7
3.2
96.4 + 0.3
68.6
1.8
81.0
IRMv2
84.3 + 0.3
76.5 + 0.7
96.8 + 0.1
70.3
2.4
82.0
Table 21: Benchmark: PACS, Model selection: test-domain validation set (oracle).
Algorithm
C
L
S
V
Avg
ERM
98.4 + 0.1
65.1 + 1.4
72.9 + 2.2
77.1 + 1.7
78.4
IRM
97.6 1.2
0.6
62.9
= 1.3
73.0 = 0.3
73.9
IRMA
98.0 0.1
6.19
71.8
= 0.8
73.9 = 1.5
77.1
IRMv2
96.3 + 1.0
6.11
70.9
= 1.3
71.9 1.5
76.5
Table 22: Benchmark: VLCS, Model selection: test-domain validation set (oracle).
25

---

Under review as a conference paper at ICLR 2022
Algorithm
ColoredMNNIST
RotatedMNISTT
PACS
VLCS
Avg
ERM
51.7 + 0.1
96.7 +: 0.0
81.1 +
78.8 + 0.4
77.0
IRMv1
51.8 + 0.2
95.2 - 0.4
1.0
76.0 - 0.5
75.4
IRMvIA
50.9 + 0.1
64.7 + 20.1
8..
77.3
- 0.2
68.4
IRMv2
50.8 + 0.4
97.1 = 0.0
82.6 12.60.99 0.9
76.5
- 0.4
76.8
Table 13: Model selection: training-domain validation set.
Algorithm
+90%
+80%
-90%
Avg
ERM
30.4 = 13.4
50.5 0.6
9.9
30.2
IRMv1
50.1 0.4
60.6 =
40.0 14.1
46.9
IRMv1A
69.5 = 14.4
49.8 0.2
10.0
43.1
IRMv2
10.0 + 0.1
36.4 + 3.0
9.9
18.8
Table 14: Benchmark: ColoredMNIST, Model selection: leave-one-domann-out cross-validation.
Algorithm
C
L
S
Avg
ERM
97.5 0.8
60.3 + 3.2
70.1
+ 4.2
75.5
75.9
IRM
93.3 + 2.7
61.8
0.4
72.9
+ 0.9
74.1
3.1
75.5
IRMA
79.2 + 12.8
66.8 - 1.4
68.3
4.1
73.5
1.3
71.9
IRMv2
98.2 = 0.1
6.1
7.
69.9
- 0.2
76.4
Table 17: Benchmark: VLCS, Model selection: leave-one-oominnouu caosiiaiiiioooooooiiaaaiioo.....
Algorithm
ColoredMNIST
RotatedMNIST
PACS
VLCCS
Avg
ERM
30.2 + 4.3
96.0 = 0,4
78,0
75.9 + 0.7
70.0
IRMv1
46.9 + 2.1
94.8 0.2
72.9 3.0
75.5
72.5
IRMvIA
43.1 + 4.8
62.4 21.6
74.0
71.9 + 2.8
62.8
IRMv2
18.8 + 1.1
97.0 + 0.1
82.5 + 1.0
76.4 1 0.3
68.7
Table 18: Model selection: leave--ne-domaii---ut-oooooo---- www.:/aaianiaoooooooooosoaaiiiia
Algorithm
+90 %
+80 %
-90% %
Avg
ERM
72.2
0.3
72.6 + 0.4
12.4
+ 1.1
52.4
IRMvl
72.5 + 0.3
72.9
+ 0.1
63.3
+ 6.6
69.6
IRMvLA
71.4 + 0.2
72.8
- 0.3
50.2
0.1
64.8
IRMv2
70.6 + 1.2
71.9
+ 0.6
20.9
+ 0.9
54.4
Table 19: Benchmark: ColoredMNIST, Model selection: test-domain validation set (oracle).
26

---

Under review as a conference paper at ICLR 2022
1e25
1e20
1e15
Test Accuracy
Condition Number
0.4
1e10
1e5
1e0
2
8
16
32
64
128
256
512
1024
2048
(a) IRMv1
1e25
1e20
1e15
1e10
Test Accuracy Condition Number
1e5
1e0
16
32
64
128
256
512
1024
2048
(b) IRMvIA
1620
0.65
0.60
1e15
0.55
0.50
Test Accuracy Condition Number
0.45
1e5
0.40
1e0
16
512 1024 2048
(c) IRMv2
Figure 2: The figure depicts the evolution of test accuracy and the condition number of I(yp) in relation to the dimension of the representation.
27


Reviews:
[{"review_id": "lvtgb_y0Ozn", "reviewer": "Reviewer_2SD7", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper is well written but the innovation is insufficient. The form of proposed invariant penalty is similar to the one of invariant penalty proposed in [1] and limited to the least-square loss. Similar theoretical result was obtained in [2] when  in the linear case. Moreover, considering that there have been many alternatives of IRM in recent literature, a detailed explanation should be given to fully illustrate the strength and weakness of the proposed method compared to the existing alternatives.", "summary_of_the_paper": "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.", "summary_of_the_review": "Specific Comments\n1）The similar theoretical result in this paper has been obtained in [2]. In other words, under the same condition, the theoretical result also holds for IRM. It confuses me whether the proposed method is essentially different from IRM.\n2）It seems to me that the proposed invariance penalty doesn’t solve the problem of IRM. This paper doesn’t guarantee that the proposed IRMv2 can recover invariant predictor at least in the theoretical scenario where IRM fails.\n3）I suggest that a detailed comparison with other alternatives for IRM should be given to better illustrate the strength and weakness of the newly proposed invariant penalty.\n4）There are a few places that are not very rigorous. (1) In 5-th line of Algorithm 1,  is computed according to Eq. equation 4. However, the distribution is usually unknown and expectation can’t be computed. (2) In 6-th line of Algorithm1,  is computed according to Eq. equation 12 instead of Eq. equation 11.\n\n[1] Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint arXiv:1907.02893, 2019.\n[2] Elan Rosenfeld, Pradeep Kumar Ravikumar, and Andrej Risteski. The risks of invariant risk minimization. In International Conference on Learning Representations, 2021. URL https://openreview.net/forum?id=BbNIbVPJ-42.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "TG2XR4bW4mU", "reviewer": "Reviewer_AZJB", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "a)\tStrength: \ni.\tThis paper analyzes the counter-example of IRM and proposes a new invariance penalty to solve the OOD Generalization problem.\nii.\tPractical insights on avoiding the failure case are provided for a better understanding of the counter-examples.\nb)\tWeaknesses:\ni.\tThe proposed method is restricted to mean squared loss and cannot be extended to more general cases.\nii.\tAlthough the role of I_e is demonstrated in Section 4.3, why the newly-proposed penalty is necessary and good is still unclear. And the derivation of the invariance penalty seems ad-hoc.\niii.\tThere lack of enough baselines in experiments. The authors only compare with ERM and IRM, while there are many other invariant learning methods, such as Rex[1], EIIL[2], HRM[3], etc.\niv.\tExperimental results are not good enough to validate the effectiveness of the proposed method.\n\n[1] D. Krueger, E. Caballero, J.-H. Jacobsen, A. Zhang, J. Binas, D. Zhang, R. Le Priol, and A. Courville, “Out-of-distribution generalization via risk extrapolation (rex),” ICML2021.\n[2] E. Creager, J.-H. Jacobsen, and R. Zemel, “Environment inference for invariant learning,”ICML2021.\n[3] J. Liu, Z. Hu, P. Cui, B. Li, and Z. Shen, “Heterogeneous risk minimization,”ICML 2021.\n", "summary_of_the_paper": "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.", "summary_of_the_review": "This paper analyzes the counter-example of IRM and proposes a new invariance penalty based on that. Although some intuitions are provided, it is unclear why the proposed invariance penalty is necessary. And baselines, as well as the empirical results, are not enough to support the method.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "T4kFeiCVdBq", "reviewer": "Reviewer_256p", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "In an earlier version of this review, I wrote that I felt that the objective proposed by this work is very similar to an existing linear regression objective from 2018, known as Anchor Regression [1]. The present paper proposes to consider linear regression *only*, which of course makes for an easier problem, but it also means there is lots of existing work in this area (for which there is insufficient discussion in the related work).\n\nHowever, after further thought, I've realized that the two are only equal in quite specific circumstances. On the other hand, I have come across an observation that I think demonstrates the extent to which this work does not analyze their objective closely enough.\n\nConsider the objective proposed by the paper:\n$$\\sum_e R_e(w^T\\varphi) + \\lambda ||I_e(\\varphi)^{1/2} (w-w_e^*)||^2.$$\nThe first term is simply ERM. The second term, given features $X = \\varphi(x)$, can be written $(w-w_e^*)^TX^TX(w-w_e^*)$. Consider what happens when we take the gradient of this objective with respect to our predictor $w$. We arrive at $2X^TX(w-w_e^*)$. Observe that this is precisely the gradient of linear regression, if we are assuming that our target is $Xw_e^*$ instead of the given target $y$. In other words, this is just linear regression where the targets are provided by a proxy classifier, as in two-stage least-squares. This tells us two things:\n\n1. The objective is simply a weighted sum of the risk when predicting the true target $y$ and the \"instrumental variable\" target $Xw_e^*$ (equivalently, $w_e^{*T}\\varphi(x)$). I believe this is nowhere mentioned in the paper and is worth exploring, if only because there doesn't seem to be any real understanding of what is going on with this objective.\n2. I think my original impression that this was similar to Anchor Regression now appears pretty close to accurate, since Anchor Regression can also be seen as a trading off between ERM and a two-stage regression with instrumental variables.\n\nBoth of these points indicate to me that there is indeed something interesting going on with this objective, but this paper does not adequately explore the implications; nor does it properly discuss the previous related work (this paper takes existing work on logistic regression and applies it to linear regression, but there is a great deal of existing work in the causal literature on linear regression which is not discussed).\n\nReturning to the paper itself, I find the empirical and theoretical results quite weak and I think this submission needs work even if it was a totally new idea. If there were some interesting new technical results to provide justification for this method in a different setting (such as classification) or under a different/more general model, it would certainly be noteworthy. However, I do not believe this to be the case: the setup Lemmas are simple linear algebra, and the primary theorem is almost identical to the related theorem in Rosenfeld et al. (furthermore, it only shows that it has the *same* linear environment complexity as IRM, rather than improving on it as the paper claims). The only new thing here is learning the feature embedder $\\varphi$, but this paper provides no additional justification for when/why this should work---that is, all of the arguments for the benefits of this approach are agnostic to the actual quality of $\\varphi$.\n\nI also don't buy the claim that this would work in the non-linear setting; the result of Rosenfeld et al. is not specific to the type of loss, just the non-linearity of $\\varphi$. The point behind that proof was that given only observations from the training environment, one can construct regions which contain all training points and on which the featurizer is perfect, and then make the featurizer arbitrarily bad elsewhere. In the absolute worst case, imagine a $\\varphi = \\varphi^*$ on all training points, and $\\varphi = 0$ elsewhere. It seems clear to me that whatever claim the authors are making here with regards to the OOD performance of their predictor is not correct or is stated unclearly, because this solution would behave just like $\\varphi^*$ at train time and give arbitrary error at test time. Even when observing the full environmental distributions, the fact that as we move away from the distribution mean the measure decays exponentially but the loss only grows quadratically means the same result would hold. I encourage the authors to revisit this section and work to clarify what precise mathematical claim they are making, and what this actually means in English.\n\nFinally, *what is it with bolding empirical results which do not meet the bar for statistically significant improvement?* In Table 2, the confidence intervals of the performances of **all compared methods** overlap. Further, why aren't the standard deviations of the individual test accuracies reported over the five runs? There is zero indication that this method outperforms the ones to which it is being compared.\n\n[1] Anchor regression: heterogeneous data meets causality. Dominik Rothenhäusler, Nicolai Meinshausen, Peter Bühlmann, Jonas Peters.\n\n\n-------------------- Update after rebuttal --------------------\n\nAfter discussion with the authors, I believe there are interesting aspects of this algorithm that are promising and worth exploring further. However, I still don't think the proposed method is analyzed or understood well enough, nor does there appear to be any substantial empirical or theoretical improvement here over existing works. I'm leaving my recommendation as is.", "summary_of_the_paper": "The paper proposes a small modification to the IRMv1 objective which leads to an improvement over the original. They propose potential theoretical explanations for the improvement.", "summary_of_the_review": "The idea behind this paper, while nice, is a repackaging of IV regression, and does not adequately analyze the implications of the objective (nor do they properly contextualize it within the great deal of similar related work in other fields). They provide no new theoretical justification (in fact their main technical result, very similar to another work, simply shows that this algorithm performs on par with the existing IRMv1). Further, their explication of performance in the non-linear setting is incomplete.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "GrLumwVQInU", "reviewer": "Reviewer_PkSM", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## Major points\n\n1. IRM penalties are most often studied alongside representations learnt by deep learning models, which are most often trained using some variant of stochastic gradient descent. Concretely, Arjovsky et al. 2019 proposed an unbiased stochastic estimator of the gradient of the original IRM penalty. However, to the best of my knowledge, the proposed approach as presented would in principle require computing the gradient of the ERM regressor $w^{*}(\\Phi) = (\\Phi^{T}\\Phi)^{-1}\\Phi^{T}y$ with respect to $\\Phi$. Similarly, the adaptive rule to adjust the penalty weight for IRMv1 presented in Section 3.1. would in principle require estimating the smallest eigenvalue of $\\Phi_{e}^{T}\\Phi_{e}$, which might be difficult in a stochastic setting.\n\n    I believe it would be ideal if the authors could address this limitations either theoretically (e.g. by proposing unbiased gradient estimators) or empirically (e.g. by proposing heuristic, possibly-biased gradient estimators and demonstrating good empirical performance nevertheless). At the very least, the issue of how to estimate the gradient of $L_{t}(\\phi_{\\theta_{t}})$ in Line 7 of Algorithm 1 in practice should be described in more detail. If it turns out that the proposed approach can only be used with (non-stochastic) gradient descent, this should at least be explicitly stated in the manuscript to ease the way for future work.\n\n2. The proposed approach can only be used alongside the MSE loss. While this is acknowledged by the authors, which is appreciated, I do believe this to be an important practical limitation, as this loss is rarely used for supervised classification. Moreover, this makes the soundness of the adaptive rule in Section 3.1. questionable whenever IRMv1 is used in conjunction with e.g. the logistic loss. \n\n3. The adaptive rule in Section 3.1. could be better justified. While the bound in Lemma 3 is to the best of my knowledge sound for the MSE loss, it is not explicitly stated how that motivates Eq. (15) exactly and what the theoretical implications are, if any. This is in my opinion important as it could be argued this adaptive rule is the contribution that has the most practical impact in performance according to the experiments in Section 5.\n\n4. The authors make a point that there might be a link between ill-conditioning of the Gram matrix $\\Phi_{e}^{T}\\Phi_{e}$ and failure modes of IRM penalties, providing theoretically that some of the most well-known counterexamples to these penalties do exhibit ill-conditioning. However, the empirical results in Figure 2 in the appendix are somewhat conflicting. In particular, it would appear that the \"phase transition\" in which test accuracy sharply drops precedes the sharp increase in ill-conditioning particularly for IRMv1 and IRMv1A.\n\n5. All in all, the experimental results do not convey a clear benefit of using IRMv2 vs IRMv1. The HealthyGutTests results in Table 2 are, as far as I can tell, well within the margin of error. Most importantly, it would appear IRMv1 (and the IRMv1A variant) dominate IRMv2 in the LinearUnitTests experiments.\n\n## Typos\n\nProofs in the appendix often refer to the $\\mathcal{l} - 2$ ball instead of $\\mathcal{l}_{2}$.", "summary_of_the_paper": "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,\n\n$$\\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},$$\n\nwhere $\\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:\n\n$$\\rho_{e}^{\\mathrm{IRMv2}}(\\Phi, w) := \\vert\\vert (\\Phi_{e}^{T}\\Phi_{e})^{1/2}(w -w_{e}^{*})\\vert\\vert_{2}^{2},$$\n\nwhich 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)$.\n\nThe 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.", "summary_of_the_review": "Out-of-distribution generalization is one of the most pressing, largely open problems in the field. Concretely, the IRM framework has recently emerged as a promising line of research in this direction, but its characterization is far from conclusive both from a theoretical and a practical standpoint. Thus, I consider the paper's topic timely and of great relevance to the field.\n\nI believe the manuscript  explores an interesting idea and is technically sound, albeit not without what I consider to be important limitations, namely,\n+ The proposed approach is only applicable to the MSE loss.\n+ It is unclear whether it would be sound to use the proposed penalty in a stochastic gradient-based optimization setting, as it requires backpropagating through the optimal linear regressor, which depends on all data.\n+ The experimental results do not provide compelling evidence that the proposed penalty is statistically significantly superior to the original formulation.\n\nAll in all, I believe IRM to be a promising yet not fully mature idea with many crucial open questions remaining to be answered. Thus, I consider exploratory ideas such as those presented in this article to be worthwhile despite potential shortcomings that might need to be revisited by future work. Hence, despite the aforementioned limitations, I lean towards supporting acceptance of the manuscript. \n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This manuscript was the object of a rich and lengthy discussion. The AC also felt compelled to read the paper in details and discussed it further with the SAC.

The authors did a thorough job at addressing some of the reviewers points. The added results on cross-entropy loss and additional discussion, as well as the points made in "Further Discussion on the Numerical Experiments" are very much appreciated. 

However, significant concerns remain on establishing connections with prior work, including related ideas on invariance from the causality literature, so as to gain deeper understanding of the implications of the proposed objective. We also strongly encourage the authors to further work on strengthening their theoretical analysis to clearly demonstrate the value of the proposed approach.

The proposed formulation is certainly thought provoking and we urge the authors to pursue their work in view of the above comments.
Decision: Reject