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: 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 label during 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 label during training.

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

Original: However, more recent studies have demonstrated that IRM-v1, a practical version of IRM, can fail in various settings.
Revised: However, more recent studies have demonstrated that IRM-v1, a practical version of IRM, can fail in various settings.

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

Original: We then introduce a complementary notion of invariance, MRI, based on conserving the label-conditioned feature expectation E e [ f ( x ) | y ] , which is free of this ﬂaw.
Revised: We then introduce a complementary notion of invariance, MRI, based on conserving the label-conditioned feature expectation E e [ f ( x ) | y ] , which is free of this ﬂaw.

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

Original: We note that this constraint is convex which confers it with an advantage over IRM-v1, which imposes non-convex constraints.

Original: We prove that in a general linear problem setting, MRI-v1 can guarantee invariant predictors given sufﬁcient  environments.
Revised: We prove that for general linear problems, MRI-v1 guarantees invariant predictors given sufﬁcient number of environments.

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

Original: MRI strongly out-performs IRM and consistently achieves near-optimal OOD generalization in image-based nonlinear problems.
Revised: IRM-v1 and consistently achieves 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 show great generalization properties when tested on the same distribution as the training dataset (indistribution generalization).
Revised: These models show great generalization properties when tested on the same distribution as the training dataset (indistribution generalization).

Original: However, these models often show catastrophic failure when tested on out-of-distribution dataset, revealing that they learnedto use easier-to-ﬁt features that are spuriously correlated to the label in the given training domains but do not generalize to the testing domains.
Revised: However, these models often show catastrophic failure when tested on out-of-distribution dataset, revealing that they learned features that are spuriously correlated 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 the training domain will use the background color as the predictive feature which is easier to learn and generalize poorly to pictures of cow with a dessert background.
Revised: For example, deep networks trained on pictures of cow with only grassy backgrounds in the training domain will use the background color as the predictive feature which is easier to learn and generalize poorly to pictures of cow with a dessert background.

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

Original: In particular, there has been a recent body of works that focus on developing algorithms that attempt to learn invariant predictors (Arjovsky et al., 2019; Ahuja et al., 2021; Peters et al., 2016; Rojas-Carulla et al., 2018; Heinze-Deml et al., 2018).
Revised: In particular, there has been a recent body of works that focus on developing algorithms that attempt to learn invariant predictors (Arjovsky et al., 2019; Ahuja et al., 2021; Peters et al., 2016; Rojas-Carulla et al., 2018; Heinze-Deml et al., 2018).

Original: Invariant Risk Minimization (IRM) and its practical version, IRM-v1,in particular, have attracted much attention as one of the ﬁrst methods that are compatible with deep learning techniques.
Revised: Invariant Risk Minimization (IRM) and its practical version, IRM-v1, have garnered signiﬁcant attention as one of the initial methods that are compatible with deep learning techniques.

Original: However, several follow-up studies have empirically demonstrated that IRM-v1 is unreliable at learning invariant representations (Kamath et al., 2021;
Revised: However, several follow-up studies have empirically demonstrated that IRM-v1 is unreliable at learning invariant representations (Kamath et al., 2021;

Original: Rosenfeld et al., 2020; Gulrajani and Lopez-Paz, 2020; Ahuja et al., 2020, 2021).
Revised: Rosenfeld et al., 2020; Gulrajani and Lopez-Paz, 2020; Ahuja et al., 2020, 2021).

Original: Here, we identify a fundamental ﬂaw of IRM formulation that causes this limitation and propose a new method that is free of this ﬂaw.
Revised: Here, we identify a fundamental ﬂaw of IRM formulation that causes this limitation and propose a new method that is free 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 kernel methods (Muandet et al., 2013; Ghifary et al., 2016;
Revised: They vary from learning domain-invariant feature representations conserving P ( f ( x )) using kernel methods (Muandet et al., 2013; Ghifary et al., 2016;

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

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

Original: ; Akuzawa et al.
Revised: ; Akuzawa et al.

Original: (2019); Albuquerque et al.
Revised: (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: (Gong et al., 2016; Li et al., 2018b) in the context of domain adaptation, which assumes access to the test distribution for adaptation.
Revised: (Gong et al., 2016; Li et al., 2018b) in the context of domain adaptation, which assumes access to the test distribution for adaptation.

Original: There is also a large body of work to learn invariant representations in the ﬁ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 in the ﬁeld of domain generalization that doesn’t assume access to test distributions.

Original: This includes imposing invariance of E e [ y | f ( x )]
Revised: This includes imposing 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 domain inference (Creager et al., 2021), model calibration (Wald et al., 2021), and others (Krueger et al., 2021;
Revised: (Arjovsky et al., 2019) with information bottleneck constraint (Ahuja et al. (2021)), imposing object-invariant condition (Mahajan et al. (2021)), using domain inference (Creager et al., 2021), model calibration (Wald et al., 2021), and others (Krueger et al., 2021;

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

Original: Our Contributions We introduce a simpler formulation of IRM that is more generalizable, and show that it can be modiﬁed to yield a new complementary notion of invariance, called MRI.
Revised: Our Contributions We introduce a variational formulation of IRM, and show that it can be modiﬁed to yield a new complementary notion of invariance, called MRI.

Original: Wealso show that IRM has a fundamental ﬂaw due to the indirect way it imposes invariance which leads to the failure of IRM-v1.
Revised: We show that IRM has a fundamental ﬂaw due to the indirect way of imposing invariance which leads to the failure of IRM-v1.

Original:  MRI is shown to be free of this ﬂaw.
Revised: In constrast, MRI is shown to be free of this ﬂaw.

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

Original: We then empirically demonstrate that MRI strongly out-performs IRM and consistently achieves near-optimal OOD generalization inimage-based nonlinear  problems.
Revised: We also show empirical demonstrations that MRI strongly out-performs IRM and consistently achieves near-optimal OOD generalization in nonlinear image-based problems.

Original:  Problem Formulation
Revised:  Problem Formulation

Original: Consider a set of environments E = { e } , each of which deﬁnes a distribution P e ( x, y ) over inputs and 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 inputs and labels, from which the dataset of the environment is drawn D e ⌘ { ( x ej 2 R d , y ej 2 R ) } .

Original: The distribution P e ( x, y ) is assumed to be generated according to the causal graph in Fig 1 (Rosenfeld et al., 2020), which includes latent features that are invariantly z i or spuriously z s correlated with the label, from which the observation x is generated.Y Z i Z s
Revised: The distribution P e ( x, y ) is assumed to be generated according to the causal graph in Fig 1 (Rosenfeld et al., 2020), which includes latent features that are invariantly z i or spuriously z s correlated with the label, from which the observation x is generated.

Original: X
Revised: [Figure]

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 is deﬁned as the population average
Revised: O in environment e is deﬁned as the population average

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: = 1 2 ( o  y ) 2 for regression ( O, Y ✓ R ) and the binary-cross-entropy (BCE) loss l log ( o, y )
Revised: = 12 ( o  y ) 2 for regression ( O, Y ✓ R ) and the binary-cross-entropy (BCE) loss l log ( o, y )

Original: =
Revised: =

Original:  (1 + y ) log( ⌘ ( o ))  (1  y ) log(1  ⌘ ( o )) for binary classiﬁcation ( O ✓ R , Y ✓ [  1 , 1] ), where ⌘ ( o ) ⌘ 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 ⌘ ( o ) ⌘ 1 / (1 + e  o ) is the sigmoid function.

Original:  IRM vs MRI Invariance 3.1 IRM Paradigm 3.1.1 Original formulation
Revised:  IRM vs MRI Invariance 3.1 IRM Paradigm 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 : X ! O , the feature representation f : X !
Revised: The feature representation f : X !

Original: Z is IRM invariant  over a set of environments E  if there exists : Z !O such that is While Fig 1 considers the causal direction Y !
Revised: Z elicits an IRM invariant predictor  f : X ! O over a set of environments E , if there exists : Z !

Original: Z i , our results on IRM and MRI are also consistent with the causal direction Z i !

Original: Y . simultaneously optimalon f for all environments in E , i.e.
Revised: O such that is simultaneously optimal for all environments in E , i.e.

Original: [Equation]
Revised: [Equation]

Revised: While Fig 1 only shows the causal direction Y !

Revised: Z i , the other direction Z i !

Revised: Y is also consistent with our analysis here as long as P ( y ) is independent from the environment index e .

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 4 , Deﬁnition 3.1 is equivalent to
Revised: For standard loss functions 4 , 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 function that depends on the loss function.
Revised: [ y ] is the feature-conditioned label expectation, and  is a monotonic 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 optimality condition on .
Revised: The original formulation above is overly complex due to the composite predictor  f and the optimality condition on .

Original: For further analysis, we introduce the following variational formulation.
Revised: For further analysis, we introduce the following variational formulation.

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 the risk L e ( f ) remains stationary under arbitrary inﬁnitesimal perturbations on the predictor’s output o = f
Revised: O is IRM invariant over a set of environments E , if the risk L e ( f ) remains stationary under arbitrary inﬁnitesimal perturbations on the predictor output o = 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: where  : O !

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

Original: Lemma 3.4.
Revised: Lemma 3.4.

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

Original: [Equation]
Revised: [Equation]

Original: which is equivalent to Lemma 3.2 withf replacing the composite predictor [Equation]
Revised: which is equivalent to Lemma 3.2 with the composite predictor  f replaced by 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 loss derivative has the form @
Revised: For standard loss functions, the loss derivative 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:  Deﬁnition 3.3 describes  the ﬁrst-order condition for the predictorf to be simultaneously optimal over all environments. This is indeed a sufﬁcient condition for optimality given a convex loss function l .
Revised: Note that even though Deﬁnition 3.3 describes only the ﬁrst-order condition for the predictor to be simultaneously optimal over all environments, this is indeed the necessary and sufﬁcient condition for optimality, since the loss function l is convex.

Original: This yields a simpler formulation of IRM without requiring a composite form for the predictor  f .
Revised: This yields a simpler formulation of IRM without requiring a composite form for the predictor  f .

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: As noted in Arjovsky et al.

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

Original: (2021) that the essense of IRM’s invariance is the conservation of the feature-conditioned label expectation , i.e.
Revised: (2021), the essence of IRM’s invariance is the 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 respect to the environment index e .
Revised: This result can be easily seen from eq (5),(7), since their RHS term  ( o ) is constant with respect to the environment index e .

Original: However, there is a critical difference: Eq (8) only requires  E e
Revised: Remark Notice an intriguing discrepancy in the number of constraints: IRM (eq (4),(5),(7)) imposes one constraint per environment, total of |E| constraints, whereas the conservation law describes |E|  1 equality relationships that E e [ y | o ] should share the same value across environments. The missing constraint is that IRM additionally requires the shared value of E e

Original: [ y | o ] toshare the same value across environments, thus imposing |E|  1 constraints, whereas eq (7) (also eq (5)) additionally imposes that value to be also  equal tothe RHS term  ( o ) . This extra constraint reﬂects the indirect way IRM attains the conservation, i.e. by equating E e [ y | o ] to  an intermediary term  ( o ) , imposing 1 constraint per environment.
Revised: [ y | o ] to also be equal to  ( o ) . This discrepancy emphasizes the fact that the equality relationships in eq (8) are not direct constraints imposed to hold between environments, but rather a byproduct — an indirect consequence of separate individual constraints all sharing a common intermediate term,  ( o ) . Furthermore, it shows that the invariance guarantee of IRM singularly depends on the constancy of this shared term across environments, which proves to be a single point of failure for IRM.

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 corresponds to restricting the output perturbations  to the space of linear functions, which, for scalar outputs, is equivalent to the identity function,  ( o )
Revised: In the variational formulation, this corresponds to restricting the output perturbations  to the space of linear functions, which, for scalar outputs, is equivalent to the identity function,  ( o )

Original: =
Revised: =

Original: o .
Revised: 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: In Arjovsky et al.
Revised: In Arjovsky et al.

Original: (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)).
Revised: (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 widely used for the reduced formulation regardless of whether hard or soft constraints are used, which we adopt here.
Revised: In the literature, the term IRM-v1 is widely used for the reduced 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 quite differently from IRM  (Kamath et al., 2021).
Revised: However, IRM-v1 has been empirically found to behave quite differently from IRM and fail even in simple problems (Kamath et al., 2021).

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

Original: Since @ o l ( o, y ) · o = o (  ( o )  y ) for standard loss functions 4 , eq (9) is equivalent to
Revised: Since @ o l ( o, y ) · o = o (  ( o )  y ) for standard loss functions 4 , eq (9) is equivalent to

Original: [Equation]
Revised: [Equation]

Original: Unlike in eq (7), however, RHS of eq (10)  is not constant across environments, and hence does not prescribe a meaningful invariance.
Revised: Note that the RHS of eq (10) originates from the RHS of eq (7). Unlike  ( o ) of eq (7), however, E e [ o  ( o )] is not constant, since it involves expectation that depends on the environment, and therefore it fails to mediate any meaningful invariance relationship.

Revised: In Supplementary Materials, we generalize this result to the wider class of perturbations that are mixtures of nonlinear basis functions (See

Original: Therefore, restricting the function space of  (or ) invalidates the indirect way ofimposing invariance in IRM involving an intermediary term.
Revised: Supplementary Materials B). Fundamental ﬂaw of IRM The above analysis shows that IRM’s indirect mechanism for attaining invariance through a shared intermediate term is in fact quite fragile, which easily breaks when the function space of  (or ) gets restricted. We identify this as the fundamental design ﬂaw of IRM .

Original: 3.2 MRI Paradigm
Revised: 3.2 MRI Paradigm

Original: We now introduce a complementary notion of invariance by considering inﬁnitesimal perturbations on label, which we call  Mirror Reﬂected IRM , or MRI .
Revised: We now introduce a complementary notion of invariance by considering inﬁnitesimal perturbations on label, which we call the 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 change in 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 change in risk due to arbitrary inﬁnitesimal label perturbations is shared across environments, i.e.

Original: [Equation]
Revised: [Equation]

Original: where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: where  : Y !
Revised: where  : Y !

Original: Y is an arbitrary perturbation that is unrestricted in the space of all measurable functions, and  L e ( f ) denotes the resulting change in risk.
Revised: Y is an arbitrary perturbation that is unrestricted in the space of all measurable functions, 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: which conserves the label-conditioned feature expectation E e
Revised: 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, the loss derivative has the form @
Revised: For standard loss functions, the loss derivative has the form @

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

Original: =
Revised: =

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

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 MRI attains invariance in a direct manner without involving any intermediary terms, thus imposing thecorrect number of constraints in eq (11) as described by the conservation law eq (13). 
Revised: Remark Note that MRI attains invariance in a direct manner without involving any intermediate term, which results in the number of constraints in eq (11) matching the conservation law eq (13). Algorithm [Equation]

Original: 3.2.1 MRI-v 
Revised: 3.2.1 MRI-v

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  L e ( f ) =
Revised: y , reduces eq (12) to  L e ( f ) =

Original: E e [ @ y l ( o, y ) · y ] .
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 describes a necessary condition for the MRI invariance eq (13).
Revised: which describes a necessary condition for the MRI invariance eq (13).

Original: Therefore,the direct way of MRI’s constraints imposing invariance continues to hold when  is restricted to the space of linear functions.Algorithm [Equation]
Revised: Therefore, MRI’s direct mechanism for attaining invariance continues to hold when  is restricted to the space of linear functions.

Original:  Methods 4.1 Constrained optimization problem
Revised:  Methods 4.1 Constrained optimization problem

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 ( f 
Revised: subject to ~ c ( f ) = ~

Original: [Equation]
Revised: [Equation]

Original: , (15) 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 w 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 .

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

Original: For example, for a training set of two environments E tr = { e 1 , e 2 } , Q ~ L = (  L e 1   L e 2 ) / p 2 , since Q = [1 ,  1] / p 2 .
Revised: For example, for a training set of two environments E tr = { e 1 , e 2 } , Q ~ L = (  L e 1   L e 2 ) / p 2 , since Q = [1 ,  1] / p 2 .

Original: The constraints are summarized in Table 1.
Revised: See Table 1.

Original: 4.2 Soft-constraint methods
Revised: 4.2 Soft-constraint methods

Original: Numerically solving IRM-v1 and MRI-v1 requires converting the hard constraints ~
Revised: Numerically solving IRM-v1 and MRI-v1 requires converting the hard constraints ~

Original: c ( f ) =
Revised: c ( f ) =

Original: ~ 0 to soft constraints, which allows the use of off-the-shelf gradient-based optimization algorithms.
Revised: ~ 0 to soft constraints, which allows the use of off-the-shelf gradient-based optimization algorithms.

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

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

Original: (2019), which adds the squared residual constraints as a penalty term to the objective,[Equation] w [Equation]
Revised: (2019), which adds the squared residual constraints as a penalty term to the objective,

Original: tr

Original: [Equation]

Original: + µ k ~

Original: [Equation]

Original: f )

Original: [Equation]
Revised: [Equation]

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

Original: 1 over training iteration t in order to approximate the exact hard-constraint, which leads to training instability and slow convergence (Bertsekas, 1976).
Revised: 1 over training iteration t in order to approximate the 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 of penalty method (Bertsekas, 1976), which adds a Lagrange multiplier term to (16)[Equation]
Revised: Augmented Lagrangian Method (ALM) ALM was introduced to overcome the limitations of penalty method (Bertsekas, 1976), which adds a Lagrange multiplier term to (16)

Original: [Equation]
Revised: [Equation]

Original:  where ~ is typically initialized at ~ 0 and updated at each training iteration t to accumulate the residual constraints ~ c ( f ( w t )) .
Revised: (17) where ~  is typically initialized at ~ 0 and updated at each training iteration t to accumulate the residual constraints ~ c ( f ( w 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:  Analytic Results 5.1 General Linear SEM
Revised:  Analytic Results 5.1 General Linear SEM

Original: In this section, we show that MRI-v1 can correctly remove allspurious features of a linear predictor given a sufﬁcient number of environments.
Revised: In this section, we demonstrate that MRI-v1 can effectively eliminate all features that are spuriously correlated with the label in a linear predictor, given a sufﬁcient number of environments.

Original: Consider a data generating process according to Fig 1, in which the observation x = g ( z i , z s ) is an injective linear function of the latent features z
Revised: Consider a data generating process according to Fig 1, in which the observation x = g ( z i , z s ) is an injective linear function of the latent features z

Original: i  R d i , z s 2 R d s .
Revised: i 2 R d i , z s 2 R d s .

Original: Note that this Structural Equation Model (SEM) (Pearl, 2009) does not require any assumptions on the generation process Y !
Revised: Note that this Structural Equation Model (SEM) (Pearl, 2009) does not require any assumptions on the generation process Y !

Original: Z i , Z s , which generalizes the SEM of Rosenfeld et al.
Revised: Z i , Z s , which generalizes the SEM of Rosenfeld et al.

Original: (2020), which additionally assumed binary labels and additive Gaussian noise for generating the latent features.
Revised: (2020), which additionally assumed binary labels and additive Gaussian noise for generating the latent features.

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

Original: O .
Revised: O .

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

Original: [Equation]
Revised: [Equation]

Original:  with  w = { w i 2 R d i , w s 2 R d s } .
Revised: 18) with parameters w ⌘ { w i 2 R d i , w s 2 R d s } .

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 positions, 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 positions, MRI-v1 will eliminate all spurious feature dimensions.

Original: Proof.
Revised: Proof.

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

Original: [Equation]
Revised: [Equation]

Original: which can be expressed in a matrix form as
Revised: which 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 {+
+} = 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: [Equation]
Revised: [Equation]

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

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

Original: =
Revised: =

Original: d s , eq (19) is equivalent to w s = 0 .
Revised: d s , eq (19) is equivalent to w s = 0 .

Original: A similar result was shown for IRM-v1, but under a more restricted setting, such as a specialized linear family of environments with binary labels and additive Gaussian noise (Rosenfeld et al., 2020).
Revised: A similar result was shown for IRM-v1, but under a more restricted setting, such as a specialized linear family of environments with binary labels and additive Gaussian noise (Rosenfeld et al., 2020).

Original: In fact, IRM-v1 has been shown to fail in more general linear problems with non-Gaussian noise (Arjovsky et al., 2019; Kamath et al., 2021).
Revised: In fact, IRM-v1 has been shown to fail in more general linear problems with non-Gaussian noise (Arjovsky et al., 2019; Kamath et al., 2021).

Original: 5.2 Minimal Example: d i = 1 , d s = 1 , |E tr | =
Revised: 5.2 Minimal Example: d i = 1 , d s = 1 , |E tr | =

Original: Here, we demonstrate a minimal case of the above general linear problem that involves one invariant feature, one spurious feature and two training environments (Shape-Texture linear regression problem in Section 6.1).
Revised: Here, we demonstrate a minimal case of the above general linear problem that involves one invariant feature, one spurious feature and two training environments (Shape-Texture linear regression problem in Section 6.1).

Original: See Supplementary Materials for the detailed experimental set up and the analytical solutions.
Revised: See Supplementary Materials for the detailed experimental set up and the analytical solutions.

Original: A similar minimal examples for linear binary classiﬁcation are also analyzed and shown in
Revised: A similar minimal examples for linear binary classiﬁcation are also analyzed and shown in

Original: Supplementary Materials (Fig 6, linear shape-texture classiﬁcation and toy-CMNISTa/b).
Revised: Supplementary Materials (Fig 6, linear shape-texture classiﬁcation and toy-CMNISTa/b).

Original: Analytic solutions (Hard-constraints, Fig 2A) IRM-v1 has two quadratic equality constraints,
Revised: Analytic solutions (Hard-constraints, Fig 2A) IRM-v1 has two quadratic equality constraints,

Original: E e 1 [ o 2  oy ] =[-
 =-]
Revised: E e 1 [ o 2  oy ] =[-
 =-]

Original: E e 2 [ o 2  oy ]  , shown as two elliptic curves.
Revised: E e 2 [ o 2  oy ] = 0 , shown as two elliptic curves.

Original: The intersection between the non-convex constraints on a 2-D feature space yields a disjoint set of 0-D points, all of which are local constrained optima, including the true invariant optimum, a zero-predictor solution, and two non-invariant solutions.
Revised: The intersection between the non-convex constraints on a 2-D feature space yields a disjoint set of 0-D points, all of which are local constrained optima, including the true invariant optimum, a zero-predictor solution, and two non-invariant solutions.

Original: Note that one of the non-invariant solutions exhibits lower train loss than the true invariant optimum solution.
Revised: Note that one of the non-invariant solutions exhibits lower train loss than the true invariant optimum solution.

Original: We also test the relaxed version of IRM-v1 by removing the extraneous constraint.
Revised: We also test the relaxed version of IRM-v1 by removing the extraneous constraint.

Original: The relaxed version has a single constraint E e 1 [ o 2  oy ] =
Revised: The relaxed version has a single constraint E e 1 [ o 2  oy ] =

Original: E e 2 [ o 2  oy ] , which describes a pair of hyperbolic curves (appears as two straight lines in Fig 2A), i.e. a non-convex constraint.
Revised: E e 2 [ o 2  oy ] , which describes a pair of hyperbolic curves (appears as two straight lines in Fig 2A), i.e. a non-convex constraint.

Original: This problem exhibits two local constrained optima, including the true invariant optimum and a non-invariant solution, with the non-invariant solution exhibiting a lower train loss.
Revised: This problem exhibits two local constrained optima, including the true invariant optimum and a non-invariant solution, with the non-invariant solution exhibiting a lower train loss.

Original: Therefore, simply relaxing the extraneous constraint does not resolve the fundamental problem of IRM-v1.
Revised: Therefore, simply relaxing the extraneous constraint does not resolve the fundamental problem of IRM-v1.

Original: In contrast, MRI-v1 has one linear equality constraint, E e 1 [ oy ] =
Revised: In contrast, MRI-v1 has one linear equality constraint, E e 1 [ oy ] =

Original: E e 2 [ oy ] , that exactly prescribes the set of all invariant solutions w s = 0 .
Revised: E e 2 [ oy ] , that exactly prescribes the set of all invariant solutions w s = 0 .

Original: That is, a solution is an invariant predictor if and only if it satisﬁes this constraint.
Revised: That is, a solution is an invariant predictor if and only if it satisﬁes this constraint.

Original: This is a convex problem, since both the objective and the constraint are convex, and thus features a unique optimum, which is the true invariant optimum solution. 
Revised: This is a convex problem, since both the objective and the constraint are convex, and thus features a unique optimum, which is the true invariant optimum solution. T r a i n R i s k T e s t

Revised: R i s k

Original:  Convergence Dynamics (Soft-constraints, Fig 2B/C)
Revised: [Figure] Convergence Dynamics (Soft-constraints, Fig 2B/C)

Original: Here, we analyze the optimization dynamics under full-batch gradient descent (penalty method eq (16) with µ = 5 ⇥ 10 4 ).
Revised: Here, we analyze the optimization dynamics under full-batch gradient descent (penalty method eq (16) with µ = 5 ⇥ 10 4 ).

Original: IRM-v1’s convergence is highly dependent on the initialization (shown with differently colored trajectories), due to the presence of multiple local minima.
Revised: IRM-v1’s convergence is highly dependent on the initialization (shown with differently colored trajectories), due to the presence of multiple local minima.

Original: Note that most trajectories do not converge to the true invariant optimum.
Revised: Note that most trajectories do not converge to the true invariant optimum.

Original: IRM-relaxed also exhibits complex dynamics due to a saddle point near the true invariant solution: Some trajectories (red and magenta)
Revised: IRM-relaxed also exhibits complex dynamics due to a saddle point near the true invariant solution: Some trajectories (red and magenta)

Original: ﬁrst approach the line of invariant solutions, but all trajectories eventually converge to the non-invariant solution.
Revised: ﬁrst approach the line of invariant solutions, but all trajectories eventually converge to the non-invariant solution.

Original: Note that the true invariant optimum solution is not even a local optimum of IRM-relaxed at this value of µ , but it would exist in the limit µ ! 1 .
Revised: Note that the true invariant optimum solution is not even a local optimum of IRM-relaxed at this value of µ , but it would exist in the limit µ ! 1 .

Original: In contrast, MRI-v1 always converges to the true invariant optimum regardless of initialization since it is a unique minimum.T r a i n R i s k T e s t
Revised: In contrast, MRI-v1 always converges to the true invariant optimum regardless of initialization since it is a unique minimum.

Original: R i s k

Original: [Figure]

Original:  Nonlinear Image-based Problems
Revised:  Nonlinear Image-based Problems

Original: Here, we explore how MRI-v1 performs in nonlinear problems. Note that unlike the linear problem, there is no guarantee for either IRM-v1 or MRI-v1 for these nonlinear  problems.
Revised: Unlike for linear problems, theoretical proof for invariance is difﬁcult to show for nonlinear problems. Here, we empirically investigate the performance of IRM-v1 and MRI-v1 in nonlinear image-based problems.

Original: 6.1 Datasets 
Revised: 6.1 Datasets

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

Original: We introduce a newversatile dataset  designed for domain generalization on linear and nonlinear image-based regression and classiﬁcation problems.
Revised: We introduce a new dataset that is designed to evaluate domain generalization algorithms across various settings, including linear regression, linear classiﬁcation, nonlinear image-based regression, and nonlinear image-based classiﬁcation.

Original: The generative process of the dataset involves an invariant feature z i =
Revised: The generative process of the dataset involves an invariant feature z i =

Original: e i ✓ i , a spurious feature z s = e i ✓ s , and a label feature e i ✓ y , each of which represents an orientation on a complex unit circle: e i ✓ 2 S 1 .
Revised: e i ✓ i , a spurious feature z s = e i ✓ s , and a label feature e i ✓ y , each of which represents an orientation on a complex unit circle: e i ✓ 2 S 1 .

Original: The angles of orientations are generated as ✓ y ⇠ U S 1 , where U S 1 is the circular uniform distribution, and for ⇤ 2 { i, s } , ✓ ⇤ = ✓ y with probability p ⇤ or ✓ ⇤ ⇠ U S 1 with probability 1  p ⇤ .
Revised: The angles of orientations are generated as ✓ y ⇠ U S 1 , where U S 1 is the circular uniform distribution, and for ⇤ 2 { i, s } , ✓ ⇤ = ✓ y with probability p ⇤ or ✓ ⇤ ⇠ U S 1 with probability 1  p ⇤ .

Original: The parameter p i = 0 .
Revised: The parameter p i = 0 .

Original: 75 is ﬁxed across environments, whereas p s varies from one environment to another.
Revised: 75 is ﬁxed across environments, whereas p s varies from one environment to another.

Original: We consider two training environments E tr = { e 1 , e 2 } with p s e 1 = 1 , p s e 1 = 0 .
Revised: We consider two training environments E tr = { e 1 , e 2 } with p s e 1 = 1 , p s e 1 = 0 .

Original: 8 and one testing environment E test =
Revised: 8 and one testing environment E test =

Original: { e 0 } with p s e 0 = 0 .
Revised: { e 0 } with p s e 0 = 0 .

Original: In the linear regression task (section 5.2), the observed input issimply the concatenated latent features x = [ e i ✓ i , e i ✓ s ] and the label is y =
Revised: In the linear regression task (section 5.2), the observed input is the concatenated latent features x = [ e i ✓ i , e i ✓ s ] and the label is y =

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

Original: In the linear classiﬁcation task, the input is[Table] x = [sin( ✓ i ) , sin( ✓ s )] and the label is y
Revised: In the linear classiﬁcation task, the input is x = [sin( ✓ i ) , sin( ✓ s )] and the label is y

Original: = H (sin( ✓ y )) , where H is the sign function.
Revised: = H (sin( ✓ y )) , where H is the sign function.

Original: In the nonlinear regression/classiﬁcation tasks, the  inputobservation is an image composed of two planar waves, in which ✓ i is the orientation of the low frequency wave ( i [Table] e. shape) and ✓ s is that of the high frequency wave ( i [Table]
Revised: In the nonlinear regression/classiﬁcation tasks, the observed input is the image composed of two planar waves, in which ✓ i is the orientation of the low frequency wave ( i.e. shape) and ✓ s is that of the high frequency wave ( i.e. texture), as shown in Fig 5.

Original: e. texture), as shown in Fig 5.

Original: Colored MNIST (CMNIST) CMNIST (Arjovsky et al., 2019) is a synthetic dataset derived from
Revised: Colored MNIST (CMNIST) CMNIST (Arjovsky et al., 2019) is a synthetic dataset derived from

Original: MNIST for binary classiﬁcation.
Revised: MNIST for binary classiﬁcation.

Original: In this dataset, the label y assigned to an image is based on the digit bit z i (1 for digits 

Revised: In this dataset, the label y assigned to an image is based on the digit bit z i (1 for digits

Original: ⇠ 4 and -1 for 5 ⇠ 9 ) such that y
Revised: ⇠ 4 and -1 for 5 ⇠ 9 ) such that y

Original: =
Revised: =

Original: z i with probability p
Revised: z i with probability p

Original: i
Revised: i

Original: or
Revised: or

Original:  z i with probability 1  p i .
Revised:  z i with label such z the testing E

Original: The color bit z s (1 for red -1 for green) is chosen based on the label such that z s = y with probability p s or  y with probability 1  p s . We consider two versions, CMNISTa and
Revised: { e 0 }

Original: CMNISTb, with two sets of environmental parameters: CMNISTa is the version from Arjovsky et al. (2019) with p i = 0 . 75 , p s e 1 = 0 . 9 , p s e 1 = 0 . 8 , and p[-s e 
-]
Revised: ; CMNISTb uses p[-s e 
-]

Original:  =
Revised: 1 =

Original: 1 for the training E tr = { e 1 , e 2 } and the testing E test = { e 0 } environments; CMNISTb uses p i = 0 . 9 , p s e 1 = 1 , p s e 1 = 0 .
Revised: p s e 1 = 0 .

Original: 8 , p s e [-
-] =
Revised: 8 , p s e [-
-]0 =

Original: In the nonlinear tasks, the input observation x is the colored MNIST image.
Revised: In the nonlinear tasks, the input observation x is the colored MNIST image.

Original: In the abstracted versions, called toy-CMNISTa/b, the input observation is the two-bit data x =
Revised: In the abstracted versions, called toy-CMNISTa/b, the input observation is the two-bit data x =

Original: [ z i , z s ] (Kamath et al., 2021).
Revised: [ z i , z s ] (Kamath et al., 2021).

Original:  Note that both Shape-Texture and CMNIST datasets can be equally understood as being generated from Fig 1 with causal directions Y !
Revised: Remark Note that both Shape-Texture and CMNIST datasets can be equally understood as being generated from Fig 1 with causal directions Y !

Original: Z i or Z i !
Revised: Z i or Z i !

Original: Y (Fig 9).
Revised: Y (Fig 9).

Original: 6.2 Result
Revised: 6.2 Result

Original: Here, we report the performance of IRM-v1, MRI-v1, as well asERM algorithm, which simply minimizes the average empirical risk without  any invariance constraint.
Revised: Here, we report the performance of IRM-v1, MRI-v1, as well as the vanilla Empirical Risk Minimization (ERM) algorithm ( i.e. without imposing any invariance constraint).

Original: For reference, the results are compared to the Oracle’s performance, which is ERM’s performance on the dataset in which the spurious features are  uncorrelated with the label in the training dataset.
Revised: For reference, the results are compared to the Oracle performance, which is obtained by applying ERM on the modiﬁed training datasets in which the spurious features are rendered uncorrelated with the label.

Original: We tested the algorithms under a wide range of hyperparameters for both the Shape-Texture (Fig 3) and the CMNIST-b (Fig 4) datasets.
Revised: We tested the algorithms under a wide range of hyperparameters for both the Shape-Texture (Fig 3) and the CMNIST-b (Fig 4) datasets.

Original: Overall, MRI-v1 consistently achieves good invariant performance close to the Oracle, whereas IRM-v1 often shows either chance-level performance or poor O.O.D generalization with large differences between the train and the test domain risks.
Revised: Overall, MRI-v1 consistently achieves good invariant performance close to the Oracle, whereas IRM-v1 often shows either chance-level performance or poor OOD generalization with large differences between the train and the test domain risks.

Original: The results for a speciﬁc hyperparameter setting is shown in Table 2.
Revised: The results for a speciﬁc hyperparameter setting is shown in Table 2.

Original: [Figure] Under the PM setting, we observe that IRM-v1 often drives the models to the zero-predictor solution, i.e. making zero output regardless of the input, even with annealing the penalty term to be applied only in the later phase of training.
Revised: Under the PM setting, we observe that IRM-v1 often drives the models to the zero-predictor solution, i.e. making zero output regardless of the input, even with annealing the penalty term to be applied only in the later phase of training.

Original: This explains IRM-v1’s identically low performance across train and test domains, consistent to the previously reported results in Gulrajani and Lopez-Paz (2020).
Revised: This explains IRM-v1’s identically low performance across train and test domains, consistent to the previously reported results in Gulrajani and Lopez-Paz (2020).

Original: In contrast, MRI-v1 never drives the models to the zero-predictor solution, consistent with the ﬁnding in
Revised: In contrast, MRI-v1 never drives the models to the zero-predictor solution, consistent with the ﬁnding in

Original: Sec 5.2 that MRI-v1 does not have an local minima at the zero-predictor in the linear problem setting.
Revised: Sec 5.2 that MRI-v1 does not have an local minima at the zero-predictor in the linear problem setting.

Original: Interestingly, the ALM setting greatly improves IRM-v1’s performance especially in terms of accuracy, while MRI-v1’s performance remains relatively unchanged between PM and ALM.
Revised: Interestingly, the ALM setting greatly improves IRM-v1’s performance especially in terms of accuracy, while MRI-v1’s performance remains relatively unchanged between PM and ALM.

Original: Under the ALM setting, IRM-v1 shows comparable or slightly higher accuracy than MRI-v1 for the CMNIST dataset.
Revised: Under the ALM setting, IRM-v1 shows comparable or slightly higher accuracy than MRI-v1 for the CMNIST dataset.

Original: In the Supplementary Materials, we also report the results for other recent domain generalization methods including MMD (Li et al., 2018b), GroupDRO (Sagawa et al., 2019), and IB-IRM (Ahuja et al., 2021) on the Shape-Texture classiﬁcation task, which exhibit signiﬁcantly worse performance than MRI-v1 (Fig 7). 
Revised: In the Supplementary Materials, we also report the results for other recent domain generalization methods including MMD (Li et al., 2018b), GroupDRO (Sagawa et al., 2019), and IB-IRM (Ahuja et al., 2021) on the Shape-Texture classiﬁcation task, which exhibit signiﬁcantly worse performance than MRI-v1 (Fig 7). Risk [Table]

Original:  Discussion
Revised:  Discussion

Original:  In this work, we introduce MRI, a novel method based on conserving the label-conditioned feature expectation E e [f ( x ) | y  ]  across environments.
Revised: Limitations In principle, IRM only requires E e [ y | z i ] to be constant across domains in order to guarantee invariance

Original: This method substantially outperforms IRM on certain classes of OOD generalization problems, in both linear and non-linear settings. Here, we conclude with a discussion on the limitations of our method and the limitations of accuracy as a metric to compare the degree of invariance achieved by various algorithms. Accuracy metric
Revised: [Table]

Original: The accuracy metric is widely used in the ﬁeld of domain generalization (Gulrajani and Lopez-Paz, 2020) to compare the performance of various algorithms. However, accuracy is not sensitive enough to reliably measure the degree of invariance of a model.
Revised: and therefore it has been previously thought to be generally applicable in a wide range of problems, even though this guarantee was only shown in the impractical limit of optimizing over the unrestricted function space of .

Original: Consider the linear classiﬁcation problem of toy-CMINSTb shown in Fig 8. For this problem, MRI-v1 has a unique
Revised: Here

Original: [Figure]
Revised: [Table]

Original: optimum, which is the optimal invariant solution, while IRM-v1 exhibits multiple local optima, including non-invariant solutions.
Revised: we showed a strong negative result that no meaningful form of invariance can be stated for IRM when the function space of is restricted.

Original: Yet, all these solutions exhibit identical accuracy (except for the zero-predictor solution of IRM-v1), since the accuracy landscape (Fig 8f) is ﬂat over a wide range ofsolutions: invariant accuracy is achieved as long as w i > | w s | .
Revised: In contrast, MRI requires a more limiting condition that the label distribution P ( y ) and P ( z i ) should be constant across domains [Table] but it can be more generally applied even for the case of restricted function space of .

Original: As a result, accuracy cannot be used as a measure of invariance in this problem. Therefore, the fact that MRI-v1 and IRM-v1 exhibit similar test-domain accuracy does not necessarily indicate that they achieve similar degree of invariance.
Revised: which yields more practicality.

Original: Limitations A common limitation of both MRI and IRM (Rosenfeld et al., 2020; Ahuja et al., 2021) is that they require signiﬁcant support overlap across domains.
Revised: A common limitation of both MRI and IRM (Rosenfeld et al., 2020; Ahuja et al., 2021) is that they require signiﬁcant support overlap across domains in order to guarantee OOD generalization.

Original: This may limit the applicability of these methods on certain domain generalization benchmarks(Gulrajani and Lopez-Paz, 2020) that consists of domains with different styles or natural backgrounds, like VLCS (Fang et al., 2013) and
Revised: This may limit the applicability of these methods on certain domain generalization benchmarks that consist of domains that lack such overlap, such as VLCS (different image stylization, Fang et al.

Original:  Terra-Incognita (Beery et al., 2018).
Revised: and Terra-Incognita (different natural backgrounds, Beery et al.

Original: In the Supplementary Materials, we report that both methods fail to improve over ERM(in terms of test accuracy) on both VLCS and Terra-Incognita datasets (Table 3). 
Revised: In the Supplementary Materials, we report that both methods do not show signiﬁcant improvement over ERM on these datasets (Table 3). Insensitivity of accuracy metric in evaluating invariance For CMINSTa/b, IRM-v1 (ALM) exhibits comparable performance to MRI-v1 in terms of test domain accuracy, despite having signiﬁcantly worse performance in term of test domain risk (See Table 2). We investigated this phenomenon by analyzing the linear version of the task, toy-CMINSTa/b.

Original: Another limitation of MRI is that it requires the label distribution P e ( y ) and  the invariant feature distribution P e ( z i ) to remain invariant across domains.
Revised: The accuracy landscape is piece-wise constant (Fig 8C,F). Especially, it exhibits identical ( i.e. invariant) accuracy between train and test domain in the region deﬁned by w i > | w s | .

Original:  In contrast, IRM does not require such assumptions.
Revised: Therefore, the accuracy metric cannot distinguish invariant solutions ( w s = 0 ) from non-invariant solutions within the region. In contrast, the train domain and the test domain risk share the same value only if the solution is invariant ( w s = 0 ) (Fig 8B,E). The constraint function of IRM-v1 (Fig 8A,D) shows that toy-CMNISTa only has invariant local optima and that toy-CMNISTb has additional non-invariant local optima, all of which satisfy w i > | w s | , thus exhibiting the same accuracy performance as the invariant optimum solution.

Original: However, the practical implementation of IRM, IRM-v1, which can behave very differently from IRM (Kamath et al., 2021), has been shown to fail for such problems (Ahuja et al., 2020, 2021).
Revised: This result illustrates that using accuracy metric alone for evaluating the degree of invariance could be insufﬁcient, and highlights the need to also consider risk for evaluations.

Original: Acknowledgments and Disclosure of Funding
Revised: Acknowledgments and Disclosure of Funding

Original: This research was done as part of Avinash Baidya’s internship at IBM.
Revised: This research was done as part of Avinash Baidya’s internship at IBM.

Original: We thank Joel Dappelo for helpful discussions.
Revised: We thank Joel Dapello for helpful discussions.

Original: We also thank the anonymous reviewers for constructive comments.
Revised: We also thank the anonymous reviewers for constructive comments.

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 on how to answer these questions.
Revised: Please read the checklist guidelines carefully for information on how 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 referencing the 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 referencing the 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 are proprietary.
Revised: [No] The code and the data are proprietary.

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

Original: For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?
Revised: For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?

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

Original: [Yes] See Section 7.
Revised: [Yes] See Section 7.

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

Original: [N/A] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
Revised: [N/A] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?

Original: [Yes] 2.
Revised: [Yes] 2.

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

Original: [Yes] See Sectionand 5.
Revised: [Yes] See Sectionand 5.

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

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

Original: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
Revised: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?

Original: [Yes] https://github.com/IBM/MRI (b) Did you specify all the training details (e
Revised: [Yes] https://github.com/IBM/MRI (b) Did you specify all the training details (e

Original: g., data splits, hyperparameters, how they were chosen)?
Revised: g., data splits, hyperparameters, how they were 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 experiments multiple times)?
Revised: (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?

Original: [Yes] See Table 2.
Revised: [Yes] See Table 2.

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

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

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

Original: [Yes] See Section 6.1 and Supplementary Materials.
Revised: [Yes] See Section 6.1 and Supplementary 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] https://github.com/IBM/MRI (d) Did you discuss whether
Revised: [Yes] https://github.com/IBM/MRI (d) Did you discuss whether

Original: and how consent was obtained from people whose data you’re using/curating?
Revised: and how consent was obtained from people whose data you’re using/curating?

Original: [N/A] (e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?
Revised: [N/A] (e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?

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

Original: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
Revised: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?

Original: [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
Revised: [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?

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

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