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: (2019) 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 ) 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 ) over inputsand 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 introduced in
Revised: Thedistribution P e ( x, y ) is assumed to be generated according to the causal graph in Fig.1 as introduced in

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

Original: (2020), which includes latent features that are invariantly z
Revised: (2020), which includes latent features that are invariantly z

Original: i or spuriously z s correlatedwith the label.
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 ) isassumed to be injective, which makes it possible to recover the latent features from the 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: 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
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

Original: ( Z i !
Revised: ( Z i !

Original: Y ).
Revised: Y ).

Original: This is the graph consideredby Arjovsky et al.
Revised: This is the graph consideredby Arjovsky et al.

Original: (2019) in Figure 3 with the difference being that they allow both the invariantfeature and label distribution to also vary across environments.
Revised: (2019) in Figure 3 with the difference being that they allow both the invariantfeature and label distribution to also vary across environments.

Original: Also, while Arjovsky et al.
Revised: Also, while Arjovsky et al.

Original: (2019)focused on the causal graph with Z i !
Revised: (2019)focused on the causal graph with Z i !

Original: Y , the IRM method is also applicable to the causal graphconsidered in Figure 1 (Rosenfeld et al., 2020).
Revised: Y , the IRM method is also applicable to the causal graphconsidered in Figure 1 (Rosenfeld et al., 2020).

Original: Y
Revised: Y

Original: Z i
Revised: Z i

Original: Z s
Revised: Z s

Original: X
Revised: X

Original: E
Revised: E

Original: The risk of a predictor f : X !
Revised: The risk of a predictor f : X !

Original: O in environment e isdeﬁned as the population average
Revised: O in environment e isdeﬁned as the population average

Original: [Equation]
Revised: [Equation]

Original: [Equation] [Equation]
Revised: [Equation] [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 l log ( o, y ) =
Revised: = 12 k o  y k l log ( o, y ) =

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

Original: =
Revised: =

Original: 1 / (1 + e o ) is the sigmoid function.
Revised: 1 / (1 + e o ) is the sigmoid function.

Original: 3.1 IRM Paradigm
Revised: 3.1 IRM Paradigm

Original: 3.1.1 Original formulation
Revised: 3.1.1 Original formulation

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

Original: (Arjovsky et al., 2019)
Revised: (Arjovsky et al., 2019)

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

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

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

Original: O such that is simultaneouslyoptimal on f for all environments in E , i.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: Lemma 3.2.
Revised: Lemma 3.2.

Original: (Kamath et al., 2021)
Revised: (Kamath et al., 2021)

Original: For standard loss functions 3 , Deﬁnition 3.1 is equivalent to
Revised: For standard loss functions 3 , Deﬁnition 3.1 is equivalent to

Original: [Equation]
Revised: [Equation]

Original: ) where E e [ y | f ( x )]
Revised: ) where E e [ y | f ( x )]

Original: ⌘ E P e ( y | f ( x ))
Revised: ⌘ E P e ( y | f ( x ))

Original: [ y ] is the feature-conditioned label expectation, and  is a monotonic 86 function that depends on the loss function.
Revised: [ y ] is the feature-conditioned label expectation, and  is a monotonic 86 function that depends on the loss function.

Original: 3.1.2 Variational formulation
Revised: 3.1.2 Variational formulation

Original: The original formulation above is overly complex due to the composite predictor  f and the 89 optimality condition on .
Revised: The original formulation above is overly complex due to the composite predictor  f and the 89 optimality condition on .

Original: For further analysis, we introduce a simpliﬁed formulation of IRM deﬁnedby the following variational principle.
Revised: For further analysis, we introduce a simpliﬁed formulation of IRM deﬁnedby the following variational principle.

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

Original: = f
Revised: = f

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

Original: [Equation]
Revised: [Equation]

Original: (6) where  : O !
Revised: (6) where  : O !

Original: 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: 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.4.
Revised: Lemma 3.4.

Original: For standard loss functions 3 , Deﬁnition 3.3 is equivalent to
Revised: For standard loss functions 3 , Deﬁnition 3.3 is equivalent to

Original: [Equation]
Revised: [Equation]

Original: e
Revised: e

Original: [Equation]
Revised: [Equation]

Original: E ,
Revised: E ,

Original: [Equation]
Revised: [Equation]

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

Original: Proof.
Revised: Proof.

Original: Eq (6) 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: For standard loss functions, the lossderivative has the form @
Revised: For standard loss functions, the lossderivative has the form @

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

Original: =
Revised: =

Original:  y +  ( o ) , which yields E e [ @ o l ( o, y )
Revised:  y +  ( o ) , which yields E e [ @ o l ( o, y )

Original: | o ]
Revised: | o ]

Original: =  E
Revised: =  E

Original: e
Revised: e

Original: [ y | o ]+  ( o ) = 0 .
Revised: [ y | o ]+  ( o ) = 0 .

Original: 3.1.3 Conservation law of IRM
Revised: 3.1.3 Conservation law of IRM

Original: It was noted in Arjovsky et al.
Revised: It was noted in Arjovsky et al.

Original: ; Kamath et al.
Revised: ; Kamath et al.

Original: that the essense of IRM’s invariance isthe conservation of the feature-conditioned label expectation , i.e.
Revised: that the essense of IRM’s invariance isthe conservation of the feature-conditioned label expectation , i.e.

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

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.
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: 3.1.4 IRM-v 
Revised: 3.1.4 IRM-v

Original: Due to the impracticality of considering the unrestricted function space of , Arjovsky et al.
Revised: Due to the impracticality of considering the unrestricted function space of , Arjovsky et al.

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

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 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: This reduces eq (6) to
Revised: This reduces eq (6) to

Original: [Equation]
Revised: [Equation]

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

Original: (2021), which imposes8 e 2 E , @
Revised: (2021), which imposes8 e 2 E , @

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

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

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

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

Original: This failure mode can be analytically understood here.
Revised: This failure mode can be analytically understood here.

Original: Note that E e [ @ o l ( o, y ) · o ] =
Revised: Note that E e [ @ o l ( o, y ) · o ] =

Original: E e
Revised: E e

Original: [ o ( y   ( o ))] for standard loss functions 3 .
Revised: [ o ( y   ( o ))] for standard loss functions 3 .

Original: Therefore, eq (9) isequivalent to
Revised: Therefore, eq (9) isequivalent to

Original: [Equation]
Revised: [Equation]

Original: Unlike eq (7), however, RHS of eq (10) does not necessarily remain constant across environments,and hence does not prescribe a meaningful invariance.
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: Therefore, restricting the function space of (or ) invalidates the indirect way of imposing invariance in IRM involving an intermediary term.
Revised: Therefore, restricting the function space of (or ) invalidates the indirect way of imposing invariance in IRM involving an intermediary term.

Original: Note that this holds true for a general causal graph and not just for the causal graph considered in
Revised: Note that this holds true for a general causal graph and not just for the causal graph considered in

Original: Figure 1.
Revised: Figure 1.

Original: 3.2 MRI Paradigm
Revised: 3.2 MRI Paradigm

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: 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.5.
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 the changein risk due to arbitrary inﬁnitesimal label perturbations is shared across environments, 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: 11) where
Revised: 11) where

Original: [Equation]
Revised: [Equation]

Original: 12) where  : Y !
Revised: 12) 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.
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.

Original: E e [ @ y l ( o, y ) | y ] ⌘ E P e
Revised: E e [ @ y l ( o, y ) | y ] ⌘ E P e

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

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

Original: Lemma 3.6.
Revised: Lemma 3.6.

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

Original: [Equation]
Revised: [Equation]

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

Original: [ o | y ] ⌘ E P e ( o | y )
Revised: [ o | y ] ⌘ E P e ( o | y )

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

Original: Proof.
Revised: Proof.

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

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

Original: =
Revised: =

Original:  o + ⇢ ( y ) 5 .
Revised:  o + ⇢ ( y ) 5 .

Original: Therefore, E e 1 [ @ y l ( o, y ) |
Revised: Therefore, E e 1 [ @ y l ( o, y ) |

Original: y ]  E e 2 [ @ y l ( o, y ) | y ] =
Revised: y ]  E e 2 [ @ y l ( o, y ) | y ] =

Original: E e 1 [ o | y ]  E e 2 [ o | y ]
Revised: E e 1 [ o | y ]  E e 2 [ o | y ]

Original: = 0 .
Revised: = 0 .

Original: Note that unlike IRM, MRI eq (11) imposes the correct number of constraints as needed by theconservation law eq (13).
Revised: Note that unlike IRM, MRI eq (11) imposes the correct number of constraints as needed by theconservation law eq (13).

Original: That is, MRI attains invariance in a direct manner without involving anyintermediary terms.
Revised: That is, MRI attains invariance in a direct manner without involving anyintermediary terms.

Original: Algorithm
Revised: Algorithm

Original: [Equation]
Revised: [Equation]

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

Original: y , reduces eq (12) to
Revised: y , reduces eq (12) to

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

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

Original: Therefore, the direct way of
Revised: Therefore, the direct way of

Original: MRI’s constraints imposing invariance continues to hold even when  is restricted to the space oflinear functions.
Revised: MRI’s constraints imposing invariance continues to hold even when  is restricted to the space oflinear functions.

Original:  Training methods
Revised:  Training methods

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

Original: [Equation]
Revised: [Equation]

Original: subject to ~ c
Revised: subject to ~ c

Original: [Equation]
Revised: [Equation]

Original: f ( ✓ )) =
Revised: f ( ✓ )) =

Original: ~
Revised: ~

Original: [Equation]
Revised: [Equation]

Original: (16) where ¯ L tr = 1 |E tr |
Revised: (16) where ¯ L tr = 1 |E tr |

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

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

Original: |E tr | is an orthonormal matrix that satisﬁes
Revised: |E tr | is an orthonormal matrix that satisﬁes

Original: Q ~ 1 = ~ 0 (See Table 1).
Revised: Q ~ 1 = ~ 0 (See Table 1).

Original: For example, in the case of a training set of two environments E tr = { e 1 , e 2 } ,
Revised: For example, in the case of a training set of two environments E tr = { e 1 , e 2 } ,

Original: Q ~ L is simply the difference of perturbed risks (  L e 1   L e 2 ) / p 2 , since Q = [1 ,  1] / p 2 .
Revised: Q ~ L is simply the difference of perturbed risks (  L e 1   L e 2 ) / p 2 , since Q = [1 ,  1] / p 2 .

Original: 4.1 Soft-constraint methods
Revised: 4.1 Soft-constraint methods

Original: Numerically solving the constrained optimization problem requires converting the hard constraints~ c ( f ( ✓ ))
Revised: Numerically solving the constrained optimization problem requires converting the hard constraints~ c ( f ( ✓ ))

Original: =
Revised: =

Original: ~ 0 to soft regularization terms, which allows the use of any off-the-shelf gradient-basedoptimization algorithms.
Revised: ~ 0 to soft regularization terms, which allows the use of any off-the-shelf gradient-basedoptimization algorithms.

Original: Penalty Method (PM)
Revised: Penalty Method (PM)

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

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

Original: [Equation]
Revised: [Equation]

Original: However, this method requires increasing µ t !
Revised: However, this method requires increasing µ t !

Original: 1 over training iterations t in order to approximatethe exact hard-constraint, which leads to training instability and slow convergence (Bertsekas, 1976).
Revised: 1 over training iterations t in order to approximatethe exact hard-constraint, which leads to training instability and slow convergence (Bertsekas, 1976).

Original: Augmented Lagrangian Method (ALM) ALM was introduced to overcome the limitations ofpenalty method (Bertsekas, 1976), which adds a Lagrange multiplier term to (17)
Revised: Augmented Lagrangian Method (ALM) ALM was introduced to overcome the limitations ofpenalty method (Bertsekas, 1976), which adds a Lagrange multiplier term to (17)

Original: [Equation]
Revised: [Equation]

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

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

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

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

Original: 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.
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.

Original: In ALM, the
Revised: In ALM, the

Original: Lagrange multiplier  is accumulated over iterations, thus imposing the model to satisfy an averagedconstraint.
Revised: Lagrange multiplier  is accumulated over iterations, thus imposing the model to satisfy an averagedconstraint.

Original: In contrast, PM penalizes the model to satisfy all of the conﬂicting constraints, and thusdecaying the model toward the zero-predictor.
Revised: In contrast, PM penalizes the model to satisfy all of the conﬂicting constraints, and thusdecaying the model toward the zero-predictor.

Original: [ELABORATE]
Revised: [ELABORATE]

Original:  Analytic Results
Revised:  Analytic Results

Original: 5.1 General Linear SEM
Revised: 5.1 General Linear SEM

Original: In this section, we show that MRI-v1 can correctly remove all spurious features given a sufﬁcientnumber of environments.
Revised: 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 Figure 1, in which x = g ( z i , z s ) is a linearfunction.
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 Structural Equation Model (SEM, Pearl (2009)) is a generalization of thelinear Gaussian model considered in Rosenfeld et al.
Revised: Note that this Structural Equation Model (SEM, Pearl (2009)) is a generalization of thelinear Gaussian model considered in Rosenfeld et al.

Original: (2020), which additionally assumed binary yand a speciﬁc noise model ( i.e. linear additive Gaussian) for the latent features.
Revised: (2020), which additionally assumed binary yand a speciﬁc noise model ( i.e. linear additive Gaussian) for the latent features.

Original: In comparison, wedon’t put any assumptions on the noise generating process or the label.
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: (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.
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 5.1.
Revised: Theorem 5.1.

Original: 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.
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 constraint eq (15) is
Revised: MRI-v1 constraint eq (15) is

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

Original: ] = const
Revised: ] = const

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

Original: [Equation]
Revised: [Equation]

Original: where M 0 = M  ¯ M · 1 with M ⌘ [ E e [ z s y ]] e 2 E
Revised: where M 0 = M  ¯ M · 1 with M ⌘ [ E e [ z s y ]] e 2 E

Original: tr 2 R d s ⇥ |E tr | and ¯ M = 1 |E tr |
Revised: tr 2 R d s ⇥ |E tr | and ¯ M = 1 |E tr |

Original: P e 2 E tr E e
Revised: P e 2 E tr E e

Original: [ z s y ] .
Revised: [ z s y ] .

Original: Sincerank ( M 0 )
Revised: Sincerank ( M 0 )

Original: =
Revised: =

Original: d s , eq (20) is equivalent to w s = 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 result that IRM-v1 can eliminate all spurious featuredimensions 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 results for MRI-v1, their result is restricted to a 199 specialized linear family of environments with binary labels and additive 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,
Revised: In fact,

Original: 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).
Revised: 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: 5.2 Minimal Example: d i = 1 , d s = 1 , |E tr | = 2
Revised: 5.2 Minimal Example: d i = 1 , d s = 1 , |E tr | = 2

Original: 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.
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: This toy example shares the same structure as the Shape-texture dataset (Section 6.1).
Revised: This toy example shares the same structure as the Shape-texture dataset (Section 6.1).

Original: It is analogousto the two-bit environment task from Kamath et al.
Revised: It is analogousto the two-bit environment task from Kamath et al.

Original: The detailed experimental setting and theanalytical solutions are presented 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 (  L e 1 ( f )
Revised: IRM-v1 IRM-v1 has two quadratic equality constraints (  L e 1 ( f )

Original: = 0 ,  L e 2 ( f ) = 0 ), shown astwo elliptic curves.
Revised: = 0 ,  L e 2 ( f ) = 0 ), shown astwo elliptic curves.

Original: 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
Revised: 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: [Figure]
Revised: [Figure]

Original: IRM-relaxed
Revised: IRM-relaxed

Original: We also consider a relaxed version of IRM-v1, whose constraint function is c =
Revised: We also consider a relaxed version of IRM-v1, whose constraint function is c =

Original: Q ~ L =
Revised: Q ~ L =

Original: (  L e 1   L e 2 ) / p 2 with  L e =
Revised: (  L e 1   L e 2 ) / p 2 with  L e =

Original: E e [ @ o l ( o, y ) · o ] , such that it has a single constriant 215 
Revised: E e [ @ o l ( o, y ) · o ] , such that it has a single constriant 215

Original: L e 1 ( f ) =
Revised: L e 1 ( f ) =

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

Original: This constraint yields a hyperbolic curve, although it appears to be two straightlines in Figure 2a.
Revised: This constraint yields a hyperbolic curve, although it appears to be two straightlines in Figure 2a.

Original: 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.
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.

Original: Therefore, simply reducing the number of constraints of
Revised: Therefore, simply reducing the number of constraints of

Original: IRM-v1 does not resolve its limitation.
Revised: IRM-v1 does not resolve its limitation.

Original: 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.
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: This problem has aconvex objective with a convex constraint, and therefore it has a unique optimum, which is the trueinvariant optimum solution.
Revised: This problem has aconvex objective with a convex constraint, and therefore it has a unique optimum, which is the trueinvariant optimum solution.

Original: In Figure 2b,c, we show the learning trajectory under gradient descent (Penalty method eq (17)with µ = 5 ⇥ 10 4 and full-batch update).
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 dependenton the initialization (shown with different trajectory colors), due to the presence of multiple localminima.
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 optimum is especially extremely difﬁcult to converge to as it hashigher training loss than the nearby non-invariant minima.
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 tothe invariant optimum regardless of initialization since 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 optimum is lost at this value of µ , 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, alltrajectories converge to the non-invariant optimum.
Revised: As a result, alltrajectories converge to the non-invariant optimum.

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

Original: -based Problems6.1 Datasets
Revised: -based Problems6.1 Datasets

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

Original: We develop a new versatile dataset designed for domain generalization,called the Shape-Texture dataset, based on the data generating model in Figure 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 isdesigned to test algorithms on linear and non-linear regression and classiﬁcation problems in multipleadjustable 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 Heaviside 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 (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 .
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 ✓ i ✓ y + ⇠ i , ✓ s ✓ y + ⇠ s , where ⇠ ⇤ ⇠ N (0 ,  ⇤ ) with probability p ⇤ , or ⇠ ⇤ ⇠ U (  ⇡ , ⇡ ) 251 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 252 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: Colored MNIST
Revised: Colored MNIST

Original: The Colored MNIST (CMNIST) dataset was introduced by Arjovsky et al.(2019) to analyze IRM-v1’s ability 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 syntheticbinary classiﬁcation task derived from MNIST.
Revised: This is a syntheticbinary classiﬁcation task derived from MNIST.

Original: In this dataset, the color of the image is determinedin correlation with the label whose correlation 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: Thecorrelation probability p i between digit 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 the spurious 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 thedata generating model in Figure 1.
Revised: This dataset follows thedata generating model in Figure 1.

Original: 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 .
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 two training environments, andp s = 0 .
Revised: 9 for the two training environments, andp s = 0 .

Original: 1 for the testing environment; CMNISTb uses 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 forthe two training environments, and p s = 0 .
Revised: 8 , 1 forthe two training environments, and p s = 0 .

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

Original: We also use an abstract version of CMNIST introduced by Kamath et al.
Revised: We also use an abstract version of CMNIST introduced by Kamath et al.

Original: (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.
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.

Original: 6.2 Result
Revised: 6.2 Result

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

Original: Forsynthetic datasets, we also report the performance of the Oracle (Table 2), which is trained using
Revised: Forsynthetic datasets, we also report the performance of the Oracle (Table 2), which is trained using

Original: ERM on the dataset in which the spurious features are rendered uncorrelated with the label.
Revised: ERM on the dataset in which the spurious features are rendered uncorrelated with the label.

Original: 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: 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: Shape-Texture Tasks
Revised: Shape-Texture Tasks

Original: For both regression and classiﬁcation tasks, we ﬁnd that IRM-v1 trainedusing PM always converges to the zero-predictor (Table 2-3; Figure5).
Revised: For both regression and classiﬁcation tasks, we ﬁnd that IRM-v1 trainedusing PM always converges to the zero-predictor (Table 2-3; Figure5).

Original: In comparison, MRI-v1trained using the penalty method achieved near-optimal OOD generalization and performed close to
Revised: In comparison, MRI-v1trained using the penalty method achieved near-optimal OOD generalization and performed close to

Original: Oracle (Table 2-3; Figure5).
Revised: Oracle (Table 2-3; Figure5).

Original: This demonstrates that MRI-v1 is far superior to IRM-v1 in terms ofboth loss and accuracy when trained using the penalty method.
Revised: This demonstrates that MRI-v1 is far superior to IRM-v1 in terms ofboth loss and accuracy when trained using the penalty method.

Original: Similar to the PM case, MRI-v1 is able to perform close to Oracle when trained using ALM (Table 23).
Revised: Similar to the PM case, MRI-v1 is able to perform close to Oracle when trained using ALM (Table 23).

Original: In comparison, while IRM was able to avoid the zero-predictor, it still achieved worse test lossand accuracy than MRI-v1.
Revised: In comparison, while IRM was able to avoid the zero-predictor, it still achieved worse test lossand accuracy than MRI-v1.

Original: Nevertheless, by using ALM, we were able to signiﬁcantly improve
Revised: Nevertheless, by using ALM, we were able to signiﬁcantly improve

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

Original: 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).
Revised: 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: This results for MRI-v1 holds true also when trainedusing ALM, albeit with a slightly improved performance (Table 2- 3; Figure 7).
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 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
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: CMNISTb datasets.
Revised: CMNISTb datasets.

Original: In contrast to the loss results, IRM-v1 achieves a similar test accuracy to MRI-v1for both the datasets (Table 3; Figure 7).
Revised: In contrast to the loss results, IRM-v1 achieves a similar test accuracy to MRI-v1for both the datasets (Table 3; Figure 7).

Original: [Table] [Table]
Revised: [Table] [Table]

Original:  Discussion
Revised:  Discussion

Original: Limitations of the accuracy metric
Revised: Limitations of the accuracy metric

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: The accuracy metric is widely used in the ﬁeld of domaingeneralization (Gulrajani and Lopez-Paz, 2020) to compare the performance of various algorithms.

Original: 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.
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: For example, we ﬁnd that for the toy-CMNISTb dataset, both invariant ( ( w i , w s ) ⇡ (3 , 0) )
Revised: For example, we ﬁnd that for the toy-CMNISTb dataset, both invariant ( ( w i , w s ) ⇡ (3 , 0) )

Original: andnon-invariant predictors ( ( w i , w s ) 2 { (3 ,  2) , (3 , 2) } ) achieve the same accuracy (Figure 8).
Revised: andnon-invariant predictors ( ( w i , w s ) 2 { (3 ,  2) , (3 , 2) } ) achieve the same accuracy (Figure 8).

Original: Incomparison, the loss metric (Figure 8) can indeed differentiate between the invariant and non-invariantrepresentations.
Revised: Incomparison, the loss metric (Figure 8) can indeed differentiate between the invariant and non-invariantrepresentations.

Original: This discrepancy between the loss and accuracy metric holds true even for non-linearimage-based datasets (Table 2-5).
Revised: This discrepancy between the loss and accuracy metric holds true even for non-linearimage-based datasets (Table 2-5).

Original: To our knowledge, such discrepancy between loss and accuracy hasnot been reported before in the domain generalization literature.
Revised: To our knowledge, such discrepancy between loss and accuracy hasnot been reported before in the domain generalization literature.

Original: 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.
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: Limitations A limitation of both MRI and IRM (Rosenfeld et al., 2020; Ahuja et al., 2021)
Revised: Limitations A limitation of both MRI and IRM (Rosenfeld et al., 2020; Ahuja et al., 2021)

Original: isthat it requires signiﬁcant overlap in the support of different domain distributions.
Revised: isthat it requires signiﬁcant overlap in the support of different domain distributions.

Original: This limits theapplicability of both MRI and IRM on most DomainBed benchmarks.
Revised: This limits theapplicability of both MRI and IRM on most DomainBed benchmarks.

Original: Rotated-MNIST, for example,expects an algorithm to generalize to an orientation of the digit that was never seen in the training set.
Revised: Rotated-MNIST, for example,expects an algorithm to generalize to an orientation of the digit that was never seen in the training set.

Original: In this case, the support of the training and testing set clearly does not overlap.
Revised: In this case, the support of the training and testing set clearly does not overlap.

Original: Similar argumentscan be made for other datasets.
Revised: Similar argumentscan be made for other datasets.

Original: For example, images with different styles across testing and trainingdomains may not satisfy the signiﬁcant support overlap condition.
Revised: For example, images with different styles across testing and trainingdomains may not satisfy the signiﬁcant support overlap condition.

Original: IRM and MRI’s domain of applicability only coincide when there is no label shift or shift invariantfeature distribution.
Revised: IRM and MRI’s domain of applicability only coincide when there is no label shift or shift invariantfeature distribution.

Original: For datasets that do not follow these conditions, MRI, but not IRM, should fail.
Revised: For datasets that do not follow these conditions, MRI, but not IRM, should fail.

Original: However, note that, while IRM is expected to work for such shifts, IRM-v1 fails for such domainshifts (Ahuja et al., 2020, 2021).
Revised: However, note that, while IRM is expected to work for such shifts, IRM-v1 fails for such domainshifts (Ahuja et al., 2020, 2021).

Original: This is because IRM-v1 and IRM are not the same, as demonstratedby us and Kamath et al. 
Revised: This is because IRM-v1 and IRM are not the same, as demonstratedby us and Kamath et al. Checklist

Revised: The checklist follows the references.

Revised: Please read the checklist guidelines carefully for information onhow to answer these questions.

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

Revised: 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.

Revised: For example:

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

Revised: [Yes] See Section ??.

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

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

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

Revised: [N/A]

Revised: Please do not modify the questions and only use the provided macros for your answers.

Revised: Note that the

Revised: Checklist section does not count towards the page limit.

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

Revised: For all authors...

Revised: (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s

Revised: contributions and scope?

Revised: [Yes]

Revised: (b) Did you describe the limitations of your work?

Revised: [Yes] See Section 7.

Revised: (c) Did you discuss any potential negative societal impacts of your work?

Revised: [N/A]

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

Revised: [Yes]

Revised: If you are including theoretical results...

Revised: (a) Did you state the full set of assumptions of all theoretical results?

Revised: [Yes] See Section 3and 5.

Revised: (b) Did you include complete proofs of all theoretical results?

Revised: [Yes] See Section 3 and 5.3.

Revised: If you ran experiments...

Revised: (a) Did you include the code, data, and instructions needed to reproduce the main

Revised: experimental results (either in the supplemental material or as a URL)?

Revised: [Yes]

Revised: https://anonymous.4open.science/r/MRI-2608/readme(b)

Revised: Did you specify all the training details (e.g., data splits, hyperparameters,

Revised: how they

Revised: were chosen)?

Revised: [Yes] See Supplementary Materials.

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

Revised: ments multiple times)?

Revised: [Yes] See Table 2 and 3.

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

Revised: [Yes] See Supplementary Materials.4.

Revised: If you are using existing assets (e

Revised: g., code, data, models) or curating/releasing new assets...

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

Revised: [Yes] See Section 6.1 andSupplementary Materials.

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

Revised: [Yes] See Supplementary Materials.

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

Revised: [Yes]

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

Revised: (d) Did you discuss whether and how consent was obtained from people whose data you’re

Revised: using/curating?

Revised: [N/A]

Revised: (e) Did you discuss whether the data you are using/curating contains personally identiﬁable

Revised: information or offensive content?

Revised: [N/A]

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

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

Revised: applicable?

Revised: [N/A]

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

Revised: Board (IRB) approvals, if applicable?

Revised: [N/A]

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

Revised: spent on participant compensation?

Revised: [N/A]