O N I NVARIANCE P ENALTIES FOR R ISK M INIMIZATION
A BSTRACT
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.
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
I NTRODUCTION
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.
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.
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.
B ACKGROUND : I NVARIANT P REDICTION
In this paper, we consider data ( X e , Y e ) collected from multiple training environments e ∈ E tr where the distribution of ( X i , Y i ) and ( X j , Y j ) may be different for i ̸ =
j with i, j ∈ E tr .
We denote by R e 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 R e ( f ) =
E X e ,Y e [ ℓ ( f ( X e ) , Y e )] .
2. I NVARIANT R ISK M INIMIZATION
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 ∈ argmin ˜ w : H→Y R e ( ˜ w ◦ φ ) for all e ∈ E .
To find such invariant predictors, Arjovsky et al.
(2019) introduce the notion of the Invariant Risk
Minimization (IRM) principle:
[Equation]
e ∈E tr
[Equation]
) subject to w ∈ argmin ˜ w : H→Y R e ( ˜ w ◦ φ ) , ∀ e ∈ E tr .
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.
2. IRM V 1: A R ELAXATION OF IRM
In order to provide an implementation of IRM, Arjovsky et al.
(2019) restrict the classifier w to linear functions, i.e.,
[Equation]
subject to
[Equation]
argmin
[Equation]
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 I e ( φ ) as
[Equation]
[Equation]
Assuming that I e ( φ ) is full rank for a fixed φ , its respective optimal classifier is unique, i.e., argmin [Equation]
=
[Equation] ⋆e ( φ ) where
[Equation]
[Equation] Hence, in this setting, the invariant constraint of IRM in equation 2 can be simplified to w = w ⋆e ( φ ) for all e ∈ E tr .
A natural relaxation of the constraint w − w ⋆e ( φ ) = 0 to a penalty is ∥ w − w ⋆e ( φ ) ∥ 2 .
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 I e ( φ ) − 1 , which appears in the computation of w ⋆e ( φ ) could improve the behavior of the invariance penalty.
That is, considering ∥I e ( φ )( w − w ⋆e ( φ ))
∥ 2 as the invariance penalty.
Moreover, for the squared loss, one can show that
[Equation]
Hence, Arjovsky et al.
(2019) propose the following invariance penalty
[Equation]
Using the penalty equation 6, the relaxation of IRM is given by
[Equation]
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 ◦ ψ −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 φ ∈ R 1 × d x .
Their relaxation of IRM, which they refer to by IRMv1 is given by
[Equation]
Although equation 5 only holds for squared loss, Arjovsky et al.
(2019, Theorem 4) justify the choice of ∥∇ w | w =1 .
R e ( 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 ⊤ Φ) ⊤ ∇ w R ( w ⊤ Φ)
=
if and only if w is optimal for all environments.
Although the penalty ρ IRMv1 e ( φ, 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 ρ IRMv1 e ( φ, 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.
IRM V 2: A N A LTERNATIVE P ENALTY
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 ∈ R d φ and w ⋆e ( φ ) as defined in equation 4.
Then,
[Equation]
Based on Lemma
[Equation]
we propose an invariance penalty that is directly comparable to risk.
[Equation]
The relaxation of IRM using the penalty equation 9 is then given by
[Equation]
We further simplify the relaxation equation 10 by finding its optimal classifier for a fixed data representation defined as
[Equation]
In the following Lemma
[Equation]
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 equationand equation 11, respectively.
Then,
[Equation]
Moreover, it holds that
[Equation]
Based on Lemmas 1 and 2, we propose the following relaxation of IRM, which we refer to by
IRMv2.
[Equation]
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 Lemmareveals 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, I e ( φ ) is incorporated differently in the invariance penalty of IRMv1 and IRMv2.
We formalize this latter observation in the Section 4.3.1.
3. IRM V 1A: A N A DAPTIVE P ENALTY C OEFFICIENT
We first bound the invariance penalty of IRMv1 in terms of the penalty of IRMv2 and the eigenvalues of I e ( φ ) .
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 ρ IRMv1 e ( φ, w ) and ρ IRMv2 e ( φ, w ) be the invariance penalties of the IRMv1 and
IRMv2 defined in Equations equation 6 and equation 9, respectively.
Then,
[Equation]
The proof of Lemma 3 directly follows from the definition of the invariance penalties ρ IRMv1 e ( φ
[Equation]
w [Equation] and ρ IRMv2 e ( φ, w ) , and the fact that for a symmetric matrix A ∈ R d × d and a vector u ∈ R d , 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 w 1 , w 2 ∈ R d φ it holds that [Equation]
Using Lemmas 3 and 4, we suggest the following rule for the penalty coefficient of IRMv1.
[Equation]
15) where λ 0 ≥ 0 is a user-specified parameter.
The role of λ 0 is to avoid numerical instability when λ min ( I e ( φ )) is small.
Note that this is an adaptive rule, as φ may change throughout training.
T HEORETICAL R ESULTS AND A NALYSIS
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.
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 I e ( φ ) , in particular, in relation to counterexamples of Arjovsky et al.
and Rosenfeld et al.
(2021), and provide practical solution to alleviate the possibility of such failure modes.
4. SEM FOR C LASSIFICATION
For each environment e , data ( X e , Y e ) is generated according to the following SEM:
[Equation]
where η ∈ [0 , 1] , and S ∈ R d × ( d e + d c ) is a left invertible matrix, i.e., there exists S † such that
S † S
= I .
In this model, Z c captures the causal variables that are invariant across environments, and
Z e captures the spurious environment dependent variables.
The variables Z c and Z e are generated as follows
[Equation]
Here, µ c ∈ R d c , µ e ∈ R d e , and N ( µ, Σ) denotes multi-variate Gaussian distribution with mean µ and covariance Σ .
We further assume that W c , W e , and Y e are independent for all environments.
4. I NVARIANT R EPRESENTATION UNDER IRM
For the setting introduced in 4.1, the invariant data representation is linear.
In particular, for any d ≥ d c , φ ( X e ) =
[Equation]
c   (cid:21)
S † X e
=
Z c 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 |E tr | > d e .
As span ( { µ e } e ∈E tr ) ≤ d e , for each e ∈ E tr there exists a set of coefficients α ei for
i ∈ E tr \ e such that
[Equation]
We say that
E tr is a non-degenerate set of environments if for all e ∈ E tr it holds that
[Equation]
where Γ e is defined as
[Equation]
The conditions of equation 20 and equation 21 specify that the span of covariance matrices of Z e ’s is
R d e .
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 Z e as in equation 21, their condition relies only on the variances σ 2 e .
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 |E tr | > d e where ( X e , Y e ) generated according to equation 16.
Consider a linear data representation Φ X =
AZ c
+ BZ e and a classifier w (Φ) on top of Φ that is invariant, i.e., w (Φ) =
w ⋆e (Φ) for all e ∈ E tr .
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. T HE R OLE OF I e ( φ )
Recall that the main difference between the invariance penalty of IRMv1 (with squared loss) and
IRMv2 is the way that I e ( φ ) is incorporated.
Although at a first glance this may not seem significant, in what follows we demonstrate the importance of I e ( φ ) .
4.3. T HE P OTENTIAL F AILURE OF I NVARIANCE P ENALIZATION
Arjovsky et al.
(2019) consider a linear Structural Equation Model (SEM) of the following form.
X 1 ∼ N (0 , σ 2 ) , Y =
X 1 + Z 1 with Z 1 ∼ N (0 , σ 2 ) , X 2 = Y + Z 2 with Z 2 ∼ (0 , 1) , and X =
[ X 1 , X 2 ] ⊤ .
The variance σ 2 is the only parameter that changes across environments, and we assume that Z 1 , Z 2 , and X 1 are independent in and across all environments.
Hence, in this setting, X 1 models the invariant features and X 2 models the environmental features as the correlation between X 2 and Y varies across environments.
Consider the data representation φ c ( x ) parameterized by a variable c ∈ R as
[Equation]
where c = 0 attains the invariant data representation, i.e., φ 0 ( X ) =
[ X 1 , 0] ⊤ .
Based on this model,
Arjovsky et al.
(2019) argue that ∥ w inv − 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,
I e ( φ c ) is ill-conditioned for both small and large c ’s. More precisely, in Appendix [-C.
-]D.1 we show that
[Equation]
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
[Equation]
where { Z e / ∈ Z ϵ } is an event with P ( Z e ∈ Z ϵ ) ≤ p e,ϵ where p e,ϵ := exp( − d e min { ϵ − 1 , ( ϵ − 1) 2 } / 8) .
The parameter ϵ ≥ 1 control the degrees to which the environmental variable Z e 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 ϵ →∞ φ ϵ ( X e ) =
φ inv ( X e ) :=
[ Z c , 0] ⊤ .
Rosenfeld et al.
(2021) show that IRMv1’s invariance penalty ∥∇ w R e ( w ⊤ φ ) ∥
= O ( p 2 e,ϵ ) , and can be made arbitrarily small, hence it is poor discrepancy as an invariance penalty.
In Appendix D.2, we show that for squared loss, it holds that | R e ( w ⋆ ( φ ϵ ) ⊤ φ ϵ )
− R e ( w ⋆ ( φ inv ) ⊤ φ inv )
| =
O ( p 2 e,ϵ ) .
In other words, although the invariance penalty is arbitrarily small, the risk is also arbitrarily close to the risk of the invariant predictor.
Moreover, we show that the matrix I e ( φ c ) is also ill-conditioned.
More specifically, we prove that lim ϵ →∞ κ ( I e ( φ ϵ ))
= ∞ .
4.3. P RACTICAL I MPLICATIONS
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.
[Equation]
[Figure]
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 ( φ )) > 10 12 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., d X = 196 .
E XPERIMENTS
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.
5. I NVARIANCE U NIT T ESTS
In this section, we evaluate the efficacy of our proposed approaches for invariance discovery on the
InvarianceUnitTests recently proposed by Aubin et al.
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 n e = 10 4 .
For all problems, the input x e ∈ R d is constructed as x e = ( x e inv , x e spu ) where x e inv ∈ R d inv and x e
spu ∈ R
d spu 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 e 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.
[Table]
1 and 1s are in 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.
5. H EALTHY G UT T ESTS
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 n e = 10 and inputs x e 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 [Table]
[Table]
C ONCLUSION
In this paper
[Table]
we have presented IRMv
[Table]
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.
R EPRODUCABILITY S TATEMENT
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.
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
B C ONSIDERATIONS FOR P RACTICAL I MPLEMENTATION
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, (cid:98) w ⋆e the empirical estimate of the optimal classifier w ⋆e is given by
[Equation]
(24) where Φ ek ∈ R m ×
d φ is a matrix each of its row is φ ( x ei ) ⊤ of the k -th mini-batch of environment e , and with some abuse of notation Y ek ∈ R m is defined as the vector containing y ie 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 ( X e )]
= w ⊤ inv φ inv ( X e ) with the error Y e − E [ Y e | φ inv ( X e )] being independent of X e and Y e , e.g., y ie = w ⊤ inv φ inv ( x ei ) + η 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 B s. Then,
[Equation]
The estimator equation 25 is unbiased if x ei are i.i.d. and y ei is independent of all other x ej conditioned on x ei . Numerical Stability: Although (cid:80) k ∈ B Φ e k ⊤ Φ e k 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 I e ( φ )
[Equation]
where λ ≥ 0 .
That is, to improve the numerical stability, we consider the ℓ 2 regularized least squares estimate
[Equation]
Then,
[Equation] We now provide an empirical estimate of the invariant penalty ρ IRMv2 e . We further partition B into B 1 , B 2 , and B 3 .
Then,
[Equation]
Note that (cid:98) ρ IRMv2 e is an unbiased estimator of ρ IRMv2 e for λ = 0 . C M ATHEMATICAL P ROOFS C. P ROOF OF L EMMA First, the risk under environment e with w = w ⋆e ( φ ) is given by
[Equation]
Then,
[Equation]
C. P ROOF OF L EMMA
Recall the definition of w ⋆ ( φ )
[Equation]
That is,
[Equation]
(cid:98) w ⋆e
[Equation]
φ
[Equation]
[Equation]
w φ (
[Equation]
Note that the objective function is a convex quadratic function of w .
Hence, using the first-order optimality condition we have that [Equation] [Equation] ⋆e ( φ
[Equation]
Then, w
[Equation]
φ ) =
[Equation]
φ
[Equation] ( φ )) (cid:33) .
Recall that w ⋆e ( φ ) =
I e ( φ ) − 1 [ φ ( X e )
Y e ] .
Then, [ φ ( X e ) Y e ] = I e ( φ ) w ⋆e ( φ ) .
Thus, w ⋆ ( φ ) = (cid:32) (cid:88)
e ∈E tr
I e ( φ ) (cid:33) − 1 (cid:32) (cid:88)
e ∈E tr
I e ( φ ) w ⋆e ( φ ) (cid:33) .
Finally, using a similar argument, we get argmin w
[Equation]
I e ( φ ) (cid:33)
[Equation]
1 (cid:32) (cid:88)
e ∈E tr
I e ( φ
[Equation] ( φ [Equation] (cid:33) .
C. P ROOF OF L EMMA
For a symmetric matrix A ∈ R d × d and a vector u ∈ R d , it holds that λ min ( A ) ∥ u ∥ 2 ≤ u ⊤ Au ≤ λ max ( A ) ∥ u ∥ 2 .
Let u
= I e ( φ ) 1 / 2 ( w − w ⋆e ( φ )) and A =
I e ( φ ) .
Then,
[Equation]
C. P ROOF OF L EMMA
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 w 1 and w 2 .
Let g ( w ) :=
log(1+exp( w ⊤ φ ( X ))) .
Then, using the multivariate mean-value theorem, there exists c ∈ (0 , 1) such that
[Equation]
We have that
[Equation]
where σ ( · ) denotes the sigmoid function.
Let ¯ w :=
(1 − c ) w 1 + cw 2 .
Then,
[Equation]
where the second equality follows from the application of mean-value theorem.
Then,
[Equation]
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.
C. P ROOF OF T HEOREM
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 (Φ) =
argmin w R e ( w ⊤ Φ) where R e ( w ⊤ Φ) is defined as
[Equation]
Then,
[Equation]
, where
[Equation]
[Equation]
Here, ¯Σ e and ¯ µ e are defined as
[Equation]
Proof of Lemma 5: We have that
[Equation]
First, we have that
[Equation]
=
[Equation]
(1 − η
[Equation]
ηµ e (1 − η ) µ e (cid:21) =
[ η ¯ µ e (1 − η )¯ µ e ] .
Thus, w ⋆e (Φ) =
[ ηβ ⋆e (Φ) (1 − η ) β
⋆e (Φ)] , (28) where β ⋆ e (Φ) =
I e (Φ) − 1 ¯ µ e . We now compute I e (Φ) − 1 in terms of ¯Σ − 1 e and ¯ µ e .
We have that
[Equation]
From the definition of Z c and Z e , it follows that
[Equation]
Thus,
I
[Equation]
Using Sherman-Morrison formula, we have that
[Equation]
Finally, using Equation equation 29 it follows that
[Equation]
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 φ ( X e ) =
AZ c + BZ e (or Φ S
=
[ A B ] ) is without loss of generality under the assumption of left-invertibilibity of S .
Then,
[Equation] [Equation]
Recall from Lemma {+
+} that β
⋆e (Φ) =
¯Σ − 1 e ¯ µ e / (1 + ¯ µ ⊤ e ¯Σ − 1 e ¯ µ e ) .
If β ⋆e (Φ) is invariant, then β ⋆ =
β ⋆e (Φ) for all e ∈ E tr .
Then, by reorganizing terms, we get
[Equation]
Thus, using Equation equation 31 and equation 32, we have that
[Equation]
Let v := − σ 2 AA ⊤ β ⋆ + (1 − µ ⊤ c Aβ ⋆ )
Aµ c .
Then,
[Equation]
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 |E tr | > d e .
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 Z e and variance of W e for this environment.
Then, we have that ˜ µ = (cid:80) d e i =
α i µ i .
By applying this linear combination to Eq.
equation 33 for this environment, we get
[Equation]
where in the last identity, we applied Eq.
equation 33 for all i = 1 , . . .
, d e .
By rearranging the terms in Eq.
equation 34, we get
[Equation]
From the non-degeneracy condition equation 20, Eq.
equation 35 is equivalent to
[Equation]
[Equation]
where Γ α is defined as
[Equation]
Note that B , β ⋆ , and v are environment independent and Γ α is an environment dependent matrix for which it holds that rank (Γ α ) =
d e from the nondegeneracy condition equation 21.
Thus, Eq.
equation 36 holds if and only v = B Γ α B ⊤ β ⋆ = 0 .
Then, Eq.
equation 33 reduces to
[Equation]
for all e ∈ E tr .
Thus, Bµ e = 0 for all e ∈ E tr , which holds if and only
if B = 0 given that the span of µ e s is R d e .
D T HE R OLE OF THE E IGENSTRUCTURE OF I e ( φ )
In this section, we elaborate on the discussions on the eigenstructure of I e ( φ ) , and in particular, its condition number in the examples of (Arjovsky et al., 2019) and (Rosenfeld et al., 2021).
D. E XAMPLE 1 OF (A RJOVSKY ET AL ., 2019)
Arjovsky et al.
(2019) consider data that is generated according to the following SEM
[Equation]
where Z 1 ∼ N (0 , σ 2 ) , Z 2 ∼ (0 , 1) , and X 1 are independent, and X = (cid:20)
X
X(cid:21) .
Consider the following data representation
[Equation]
[Equation]
Then,
[Equation]
We now find the eigenvalues of I ( φ c ) .
That is, the solutions to det ( I ( φ c ) − λI ) = 0 .
[Equation]
Then,
[Equation]
Hence,
[Equation]
Finally,
[Equation]
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 .
D. E XAMPLE OF (R OSENFELD 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 is of the same order as the invariance penalty.
[Equation]
Here, Z ϵ is defined as
[Equation]
where r := (cid:112) ϵσ 2 e d e and B r ( µ ) denotes the ℓ 2 ball of radius r centered at µ .
Note that the invariant representation is given by
[Equation]
We compare the risk of the optimal predictor given the invariant and the Rosenfeld et al.
(2021) representations.
D.2. I NVARIANT R EPRESENTATION
[Equation]
Given that I e ( φ inv ) is singular, we consider ℓ 2 regularized LSE with regularization parameter γ > 0 , i.e.,
[Equation]
Then,
[Equation]
Recall from the proof of Lemma 1 that
[Equation]
Then,
[Equation]
D.2.
R OSENFELD ET AL .
R EPRESENTATION
We have that
I e
[Equation]
Recall that the first term is equal to I e ( φ inv ) .
For the second term, we have that
[Equation]
Now define matrix A and parameter p as
[Equation]
Then,
I
[Equation]
Also, w ∗
[Equation]
Let a = E
[ Z e Y | Z e / ∈ Z ϵ ] .
Then,
[Equation]
For an invertible matrix S , we have that
[Equation]
Thus,
[Equation]
Then,
[Equation]
That is,
[Equation]
Recall that
[Equation]
Here, Z ϵ is defined as
[Equation]
where r := (cid:112) ϵσ 2 e d e and B r ( µ ) denotes the ℓ 2 ball of radius r centered at µ .
Then,
[Equation]
where I c and I e are defined as
[Equation]
(cid:35)
[Equation]
ϵ ) . Here, we establish a lower bound on the condition number of I e ( φ ϵ ) in terms of the probability of event 1 { Z e / ∈Z ϵ } .
Using Weyl’s inequality, we have that
[Equation]
As I e is positive semidefinite, λ min ( I e ) ≥ 0 .
Moreover, λ min ( I c ) = 0 .
Then, λ max ( I e ( φ ϵ [Equation] max ( I c ) ,
[Equation]
I e ) .
For the first term, we have
[Equation]
[Equation] where the last identity follows from the fact that Y 2 = 1 almost surely, and the fact that W c is independent of Y and W e , and hence is independent of the event 1 { Z e / ∈Z ϵ } .
For the second term, we have
[Equation]
, where the last identity follows similarly as Equation equation 41.
Then,
[Equation]
We have that
[Equation]
Moreover, using Cauchy-Schwarz inequality
[Equation]
we get
[Equation]
Hence,
[Equation]
Thus
[Equation]
κ
[Equation]
Note that (Rosenfeld et al., 2021, Lemma F.3.)
show that
[Equation]
where
[Equation]
Then,
[Equation]
Rosenfeld et al.
(2021) show that the invariance penalty of (Arjovsky et al., 2019) is no greater than
O ( p 2 ϵ,e ) , which can be made arbitrarily small by choosing appropriately large ϵ .
However, for such choices of ϵ , matrix I e ( φ ϵ ) is ill-conditioned.
In particular,
[Equation]
D. E XTENDED C OLORED MNIST E XPERIMENTS
The model and all its respective parameters are replicated from Arjovsky et al.
(2019)’s
GitHub repository 2 .
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
E F ULL L INEAR U NIT T ESTS E XPERIMENTS
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.
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.
F D ETAILS OF THE H EALTHY G UT T ESTS E XPERIMENTS
We summarize the details of each data set of the HealthyGutTests in Table 7.
[Table]
[Table] [Table] [Table]
[Table]
F. E XPERIMENT D ETAILS F.1. N EURAL N ETWORK A RCHITECTURE AND H YPERPARAMETER T UNING
The neural network considered consists of 5 layers with input layer
d X × d h , hidden layers of size d h × d h / 2 , d h / 2 × d h / 2 , d h / 2 × d h / 4 , respectively, and the output layer d h / 4 × d φ .
We use the
ReLU action function and the dropout parameter of each layer is set as p .
We treat d h , 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.
G F ULL D OMAIN B ED E XPERIMENTS G. E XPERIMENT SUMMARY
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.
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.
[Table]
[Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table] [Table]
[Table] [Table]
[Table] [Table] [Figure]