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 ﬂaw of IRM formulation that causes the failure.

Original: We then introduce a complementary notion of invariance,
Revised: We then introduce a complementary notion of invariance,

Original: MRI, based on conserving the label-conditioned feature expectation E e [ f ( x ) | y ] , which is free of this ﬂaw.
Revised: 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.
Revised: 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 in a general linear problem setting, MRI-v1 can guarantee invariant predictors given sufﬁcient environments.

Original: We also empirically demonstrate that
Revised: We also empirically demonstrate that

Original: MRI strongly out-performs IRM and consistently achieves near-optimal OOD generalization in image-based nonlinear problems.
Revised: MRI strongly out-performs IRM 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 learned to 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 to 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.

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

Original: 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) 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, in particular, have attracted much attention as one of the ﬁrst 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 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.

Original: We also 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 also show that IRM has a fundamental ﬂaw due to the indirect way it imposes invariance which leads to the failure of IRM-v1.

Original: MRI is shown to be free 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 invariant predictors given sufﬁcient environments.
Revised: We prove that in a general linear problem setting, MRI-v1 can guarantee invariant predictors given sufﬁcient environments.

Original: We then empirically demonstrate that MRI strongly out-performs IRM and consistently achieves near-optimal
Revised: We then empirically demonstrate that MRI strongly out-performs IRM and consistently achieves near-optimal

Original: OOD generalization in image-based nonlinear problems.
Revised: OOD generalization in image-based nonlinear 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.
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: 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 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: = 1 2 ( 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 !
Revised: Given a predictor  f : X !

Original: O , the feature representation f : X !
Revised: O , the feature representation f : X !

Original: Z is IRM invariant over a set of environments E if there exists : Z !
Revised: Z is IRM invariant over a set of environments E if there exists : Z !

Original: O such that is
Revised: O such that is

Original: While Fig 1 considers the causal direction Y !
Revised: While Fig 1 considers the causal direction Y !

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

Original: Y .
Revised: Y .

Original: simultaneously optimal on f for all environments in E , i.e.
Revised: simultaneously optimal 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 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’s 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: (6) 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 with f replacing the composite predictor
Revised: which is equivalent to Lemma 3.2 with f replacing the composite predictor

Original: [Equation]
Revised: [Equation]

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 predictor f to be simultaneously optimal over all environments.
Revised: Deﬁnition 3.3 describes the ﬁrst-order condition for the predictor f to be simultaneously optimal over all environments.

Original: This is indeed a sufﬁcient condition for optimality given a convex loss function l .
Revised: This is indeed a sufﬁcient condition for optimality given a convex loss function l .

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: It was 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 ,
Revised: (2021) that the essense of IRM’s invariance is the conservation of the feature-conditioned label expectation ,

Original: i.e.
Revised: 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: IRM eq (7) (also eq (5)) indeed implies the conservation law, since its RHS is constant with respect to the environment index e .

Original: However, there is a critical difference: Eq (8) only requires E e
Revised: However, there is a critical difference: Eq (8) only requires E e

Original: [ y | o ] to share the same value across environments, thus imposing |E|  1 constraints, whereas eq (7) (also eq (5)) additionally imposes that value to be also equal to the RHS term  ( o ) .
Revised: [ y | o ] to share the same value across environments, thus imposing |E|  1 constraints, whereas eq (7) (also eq (5)) additionally imposes that value to be also equal to the RHS term  ( o ) .

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

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: Despite the practical applicability, it has been empirically shown that IRM-v1 can behave quite differently 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: 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: Unlike in eq (7), however, RHS of eq (10) is not 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: 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 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 the correct number of constraints in eq (11) as described by the conservation law eq (13).
Revised: Note that MRI attains invariance in a direct manner without involving any intermediary terms, thus imposing the correct number of constraints in eq (11) as described by the conservation law eq (13).

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
Revised: Therefore, the direct way of

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

Original: Algorithm
Revised: Algorithm

Original: [Equation]
Revised: [Equation]

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 , and w is the parameters of the predictor f .

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

Original: 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: 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
Revised: For example, for a training set of two environments

Original: E tr = { e 1 , e 2 } , Q ~ L = (  L e 1   L e 2 ) / p 2 , since Q = [1 ,  1] / p 2 .
Revised: 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: The constraints are summarized in 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,
Revised: (2019), which adds the squared residual constraints as a penalty term to the objective,

Original: [Equation]
Revised: [Equation]

Original: w
Revised: w

Original: [Equation]
Revised: [Equation]

Original: tr
Revised: tr

Original: [Equation]
Revised: [Equation]

Original: + µ k ~
Revised: + µ k ~

Original: [Equation]
Revised: [Equation]

Original: f )
Revised: 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)
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: [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: 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 all spurious features of a linear predictor given a sufﬁcient number of environments.
Revised: In this section, we show that MRI-v1 can correctly remove all spurious features of 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
Revised: i

Original: R d i , z s 2 R d s .
Revised: 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: with 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 0 = M  ¯ M · 1 with M ⌘ [ E e [ z s y ]] e 2 E

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

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

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

Original: E
Revised: E

Original: e 2 [ o 2  oy ] , shown as two elliptic curves.
Revised: e 2 [ o 2  oy ] , 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.

Original: Convergence Dynamics (Soft-constraints, Fig 2B/C)
Revised: 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.
Revised: In contrast, MRI-v1 always converges to the true invariant optimum regardless of initialization since it is a unique minimum.

Original: T r a i n
Revised: T r a i n

Original: R i s k
Revised: R i s k

Original: T e s t
Revised: T e s t

Original: R i s k
Revised: R i s k

Original: [Figure]
Revised: [Figure]

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

Original: Here, we explore how MRI-v1 performs in nonlinear problems.
Revised: Here, we explore how MRI-v1 performs in nonlinear problems.

Original: Note that unlike the linear problem, there is no guarantee for either IRM-v1 or MRI-v1 for these nonlinear problems.
Revised: Note that unlike the linear problem, there is no guarantee for either IRM-v1 or MRI-v1 for these nonlinear problems.

Original: 6.1 Datasets
Revised: 6.1 Datasets

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

Original: We introduce a new versatile dataset designed for domain generalization on linear and nonlinear image-based regression and classiﬁcation problems.
Revised: We introduce a new versatile dataset designed for domain generalization on linear and nonlinear image-based regression and classiﬁcation problems.

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 is simply 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 simply 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
Revised: In the linear classiﬁcation task, the input is

Original: [Table]
Revised: [Table]

Original: x = [sin( ✓ i ) , sin( ✓ s )] and the label is y
Revised: 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 input observation is an image composed of two planar waves, in which ✓ i is the orientation of the low frequency wave ( i
Revised: In the nonlinear regression/classiﬁcation tasks, the input observation is an image composed of two planar waves, in which ✓ i is the orientation of the low frequency wave ( i

Original: [Table]
Revised: [Table]

Original: e. shape) and ✓ s is that of the high frequency wave ( i
Revised: e. shape) and ✓ s is that of the high frequency wave ( i

Original: [Table]
Revised: [Table]

Original: e. texture), as shown in Fig 5.
Revised: 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 probability 1  p i .

Original: The color bit z s (1 for red -1 for green) is chosen based on the label such that z s =
Revised: The color bit z s (1 for red -1 for green) is chosen based on the label such that z s =

Original: y with probability p s or  y with probability 1  p s .
Revised: y with probability p s or  y with probability 1  p s .

Original: We consider two versions, CMNISTa and
Revised: We consider two versions, CMNISTa and

Original: CMNISTb, with two sets of environmental parameters: CMNISTa is the version from Arjovsky et al.
Revised: CMNISTb, with two sets of environmental parameters: CMNISTa is the version from Arjovsky et al.

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

Original: 75 , p s e 1 = 0 .
Revised: 75 , p s e 1 = 0 .

Original: 9 , p s e 1 = 0 . 8 , and p s e 

Revised: 9 , p s e 1 = 0 . 8 , and p s e

Original: =
Revised: =

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

Original: { e 0 } environments; CMNISTb uses p i = 0 .
Revised: { e 0 } environments; CMNISTb uses p i = 0 .

Original: 9 , p s e 1 = 1 , p s e 1 = 0 .
Revised: 9 , p s e 1 = 1 , p s e 1 = 0 .

Original: 8 , p s e 

Revised: 8 , p s e

Original: =
Revised: =

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: 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 as ERM 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 ERM algorithm, which simply minimizes the average empirical risk without 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 ’s performance, which is ERM’s performance on the dataset in which the spurious features are uncorrelated with the label in the training dataset.

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 O.O.D 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]
Revised: [Figure]

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

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: In this work, we introduce MRI, a novel method based on conserving the label-conditioned feature expectation E e [ f ( x ) | y ] across environments.

Original: This method substantially outperforms IRM on certain classes of OOD generalization problems, in both linear and non-linear settings.
Revised: This method substantially outperforms IRM on certain classes of OOD generalization problems, in both linear and non-linear settings.

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

Original: Accuracy metric
Revised: Accuracy metric

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

Original: However, accuracy is not sensitive enough to reliably measure the degree of invariance of a model.
Revised: However, accuracy is not sensitive enough to reliably measure the degree of invariance of a model.

Original: Consider the linear classiﬁcation problem of toy-CMINSTb shown in Fig 8.
Revised: Consider the linear classiﬁcation problem of toy-CMINSTb shown in Fig 8.

Original: For this problem, MRI-v1 has a unique
Revised: For this problem, MRI-v1 has a unique

Original: [Figure]
Revised: [Figure]

Original: optimum, which is the optimal invariant solution, while IRM-v1 exhibits multiple local optima, including non-invariant solutions.
Revised: optimum, which is the optimal invariant solution, while IRM-v1 exhibits multiple local optima, including non-invariant solutions.

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 of solutions: invariant accuracy is achieved as long as w i > | w s | .
Revised: 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 of solutions: invariant accuracy is achieved as long as w i > | w s | .

Original: As a result, accuracy cannot be used as a measure of invariance in this problem.
Revised: As a result, accuracy cannot be used as a measure of invariance in this problem.

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

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

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 (Gulrajani and Lopez-Paz, 2020) that consists of domains with different styles or natural backgrounds, like VLCS (Fang et al., 2013) and

Original: Terra-Incognita (Beery et al., 2018).
Revised: Terra-Incognita (Beery et al., 2018).

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 fail to improve over ERM (in terms of test accuracy) on both VLCS and Terra-Incognita datasets (Table 3).

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

Original: In contrast, IRM does not require such assumptions.
Revised: In contrast, IRM does not require such assumptions.

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

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

Revised: A Analytical Solution for Toy Shape-Texture Regression

Revised: The Shape-Texture dataset involves an invariant latent feature ✓ i 2 R , a spurious latent feature ✓ s 2 R , and the label ✓ y 2 R , each of which represents an angle of orientation, and therefore is naturally represented on a unit circle in the complex plane as e i ✓ 2 S 1 .

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

Revised: e i ✓ y .

Revised: For all task settings, the latent angles and the label angle are linearly related as ✓ i ✓ y + ⇠ i , ✓ s ✓ y + ⇠ s , where ⇠ ⇤ ⇠ N (0 ,  ⇤ ) with probability p ⇤ , or ⇠ ⇤ ⇠ U (  ⇡ , ⇡ ) probability 1  p

Revised: ⇤ for ⇤ 2 { i, s } .

Revised: We used p i = 0 .

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

Original: {+8 , p s 2 = 1 for the two training environments, p s test = 
+}
Revised: {+8 , p s 2 = 1 for the two training environments, p s test = 
+}

Revised: for the testing environment, and  i

Revised: =

Revised:  s e = 0 for all environments.

Revised: Note that

Revised: E e

Revised: [Equation]

Revised: = ⇢ 1 , with prob.

Revised: p i 0 , with prob.

Revised: 1  p i

Revised: E e [ e

Revised: [Equation]

Revised: ( ✓ s  ✓ y

Revised: [Equation]

Revised: ] = ⇢ 1 , with prob.

Revised: p s 0 , with prob.

Revised: 1  p s

Revised: E e [ e i ( ✓ i  ✓ s ) ]

Revised: = ⇢ 1 , with prob.

Revised: p i p s 0 , with prob.

Revised: 1  p i p s which yields

Revised: [Equation]

Revised: We consider real weights w =

Revised: [ w i w s ] .

Revised: Since o = wx ,

Revised: E e [ oo

Revised: [Equation]

Revised: [Equation]

Revised: xx

Revised: [Equation]

Revised: E e

Revised: [Equation]

Revised: e [

Revised: [Equation]

Revised: [Equation]

Revised: i p i

Revised: [Equation]

Revised: s

Revised: [Equation]

Revised: e

Revised: The risk of environment e is

Revised: [Equation]

Revised: Therefore, the average risk for training is

Revised: [Equation]

Revised: MRI constraint

Revised: [Equation]

Revised: Minimizing ¯ L tr ( f ) subject to the above constraint yields ( w ⇤ i , w ⇤ s )

Revised: = ( p i , 0) as the constrained optimum.

Revised: IRM-relaxed constraint

Revised: [Equation]

Revised: [Equation]

Revised: [Figure]

Revised: Minimizing ¯ L tr ( f ) subject to the above constraint yields

Revised: [Equation]

Revised: as the constrained optima.

Revised: IRM constraint

Revised: [Equation]

Revised: which yields

Revised: [Equation]

Revised: as constraints, which are also the constrained optima of the problem.

Revised: B Training Procedure

Revised: PyTorch version 1.10.0 was used.

Revised: For all tasks, Tesla V-100 with a memory of 16GB was used.

Revised: For all our simulations, we utilized the DomainBed suite (Gulrajani and Lopez-Paz (2020); available under MIT license).

Revised: B.1 Linear Datasets

Revised: For Fig 2, 6, and 8, we used a linear predictor f ( x ) =

Revised: w i

Revised: z i + w s z s .

Revised: In all cases, we used full-batch training with a batch size of 40000.

Revised: We optimized using the penalty method with µ = 5 ⇥ 10 4 .

Revised: For

Revised: Fig 2, we used the SGD optimizer with no momentum.

Revised: For Fig 6 and 8, we used the Adam optimizer with betas of (0.975, 0.999).

Revised: The learning rate was set to 0.005, the total number of gradient steps were set to 2000, and the gradient norms were clipped at 2 for all cases.

Revised: [Figure]

Revised: B.2 Non-Linear Datasets

Revised: The results for the nonlinear datasets are presented in Table 2-3 and Fig 3, 4, and, 7.

Revised: The results in

Revised: Table 2-3 were produced from 10 different model initialization using random seeds.

Revised: Shape-Texture Datasets For the shape-texture datasets

Revised: , we used a single convolutional layer with 12 channels, a kernel size of 5, and stride of 1.

Revised: The convolutional layer was followed by a 2D adaptive average pooling layer with output size of 1x1 followed by a fully-connected linear layer.

Revised: For results in Table 2, for optimization using the ALM method we used µ = 10 and  = 10 and for optimization

Revised: [Figure]

Revised: using PM we used µ = 10000 .

Revised: We used mini-batch training with a mini-batch size 8000 and total dataset size 50000.

Revised: We used the Adam optimizer with betas of (0.9, 0.999).

Revised: The learning rate was set to 0.01.

Revised: The total number of gradient steps were set to 2000.

Revised: Gradient clipping was only used when training using PM such that gradient norms were clipped at 2.

Revised: We used a weight decay of 0.003 for the classiﬁcation setting and 0.005 for the regression setting.

Revised: For Fig 3 and Fig 7, we used the same training procedure as above for the penalty method except that the seeds and algorithm-speciﬁc parameters were chosen at random to generate 50 random hyperparameter combinations for each algorithm.

Revised: For IRM-v1 and MRI-v1, µ was chosen randomly from the distribution 10 Uniform [2 , 5] .

Revised: We used the DomainBed (Gulrajani and Lopez-Paz, 2020) implementation of IB-IRM, GroupDRO, and MMD.

Revised: For IB-IRM,  IRM and  IB were both chosen randomly from the distribution 10 Uniform [2 , 5] .

Revised: For GroupDRO, ⌘ GroupDRO was chosen from the distribution 10 Uniform [  3 , 1] .

Revised: For MMD,  MMD was chosen from the distribution 10 Uniform [  1 , 1] .

Revised: CMNIST For the CMNIST datasets, we used the LeNet architecture (LeCun et al., 1998).

Revised: We used mini-batch training with a mini-batch size of 256.

Revised: For optimization using the ALM method with  = 10 for both datasets, but µ = 1 for CMNISTa and µ = 10 for CMNISTb datasets for results in

Revised: Table 2.

Revised: For optimization using the PM method, we used µ = 10000 for both datasets.

Revised: We used the

Revised: Adam optimizer with betas of (0.975, 0.999).

Revised: The learning rate was set to 0.001.

Revised: The total number of gradient steps were set to 10000.

Revised: The weight decay was set to 0.001.

Revised: Gradient clipping was only used when training using PM such that gradient norms were clipped at 2.

Revised: For the CMNIST datasets, we also used data augmentation during training.

Revised: This was done using Torchvision’s Random afﬁne transformation class with degrees=30, translate=(0.1, 0.1), and scale=(0.8, 1.2).

Revised: For Fig 4, we used the same training procedure as above for training except that the seeds and algorithm-speciﬁc parameters were chosen at random to generate 50 random hyperparameter combinations for each algorithm.

Revised: For training with PM, µ was chosen independently and randomly from the distribution 10 Uniform [2 , 5] .

Revised: For training with ALM,  and µ were chosen randomly from the distribution Uniform [1 , 10] .

Revised: The total number of gradient steps were set to 5000.

Revised: VLCS/Terra-Incognita VLCS (Fang et al., 2013) and Terra-Incognita (Beery et al., 2018) are popular domain generalization datasets that contain realistic images.

Revised: For the binary classiﬁcation tasks, we only used two classes from each dataset (VLCS: bird and person, Terra-Incognita: opossum and raccoon).

Revised: We used the LeNet architecture (LeCun et al., 1998) for both the datasets.

Revised: We used mini-batch training with a mini-batch size of 256.

Revised: We trained the models using the ALM method with µ = 10 and  = 10 .

Revised: We used the Adam optimizer with betas of (0.975, 0.999).

Revised: The learning rate was set to 0.001.

Revised: The total number of gradient steps were set to 2000.

Revised: The weight decay was set to 0.001.

Revised: No gradient clipping was used.

Revised: The images were resized to the shape (64,64,3) and normalized with the mean (0.485, 0.456, 0.406) and standard deviation (0.229, 0.224, 0.225).

Revised: For data augmentation, we used Torchvision’s RandomHorizontalFlip transform, RandomGrayscale transform,

Revised: RandomResizedCrop transform with scale = (0.7, 1.0), and ColorJitter transform with brightness, contrast, saturation, and

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Figure]