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 conserving thefeature-conditioned label 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-v1, a practical version of IRM,can fail in various task settings.

Original: Here, we identify a fundamental ﬂaw of IRM formulation that causesthe failure.
Revised: Here, we identify a fundamental ﬂaw of IRMformulation that causes the 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 ofinvariance, MRI, based on conserving the label-conditioned feature expectation E e [ f ( x ) | y ] across environments, which is free of this ﬂaw.

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

Original: We note that this constraint is convex which confers it with an advantage overthe practical version of IRM, IRM-v1, which imposes non-convex constraints.
Revised: We note thatthis constraint is convex which confers it with an advantage over IRM-v1, whichimposes 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 general linear problem setting, MRI-v1 can guarantee invariant predictors given sufﬁcient environments.

Original: We also empirically demonstrate that MRI strongly out-performs IRM and consistently achieves near-optimal OOD generalization in image-basednonlinear problems.
Revised: We alsoempirically demonstrate that MRI strongly out-performs IRM and consistentlyachieves near-optimal OOD generalization in image-based nonlinear 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 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.

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

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., 2016;

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

Original: Invariant Risk Minimization (IRM) ,in particular, has attracted a lot of attention asit is one of the ﬁrst methods to be compatible withdeep learning techniques.
Revised: 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.

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

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 IRM formulation has been overlooked andwe propose a  method to eliminates this ﬂaw.
Revised: Here, we identify afundamental ﬂaw of IRM formulation that causes this limitation and propose a new method that isfree of 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 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.

Original: We also show that IRMhas a fundamental ﬂaw due to an extraneous constraint which causes it to fail.
Revised: Wealso show that IRM has a fundamental ﬂaw due to the indirect way it imposes invariance which leadsto the failure of IRM-v1.

Original: MRI is shown to befree of this ﬂaw.
Revised: MRI is shown to be free 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 linearproblem setting, MRI-v1 can guarantee invariant predictors 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 thenempirically demonstrate that MRI strongly out-performs IRM and consistently achieves near-optimal OOD generalization in image-based nonlinear problems.

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: 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]

Original: [Equation]
Revised: [Equation]

Original: where o
Revised: where o

Original: = f
Revised: = f

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

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

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

Original: =12 k o  y k l log ( o, y ) =
Revised: =

Original:  (1 + y ) log( ⌘ ( o))  (1  y )log(1  ⌘ ( o ))for binary classiﬁcation 2  ( O ✓ R , Y ✓ [  1 , 1] ), where 76 ⌘ ( o )
Revised: 12 ( o  y ) 2 for regression ( O, Y ✓ R ) and the binary-cross-entropy (BCE) lossl log ( o, y )

Original: =
Revised: =

Original:  1 / (1 + e  o ) is the sigmoid function.
Revised:  (1 + y ) log( ⌘ ( o ))  (1  y ) log(1  ⌘ ( o ))for binary classiﬁcation ( O ✓ R , Y ✓ [  1 , 1] ), where 75 ⌘ ( o ) ⌘ 1 / (1 + e  o ) is the sigmoid function.

Original:  3.1 IRM Paradigm
Revised:  3.1 IRM Paradigm3.1.1 Original formulation

Original: 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 80 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 2 , Deﬁnition 3.1 is equivalent to

Original: [Equation]
Revised: [Equation]

Original: ) where E e [ y | f ( x )] ⌘ E P e ( y | f ( x ))
Revised: where E e [ y | f ( x )] ⌘ 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 85 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 88 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]

Revised: [Equation]

Original: (6) where  : O !
Revised: 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 [Equation] e 
Revised: For standard loss functions 2 , Deﬁnition 3.3 is equivalent to8 e 2 E ,

Original: [Equation] E ,
Revised: E e [ y | o ] =

Original: [Equation] (7) which is equivalent to Lemma 3.2 with f replacing the composite predictor  f .
Revised:  ( o ) , (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 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.
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 3 .

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 (16)).

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: Since E e [ @ o l ( o, y ) · o ] =

Original: E e
Revised: E e

Original: [ o (y   ( o ))] for standard loss functions 3 . Therefore, eq (9) isequivalent to 
Revised: [ o (  ( o )  y )] for standard loss functions 2 , eq (9) is equivalent to [Equation]

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

Original: 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 [Equation] 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 ) 4 .

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 ] = E e 1 [ o | y ]  E e 2 [ o | y ]
Revised: y ]  E e 2 [ @ y l ( o, y ) | y ] = [Equation]

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: 3.2.1 MRI-v

Original: [Equation] 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  L e ( f ) =

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

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: (14) 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. Algorithm [Equation]

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 ( f ( ✓

Original: [Equation]
Revised: [Equation]

Original: f ( ✓ )) =

Original: ~
Revised: =

Original: [Equation]
Revised: [Equation]

Original: (16) where ¯ L tr = 1 |E tr |
Revised: (15) 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) ⇥ |E tr | is an orthonormal matrix that satisﬁes Q ~ 1 = ~ 0 (See Table 1).
Revised: The constraint functions are ~ c = ~ L tr 2 R |E tr | for IRM-v1 and~ c = Q ~ L tr 2 R |E tr |  1 for MRI-v1, where ~ L tr 2 R |E tr | is a vector of perturbed risks  L e for e 2 E tr ,and Q 2 R ( |E tr |  1) ⇥ |E tr | is an orthonormal matrix that satisﬁes Q ~ 1 = ~ 0 .

Original: For example, in the case of a training set of two environments E tr = { e 1 , e 2 } , Q ~ L  is simply the difference of perturbed risks (  L e 1   L e 2 ) / p 2 , since Q = [1 ,  1] / p 2 . 
Revised: For example, in the case of atraining set of two environments E tr = { e 1 , e 2 } , Q ~ L tr is simply the difference of perturbed risks(  L e 1   L e 2 ) / p 2 , since Q = [1 ,  1] / p 2 . The constraints are summarized in Table 1.

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

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

Original: =
Revised: =

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

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 thesum of squared  constraints as a penalty term to the objective,
Revised: (2019), which adds the squared residual constraints as a penalty term to the objective,min ✓ ¯

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 (16)

Original: [Equation]
Revised: [Equation]

Original: (18) where ~  gets updated as ~ t +1 =
Revised: 17) where ~ is typically initialized at ~ 0 and 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 theresidual constraints at each training iteration 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.

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.

Original: In ALM, the

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

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

Revised: [Figure]

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: (18) 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, 182 MRI-v1 will eliminate all spurious feature dimensions.

Original: Proof.
Revised: Proof.

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

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]

Original: ] = 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 tr 2 R d s ⇥ |E tr | and ¯ M = 1 |E tr |
Revised: (19) where M 0 = M  ¯ M · 1 with M ⌘ [ E e [ z s y ]] e 2 E tr 2 R d s ⇥ |E tr | and ¯ M = 1 |E tr |

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 (19) 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 190 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 | =

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: Here, we demonstrate the result for the general linear SEM above with a minimal linear regressionproblem 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 analogous tothe 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: Note that one of thenon-invariant solutions exhibits lower training loss than the true invariant optimum solution.

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 206

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 is a set of hyperbolic curves, which 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 is a convexproblem, since both the objective and the constraint are convex, and thus features a unique optimum,which is the true invariant 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 full-batch gradient descent (penalty methodeq (16) with µ = 5 ⇥ 10 4 used).

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 dependent on theinitialization (shown with different trajectory colors), due to the presence of multiple local minima.

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: Inthis case, the invariant optimum is especially extremely difﬁcult to converge to as it has higher trainingloss 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 to the invariantoptimum 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 invariantoptimum 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, all trajectoriesconverge 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 Image-based Problems

Original: -based Problems6.1 Datasets
Revised: 6.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 (  ⇡ , ⇡ ) 242 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 243 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: ColoredMNIST
Revised: Colored

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: MNIST 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).[Table] [Table]
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:  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 has [Table] 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 (Gulrajani and Lopez-Paz, 2020) 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 thatwas 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 doesnot overlap.

Original: Similar argumentscan be made for other datasets.
Revised: Similar arguments can 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 differentstyles across testing and training domains 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.

Original: Checklist
Revised: Checklist

Original: The checklist follows the references.
Revised: The checklist follows the references.

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

Original: For each question, change the default [TODO] to [Yes] , [No] , or[N/A] .
Revised: 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.
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.

Original: For example:
Revised: For example:

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

Original: [Yes] See Section ??.
Revised: [Yes] See Section ??.•

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

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

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

Original: [N/A]
Revised: [N/A]

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

Original: Note that the
Revised: Note that the

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

Original: In your paper, please delete this instructionsblock and only keep the Checklist section heading above along with the questions/answers below.
Revised: 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...
Revised: For all authors...

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

Original: contributions and scope?
Revised: contributions and scope?

Original: [Yes]
Revised: [Yes]

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

Original: [Yes] See Section 7.
Revised: [Yes] See Section 7.(c)

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

Original: [N/A]
Revised: [N/A](d)

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

Original: [Yes]
Revised: [Yes]

Original: If you are including theoretical results...
Revised: If you are including theoretical results...

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

Original: [Yes] See Section 3and 5.
Revised: [Yes] See Section 3and 5.(b)

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

Original: [Yes] See Section 3 and 5.3.
Revised: [Yes] See Section 3 and 5.3.

Original: If you ran experiments...
Revised: If you ran experiments...

Original: (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: (a) Did you include the code, data, and instructions needed to reproduce the mainexperimental results (either in the supplemental material or as a URL)?

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

Original: [Yes] See Supplementary Materials.
Revised: [Yes] See Supplementary Materials.

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

Original: ments multiple times)?
Revised: ments multiple times)?

Original: [Yes] See Table 2 and 3.
Revised: [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)?
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)?

Original: [Yes] See Supplementary Materials.4.
Revised: [Yes] See Supplementary Materials.

Original: If you are using existing assets (e g., code, data, models) or curating/releasing new assets...
Revised: 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?
Revised: (a) If your work uses existing assets, did you cite the creators?

Original: [Yes] See Section 6.1 andSupplementary Materials.
Revised: [Yes] See Section 6.1 andSupplementary Materials.

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

Original: [Yes] See Supplementary Materials.
Revised: [Yes] See Supplementary Materials.

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

Original: [Yes]
Revised: [Yes]

Original: https://anonymous.4open.science/r/MRI-2608/readme
Revised: 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
Revised: (d) Did you discuss whether and how consent was obtained from people whose data you’re

Original: using/curating?
Revised: using/curating?

Original: [N/A]
Revised: [N/A]

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

Original: information or offensive content?
Revised: information or offensive content?

Original: [N/A]
Revised: [N/A]

Original: If you used crowdsourcing or conducted research with human subjects...
Revised: 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
Revised: (a) Did you include the full text of instructions given to participants and screenshots, if

Original: applicable?
Revised: applicable?

Original: [N/A]
Revised: [N/A]

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

Original: Board (IRB) approvals, if applicable?
Revised: Board (IRB) approvals, if applicable?

Original: [N/A]
Revised: [N/A]

Original: (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
Revised: (c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation?

Original: [N/A]
Revised: [N/A]