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 conserving thefeature-conditioned label 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-v1, a practical version of IRM,can fail in various task settings.
Revised: However, more recent studies have demonstrated that IRM can fail in various tasksettings.

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

Original: We then introduce a complementary notion ofinvariance, MRI, based on conserving the label-conditioned feature expectation E e [ f ( x ) | y ] across environments, which is free of this ﬂaw.
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 asimpliﬁed, practical version of the MRI formulation called  MRI-v1.
Revised: Further, we introduce a simpliﬁed,practical version of the MRI formulation called as MRI-v1.

Original: We note thatthis constraint is convex which confers it with an advantage over  IRM-v1, whichimposes non-convex constraints.
Revised: We note that this constraint is convex which confers it with an advantage over the practical version of IRM, IRM-v1, which imposes non-convex constraints.

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

Original: We alsoempirically demonstrate that MRI strongly out-performs IRM and consistentlyachieves near-optimal OOD generalization in image-based nonlinear problems.
Revised: We also empirically demonstrate that MRI strongly out-performs 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 learned to use easier-to-ﬁt features that are spuriouslycorrelated to the label in the given training domains but do not generalize to the testing domains.
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 the trainingdomain will use the background color as the predictive feature which is easier to learn and generalizepoorly 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., 2016;
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:  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) and itspractical version, IRM-v1, in particular, have attracted much attention as  one of the ﬁrst methods thatare compatible with deep 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 empiricallydemonstrated that IRM-v1 is unreliable at learning invariant representations (Kamath et al., 2021; Rosenfeld et al., 2020; Gulrajani and 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).

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: Here, we identify afundamental ﬂaw of IRM formulation that causes this limitation and  propose anew method that isfree of 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: (2019) to learning invariant class-conditional features P ( f ( x ) | y )
Revised: (2020) 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 introduce a simpler formulation of IRM that is more generalizable, andshow that it can be 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: Wealso show that IRM has a fundamental ﬂaw due to the indirect way it imposes invariance which leadsto the failure of IRM-v1.
Revised: We also show that IRMhas a fundamental ﬂaw due to an extraneous constraint which causes it to fail.

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

Original: We prove that in a general linearproblem setting, MRI-v1 can guarantee invariant predictors 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 thenempirically demonstrate that MRI strongly out-performs IRM and consistently achieves near-optimal OOD generalization in image-based nonlinear problems.
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: 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 ) } .
Revised: 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 ) } .

Original: Thedistribution P e ( x, y ) is assumed to be generated according to the causal graph in Fig.1 as introducedin Rosenfeld et al. (2020), which includes latent features that are invariantly z 
Revised: 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 z i or spuriously

Original: i or spuriously z scorrelated with the label, from which the observation x is generated.
Revised: z scorrelated with the environment. The observation x is generated as a function of the latent features:x =

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

Original: 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 labely ( Z i ! Y ). This is  the graph consideredby Arjovsky et al.
Revised: In this formulation, we have using d-separation that the label is independent from the environmentgiven the invariant latent feature: y ??

Original: (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: e | z i and the invariant latent feature is independent from the 65 environment given the label:

Original: Y

Original: Z i
Revised: z i

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

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

Original: [Equation]
Revised: [Equation]

Original: [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 k 2 , and thebinary-cross-entropy loss l BCE ( o, y ) =  (1 + y ) log((
 +tanh( o/ 2)) / 2)  (1  y ) log((
  tanh( o/ 2)) / 2) (for+}
Revised: {+= 12 k o  y k 2 , and thebinary-cross-entropy loss l BCE ( o, y ) =  (1 + y ) log((
 +tanh( o/ 2)) / 2)  (1  y ) log((
  tanh( o/ 2)) / 2) (for+}

Original: ⌘ is
Revised: Y =

Original:  3.1 IRM Paradigm3.1.1 Original formulation
Revised: {  1 , 1 } ) 2 . 3.1 IRM paradigm

Original: Deﬁnition 3.1.
Revised: 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]

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

Original: Given a predictor  f , the feature representation f : X !
Revised: Deﬁnition 3.1. A representation f : X !

Original: Z80 is IRM invariant over a set of environments 
Revised: Z is IRM invariant over a set of environments E

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

Original: O such that is simultaneouslyoptimal on f for all environments in E, i.e.
Revised: O such that is simultaneously optimal on f for all environments [Equation] 2 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: Lemma 3.2. (Kamath et al., 2021) For standard loss functions 2 , Deﬁnition 3.1 is equivalent to
Revised: This deﬁnition has shown to be equivalent to the following condition (Kamath et al., 2021),

Original: [Equation]
Revised: [Equation]

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

Original: [ y ] is the feature-conditioned label expectation, and  is a monotonic 85 function that depends on the loss function.
Revised: o for square loss,  ( o ) =

Original: 3.1.2 Variational formulation
Revised: 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.

Original: The original formulation above is overly complex due to the composite predictor  f and the 88 optimality condition on .
Revised: 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.

Original: For further analysis, we introduce a simpliﬁed formulation of IRM deﬁnedby the following variational principle.
Revised: 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.

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

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 ) remains stationary under arbitrary inﬁnitesimal perturbations on the predictor’s outputo
Revised: 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

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

Original: [Equation]
Revised: [Equation]

Original: [Equation] 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.
Revised: where  ( · ) is an arbitrary perturbation that is unrestricted in the space of all measurable functions.

Original: Lemma 3.4.
Revised: Lemma 3.3.

Original: For standard loss functions 2 , Deﬁnition 3.3 is equivalent to8 e 2 E , E e [ y | o ] =  ( o ) , (7) which is equivalent to Lemma 3.2 with f replacing the composite predictor  f .
Revised: This reformulated deﬁnition 3.2 is equivalent to eq (6) .

Original: Proof.
Revised: Proof.

Original: Eq (6) is satisﬁed if and only if E P e
Revised: Eq (7) 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: For standard loss functions, the lossderivative has the form @
Revised: Note that @

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

Original: =
Revised: =

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

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

Original: 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.
Revised: Therefore, we arrive at

Original: [Equation]
Revised: [Equation]

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

Original: 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 106 eq (5)) additionally requires that value to be also equal to the RHS term  ( o ) .

Original: 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:  Due to the impracticality of considering  theunrestricted function space of , Arjovsky et al.
Revised: IRM-v1 Due to the impracticality of considering optimal classiﬁers in the space of unrestrictedfunctions, Arjovsky et al.

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

Original: 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 .
Revised: 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 ) =

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

Original: [Equation]
Revised: [Equation]

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

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

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

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

Original: In Arjovsky et al.(2019), this reduced formulation is termed IRM-v1when the constraints are imposed in a soft manner ,i.e. as squared penalty terms (See eq (16)). 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.
Revised: [ @ o [Equation] This reduced formulation is called IRM-v1 in Arjovsky et al.

Original: Despite the practical applicability, it has been empirically shown that IRM-v1 can behave quitedifferently from IRM (Kamath et al., 2021).
Revised: 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).

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

Original: E  e
Revised: E P e

Original: [ o (  ( o )  y )] for standard loss functions 2 ,  eq (9) is equivalent to [Equation]
Revised: o ) [ E e [ y | o ] · o ] , a necessary condition forthe conservation of E e [ y | o ] in eq (4).

Original: Unlike eq (7), however, RHS of  eq (10) does not necessarilyremain constant across environments,and hence does not prescribe a meaningful invariance. 
Revised: 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. Although this inconsistency was empirically shown in Kamathet al.

Original: Therefore, restricting the function space of(or ) invalidates the indirect way of imposing invariance in IRM involving an intermediary term.
Revised: , the actual cause of the problem was never discussed. Our analysis reveals that the rootcause originates from the extraneous constraint of IRM. Square loss [Equation]

Original: 3.2 MRI Paradigm
Revised: BCE loss [Equation] Constraint

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

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

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

Original: O is MRI invariant over a set of environments E , if the changein risk due to arbitrary inﬁnitesimallabel perturbations is shared across environments, i.e.
Revised: 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.

Original: [Equation]
Revised: [Equation]

Original: 11) where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: where  : Y !

Original: 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 )] .
Revised: where  ( · ) is an arbitrary perturbation that is unrestricted in the space of all measurable functions.

Original: Lemma 3.6.
Revised: Lemma 3.5.

Original: For standard loss functions, Deﬁnition 3.5 is equivalent to

Original: [Equation] 13) which conserves the label-conditioned feature expectation E e [ o | y ] ⌘ E P 
Revised: Deﬁnition 3.4 implies the conservation of E e [ o | y ] ⌘ E P e

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

Original: [ o ] across environments.
Revised: [ o ] across environments,

Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: Eq (11) is satisﬁed
Revised: Eq (10) 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: For standard loss functions, theloss derivative has the form @

Original:  y l  ( o, y ) =
Revised: Since @ y l sq ( o, y ) =

Original:  o + ⇢ ( y ) 4 . Therefore, E e 1 [ @ y l ( o, y) | y ]  E e 2 [ @ y l  ( o, y )| y ] =[Equation]
Revised: y  o and@ y l BCE ( o, y ) =

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

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

Original: y , reduceseq (12) to  L e ( f ) = E e [ @ y l ( o, y ) · y ] . This reduces the conservation law eq (13)to
Revised: y , reduces the MRI formulation eq (11) to

Original: [Equation]
Revised: [Equation]

Original: e [Equation]
Revised: (13) which yields the following conservation law

Original: [Equation]
Revised: [Equation]

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

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

Original: MRI’s constraints imposing invariance continues to hold even when  is restricted to the space oflinear functions. Algorithm [Equation]
Revised: 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.

Original: Methods
Revised: 3.3 Objective

Original: 4.1 Constrained optimization problem

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

Original: [Equation]
Revised: [Equation]

Original: subject to ~ c ( f ) = ~
Revised: where the objective ¯ L tr ( f ) = |E

Original: [Equation]
Revised: [Equation]

Original: (15) 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 , andw is the parameters of the predictor f . The constraint functions are ~ c = ~ L 2 R |E tr | for IRM-v1and ~ c = Q ~ L 2 R |E tr |  1 for MRI-v1, where ~ L 2 R |E tr | is a vector of perturbed risks  L e fore 2 E tr . Q 2 R ( |E tr |  1) ⇥ |E tr | is an orthonormal matrix that satisﬁes Q ~ 1 = ~ 0 , such that Q ~ L computesthe differences of  L e between environments.
Revised: |

Original: For example, for a training set of two environments E tr = { e 1 , e 2 } , Q ~ L = (  L e 1   L e 2 ) / p 2 , since Q = [1 ,  1] / p 2 . The constraints are summarized 155 in Table 1. 4.2 Soft-constraint methods
Revised: P e 2 E tr L e ( f ) is the average risk over the set of training environments E tr ⇢ E , ✓ is the parameters of the predictor f , and the vector-valued constraint function 135 contains the perturbed risks

Original: Numerically solving IRM-v1 and MRI-v1 requires converting the hard constraints ~ c ( f ) = ~ 0 to softconstraints, which allows the use of off-the-shelf gradient-based optimization algorithms. Penalty Method (PM)
Revised: 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).

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

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

Original: [Equation]
Revised: [Figure]

Original:  However, this method requires increasing µ t ! 1 over training iteration 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 (16)
Revised:  Linear setting

Original: [Equation]
Revised: [Figure]

Original: 17) where ~ is typically initialized at ~ 0 and updated at each training iteration t to accumulate the residualconstraints ~ c ( f ( w t )) .

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

Original: Analytic Results

Original: 5.1 General Linear SEM
Revised: 4.1 General linear SEM

Original:  In this section, we show that MRI-v1 can correctly remove all spurious features of a linear predictorgiven a sufﬁcient number of environments.
Revised: [Figure] In this section, we show that MRI-v1 can correctly remove all spurious features when enoughenvironments are given.

Original: Consider a data generating process according to Figure 1,in whichthe observation x = g ( z i , z s ) is an injective linear [-function of the latent features z i 
-]
Revised: We consider a data generating process according to 1, in which x = g ( z i , z s ) is a linear [-function of the latent features z i 
-]function.

Original: R d i , z s 2 R d s . Note that thisStructural Equation Model (SEM) (Pearl, 2009) does not require anyassumptions on the generation process Y ! Z i , Z s , which generalizes the SEM  of  Rosenfeld et al.(2020), which additionally assumed binary labels and additive Gaussian noise for generating thelatent features.
Revised: 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.

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: [Equation] with w = { w i 2 R d i , w s 2 R d s } .
Revised: (16) with ✓ = { w i 2 R d i , w s 2 R d s } .

Original: Theorem 5.1.
Revised: Theorem 4.1.

Original: Given |E tr | > d s training environments,  and that E e [ z s y ] are in general linear 181 positions, MRI-v1 will eliminate all spurious feature dimensions.
Revised: 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.

Original: Proof.
Revised: Proof.

Original: MRI-v1’s constraint eq (14) yields
Revised: MRI-v1 condition:

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

Original: [Equation] where M0 = 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
Revised: ) 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 ] .

Original: [Equation]

Original: e [ z s y ] .

Original: Since rank ( M0 ) =
Revised: Since rank ( M ) =

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

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

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

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

Original: In fact, IRM-v1 has been shown to fail in more general linear problems in which thenoise is not additive Gaussian (Arjovsky et al., 2019; Kamath et al., 2021).5.2 Minimal Example: d i = 1 , d s = 1 , |E tr | =
Revised: 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).

Original: Here, we demonstrate aminimal case of the above general linear  problem that involves one invariantfeature, one spurious feature and two training environments (Shape-Texture linear regression in Section 6.1).
Revised: Here, we consider a linear regression problem on the Shape-texture dataset which has 2 trainingenvironments and include 1 invariant and 1 spurious features.

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

Original: Figure 2A shows the results for the hard-constrained problems. IRM-v1 IRM-v1 has two quadratic equality constraints, E e 1 [ o 2  oy ] =
Revised: IRM-v1 IRM-v1 has two quadratic equality constraints, shown as two ellipses in Fig 2a.

Original: E e 2 [ o 2  oy ] = 0 , shownas two 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. Note that  one of thenon-invariant solutions exhibits lower train loss than the true invariant optimum solution.
Revised: 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 which is 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).

Original: IRM-relaxed The extraneous constraint of [-IRM-v1 can be relaxed to yield E e 1 [ o 2  oy ]  E e 
 [ o 2  202 oy ]-] = 0 .
Revised: 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 [-IRM-v1 can be relaxed to yield E e 1 [ o 2  oy ]  E e 
 [ o 2  202 oy ]-]invarian solutions w s = 0 .

Original: This constraint features a pair of hyperbolic curves, that appears as two straight lines. Dueto the non-convex constraint, this problem exhibits two local constrained optima, including the trueinvariant optimum and a non-invariant solution, with the non-invariant solution exhibiting lower trainloss.
Revised: That is, every point on theconstraint line is a invariant predictor. 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).

Original: Therefore, the failure mode ofIRM-v1 cannot be resolved by simply relaxing the extraneousconstraint without ﬁxing the indirect way IRM imposes invariance.
Revised: 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 ) =  L e 2 ( f ) .

Original: MRI-v1 MRI-v1 has one linear equality constraint, [-E e 1 [ oy ]  E e 
 [ oy ] = 0 , that exactly prescribesthe-] set of invariant solutions w s = 0 .
Revised: Since it has only a single constraint, [-E e 1 [ oy ]  E e 
 [ oy ] = 0 , that exactly prescribesthe-]it has 1dcontinuous set of points as possible solutions (Fig 2a).

Original: That is, any point on the constraint is  a invariant predictor.
Revised: However, it is still a nonconvex constraint.

Original: This is a convex problem, since both the objective and the constraint are convex, and thus features aunique optimum, which isthe true invariant optimum solution. Convergence Dynamics Figure 2B/C show the convergence dynamics ofthe soft-constrainedoptimization problems under full-batch gradient descent (penalty method, µ = 5 ⇥ 10 4 , eq (16)).
Revised: Inaddition to the w s = 0 line, it also exhibits w i = 2 / 3 line as the solution (Fig 2a). As a result, the¯ L tr ( f ) objective exhibits two constrained optimum solutions, only one of which is invariant (Fig 2a). Learning Dynamics In Fig 2(b,c), we show simulation of gradient descent learning dynamics ofthe objective eq (19) ( µ = 5 ⇥ 10 4 ).

Original: IRM-v1’s convergence dynamics is highly dependent on the initialization (shown with differentlycolored trajectories), due to the presence of multiple local minima.
Revised: The convergence pattern for IRM-v1 is highly dependent onthe initialization, due to the presence of multiple local minima.

Original: Note that most trajectories do notconverge to the true invariant optimum. IRM-relaxed also exhibits complex dynamics due to a saddlepoint near thetrue invariant solution: The trajectories seem toapproach the invariant solution at ﬁrst,but eventually they all converge to  the non-invariant solution.
Revised: In this case, the invariant optimumis especially extremely difﬁcult to converge to as it has higher risk than the nearby non-invariantminima.

Original: Note that thetrue invariant optimum islost at this value of µ , which exists in the limit µ ! 1 .In contrast, MRI-v1 always converges to thetrue invariant optimum regardless of initialization since it is a unique minimum. T r a i n
Revised: In contrast, MRI-v1 always converges to the invariant optimum regardless of initializationsince it is a unique minimum. For IRM-relaxed, the invariant optimum point is lost at this valueof µ , which exists in the limit µ ! 1 .

Original: R i s k
Revised: As a result, all trajectories converge to the non-invariantoptimum, which has lower training risk. We show similar results for other linear problems includingshape-texture classiﬁcation and toy-CMNIST in Supplementary Materials. 5.1 Methods: Constrained Optimization Solving constrained optimization numerically requires converting the problem into an unconstrainedoptimization problem such that the constraint is imposed as a soft regularization term.

Original: T e s t
Revised: This allowsthe use of any off-the-shelf gradient-based optimization algorithm.

Original: R i s k [Figure]
Revised: Penalty Method (PM)

Original: Minimal examples for linear binary classiﬁcation on the shape-texture and the toy-CMNISTa/bdatasets are also analyzed in Supplementary Materials (Figure 3).
Revised: Penalty method is the most commonly used approach, including in Arjovskyet al.

Original: Nonlinear Image-based Problems
Revised: (2019), which adds the squared constraint as a penalty term to the objective, [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)

Revised: [Equation]

Original: Here, we explore how MRI-v1 performs in nonlinear problems.
Revised: 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: ~ is typically initialized at ~ 0 , which has an annealing effect on ALM.

Original: Note that unlike the linear problem,there is no guarantee for either IRM-v1 or MRI-v1 for these nonlinear problems.
Revised: Moreover, ALM has an interesting effect when used with mini-batches, in which the constraints haveslight mismatches across iterations. Since the Lagrange multiplier  is accumulated over iterations, itimposes the model to satisfy an averaged constraint. In contrast, PM penalizes the model to satisfy allof the conﬂicting constraints, and thus decaying the model toward the zero-predictor. Unless notedotherwise, we used ALM by default for our optimization.

Original: 6.1 Datasets
Revised: 5.2 Datasets

Original:  Shape-Texture Dataset
Revised: 5.2.1 Shape-Texture Dataset

Original: We introduce a new versatile dataset designed for domain generalizationon linear and non-linearimage-based regression and classiﬁcation problems. The generative processof the dataset involves an invariant feature z i
Revised: We develop a new versatile dataset designed for domain generalization, called the Shape-Texturedataset, based on the data generating model in Fig 1. This dataset is designed to test algorithms onlinear and non-linear regression and classiﬁcation problems in multiple adjustable settings.

Original: = e i ✓ i  , a spurious  featurez s = e i ✓ s  , and a label featuree i ✓ y, each of which represents an orientation on a complex unit circle: e i ✓ 2 S 1 . The angles oforientations are generated as ✓ y ⇠ U S 1 , where U S 1 is the circular uniform distribution, and for⇤ 2 { i, s } , ✓ ⇤ = ✓ y with probability p ⇤ or ✓ ⇤ ⇠ U S 1 with probability 1  p ⇤ . The parameter p i = 0 .
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: 75 is ﬁxed across environments, whereas p s varies from one environment toanother. We consider two training environments E tr = { e 1 , e  2 } with p s e 1 = 1 , p s e 1= 0 . 8 and onetesting environment E test = e 0 } with p s e 0 = 0 .
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 task (section 5.2), the observed input is simply the concatenated latentfeatures 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: In the linear classiﬁcation task, the input isx = [sin( ✓ i ) , sin( ✓ s )]  and the label is 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 )), where H is the sign function.
Revised: = H (sin( ✓ y )) where H denotes the Heavyside step function.

Original: In the nonlinear regression/classiﬁcation tasks, the input observation is an image  composed of two planarwaves, in which ✓ i is the orientation of the low frequency wave ( i.e. shape) and ✓ s is that of the highfrequency wave ( i.e. texture), as shown in Figure 4.
Revised: 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 . 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 } . p i ,  i are ﬁxed invariant parameters whereas p s ,  s can vary from 228 one environment to another. In this work, for all task settings, we use two training environmentsand one testing environment with p i = 0 .

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

Original: Colored MNIST (CMNIST) CMNIST (Arjovsky et al., 2019) is a synthetic dataset derived from
Revised: 8 , 1 for the two trainingenvironments, p s = 0 for the testing environment, and  i =  s = 0 for all environments. Result We use a small convolutional network for the image-based regression and classiﬁcationdataset. 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).

Revised: 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: MNIST for binary classiﬁcation.
Revised: 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). We ﬁnd similar results for accuracy (see Table 3). However, MRI-v1 gives almost identicalperformance across training and testing environments irrespective of the network initialization(performs as well as Oracle). As a result, MRI-v1 achieve a lower testing loss than IRM-v1 and ERMin both task settings (Table 2).

Original:  The  dataset features a ﬁxed correlation probability p
Revised: We ﬁnd similar results for accuracy (see Table 3). 5.2.2 Colored MNIST The Colored MNIST (CMNIST) dataset was introduced by Arjovsky et al. (2019) to analyze IRM-v1’sability to learn non-linear invariant predictors.

Original: i betweenthe digit and the label, and p s  between the color and  the label, which changes across  environments. 
Revised: This is a synthetic binary classiﬁcation task derivedfrom MNIST. In this dataset, the color of the image is determined in correlation with the label whosecorrelation probability p s changes between environments. The correlation probability p i betweendigit and label remains the same across all environments. Therefore, in this dataset color is thespurious feature and digit is the invariant feature.

Original: We consider two versions, CMNISTa and CMNISTb, with two sets of environmental parameters:
Revised: This dataset follows the data generating model in Fig 1. In this work, we simulate two versions of the Colored MNIST dataset corresponding to two differentsets of correlation probabilities p s and p i .

Original: CMNISTa is the version from Arjovsky et al.
Revised: The ﬁrst version, CMNISTa, uses the same correlationprobabilities as Arjovsky et al.

Original: (2019) with p i = 0 .
Revised: (2019) with p i = 0 .

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

Original: 9 , p se 1 = 0 . 8 , andp s e 0 = 0 .
Revised: 9 for the twotraining environments, and p s = 0 .

Original: 1 for thetraining E tr = { e 1 , e 2 } and the testing E test =
Revised: 1 for the testing environment.

Original: { e 0 } environments; CMNISTb 247 uses p i = 0 . 
Revised: For the other version, CMNISTb,we use p i = 0 . 95 for all environments, p s = 0 .

Original: 9 ,p s e 1 = 1 , p se 1 = 0 .[-8 , p s e 
-]
Revised: 8 , 1 for the two training environments, and p s = 0 .[-8 , p s e 
-]

Original: =
Revised: 1for the testing environment.

Original: We also use the abstracted versions, calledtoy-CMNISTa/b, in which theinput image  is replaced by the two-bit observation data x = [ x d , x c ],where x d , x c 2 {  1 , 1 } correspond to the  digitbit and  the color bit, respectively (Kamath et al., 250 2021). Note that both Shape-Texture and CMNIST datasets can be equally understood as being generatedfrom Fig 1 with causal directions Y !
Revised: We also simulate the linear toy versions, toy-CMNIST (also known as thetwo-bit environment Kamath et al. , 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. Result For both the CMNIST datasets, we use the LeNet architecture (LeCun et al., 1998).

Original: Z i or Z i ! Y (Figure 9). 6.2 Result 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).
Revised: We ﬁndthat both MRI-v1 and IRM-v1 perform better than ERM for both datasets.

Original: Forsynthetic datasets, we also report the performance of the Oracle  (Table 2), which is trained using
Revised: MRI-v1 again performsidentically across all environments and almost as well the Oracle for both datasets (Table 2).

Original: 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).
Revised: IRM-v1,however, has slightly worse OOD generalization performance than MRI-v1 for CMNISTa, but it ismuch worse for the CMNISTb dataset (Table 2). This result is consistent with the results for thelinear toy-CMNIST datasets.

Original: Shape-Texture Tasks For both regression and classiﬁcation tasks, we ﬁnd thatIRM-v1 trainedusing PM always converges to the zero-predictor (Table 2-3; Figure5).
Revised: We ﬁnd that the IRM-v1 constraint does not contain a non-invariantsolution for toy-CMNISTa, whereas for toy-CMNISTb it does contain non-invariant solutions (Fig 6).

Original: In comparison, MRI-v1 trainedusing PM achieves near-optimal OOD generalization, performing close to Oracle (Table 2-3; Figure5). This demonstrates that MRI-v1 is far superior to IRM-v1 in terms of both loss and accuracy whentrained using PM. This is an important distinction between MRI-v1 and IRM-v1 and is a consequenceof the zero-predictor being a constrained local optimum of IRM-v1, but not of MRI-v1.
Revised: However, MRI-v1 constraints the predictor to ignore 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.

Original: Similar to the PM case, MRI-v1 is able to perform close to Oracle when trained using ALM (Table 23).
Revised: Results for accuracy are consistent with the results of the linear toy version as well (see Table 3)).

Original: 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 IRM-v1’s performance by avoiding the zero-predictor. In comparison, MRI-v1 achieved near-optimal
Revised: 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.

Original:  OOD generalization irrespective of the training method. Colored MNIST Similar to the shape-texture tasks, IRM-v1 always converges to the null-predictorwhen trained using PM, while MRI-v1 achieves a signiﬁcantly better test loss and accuracy thanboth IRM-v1 and ERM (Table 2- 3; Figure 6). This result for MRI-v1 holds true also when trainedusing ALM (Table 2- 3; Figure 7).
Revised: MRI-v1 provides best OOD generalization across all tasks.

Original: In fact,  the ALM version of MRI-v1 has better performance.
Revised: In fact, it performs almost as well as the Oracle.

Original: In comparison, similar to the shape-texture tasks, IRM-v1 is able to avoid the null predictor whentrained with ALM, but it achieves worse test loss than MRI-v1 on both datasets (Table 2; Figure 7),especially on the CMNISTb datasets. In contrast to the loss results, the ALM version of IRM-v1achieves a similar test accuracy to MRI-v1 on both datasets (Table 3; Figure 7).
Revised: [Table]

Original:  Discussion
Revised:  Discussion

Revised: In this work

Original: [Table]
Revised: [Table]

Original: Limitations of the accuracy metric
Revised: 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. 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.

Revised: In fact, it performsnearly as well as the Oracle.

Revised: Disadvantage of IRM formulation.

Original: The accuracy metric is widely used in the ﬁeld of domaingeneralization (Gulrajani and Lopez-Paz, 2020) to compare the performance of various algorithms.
Revised: Both the MRI and the IRM formulation are based on seeminglyintuitive conservation laws to attain invariance. Why is it then that the MRI formulation is so muchmore successful at preserving invariance in practice? In this work, we note that the failure of IRMis a side-effect of the extra constraint that requires E e [ y | f ( x )]

Original: However, accuracy is not sensitive enough to reliably measure the degree of invariance of a model.
Revised: = f ( x ) , in addition to conserving E e [ y | f ( x )] across environments. Without the extra constraint, the practical linear versions of boththe IRM and the MRI formulations match.

Original: For example, consider the linear classiﬁcation problem of toy-CMINSTb shown in Figure 8.
Revised: To our knowledge, this is the ﬁrst study to identify thismajor ﬂaw in the IRM formulation. Convexity matters. 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.

Original: Forthis problem, MRI-v1 has a unique optimum, which is  the optimal invariant solution, while IRM-v1exhibits multiple  local optima, including non-invariant solutions.
Revised: ,2018; Garipov et al., 2018). 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 the network’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.

Original: Yet, all these solutions exhibitidentical accuracy (except for the zero-predictor solution of IRM-v1), since the accuracy landscape(Figure 8f) is ﬂat over a wide range of solutions and remains invariant across domains: invariantaccuracy is achieved as long as w i > | w s | .
Revised: 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.

Original: This shows that accuracy cannot be used as a measure ofinvariance in this problem.
Revised: To our knowledge, this is the ﬁrst study to discuss the limitations ofusing constraints that are non-convex in the output.

Original: Therefore, the fact that MRI-v1 and IRM-v1 exhibit similar test-domainaccuracy does not necessarily indicate that they achieve similar degree of invariance. Limitations A limitation of both MRI and IRM (Rosenfeld et al., 2020; Ahuja et al., 2021)
Revised: Finite samples. 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.

Original: isthat it requires signiﬁcant overlap in the support of different domain distributions.
Revised: Here, we discuss two additional disadvantages of training with ﬁnite sample sizes.

Original: This limits theapplicability of both MRI and IRM on most DomainBed (Gulrajani and Lopez-Paz, 2020) benchmarks.
Revised: 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.

Original: Rotated-MNIST, for example, expects an algorithm to generalize to an orientation of the digit thatwas never seen in the training set. In this case, the support of the training and testing set clearly doesnot overlap.
Revised: 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)

Original: Similar arguments can be made for other datasets.
Revised: but then starts memorizing the dataset. 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.

Original: For example, images with differentstyles across testing and training domains may not satisfy the signiﬁcant support overlap condition. IRM and MRI’s domain of applicability only coincide when there is no label shift or shift invariantfeature distribution.
Revised: ; Fig 4) can also reduce overﬁtting. 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).

Original: For datasets that do not follow these conditions, MRI, but not IRM, should fail. However, note that, while IRM is expected to work for such shifts, IRM-v1 fails for such domainshifts (Ahuja et al., 2020, 2021).
Revised: To our knowledge, this is the ﬁrst study to discuss these two disadvantages oflearning invariant representation with ﬁnite sample sizes.

Original: This is because IRM-v1 and IRM are not the same, as demonstratedby us and Kamath et al. Checklist The checklist follows the references. Please read the checklist guidelines carefully for information onhow to answer these questions.
Revised: Limitations.

Original: For each question, change the default [TODO] to [Yes] , [No] , or[N/A] .

Original: You are strongly encouraged to include a justiﬁcation to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description.

Original: For example:

Original: • Did you include the license to the code and datasets?

Original: [Yes] See Section ??.•

Original: Did you include the license to the code and datasets?

Original: [No] The code and the data areproprietary.•

Original: Did you include the license to the code and datasets?

Original: [N/A] Please do not modify the questions and only use the provided macros for your answers. Note that the Checklist section does not count towards the page limit.
Revised: A major limitation of our proposed method is the assumption of no label shift acrossdomains.

Original: In your paper, please delete this instructionsblock and only keep the Checklist section heading above along with the questions/answers below.

Original: For all authors...

Original: (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’scontributions and scope?

Original: [Yes] (b) Did you describe the limitations of your work? [Yes] See Section 7.(c) Did you discuss any potential negative societal impacts of your work?
Revised: However, label shift can be easily corrected for.

Original: [N/A]

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?

Original: [Yes]2.

Original: If you are including theoretical results...

Original: (a) Did you state the full set of assumptions of all theoretical results? [Yes] See Section 3and 5.(b) Did you include complete proofs of all theoretical results? [Yes] See Section 3 and 5.3.
Revised: Mahajan et al.

Original: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)? [Yes]
Revised: (2021) show that classconditional invariance can fail if P ( z c | y ) does not remain the same across domains, which alsoapplies to our method.

Original: https://anonymous.4open.science/r/MRI-2608/readme (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
Revised: 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.

Original: [Yes] See Supplementary Materials.

Original: (c) Did you report error bars (e.g., with respect to the random seed after running experi-

Original: ments multiple times)?

Original: [Yes] See Table 2 and 3.

Original: (d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?

Original: [Yes] See Supplementary Materials.

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...

Original: (a) If your work uses existing assets, did you cite the creators?

Original: [Yes] See Section 6.1 andSupplementary Materials.

Original: (b) Did you mention the license of the assets?

Original: [Yes] See Supplementary Materials.

Original: (c) Did you include any new assets either in the supplemental material or as a URL?

Original: [Yes]

Original: https://anonymous.4open.science/r/MRI-2608/readme

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating? [N/A] (e) Did you discuss whether the data you are using/curating contains personally identiﬁable
Revised: Also,like IRM (Ahuja et al., 2021), our method requires that there is signiﬁcant overlap in the support ofdifferent environment distributions.

Original: information or offensive content?

Original: [N/A]

Original: If you used crowdsourcing or conducted research with human subjects...

Original: (a) Did you include the full text of instructions given to participants and screenshots, if

Original: applicable?

Original: [N/A]

Original: (b) Did you describe any potential participant risks, with links to Institutional Review

Original: Board (IRB) approvals, if applicable?

Original: [N/A]

Original: (c) Did you include the estimated hourly wage paid to participants and the total amount

Original: spent on participant compensation?

Original: [N/A]