Original: The Missing Invariance Principle found – the Reciprocal Twin of Invariant Risk Minimization
Revised: The Missing Invariance Principle Found – the Reciprocal Twin of Invariant Risk Minimization

Original: Anonymous Author(s)
Revised: Anonymous Author(s)

Original: AfﬁliationAddress email
Revised: AfﬁliationAddress email

Original: Abstract
Revised: Abstract

Original: Machine learning models often generalize poorly to out-of-distribution (OOD)data as a result of relying on features that are spuriously correlated with the labelduring training.
Revised: Machine learning models often generalize poorly to out-of-distribution (OOD)data as a result of relying on features that are spuriously correlated with the labelduring training.

Original: Recently, the technique of Invariant Risk Minimization (IRM)was proposed to learn predictors that only use invariant features by conservingthe feature-conditioned class expectation E e [ y | f ( x )] across environments.
Revised: Recently, the technique of Invariant Risk Minimization (IRM)was proposed to learn predictors that only use invariant features by conservingthe feature-conditioned class expectation E e [ y | f ( x )] across environments.

Original: However, more recent studies have demonstrated that IRM can fail in various tasksettings.
Revised: However, more recent studies have demonstrated that IRM can fail in various tasksettings.

Original: Here, we identify a fundamental ﬂaw of IRM formulation that causesthe failure.
Revised: Here, we identify a fundamental ﬂaw of IRM formulation that causesthe failure.

Original: We then introduce a complementary notion of invariance, MRI, that isbased on conserving the class-conditioned feature expectation E e [ f ( x ) | y ] acrossenvironments, that corrects for the ﬂaw in IRM.
Revised: We then introduce a complementary notion of invariance, MRI, that isbased on conserving the class-conditioned feature expectation E e [ f ( x ) | y ] acrossenvironments, that corrects for the ﬂaw in IRM.

Original: Further, we introduce a simpliﬁed,practical version of the MRI formulation called as MRI-v1.
Revised: Further, we introduce a simpliﬁed,practical version of the MRI formulation called as MRI-v1.

Original: We note that this constraint is convex which confers it with an advantage over the practical version of
Revised: We note that this constraint is convex which confers it with an advantage over the practical version of

Original: IRM, IRM-v1, which imposes non-convex constraints.
Revised: IRM, IRM-v1, which imposes non-convex constraints.

Original: We prove that in a generallinear problem setting, MRI-v1 can guarantee invariant predictors given sufﬁcientenvironments.
Revised: We prove that in a generallinear problem setting, MRI-v1 can guarantee invariant predictors given sufﬁcientenvironments.

Original: We also empirically demonstrate that MRI strongly out-performs
Revised: We also empirically demonstrate that MRI strongly out-performs

Original: IRM and consistently achieves near-optimal OOD generalization in image-basednonlinear problems.
Revised: IRM and consistently achieves near-optimal OOD generalization in image-basednonlinear problems.

Original:  Introduction
Revised:  Introduction

Original: Deep learning models have shown tremendous success over the past decade.
Revised: Deep learning models have shown tremendous success over the past decade.

Original: These models showgreat generalization properties when tested on the same distribution as the training dataset (indistribution generalization).
Revised: These models showgreat generalization properties when tested on the same distribution as the training dataset (indistribution generalization).

Original: However, these models often show catastrophic failure when tested onout-of-distribution dataset, revealing that they learn the shortcut to use easier-to-ﬁt features that arespuriously correlated to the label in the given training domains but don’t generalize to the testingdomains.
Revised: However, these models often show catastrophic failure when tested onout-of-distribution dataset, revealing that they learn the shortcut to use easier-to-ﬁt features that arespuriously correlated to the label in the given training domains but don’t generalize to the testingdomains.

Original: For example, deep networks trained on pictures of cow with only grassy backgrounds in thetraining domain will use the background color as the predictive feature which is easier to learn andgeneralize poorly to pictures of cow, for example, with a dessert background.
Revised: For example, deep networks trained on pictures of cow with only grassy backgrounds in thetraining domain will use the background color as the predictive feature which is easier to learn andgeneralize poorly to pictures of cow, for example, with a dessert background.

Original: Recently, there has been a growing interest in developing models that generalize well across multipledomains.
Revised: Recently, there has been a growing interest in developing models that generalize well across multipledomains.

Original: In particular, there has been a recent body of works that focus on developing algorithmsthat attempt to learn invariant predictors (Arjovsky et al., 2019; Ahuja et al., 2021; Peters et al.
Revised: In particular, there has been a recent body of works that focus on developing algorithmsthat attempt to learn invariant predictors (Arjovsky et al., 2019; Ahuja et al., 2021; Peters et al.

Original: ,2016; Rojas-Carulla et al., 2018; Heinze-Deml et al., 2018).
Revised: ,2016; Rojas-Carulla et al., 2018; Heinze-Deml et al., 2018).

Original: Invariant Risk Minimization (IRM) ,in particular, has attracted a lot of attention as it is one of the ﬁrst methods to be compatible withdeep learning techniques.
Revised: Invariant Risk Minimization (IRM) ,in particular, has attracted a lot of attention as it is one of the ﬁrst methods to be compatible withdeep learning techniques.

Original: However, several follow-up studies have demonstrated the unreliabilityof IRM at learning invariant representations (Kamath et al., 2021; Rosenfeld et al., 2020; Gulrajaniand Lopez-Paz, 2020; Ahuja et al., 2020, 2021).
Revised: However, several follow-up studies have demonstrated the unreliabilityof IRM at learning invariant representations (Kamath et al., 2021; Rosenfeld et al., 2020; Gulrajaniand Lopez-Paz, 2020; Ahuja et al., 2020, 2021).

Original: While many of these studies have analyzed variouslimitations of IRM and IRM-v1 (practical version of IRM), we argue here that a major ﬂaw in the
Revised: While many of these studies have analyzed variouslimitations of IRM and IRM-v1 (practical version of IRM), we argue here that a major ﬂaw in the

Original: IRM formulation has been overlooked and we propose a method to eliminates this ﬂaw.
Revised: IRM formulation has been overlooked and we propose a method to eliminates this ﬂaw.

Original: Related works There has been considerable work in the ﬁeld of learning invariant representations.
Revised: Related Works There has been considerable work in the ﬁeld of learning invariant representations.

Original: They vary from learning domain-invariant feature representations conserving P ( f ( x )) using kernelmethods (Muandet et al., 2013; Ghifary et al., 2016;
Revised: They vary from learning domain-invariant feature representations conserving P ( f ( x )) using kernelmethods (Muandet et al., 2013; Ghifary et al., 2016;

Original: Hu et al., 2020), variational autoencoder (Ilseet al., 2020), and adversarial networks Ganin et al.
Revised: Hu et al., 2020), variational autoencoder (Ilseet al., 2020), and adversarial networks Ganin et al.

Original: ; Long et al.
Revised: ; Long et al.

Original: ; Akuzawa et al.(2019); Albuquerque et al.
Revised: ; Akuzawa et al.(2019); Albuquerque et al.

Original: (2020) to learning invariant class-conditional features P ( f ( x ) | y )
Revised: (2019) to learning invariant class-conditional features P ( f ( x ) | y )

Original: (Gonget al., 2016; Li et al., 2018b) in the context of domain adaptation, which assumes access to the testdistribution for adaptation.
Revised: (Gonget al., 2016; Li et al., 2018b) in the context of domain adaptation, which assumes access to the testdistribution for adaptation.

Original: There is also a large body of work to learn invariant representations inthe ﬁeld of domain generalization that doesn’t assume access to test distributions.
Revised: There is also a large body of work to learn invariant representations inthe ﬁeld of domain generalization that doesn’t assume access to test distributions.

Original: This includesimposing invariance of E e [ y | f ( x )]
Revised: This includesimposing invariance of E e [ y | f ( x )]

Original: (Arjovsky et al., 2019) with information bottleneck constraint(Ahuja et al. (2021)), imposing object-invariant condition (Mahajan et al. (2021)), using domaininference (Creager et al., 2021), model calibration (Wald et al., 2021), and others (Krueger et al.
Revised: (Arjovsky et al., 2019) with information bottleneck constraint(Ahuja et al. (2021)), imposing object-invariant condition (Mahajan et al. (2021)), using domaininference (Creager et al., 2021), model calibration (Wald et al., 2021), and others (Krueger et al.

Original: Li et al., 2018a; Shankar et al., 2018).
Revised: Li et al., 2018a; Shankar et al., 2018).

Original: Our Contributions. We reformulate IRM in a simple more generalizable way, and show that it canbe modiﬁed to yield a new complementary notion of invariance, called MRI.
Revised: Our Contributions We reformulate IRM in a simple more generalizable way, and show that it canbe modiﬁed to yield a new complementary notion of invariance, called MRI.

Original: We also show that IRMhas a fundamental ﬂaw due to an extraneous constraint which causes it to fail.
Revised: We also show that IRMhas a fundamental ﬂaw due to an extraneous constraint which causes it to fail.

Original: MRI is shown to befree of this ﬂaw.
Revised: MRI is shown to befree of this ﬂaw.

Original: We prove that in a general linear problem setting, MRI-v1 can guarantee invariantpredictors given sufﬁcient environments.
Revised: We prove that in a general linear problem setting, MRI-v1 can guarantee invariantpredictors given sufﬁcient environments.

Original: We then empirically demonstrate that MRI strongly outperforms IRM and consistently achieves near-optimal OOD generalization in image-based nonlinearproblems.
Revised: We then empirically demonstrate that MRI strongly outperforms IRM and consistently achieves near-optimal OOD generalization in image-based nonlinearproblems.

Original:  Problem Formulation
Revised:  Problem Formulation

Original: We consider a set of environments E = { e } , each of which deﬁnes a distribution P e ( x, y ) overinputs and labels, from which the dataset of the environment is drawn D e ⌘ { ( x ej 2 R d , y ej 2 R ) } .
Revised: We consider a set of environments E = { e } , each of which deﬁnes a distribution P e ( x, y ) over inputsand labels, from which the dataset of the environment is drawn D e ⌘ { ( x ej 2 R d , y ej 2 R ) } .

Original: The distribution P e ( x, y ) is assumed to be generated according to the causal graph  introduced in(Rosenfeld et al., 2020) (Fig.1), which includes latent features that are invariantly zi or spuriously
Revised: Thedistribution P e ( x, y ) is assumed to be generated according to the causal graph in Fig.1 as introduced in Rosenfeld et al. (2020), which includes latent features that are invariantly z

Original:  z scorrelated with the environment.
Revised: i or spuriously z s correlatedwith the label.

Original: The observation x is generated as a function of the latent features:x = g ( z i , z s ) is assumed to be injective, which makes it possible to recover the latent features fromthe observation.
Revised: The observation x is generated as a function of the latent features: x = g ( z i , z s ) isassumed to be injective, which makes it possible to recover the latent features from the observation.

Original: In this formulation, we have using d-separation that  the label is independent from the environmentgiven the invariant latent feature: y ??
Revised: Note that while this work focuses on the causal graph in Figure 1, our results are also applicable tothe graph in which the invariant feature z i causes the label y ( Z i ! Y ). This is the graph consideredby Arjovsky et al.

Original: e | z i and  the invariant latent feature is independent from the 65 environment given the label:
Revised: (2019) in Figure 3 with the difference being that they allow both the invariantfeature and label distribution to also vary across environments. Also, while Arjovsky et al. (2019)focused on the causal graph with Z i ! Y , the IRM method is also applicable to the causal graphconsidered in Figure 1 (Rosenfeld et al., 2020).

Revised: Y

Original: z i
Revised: Z i

Original: [Equation] [Equation]
Revised: Z s X

Original:  The risk of  f  in environment e is deﬁned as the populationaverage
Revised: E The risk of a predictor f : X ! O in environment e isdeﬁned as the population average

Original: [Equation]
Revised: [Equation]

Revised: [Equation]

Revised: [Equation]

Original: where o
Revised: where o

Original: = f
Revised: = f

Original: ( x ) is the predictor’s output.
Revised: ( x ) is the predictor’s output.

Original: Here, we consider standard convex loss functions l : O ⇥ Y !
Revised: Here, we consider standard convex loss functions l : O ⇥ Y !

Original: R  0 ,including the square loss l sq ( o, y )
Revised: R  0 ,including the square loss l sq ( o, y )

Original: = 12 k o  y k2 , and thebinary-cross-entropy loss l BCE ( o, y ) =
Revised: = 12 k o  y k l log ( o, y ) =

Original:  (1 + y ) [-log((
 +tanh( o/ 2)) / 2)-]  (1  y ) [-log((
-]  tanh( o/ 2)) / 2) (for
Revised:  (1 + y ) [-log((
 +tanh( o/ 2)) / 2)-]log( ⌘ ( o ))  (1  y ) [-log((
-]log(1  ⌘ ( o ))for binary classiﬁcation 2 ( O ✓ R , Y ✓ [  1 , 1] ), where 76 ⌘ ( o )

Original: Y =
Revised: =

Original: {  1 , 1 } ) 2 .
Revised: 1 / (1 + e o ) is the sigmoid function.

Original: 3.1 IRM paradigm
Revised: 3.1 IRM Paradigm

Original: The core idea underlying IRM is the invariance of feature-conditioned label expectation E e [ y | o ] ⌘ E P e ( y | o ) [ y ] across environments, i.e. [Equation]
Revised: 3.1.1 Original formulation

Original: The IRM paradigm (Arjovsky et al., 2019) proposes to impose this invariance via the followingformulation:

Original: Deﬁnition 3.1.
Revised: Deﬁnition 3.1.

Revised: (Arjovsky et al., 2019)

Original: A representation f : X !
Revised: Given a predictor  f , the feature representation f : X !

Original: Z  is IRM invariant over a set of environmentsE
Revised: Z 81 is IRM invariant over a set of environments

Original:  if thereexists : Z !
Revised: E if there exists : Z !

Original: O such that is simultaneously optimal on f for all environments [Equation] 2 E 
Revised: O such that is simultaneouslyoptimal on f for all environments in E , i.e.

Original: [Equation]
Revised: [Equation]

Original: where is assumed to be unrestricted in the space of all measurable functions.
Revised: where is assumed to be unrestricted in the space of all measurable functions.

Original: This deﬁnition has shown to be equivalent to the following condition (Kamath et al., 2021),
Revised: Lemma 3.2. (Kamath et al., 2021) For standard loss functions 3 , Deﬁnition 3.1 is equivalent to

Original: [Equation]
Revised: [Equation]

Original: where  ( · ) is a monotonic function such that  ( o ) =
Revised: ) where E e [ y | f ( x )]

Original: o for square loss,  ( o ) =
Revised: ⌘ E P e ( y | f ( x ))

Original: tanh( o/ 2) for BCEloss. Since the RHS is constant with respect to environments, eq (6) indeed implies eq (4). However,there is a critical difference: Eq (4) only requires the conditional expectation E e [ y| o ] to share thesame value across environments, i.e. , it imposes only |E|  1 number of constraints, whereas eq (6) 88 requires that value to be also equal to the RHS term imposing |E| constraints.
Revised: [ y ] is the feature-conditioned label expectation, and  is a monotonic 86 function that depends on the loss function.

Original: This extra constraint 89 originates from the fact that IRM Eq (5) imposes 1 constraint for each environment. Perturbation-based IRM Deﬁnition 3.1 is overly complex due to the presence of and itsoptimality condition.
Revised: The original formulation above is overly complex due to the composite predictor  f and the 89 optimality condition on .

Original: For further analysis, it helps to simplify the deﬁnition by absorbing intothe predictor f and considering arbitrary perturbations on the output, which yields the followingreformulation of IRM.
Revised: For further analysis, we introduce a simpliﬁed formulation of IRM deﬁnedby the following variational principle.

Original: Deﬁnition 3.2.
Revised: Deﬁnition 3.3.

Original: A predictor f : X !
Revised: A predictor f : X !

Original: O is IRM invariant over a set of environments E , if therisk L e ( f ) is preserved under arbitrary inﬁnitesimal perturbations on its output o
Revised: O is IRM invariant over a set of environments E , if therisk L e ( f ) remains stationary under arbitrary inﬁnitesimal perturbations on the predictor’s outputo

Original: = f ( x ) for allenvironments, i.e.
Revised: = f ( x ) for all environments in E , i.e.

Original: [Equation]
Revised: [Equation]

Original:  where  ( · ) is an arbitrary perturbation that is unrestricted in the space of all measurable functions.
Revised: (6) where  : O ! O is an arbitrary perturbation that is unrestricted in the space of all measurablefunctions, and  L e ( f ) denotes the resulting change in risk.

Original: Lemma 3.3.
Revised: Lemma 3.4.

Original: This reformulated deﬁnition 3.2 is equivalent to eq (6) .
Revised: For standard loss functions 3 , Deﬁnition 3.3 is equivalent to [Equation] e [Equation]

Revised: E ,

Revised: [Equation]

Revised: (7) which is equivalent to Lemma 3.2 with f replacing the composite predictor  f .

Original: Proof.
Revised: Proof.

Original: Eq (7) is satisﬁed if and only if E P e
Revised: Eq (6) is satisﬁed if and only if E P e

Original: ( y | o )
Revised: ( y | o )

Original: [ @ o l ( o, y )]
Revised: [ @ o l ( o, y )]

Original: = 0 .
Revised: = 0 .

Original: Note that @
Revised: For standard loss functions, the lossderivative has the form @

Original: o lsq ( o, y )
Revised: o l ( o, y )

Original: =
Revised: =

Original: o  y and@ o lBCE ( o, y )  = 
Revised:  y +  ( o ) , which yields E e [ @ o l ( o, y ) | o ] =  E

Original: tanh( o/ 2)  y  .
Revised: e [ y | o ]+  ( o ) = 0 .

Original: Therefore, we arrive at
Revised: It was noted in Arjovsky et al. ; Kamath et al. that the essense of IRM’s invariance isthe conservation of the feature-conditioned label expectation , i.e.

Original: [Equation]
Revised: [Equation]

Revised: IRM eq (7) (also eq (5)) indeed implies the conservation law, since its RHS is constant with respectto the environment index e .

Revised: However, there is a critical difference: Eq (8) only requires E e [ y | o ] toshare the same value across environments, thus imposing |E|  1 constraints, whereas eq (7) (also 107 eq (5)) additionally requires that value to be also equal to the RHS term  ( o ) .

Revised: This extra constraint isdue to the indirect way IRM attains the conservation law, i.e. by equaling E e [ y | o ] to an intermediaryterm  ( o ) , which imposes 1 constraint per environment.

Original: IRM-v1 Due to the impracticality of consideringoptimal classiﬁers in the  space of unrestrictedfunctions, Arjovsky et al.
Revised: Due to the impracticality of considering the unrestricted function space of , Arjovsky et al.

Original: (2019) suggested restrictingthe functional class of to the space of linearfunctions.
Revised: (2019)suggested restricting to the space of linear functions.

Original: In deﬁnition 3.2, this corresponds to restricting the output perturbations  to the space oflinear functions, or equivalently, an identity function: i.e.  ( o ) = 
Revised: In the variational formulation, this correspondsto restricting the output perturbations  to the space of linear functions, which, for scalar outputs, isequivalent to the identity function,  ( o ) = o .

Original: o , which reduces eq (7) to L e (
Revised: This reduces eq (6) to

Original: [Equation]
Revised: [Equation]

Original: P
Revised: The reduced constraints in eq (9) are identical to IRM-s in Kamath et al.

Original: [Equation] ( o [Equation] [Equation]
Revised: (2021), which imposes8 e 2 E , @

Original: P [Equation] y |
Revised: E e [ l ( · o, y )]

Original: [Equation]
Revised: = 0 , where is simply a scalar multiplicative factor 4 .

Original: [ @ o [Equation] This reduced formulation is called IRM-v1  in Arjovsky et al.
Revised: In Arjovsky et al.(2019), this reduced formulation is termed IRM-v1 when the constraints are imposed in a soft manner ,i.e. as squared penalty terms (See eq (17)). In the literature, the term IRM-v1 is widly used for thereduced formulation regardless of whether hard or soft constraints are used, which we adopt here.

Original: or IRM-s in Kamath et al. Note that the ﬁrst term E e [ o  ( o )] of eq (9) originates from the RHS of eq (6).
Revised: Despite the practical applicability, it has been empirically shown that IRM-v1 can behave quitedifferently from IRM (Kamath et al., 2021).

Original: If this term wereconstant, it would yield the conservation of E e [ oy ] =
Revised: This failure mode can be analytically understood here. Note that E e [ @ o l ( o, y ) · o ] =

Original: EP e
Revised: E e

Original: o ) [E e [ y | o ] · o ] , a necessary condition forthe conservation of E e [ y | o ] in eq (4). However, since the constancy of this term is generally notguaranteed, eq (9) does not necessarily prescribe any meaningful invariance. Thus, IRM-v1 is notconsistent with IRM’s invariance.
Revised: [ o ( y   ( o ))] for standard loss functions 3 .

Original: Although this inconsistency was empirically shown in Kamathet al. , the actual cause of the problem was never discussed. Our analysis reveals that the rootcause originates from the extraneous constraint of IRM. Square loss
Revised: Therefore, eq (9) isequivalent to

Original: [Equation]
Revised: [Equation]

Original: BCE loss
Revised: Unlike eq (7), however, RHS of eq (10) does not necessarily remain constant across environments,and hence does not prescribe a meaningful invariance.

Original: [Equation] Constraint
Revised: Therefore, restricting the function space of (or ) invalidates the indirect way of imposing invariance in IRM involving an intermediary term. Note that this holds true for a general causal graph and not just for the causal graph considered in

Original: [Equation]
Revised: Figure 1.

Original: 3.2 MRI paradigm
Revised: 3.2 MRI Paradigm

Original: We now introduce a complementary notion of invariance by slightly modifying deﬁnition 3.2, whichwe callthe Mirror Reﬂected IRM , or MRI.
Revised: We now introduce a complementary notion of invariance by considering inﬁnitesimal perturbationson label, which we call Mirror Reﬂected IRM , or MRI.

Original: Deﬁnition 3.4.
Revised: Deﬁnition 3.5.

Original: A predictor f : X !
Revised: A predictor f : X !

Original: O is MRI invariant over a set of environments E , if for allenvironments L e ( f ) gets identically affected by arbitrary inﬁnitesimal  perturbations on label, i.e.
Revised: O is MRI invariant over a set of environments E , if the changein risk due to arbitrary inﬁnitesimal label perturbations is shared across environments, i.e.

Original: [Equation]
Revised: [Equation]

Original:  where
Revised: 11) where

Original: [Equation]
Revised: [Equation]

Revised: 12) where  : Y !

Original: where  ( · ) is an arbitrary perturbation that is unrestricted in the space of all measurable functions.
Revised: Y is an arbitrary perturbation that is unrestricted in the space of all measurablefunctions, and  L e ( f ) denotes the resulting change in risk. E e [ @ y l ( o, y ) | y ] ⌘ E P e ( o | y ) [ @ y l ( o, y )] .

Original: Lemma 3.5.
Revised: Lemma 3.6.

Original:  Deﬁnition 3.4 implies the conservation of E e [ o | y ] ⌘ E P e
Revised: For standard loss functions, Deﬁnition 3.5 is equivalent to

Original: ( o | y )
Revised: [Equation]

Revised: 13) which conserves the label-conditioned feature expectation E e

Original: [ o  ] across environments, [Equation]
Revised: [ o | y ] ⌘ E P e ( o | y ) [ o ] across environments.

Original: Proof.
Revised: Proof.

Original: Eq (10) is satisﬁed
Revised: Eq (11) is satisﬁed

Original: if and only if E e [ @ y l ( o, y ) | y ] is conserved.
Revised: if and only if E e [ @ y l ( o, y ) | y ] is conserved.

Original: Since @ y lsq ( o, y ) = 
Revised: For standard loss functions, theloss derivative has the form @ y l ( o, y ) =  o + ⇢ ( y ) 5 .

Original:  y   o and@ y lBCE ( o, y )  = 
Revised: Therefore, E e 1 [ @ y l ( o, y ) | y ]  E e 2 [ @ y l ( o, y ) | y ] = E e 1 [ o | y ]  E e 2 [ o | y ] = 0 .

Original:  o , this yields eq (12). Note that unlike IRM, MRI prescribes E  1 constraints, which is the right number of constraints.
Revised: Note that unlike IRM, MRI eq (11) imposes the correct number of constraints as needed by theconservation law eq (13). That is, MRI attains invariance in a direct manner without involving anyintermediary terms. Algorithm

Original: MRI-v1 Restricting the label perturbations to the space of linear functions, or equivalently, anidentity function  ( y ) =
Revised: [Equation] Restricting the label perturbations to the space of linear functions, or equivalently, an identity function ( y ) =

Original: y , reducesthe MRI formulation eq (11) to
Revised: y , reduces eq (12) to

Original: [Equation]
Revised: [Equation]

Original: (13) which yields thefollowing conservation law 
Revised: This reduces the conservation law eq (13) to

Original: [Equation]
Revised: [Equation]

Original: which is indeed a necessary condition for  MRI invariance eq (12).
Revised: which is in fact a necessary condition for the MRI invariance eq (13).

Original: Moreover, MRI-v1 eq (14) imposeslinear equality constraint with respect to the output o .
Revised: Therefore, the direct way of

Original: The consistency with MRI invariance, as wellas the convexity of constraint, provides MRI-v1 with great practical advantages over IRM-v1, whichwe discuss in following sections.
Revised: MRI’s constraints imposing invariance continues to hold even when  is restricted to the space oflinear functions.

Original:  3.3 Objective
Revised:  Training methods

Original: The full methods of IRM-v1 and MRI-v1 involve solving the following constrained optimization 
Revised: The full methods of IRM-v1 and MRI-v1 can be formalized as a constrained optimization problem

Original: [Equation]
Revised: [Equation]

Original: where the objective ¯ L tr ( f ) =
Revised: subject to ~ c [Equation] f ( ✓ )) =

Original: |E
Revised: ~

Original: [Equation]
Revised: [Equation]

Original:  | 
Revised: (16) where ¯ L tr = 1 |E tr | P e 2 E tr L e is the average risk over the set of training environments E tr ⇢ E , and ✓is the parameters of the predictor f . The constraint functions are ~ c = ~ L 2 R |E tr | for IRM-v1 and~ c = Q ~ L 2 R |E tr |  1 for MRI-v1, where Q 2 R ( |E tr |  1) ⇥ |E tr | is an orthonormal matrix that satisﬁes

Original: P e 2 E tr L e ( f ) is the average risk over the set of  training  environments E tr ⇢ E , ✓ is  the parameters ofthe predictor f , and the vector-valued constraint function 135 contains the perturbed risks   
Revised: Q ~ 1 = ~ 0 (See Table 1). For example, in the case of a training set of two environments E tr = { e 1 , e 2 } , Q ~ L is simply the difference of perturbed risks (  L e 1   L e 2 ) / p 2 , since Q = [1 ,  1] / p 2 .

Original: L e ( f ) , 8 e 2 E tr for the case of IRM-v1 as the perturbed risk itself, or 136 the difference between pairs of perturbed risks for the case of MRI-v1 (See Table 1).
Revised: Numerically solving the constrained optimization problem requires converting the hard constraints~ c ( f ( ✓ )) = ~ 0 to soft regularization terms, which allows the use of any off-the-shelf gradient-basedoptimization algorithms.

Revised: Penalty Method (PM)

Revised: Penalty method is the most commonly used approach, including in Arjovskyet al.

Revised: (2019), which adds the sum of squared constraints as a penalty term to the objective,

Original: [Figure]
Revised: [Equation]

Original: Linear setting
Revised: However, this method requires increasing µ t ! 1 over training iterations t in order to approximatethe exact hard-constraint, which leads to training instability and slow convergence (Bertsekas, 1976). Augmented Lagrangian Method (ALM) ALM was introduced to overcome the limitations ofpenalty method (Bertsekas, 1976), which adds a Lagrange multiplier term to (17)

Original: [Figure]
Revised: [Equation]

Revised: (18) where ~ gets updated as ~ t +1 =

Revised: ~ t + µ ~ c ( f ( ✓ t )) to accumulate the residual constraint at each trainingiteration t .

Revised: In practice, ALM can operate with moderate values of µ ( ⇠ 10 ) without ﬁne tuning, and thus exhibits fast and stable convergence (Bertsekas, 1976).

Revised: ~ is typically initialized at ~ 0 , which hasan annealing effect on ALM.

Original: 4.1 General linear SEM
Revised: ALM has an additional advantage when mini-batch updates are used, such as stochastic gradientdescent (SGD), in which the constraints have slight mismatches across iterations. In ALM, the Lagrange multiplier  is accumulated over iterations, thus imposing the model to satisfy an averagedconstraint. In contrast, PM penalizes the model to satisfy all of the conﬂicting constraints, and thusdecaying the model toward the zero-predictor.

Original: [Figure]
Revised: [ELABORATE]

Original:  In this section, we show that MRI-v1 can correctly remove all spurious features when enoughenvironments are given.
Revised: 5.1 General Linear SEM In this section, we show that MRI-v1 can correctly remove all spurious features given a sufﬁcientnumber of environments.

Original: We consider a data generating process according to  1, in which x = g ( z i , z s ) is a linear function.
Revised: We consider a data generating process according to Figure 1, in which x = g ( z i , z s ) is a linearfunction.

Original: Note that this SEM is a generalization of the linear Gaussian model considered in Rosenfeld et al.(2020), which additionally assumed binary y and a speciﬁc noise model (linear additive Gaussian)for the latent features.
Revised: Note that this Structural Equation Model (SEM, Pearl (2009)) is a generalization of thelinear Gaussian model considered in Rosenfeld et al. (2020), which additionally assumed binary yand a speciﬁc noise model ( i.e. linear additive Gaussian) for the latent features.

Revised: In comparison, wedon’t put any assumptions on the noise generating process or the label.

Original: We consider a linear predictor f : X !
Revised: We consider a linear predictor f : X !

Original: O .
Revised: O .

Original: Since g is injective and has an inverse over its range,without loss of generality, we can deﬁne f as a linear function directly over the latents as
Revised: Since g is injective and has an inverse over its range,without loss of generality, we can deﬁne f as a linear function directly over the latents as

Original: [Equation]
Revised: [Equation]

Original: (16) with ✓ = { w i 2 R d i , w s 2 R d s } .
Revised: (19) with ✓ = { w i 2 R d i , w s 2 R d s } , where d i is the dimension of invariant features, d s is the dimensionof spurious features.

Original: Theorem 4.1.
Revised: Theorem 5.1.

Original: Given |E tr | > d s training environments,where d s is the dimension of the spurious 149 features, and that E e [ z s y ] are in general linear position,  MRI-v1 will eliminate all spurious featuredimensions.
Revised: Given |E tr | > d s training environments, and that E e [ z s y ] are in general linear position, 191 MRI-v1 will eliminate all spurious feature dimensions.

Original: Proof.
Revised: Proof.

Original: MRI-v1 condition:
Revised: MRI-v1 constraint eq (15) is

Original: [Equation]
Revised: [Equation]

Revised: e [ oy ] = w i · E e [ z i y ]

Revised: [Equation]

Revised: ] = const

Original: This can be written in a matrix form as
Revised: This can be expressed in a matrix form as

Original: [Equation]
Revised: [Equation]

Original: ) where M  ⌘ [ E e [ z s y ]] e 2 E tr 2 R d s ⇥ |E tr |, and c 0 ⌘ const  w i · E e [ z i y ] .
Revised: where M 0 = M  ¯ M · 1 with M ⌘ [ E e [ z s y ]] e 2 E tr 2 R d s ⇥ |E tr | and ¯ M = 1 |E tr | P e 2 E tr E e [ z s y ] .

Original: Since rank ( M  ) =
Revised: Sincerank ( M 0 ) =

Original: d s , 154 Eq (18) is true if and only if w s = 0and c 0 = 0 .
Revised: d s , eq (20) is equivalent to w s = 0 .

Original: Rosenfeld et al.
Revised: Rosenfeld et al.

Original: (2020) derives a similar resultto that of our Theorem 4.1 stating that IRM-v1constraint can eliminate all spurious feature dimensions given |E tr | > d s .
Revised: (2020) derives a similar result that IRM-v1 can eliminate all spurious featuredimensions given |E tr | > d s .

Original: However, unlike our results158 for MRI, their result is restricted to a  specialized linear family of environments with binary labels andadditive gaussian noise.
Revised: However, unlike our results for MRI-v1, their result is restricted to a 199 specialized linear family of environments with binary labels and additive gaussian noise.

Original: In fact, IRM-v1 is known to fail to eliminate spurious feature dimensions ina more general linear problems in which the noise is not additive Gaussian (Arjovsky et al., 2019; Kamath et al., 2021).
Revised: In fact, IRM-v1 has been shown to fail to eliminate spurious features in a more general linear problems inwhich the noise is not additive Gaussian (Arjovsky et al., 2019; Kamath et al., 2021).

Original: 4.2 Shape-texture Linear Regression
Revised: 5.2 Minimal Example: d i = 1 , d s = 1 , |E tr | = 2

Revised: In Figure 2a, we demostrate the result for the general linear SEM above with a minimal linear regression problem that involves one invariant feature, one spurious feature and two training environments.

Original: Here, we consider a linear regression problem on the Shape-texture dataset which has 2 trainingenvironments and include 1 invariant and 1 spurious features.
Revised: This toy example shares the same structure as the Shape-texture dataset (Section 6.1). It is analogousto the two-bit environment task from Kamath et al.

Original: The full analytical solutions arepresented in the Supplementary Materials.
Revised: The detailed experimental setting and theanalytical solutions are presented in the Supplementary Materials.

Original: IRM-v1 IRM-v1 has two quadratic equality constraints, shown as two ellipses in Fig 2a.
Revised: IRM-v1 IRM-v1 has two quadratic equality constraints (  L e 1 ( f )

Original: Theintersection between 2  non-convex constraints on a 2 dimensional problem leaves the solution spaceto be a disjoint set of (0-d) points, shown as 4 stars. Note that among the 4 solutions, only two ofthem correspond to an invariant ( w s = 0 ) predictor, one of whichis the trivial zero predictor, whilethe other two are non-invariant, which conﬁrms the inconsistency of IRM-v1 at identifying invariantpredictors (Kamath et al., 2021; Arjovsky et al., 2019).
Revised: = 0 ,  L e 2 ( f ) = 0 ), shown astwo elliptic curves. The intersection between the non-convex constraints on a 2-D feature spaceyields a disjoint set of 0-D points, all of which are local constrained optima, including the trueinvariant optimum, a zero-predictor solution, and two non-invariant solutions

Original: MRI-v1 In contrast, consistent with Theorem 4.1, MRI-v1 has a single linear equality constraint(Fig 2a), which exactly yields the set of invarian solutions w s =0 .
Revised: [Figure] IRM-relaxed We also consider a relaxed version of IRM-v1, whose constraint function is c =

Original: That is, every point on theconstraint line is a invariant predictor.
Revised: Q ~ L =

Original: Since this is a convex constraint, the convex objective ¯ L tr ( f )has a unique constrained optimum, which yields the invariant optimum predictor (Fig 2a).
Revised: (  L e 1   L e 2 ) / p 2 with  L e = E e [ @ o l ( o, y ) · o ] , such that it has a single constriant 215

Original: IRM-relaxed To isolate the non-convexity issue of the IRM-v1 constraint, we consider a relaxedversion of IRM-v1 such that  L e 1 ( f ) =
Revised: L e 1 ( f ) =

Original:  L e 2 ( f ) .
Revised:  L e 2 ( f ) .

Original: Since it has only a single constraint, it has 1dcontinuous set of points as possible solutions (Fig 2a).
Revised: This constraint yields a hyperbolic curve, although it appears to be two straightlines in Figure 2a.

Original: However, it is still a nonconvex constraint.
Revised: Due to the non-convex constraint, this problem exhibits two local constrainedoptima, including the true invariant optimum and a non-invariant solution, with the non-invariantsolution exhibiting lower training loss. Therefore, simply reducing the number of constraints of IRM-v1 does not resolve its limitation.

Original: Inaddition to the  w s = 0 line, it also exhibits w i = 2 / 3 line as the solution (Fig 2a).
Revised: MRI-v1 MRI-v1 has a single linear equality constraint, which is exactly the set of invariantsolutions ( w s = 0 ): That is, any point on the constraint is a invariant predictor.

Original: As a result, the¯ L tr ( f ) objective exhibits two constrained optimum solutions, only one of which is invariant (Fig 2a).
Revised: This problem has aconvex objective with a convex constraint, and therefore it has a unique optimum, which is the trueinvariant optimum solution.

Original: Learning Dynamics In Fig 2(b,c), we show simulation of gradient descent learning dynamics ofthe objective eq (19) ( µ = 5 ⇥ 10 4 ).
Revised: In Figure 2b,c, we show the learning trajectory under gradient descent (Penalty method eq (17)with µ = 5 ⇥ 10 4 and full-batch update).

Original: The convergence pattern for IRM-v1 is highly dependent onthe initialization, due to the presence of multiple local minima.
Revised: The convergence pattern for IRM-v1 is highly dependenton the initialization (shown with different trajectory colors), due to the presence of multiple localminima.

Original: In this case, the invariant optimumis especially extremely difﬁcult to converge to as it has higher risk than the nearby non-invariantminima.
Revised: In this case, the invariant optimum is especially extremely difﬁcult to converge to as it hashigher training loss than the nearby non-invariant minima.

Original: In contrast, MRI-v1 always converges to the invariant optimum regardless of initializationsince it is a unique minimum.
Revised: In contrast, MRI-v1 always converges tothe invariant optimum regardless of initialization since it is a unique minimum.

Original: For IRM-relaxed, the invariant optimumpoint is lost at this valueof µ , which exists in the limit µ ! 1 .
Revised: For IRM-relaxed,the invariant optimum is lost at this value of µ , which exists in the limit µ ! 1 .

Original: As a result, all trajectories converge to the non-invariantoptimum, which has lower training risk.
Revised: As a result, alltrajectories converge to the non-invariant optimum.

Original: We show similar results for other linear problems includingshape-texture classiﬁcation and toy-CMNIST in Supplementary Materials. 5.1 Methods: Constrained Optimization
Revised: Minimal examples for binary classiﬁcation, based on the shape-texture classiﬁcation and the toyCMNISTa/b datasets, are also analyzed in Supplementary Materials (Figure 3).

Original:  Solving constrained optimization numerically requires converting the problem into an unconstrainedoptimization problem such that the constraint is imposed as a soft regularization term. This allowsthe use of any off-the-shelf gradient-based optimization algorithm. Penalty Method (PM) Penalty method is the most commonly used approach, including in Arjovskyet al.
Revised:  Nonlinear

Original: (2019), which adds the squared constraint as a penalty term to the objective,

Original: [Equation] However, this method requires large values for µ to approximate the exact constraint, which leads totraining instability and slow convergence Bertsekas (1976). Augmented Lagrangian Method (ALM) ALM was introduced to overcome the limitations ofpenalty method (Bertsekas, 1976), which adds a Lagrange multiplier term ~ to (19) [Equation]
Revised: -based Problems6.1 Datasets

Original: This method involves updating  i at each training iteration t to accumulate the residual constraint,i.e. ~ t +1 e =  te + µ ~ c ( ✓ t ) , and then updating ✓ t according to a gradient descent step with respectto the combined objective in eq (20). ALM quickly converges to exactly satisfy the constraintswith moderate values of µ ( ⇠ 10 ) without ﬁne tuning, and thus exhibits much faster and stableconvergence Bertsekas (1976).
Revised: Shape-Texture Dataset

Original: ~ is typically initialized at ~ 0 , which has an annealing effect on ALM.

Original: Moreover, ALM has an interesting effect when used with mini-batches, in which the constraints haveslight mismatches across iterations.

Original: Since the Lagrange multiplier  is accumulated over iterations, itimposes the model to satisfy an averaged constraint.

Original: In contrast, PM penalizes the model to satisfy allof the conﬂicting constraints, and thus decaying the model toward the zero-predictor.

Original: Unless notedotherwise, we used ALM by default for our optimization.

Original: 5.2.1 Shape-Texture Dataset We develop a new versatile dataset designed for domain generalization, called the Shape-Texturedataset, based on the data generating model in Fig 1.
Revised: We develop a new versatile dataset designed for domain generalization,called the Shape-Texture dataset, based on the data generating model in Figure 1.

Original: This dataset is designed to test algorithms onlinear and non-linear regression and classiﬁcation problems in multiple adjustable settings.
Revised: This dataset isdesigned to test algorithms on linear and non-linear regression and classiﬁcation problems in multipleadjustable settings.

Original: The dataset involves an invariant latent feature ✓ i 2 R , a spurious latent feature ✓ s 2 R , and thelatent feature ✓ y 2 R for the label.
Revised: The dataset involves an invariant latent feature ✓ i 2 R , a spurious latent feature ✓ s 2 R , and thelatent feature ✓ y 2 R for the label.

Original: Each of them represents an angle of orientation, and therefore isnaturally represented on a unit circle in the complex plane as e i ✓ 2 S 1 .
Revised: Each of them represents an angle of orientation, and therefore isnaturally represented on a unit circle in the complex plane as e i ✓ 2 S 1 .

Original: In the linear regression setting,the observed input is simply the concatenated features x = [ e i ✓ i , e i ✓ s ] , and the label is y =
Revised: In the linear regression setting,the observed input is simply the concatenated features x = [ e i ✓ i , e i ✓ s ] , and the label is y =

Original: e i ✓ y .
Revised: e i ✓ y .

Original: Whereas in the linear classiﬁcation setting, the input is given by x = [sin( ✓ i ) , sin( ✓ s )] , and the labelis y
Revised: Whereas in the linear classiﬁcation setting, the input is given by x = [sin( ✓ i ) , sin( ✓ s )] , and the labelis y

Original: = H (sin( ✓ y ))
Revised: = H (sin( ✓ y ))

Original: where H denotes the Heavyside step function.
Revised: where H denotes the Heaviside step function.

Original: In the non-linear regression andclassiﬁcation setting, the label remains the same as their respective linear versions but the input isan image (Fig 7) composed of two planar waves of differing frequencies; The orientation of thelow frequency wave (a.k.a shape) is given by ✓ i / 2 , and the orientation of the high frequency wave(a.k.a texture) is given by ✓ s / 2 .
Revised: In the non-linear regression andclassiﬁcation setting, the label remains the same as their respective linear versions but the input isan image (Figure 4) composed of two planar waves of differing frequencies; The orientation of thelow frequency wave (a.k.a shape) is given by ✓ i / 2 , and the orientation of the high frequency wave(a.k.a texture) is given by ✓ s / 2 .

Original: For all task settings, the latent angles and the label angle are linearlyrelated as ✓y ✓ i  + ⇠ o , ✓ s ✓ y + ⇠ e , where ⇠ ⇤ ⇠ N (0 ,  ⇤ ) with probability p ⇤ , or ⇠ ⇤ ⇠ U (  ⇡ , ⇡ ) 227 probability p ⇤ for ⇤ 2 { i, s } .
Revised: For all task settings, the latent angles and the label angle are linearlyrelated as ✓ i ✓ y + ⇠ i , ✓ s ✓ y + ⇠ s , where ⇠ ⇤ ⇠ N (0 ,  ⇤ ) with probability p ⇤ , or ⇠ ⇤ ⇠ U (  ⇡ , ⇡ ) 251 probability p ⇤ for ⇤ 2 { i, s } .

Original: p i ,  i are ﬁxed invariant parameters whereas p s ,  s can vary from 228 one environment to another.
Revised: p i ,  i are ﬁxed invariant parameters whereas p s ,  s can vary from 252 one environment to another.

Original: In this work, for all task settings, we use two training environmentsand one testing environment with p i = 0 .
Revised: In this work, for all task settings, we use two training environmentsand one testing environment with p i = 0 .

Original: 75 for all environments, p s = 0 .
Revised: 75 for all environments, p s = 0 .

Original: 8 , 1 for the two trainingenvironments, p s = 0 for the testing environment, and  i =
Revised: 8 , 1 for the two trainingenvironments, p s = 0 for the testing environment, and  i =

Original:  s = 0 for all environments.
Revised:  s = 0 for all environments.

Original: Result We use a small convolutional network for the image-based regression and classiﬁcationdataset.

Original: In addition to reporting the results for IRM-v1 and MRI-v1 algorithms, we report theperformance of the ERM algorithm, which minimizes the average risk across environments (Table 2).

Original: We also report the performance of the Oracle (Table 2), where the Oracle is trained on environmentsin which the spurious features are not correlated with the label.

Original: For both regression and classiﬁcationsetting, we ﬁnd that IRM-v1 it is better than ERM although there is a signiﬁcant generalization error(Table 2).

Original: We ﬁnd similar results for accuracy (see Table 3).

Original: However, MRI-v1 gives almost identicalperformance across training and testing environments irrespective of the network initialization(performs as well as Oracle).

Original: As a result, MRI-v1 achieve a lower testing loss than IRM-v1 and ERMin both task settings (Table 2).

Original: We ﬁnd similar results for accuracy (see Table 3).

Original: 5.2.2 Colored MNIST
Revised: Colored MNIST

Original: The Colored MNIST (CMNIST) dataset was introduced by Arjovsky et al. (2019) to analyze IRM-v1’sability to learn non-linear invariant predictors.
Revised: The Colored MNIST (CMNIST) dataset was introduced by Arjovsky et al.(2019) to analyze IRM-v1’s ability to learn non-linear invariant predictors.

Original: This is a synthetic binary classiﬁcation task derivedfrom MNIST.
Revised: This is a syntheticbinary classiﬁcation task derived from MNIST.

Original: In this dataset, the color of the image is determined in correlation with the label whosecorrelation probability p s changes between environments.
Revised: In this dataset, the color of the image is determinedin correlation with the label whose correlation probability p s changes between environments.

Original: The correlation probability p i betweendigit and label remains the same across all environments.
Revised: Thecorrelation probability p i between digit and label remains the same across all environments.

Original: Therefore, in this dataset color is thespurious feature and digit is the invariant feature.
Revised: Therefore,in this dataset color is the spurious feature and digit is the invariant feature.

Original: This dataset follows the data generating model in Fig 1.
Revised: This dataset follows thedata generating model in Figure 1.

Original: In this work, we simulate two versions of the Colored MNIST dataset corresponding to two differentsets of correlation probabilities p s and p i . The ﬁrst version, CMNISTa, uses the same correlationprobabilities as Arjovsky et al. (2019) with p i = 0 .
Revised: We use two versions of CMNIST, CMNISTa and CMNISTb, that are generated with two differentsets of correlation probabilities: CMNISTa uses the same correlation probabilities as Arjovsky et al.(2019) with p i = 0 .

Original: 75 for all environments, p s = 0 .
Revised: 75 for all environments, p s = 0 .

Original: 9 for the twotraining environments, and p s = 0 .
Revised: 9 for the two training environments, andp s = 0 .

Original: 1 for the testing environment. For the other version, CMNISTb,we use p i = 0 .
Revised: 1 for the testing environment; CMNISTb uses p i = 0 .

Original: 95 for all environments, p s = 0 .
Revised: 95 for all environments, p s = 0 .

Original: 8 , 1 for the two training environments, and p s = 0 .
Revised: 8 , 1 forthe two training environments, and p s = 0 .

Original: 1for the testing environment.
Revised: 1 for the testing environment.

Original: We also simulate the linear toy versions, toy-CMNIST (also known as thetwo-bit environment Kamath et al.
Revised: We also use an abstract version of CMNIST introduced by Kamath et al.

Original: , of both these datasets in which the image input is replacedby x = [ X 1 , X 2 ] , where X 1 2 {  1 , 1 } corresponds to the label of the digit  and X 2 2 {  1 , 1 }257 corresponds to the color. 
Revised: (2021), called toyCMNISTa/b, in which the image data is replaced by a two-bit data x = [ x d , x c ] , where x d 2 {  1 , 1 }corresponds to the label of the digit (without label noise) and x c 2 {  1 , 1 } corresponds to the color. 6.2 Result

Original: Result For both the CMNIST datasets, we use the LeNet architecture (LeCun et al., 1998).
Revised: In addition to IRM-v1 and MRI-v1, we report the performance of the ERM algorithm, whichminimizes the average risk across environments without any invariance constraints (Table 2). Forsynthetic datasets, we also report the performance of the Oracle (Table 2), which is trained using ERM on the dataset in which the spurious features are rendered uncorrelated with the label. For theshape-texture datasets, we use a small convolutional network (single convolution layer), while for theother image-based datasets, we use the LeNet architecture (LeCun et al., 1998).

Original: We ﬁndthat both MRI-v1 and  IRM-v1 perform better than ERM for both datasets.
Revised: Shape-Texture Tasks For both regression and classiﬁcation tasks, we ﬁnd that IRM-v1 trainedusing PM always converges to the zero-predictor (Table 2-3; Figure5). In comparison, MRI-v1trained using the penalty method achieved near-optimal OOD generalization and performed close to

Original:  MRI-v1 again performsidentically across all environments and almost as well the  Oracle for both datasets (Table 2).
Revised: Oracle (Table 2-3; Figure5). This demonstrates that MRI-v1 is far superior to IRM-v1 in terms ofboth loss and accuracy when trained using the penalty method. Similar to the PM case, MRI-v1 is able to perform close to Oracle when trained using ALM (Table 23).

Original: IRM-v1,however, has slightly worse OOD generalization performance than MRI-v1 for CMNISTa, but it ismuch worse for the CMNISTb dataset (Table 2).
Revised: In comparison, while IRM was able to avoid the zero-predictor, it still achieved worse test lossand accuracy than MRI-v1. Nevertheless, by using ALM, we were able to signiﬁcantly improve

Original: This result is consistent with the results for thelinear toy-CMNIST datasets. We ﬁnd that the IRM-v1 constraint does notcontain a non-invariantsolution for toy-CMNISTa, whereas for toy-CMNISTb it does contain non-invariant solutions (Fig 6). However, MRI-v1 constraints the predictor toignore the spurious feature ( w s = 0 ) in the toy versionof both the datasets (Fig 6) and therefore, it should converge to the invariant solution in both cases.
Revised: IRM-v1’s performance by avoiding the zero-predictor, while MRI-v1 performed similarly irrespectiveof the training method as the the zero-predictor is not a local constrained optimum of MRI-v1 (ALMslightly improves MRI-v1’s performance). Colored MNIST Similar to the shape-texture tasks, IRM-v1 always converges to the null-predictorwhen trained using the penalty method while MRI-v1 achieves a much better test loss and accuracythat IRM-v1 and ERM (Table 2- 3; Figure 6).

Original: Results for accuracy are consistent with the results of the linear toy version as well (see Table 3)).
Revised: This results for MRI-v1 holds true also when trainedusing ALM, albeit with a slightly improved performance (Table 2- 3; Figure 7).

Original: In conclusion, we conﬁrm that the results from the linear settings also apply to practical training ofdeep networks in non-linear image-based regression and classiﬁcation tasks. MRI-v1 provides best
Revised: In comparison,similar to the shape-texture tasks, IRM-v1 is able to avoid the null predictor when trained with ALM,but it achieves worse test loss than MRI-v1 on both the datasets (Table 2; Figure 7), especially for the

Original: OOD generalization across all tasks.
Revised: CMNISTb datasets.

Original: In fact, it performs almost as well as the Oracle.
Revised: In contrast to the loss results, IRM-v1 achieves a similar test accuracy to MRI-v1for both the datasets (Table 3; Figure 7). [Table]

Original: [Table]
Revised: [Table]

Original:  Discussion
Revised:  Discussion

Original: In this work [Table] we introduce the MRI formulation that learns invariant representations using a newnotion of invariance based on conserving the class-conditional feature expectation E e [ f ( x ) | y ] across environments. We derive a simpliﬁed version of the MRI formulation, MRI-v1, that allows us toapply this new notion of invariance to practical training of deep networks.
Revised: Limitations of the accuracy metric

Original: We provide analyticalguarantees that MRI will learn invariant predictors under certain conditions. We also demonstratenumerically in multiple task settings, including practical image-based tasks, that MRI provides muchbetter OOD generalization, than IRM (which conserves E e [ y | f ( x )] ) and ERM. In fact, it performsnearly as well as the Oracle.
Revised: The accuracy metric is widely used in the ﬁeld of domaingeneralization (Gulrajani and Lopez-Paz, 2020) to compare the performance of various algorithms.

Original: Disadvantage of IRM formulation. Both the MRI and the IRM formulation are based on seeminglyintuitive conservation laws to attain invariance. Why is itthen that the MRI formulation is so muchmore successful at preserving invariance in practice?
Revised: While accuracy does provide some insight into the performance of an algorithm, it is not always asensitive enough metric to discriminate between the representations learned by different algorithms.

Original: In this work, we note that  the failure of IRMis a side-effect of the extra constraint that requires E e [ y | f (x )] = f ( x )  , in addition to conserving
Revised: For example, we ﬁnd that for the toy-CMNISTb dataset, both invariant ( ( w i , w s ) ⇡ (3 , 0) )

Original: E e [ y | f ( x )] across environments. Without the extra constraint, the practical linear versions of boththe IRM and the MRI formulations match. To our knowledge, this is the ﬁrst study to identify thismajor ﬂaw in the IRM formulation. Convexity matters.
Revised: andnon-invariant predictors ( ( w i , w s ) 2 { (3 ,  2) , (3 , 2) } ) achieve the same accuracy (Figure 8).

Original: Deep learning in general is a highly non-convex optimization problem, yet itdoesn’t seem to have problem with convergence. It is understood that all minima is just as good as(equivalent to) the global minimum in the highly over-parameterized regime and central to this is theidea of mode-connectivity that suggests that all minima are connected to each other (Draxler et al. ,2018; Garipov et al., 2018).
Revised: Incomparison, the loss metric (Figure 8) can indeed differentiate between the invariant and non-invariantrepresentations.

Original: In this paper, we suggest that this property is highly dependent on thefact that most loss functions used in deep learning are convex with respect to thenetwork’s output. We argue that when loss functions have multiple disconnected local minima in the output space, theoverall learning problem with respect to the network parameters can also exhibit disconnected localminima, i.e., broken mode connectivity. To support our argument, we demonstrate that the IRMconstraint, which is non-convex in output, leads to multiple disconnected minima. However, the MRImethod, which uses a convex constraint, has no problem learning the invariant optima even in highlyover-parametrized networks.
Revised: This discrepancy between the loss and accuracy metric holds true even for non-linearimage-based datasets (Table 2-5).

Original: To our knowledge, this is the ﬁrst study to discuss the limitations ofusing constraints that are non-convex in the output.
Revised: To our knowledge, such discrepancy between loss and accuracy hasnot been reported before in the domain generalization literature.

Original: Finite samples.

Original: Although we have deﬁned the algorithms (MRI-v1 and IRM-v1) in terms ofpopulation losses, in practice these algorithms are trained with ﬁnite sample sizes. Kamath et al.(2021) pointed out that IRM-v1 learns the trivial 0 predictor for high penalty with ﬁnite sample sizes. Here, we discuss two additional disadvantages of training with ﬁnite sample sizes. Firstly, we ﬁndthat both MRI-v1 and IRM-v1 are prone to the property of deep networks of memorizing the dataset(Goodfellow et al., 2016) instead of learning patterns.
Revised: Based on this, we suggest that thereliance over accuracy as the sole metric of comparison between domain generalization algorithmscan be misleading, but loss and accuracy taken together can provide a meaningful comparison.

Original: We demonstrate this using a one-hiddenlayer ReLU network that learns the linear shape-texture regression dataset. Both the algorithmsﬁrst learns the correct solution (i.e, invariant solution for MRI-v1 and non-invariant solution for IRM-v1) but then starts memorizing the dataset.
Revised: Limitations IRM and MRI’s domain of applicability only coincide when there is no label shift orshift invariant feature distribution.

Original: However, since this overﬁtting issue arises as resultof ﬁnite sample size we can reduce it by increasing the sample size (Fig 4). Additionally, commonregularization techniques like mini-batch training (Smith et al. (2021); Fig 4) and weight decay(Goodfellow et al. ; Fig 4) can also reduce overﬁtting.
Revised: For datasets that do not follow these conditions, MRI, but not

Original: Secondly, we ﬁnd that due to ﬁnitesample size noise in the joint input-label distribution, MRI-v1 learns ’imprecise’ invariant predictor(i.e. , the training and testing loss don’t exactly match) and this precision improves with increase insample size (Fig 4). To our knowledge, this is the ﬁrst study to discuss these two disadvantages oflearning invariant representation with ﬁnite sample sizes. Limitations.
Revised: IRM, should fail.

Original: A major limitation of our proposed method is the assumption of no label shift acrossdomains. However, label shift can be easily corrected for.
Revised: However, note that, while IRM is expected to work for such shifts, IRM-v1 fails forsuch domain shifts ((Ahuja et al., 2020, 2021)).

Original: Mahajan et al. (2021) show that classconditional invariance can fail if P ( z c | y ) does notremain the same across domains, which alsoapplies to our method. Like many other invariant representation learning algorithms, including IRM(Ahuja et al., 2020, 2021), our method will be unable to handle covariate shift across domains. Also,like IRM (Ahuja et al., 2021), our method requires that there is signiﬁcant overlap in the support ofdifferent environment distributions.
Revised: This is because IRM-v1 and IRM are not the same,as demonstrated by us and Kamath et al.