Original: The Missing Invariance Principle Found – the Reciprocal Twin of Invariant Risk Minimization
Revised: The Missing Invariance Principle Found – the Reciprocal Twin of Invariant Risk Minimization

Original: Anonymous Author(s)
Revised: Anonymous Author(s)

Original: AfﬁliationAddress email
Revised: AfﬁliationAddress email

Original: Abstract
Revised: Abstract

Original: Machine learning models often generalize poorly to out-of-distribution (OOD)data as a result of relying on features that are spuriously correlated with the labelduring training.
Revised: Machine learning models often generalize poorly to out-of-distribution (OOD)data as a result of relying on features that are spuriously correlated with the labelduring training.

Original: Recently, the technique of Invariant Risk Minimization (IRM)was proposed to learn predictors that only use invariant features by conserving thefeature-conditioned label expectation E e [ y | f ( x )] across environments.
Revised: Recently, the technique of Invariant Risk Minimization (IRM)was proposed to learn predictors that only use invariant features by conserving thefeature-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 task settings.
Revised: However,more recent studies have demonstrated that IRM-v1, a practical version of IRM,can fail in various task settings.

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

Original: We then introduce a complementary notion ofinvariance, MRI, based on conserving the label-conditioned feature expectation
Revised: We then introduce a complementary notion ofinvariance, MRI, based on conserving the label-conditioned feature expectation

Original: E e [ f ( x ) | y ] across environments, which is free of this ﬂaw.
Revised: E e [ f ( x ) | y ] across environments, which is free of this ﬂaw.

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

Original: We note thatthis constraint is convex which confers it with an advantage over IRM-v1, whichimposes non-convex constraints.
Revised: We note thatthis constraint is convex which confers it with an advantage over IRM-v1, whichimposes non-convex constraints.

Original: We prove that in a general linear problem setting,
Revised: We prove that in a general linear problem setting,

Original: MRI-v1 can guarantee invariant predictors given sufﬁcient environments.
Revised: MRI-v1 can guarantee invariant predictors given sufﬁcient environments.

Original: We alsoempirically demonstrate that MRI strongly out-performs IRM and consistentlyachieves near-optimal OOD generalization in image-based nonlinear problems.
Revised: We alsoempirically demonstrate that MRI strongly out-performs IRM and consistentlyachieves near-optimal OOD generalization in image-based nonlinear problems.

Original:  Introduction
Revised:  Introduction

Original: Deep learning models have shown tremendous success over the past decade.
Revised: Deep learning models have shown tremendous success over the past decade.

Original: These models showgreat generalization properties when tested on the same distribution as the training dataset (indistribution generalization).
Revised: These models showgreat generalization properties when tested on the same distribution as the training dataset (indistribution generalization).

Original: However, these models often show catastrophic failure when tested onout-of-distribution dataset, revealing that they learned to use easier-to-ﬁt features that are spuriouslycorrelated to the label in the given training domains but do not generalize to the testing domains.
Revised: However, these models often show catastrophic failure when tested onout-of-distribution dataset, revealing that they learned to use easier-to-ﬁt features that are spuriouslycorrelated to the label in the given training domains but do not generalize to the testing domains.

Original: For example, deep networks trained on pictures of cow with only grassy backgrounds in the trainingdomain will use the background color as the predictive feature which is easier to learn and generalizepoorly to pictures of cow, for example, with a dessert background.
Revised: For example, deep networks trained on pictures of cow with only grassy backgrounds in the trainingdomain will use the background color as the predictive feature which is easier to learn and generalizepoorly to pictures of cow, for example, with a dessert background.

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

Original: In particular, there has been a recent body of works that focus on developing algorithmsthat attempt to learn invariant predictors (Arjovsky et al., 2019; Ahuja et al., 2021; Peters et al., 2016;
Revised: In particular, there has been a recent body of works that focus on developing algorithmsthat attempt to learn invariant predictors (Arjovsky et al., 2019; Ahuja et al., 2021; Peters et al., 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 itspractical version, IRM-v1, in particular, have attracted much attention as one of the ﬁrst methods thatare compatible with deep learning techniques.
Revised: Invariant Risk Minimization (IRM) and itspractical version, IRM-v1, in particular, have attracted much attention as one of the ﬁrst methods thatare compatible with deep learning techniques.

Original: However, several follow-up studies have empiricallydemonstrated that IRM-v1 is unreliable at learning invariant representations (Kamath et al., 2021;
Revised: However, several follow-up studies have empiricallydemonstrated 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 afundamental ﬂaw of IRM formulation that causes this limitation and propose a new method that isfree of this ﬂaw.
Revised: Here, we identify afundamental ﬂaw of IRM formulation that causes this limitation and propose a new method that isfree of this ﬂaw.

Original: Related Works There has been considerable work in the ﬁeld of learning invariant representations.
Revised: Related Works There has been considerable work in the ﬁeld of learning invariant representations.

Original: They vary from learning domain-invariant feature representations conserving P ( f ( x )) using kernelmethods (Muandet et al., 2013; Ghifary et al., 2016;
Revised: They vary from learning domain-invariant feature representations conserving P ( f ( x )) using kernelmethods (Muandet et al., 2013; Ghifary et al., 2016;

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

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

Original: ; Akuzawa et al.(2019); Albuquerque et al.
Revised: ; Akuzawa et al.(2019); Albuquerque et al.

Original: (2019) to learning invariant class-conditional features P ( f ( x ) | y )
Revised: (2019) to learning invariant class-conditional features P ( f ( x ) | y )

Original: (Gonget al., 2016; Li et al., 2018b) in the context of domain adaptation, which assumes access to the testdistribution for adaptation.
Revised: (Gonget al., 2016; Li et al., 2018b) in the context of domain adaptation, which assumes access to the testdistribution for adaptation.

Original: There is also a large body of work to learn invariant representations inthe ﬁeld of domain generalization that doesn’t assume access to test distributions.
Revised: There is also a large body of work to learn invariant representations inthe ﬁeld of domain generalization that doesn’t assume access to test distributions.

Original: This includesimposing invariance of E e [ y | f ( x )]
Revised: This includesimposing invariance of E e [ y | f ( x )]

Original: (Arjovsky et al., 2019) with information bottleneck constraint(Ahuja et al. (2021)), imposing object-invariant condition (Mahajan et al. (2021)), using domaininference (Creager et al., 2021), model calibration (Wald et al., 2021), and others (Krueger et al.
Revised: (Arjovsky et al., 2019) with information bottleneck constraint(Ahuja et al. (2021)), imposing object-invariant condition (Mahajan et al. (2021)), using domaininference (Creager et al., 2021), model calibration (Wald et al., 2021), and others (Krueger et al.

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

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

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

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

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

Original: We thenempirically demonstrate that MRI strongly out-performs IRM and consistently achieves near-optimal
Revised: We thenempirically 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 inputsand labels, from which the dataset of the environment is drawn D e ⌘ { ( x ej 2 R d , y ej 2 R ) } .
Revised: Consider a set of environments E = { e } , each of which deﬁnes a distribution P e ( x, y ) over inputsand labels, from which the dataset of the environment is drawn D e ⌘ { ( x ej 2 R d , y ej 2 R ) } .

Original: Thedistribution P e ( x, y ) is assumed to be generated according to the causal graph in Fig.1 as introduced in
Revised: Thedistribution P e ( x, y ) is assumed to be generated according to the causal graph in Fig.1 as introduced in

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

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

Original: i or spuriously z s correlatedwith the label.
Revised: i or spuriously z s correlatedwith the label.

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

Original: Note that while this work focuses on the causal graph in Figure 1, our results are also applicable tothe graph in which the invariant feature z i causes the label y
Revised: Note that while this work focuses on the causal graph in Figure 1, our results are also applicable tothe graph in which the invariant feature z i causes the label y

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

Original: Y ).
Revised: Y ).

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

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

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

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

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

Original: Y
Revised: Y

Original: Z i
Revised: Z i

Original: Z s
Revised: Z s

Original: X
Revised: X

Original: E
Revised: E

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

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

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

Original: where o
Revised: where o

Original: = f
Revised: = f

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

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

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

Original: =
Revised: =

Original: 12 ( o  y ) 2 for regression ( O, Y ✓ R ) and the binary-cross-entropy (BCE) lossl log ( o, y )
Revised: 12 ( o  y ) 2 for regression ( O, Y ✓ R ) and the binary-cross-entropy (BCE) lossl 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 75 ⌘ ( 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 75 ⌘ ( o ) ⌘ 1 / (1 + e  o ) is the sigmoid function.

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

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

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

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

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

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

Original: O such that is simultaneouslyoptimal on f for all environments in E , i.e.
Revised: O such that is simultaneouslyoptimal on f for all environments in E , i.e.

Original: [Equation]
Revised: [Equation]

Original: where is assumed to be unrestricted in the space of all measurable functions.
Revised: where is assumed to be unrestricted in the space of all measurable functions.

Original: Lemma 3.2.
Revised: Lemma 3.2.

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

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

Original: [Equation]
Revised: [Equation]

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

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

Original: 3.1.2 Variational formulation
Revised: 3.1.2 Variational formulation

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

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

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

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

Original: O is IRM invariant over a set of environments E , if therisk L e ( f ) remains stationary under arbitrary inﬁnitesimal perturbations on the predictor’s outputo
Revised: O is IRM invariant over a set of environments E , if therisk L e ( f ) remains stationary under arbitrary inﬁnitesimal perturbations on the predictor’s outputo

Original: = f
Revised: = f

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

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

Original: where  : O !
Revised: where  : O !

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

Original: Lemma 3.4.
Revised: Lemma 3.4.

Original: For standard loss functions 2 , Deﬁnition 3.3 is equivalent to8 e 2 E ,
Revised: For standard loss functions 2 , Deﬁnition 3.3 is equivalent to8 e 2 E ,

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

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

Original: Proof.
Revised: Proof.

Original: Eq (6) is satisﬁed if and only if E P e
Revised: Eq (6) is satisﬁed if and only if E P e

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

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

Original: = 0 .
Revised: = 0 .

Original: For standard loss functions, the lossderivative has the form @
Revised: For standard loss functions, the lossderivative has the form @

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

Original: =
Revised: =

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

Original: | o ]
Revised: | o ]

Original: =  E
Revised: =  E

Original: e
Revised: e

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

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

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

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

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

Original: [Equation]
Revised: [Equation]

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

Original: However, there is a critical difference: Eq (8) only requires E e [ y | o ] toshare the same value across environments, thus imposing |E|  1 constraints, whereas eq (7) (also 106 eq (5)) additionally requires that value to be also equal to the RHS term  ( o ) .
Revised: However, there is a critical difference: Eq (8) only requires E e [ y | o ] toshare the same value across environments, thus imposing |E|  1 constraints, whereas eq (7) (also 106 eq (5)) additionally requires that value to be also equal to the RHS term  ( o ) .

Original: This extra constraint isdue to the indirect way IRM attains the conservation law, i.e. by equaling E e [ y | o ] to an intermediaryterm  ( o ) , which imposes 1 constraint per environment.
Revised: This extra constraint isdue to the indirect way IRM attains the conservation law, i.e. by equaling E e [ y | o ] to an intermediaryterm  ( o ) , which imposes 1 constraint per environment.

Original: 3.1.4 IRM-v 
Revised: 3.1.4 IRM-v

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

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

Original: In the variational formulation, this correspondsto restricting the output perturbations  to the space of linear functions, which, for scalar outputs, isequivalent to the identity function,  ( o ) = o .
Revised: In the variational formulation, this correspondsto restricting the output perturbations  to the space of linear functions, which, for scalar outputs, isequivalent to the identity function,  ( o ) = o .

Original: This reduces eq (6) to
Revised: This reduces eq (6) to

Original: [Equation]
Revised: [Equation]

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

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

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

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

Original: In Arjovsky et al.(2019), this reduced formulation is termed IRM-v1 when the constraints are imposed in a soft manner ,i.e. as squared penalty terms (See eq (16)).
Revised: In Arjovsky et al.(2019), this reduced formulation is termed IRM-v1 when the constraints are imposed in a soft manner ,i.e. as squared penalty terms (See eq (16)).

Original: In the literature, the term IRM-v1 is widly used for thereduced formulation regardless of whether hard or soft constraints are used, which we adopt here.
Revised: In the literature, the term IRM-v1 is widly used for thereduced formulation regardless of whether hard or soft constraints are used, which we adopt here.

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

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

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

Original: E e
Revised: E e

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

Original: [Equation]
Revised: [Equation]

Original: Unlike eq (7), however, RHS of eq (10) does not necessarily remain constant across environments,and hence does not prescribe a meaningful invariance.
Revised: Unlike eq (7), however, RHS of eq (10) does not necessarily remain constant across environments,and hence does not prescribe a meaningful invariance.

Original: Therefore, restricting the function space of (or ) invalidates the indirect way of imposing invariance in IRM involving an intermediary term.
Revised: Therefore, restricting the function space of (or ) invalidates the indirect way of imposing invariance in IRM involving an intermediary term.

Original: 3.2 MRI Paradigm
Revised: 3.2 MRI Paradigm

Original: We now introduce a complementary notion of invariance by considering inﬁnitesimal perturbationson label, which we call Mirror Reﬂected IRM , or MRI.
Revised: We now introduce a complementary notion of invariance by considering inﬁnitesimal perturbationson label, which we call Mirror Reﬂected IRM , or MRI.

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

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

Original: O is MRI invariant over a set of environments E , if the changein risk due to arbitrary inﬁnitesimal label perturbations is shared across environments, i.e.
Revised: O is MRI invariant over a set of environments E , if the changein risk due to arbitrary inﬁnitesimal label perturbations is shared across environments, i.e.

Original: [Equation]
Revised: [Equation]

Original: 11) where
Revised: 11) where

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

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

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

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

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

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

Original: Lemma 3.6.
Revised: Lemma 3.6.

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

Original: [Equation]
Revised: [Equation]

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

Original: [ o | y ] ⌘ E P
Revised: [ o | y ] ⌘ E P

Original: [Equation]
Revised: [Equation]

Original: y )
Revised: y )

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

Original: Proof.
Revised: Proof.

Original: Eq (11) is satisﬁed
Revised: Eq (11) is satisﬁed

Original: if and only if E e [ @ y l ( o, y ) | y ] is conserved.
Revised: if and only if E e [ @ y l ( o, y ) | y ] is conserved.

Original: For standard loss functions, theloss derivative has the form @
Revised: For standard loss functions, theloss derivative has the form @

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

Original: =
Revised: =

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

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

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

Original: [Equation]
Revised: [Equation]

Original: Note that MRI eq (11) imposes the correct number of constraints as needed by the conservation laweq (13), as it attains invariance in a direct manner without involving any intermediary terms.
Revised: Note that MRI eq (11) imposes the correct number of constraints as needed by the conservation laweq (13), as it attains invariance in a direct manner without involving any intermediary terms.

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

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

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

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

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

Original: Algorithm
Revised: Algorithm

Original: [Equation]
Revised: [Equation]

Original:  Training methods
Revised:  Training methods

Original: 4.1 Constrained optimization 
Revised: 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: =
Revised: =

Original: ~ 0 , (15) where L tr = 1 |E tr |
Revised: ~ 0 , (15) where L tr = 1 |E tr |

Original: P e 2 E tr L e is the average risk over the set of training environments E tr ⇢ E , and ✓is the parameters of the predictor f .
Revised: P e 2 E tr L e is the average risk over the set of training environments E tr ⇢ E , and ✓is the parameters of the predictor f .

Original: The constraint functions are ~ c = ~ Ltr 2 R |E tr | for IRM-v1 and~ c = Q ~ Ltr 2 R |E tr |  1 for MRI-v1, where ~ Ltr 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 .
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 .

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

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

Original: ~ 0 tosoft constraints, which allows the use of off-the-shelf gradient-based optimization algorithms.
Revised: ~ 0 tosoft 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 Arjovskyet al.
Revised: Penalty method is the most commonly used approach, including in Arjovskyet al.

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

Original: [Equation]
Revised: [Equation]

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

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

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

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: where ~ is typically initialized at ~ 0 and gets updated as ~ t +1 =
Revised: where ~ is typically initialized at ~ 0 and gets updated as ~ t +1 =

Original: ~ t + µ ~ c ( f ( ✓ t )) to accumulate theresidual constraints at each training iteration t .
Revised: ~ t + µ ~ c ( f ( ✓ t )) to accumulate theresidual constraints at each training iteration t .

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

Original:  Analytic Results
Revised:  Analytic Results

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

Original: In this section, we show that MRI-v1 can correctly remove all spurious features given a sufﬁcientnumber of environments.
Revised: In this section, we show that MRI-v1 can correctly remove all spurious features given a sufﬁcientnumber of environments.

Original: We consider a data generating process according to Figure 1, in which x = g ( z i , z s ) is a linearfunction.
Revised: We consider a data generating process according to Figure 1, in which x = g ( z i , z s ) is a linearfunction.

Original: Note that this Structural Equation Model (SEM, Pearl (2009)) is a generalization of thelinear Gaussian model considered in Rosenfeld et al.
Revised: Note that this Structural Equation Model (SEM, Pearl (2009)) is a generalization of thelinear Gaussian model considered in Rosenfeld et al.

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

Original: In comparison, wedon’t put any assumptions on the noise generating process or the label.
Revised: In comparison, wedon’t put any assumptions on the noise generating process or the label.

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

Original: O .
Revised: O .

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

Original: [Equation]
Revised: [Equation]

Original: ) with ✓ = { w i 2 R d i , w s 2 R d s } , where d i is the dimension of invariant features, d s is the dimensionof spurious features.
Revised: ) with ✓ = { w i 2 R d i , w s 2 R d s } , where d i is the dimension of invariant features, d s is the dimensionof spurious features.

Original: Theorem 5.1.
Revised: Theorem 5.1.

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

Original: Proof.
Revised: Proof.

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

Original: [Equation]
Revised: [Equation]

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

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

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

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

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

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

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

Original: =
Revised: =

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

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

Original: (2020) derives a similar result that IRM-v1 can eliminate all spurious featuredimensions given |E tr | > d s .
Revised: (2020) derives a similar result that IRM-v1 can eliminate all spurious featuredimensions given |E tr | > d s .

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

Original: In fact,
Revised: In fact,

Original: IRM-v1 has been shown to fail to eliminate spurious features in a more general linear problems inwhich the noise is not additive Gaussian (Arjovsky et al., 2019; Kamath et al., 2021).
Revised: IRM-v1 has been shown to fail to eliminate spurious features in a more general linear problems inwhich the noise is not additive Gaussian (Arjovsky et al., 2019; Kamath et al., 2021).

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

Original: Here, we demonstrate the result for the general linear SEM above with a minimal linear regressionproblem that involves one invariant feature, one spurious feature and two training environments. Thistoy example shares the same structure as the Shape-texture dataset (Section 6.1).It is analogous tothe two-bit environment task from Kamath et al.
Revised: Here, we demonstrate the result for the general linear SEM above with a minimal linear regressionproblem that involves one invariant feature, one spurious feature and two training environments(Shape-Texture linear regression in Section 6.1).

Original: The detailed experimental setting and theanalytical solutions are presented in the Supplementary Materials.
Revised: The detailed experimental setting and the analyticalsolutions are presented in the Supplementary Materials.

Original: IRM-v1 IRM-v1 has two quadratic equality constraints (  L e 1 ( f )
Revised: IRM-v1 IRM-v1 has two quadratic equality constraints (  L e 1 ( f )

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

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

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

Original: IRM-relaxed
Revised: IRM-relaxed

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

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

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

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

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

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

Original: This constraint is a set of hyperbolic curves, which appears to be two straightlines in Figure 2a.
Revised: This constraint is a set of hyperbolic curves, which appears to be two straightlines in Figure 2a.

Original: Due to the non-convex constraint, this problem exhibits two local constrainedoptima, including the true invariant optimum and a non-invariant solution, with the non-invariantsolution exhibiting lower training loss.
Revised: Due to the non-convex constraint, this problem exhibits two local constrainedoptima, including the true invariant optimum and a non-invariant solution, with the non-invariantsolution exhibiting lower training loss.

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

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

Original: MRI-v1 MRI-v1 has a single linear equality constraint, which is exactly the set of invariantsolutions ( w s = 0 ): That is, any point on the constraint is a invariant predictor.
Revised: MRI-v1 MRI-v1 has a single linear equality constraint, which is exactly the set of invariantsolutions ( w s = 0 ): That is, any point on the constraint is a invariant predictor.

Original: This is a convexproblem, since both the objective and the constraint are convex, and thus features a unique optimum,which is the true invariant optimum solution.
Revised: This is a convexproblem, since both the objective and the constraint are convex, and thus features a unique optimum,which is the true invariant optimum solution.

Original: In Figure 2b,c, we show the learning trajectory under full-batch gradient descent (penalty methodeq (16) with µ = 5 ⇥ 10 4 used).
Revised: In Figure 2b,c, we show the learning trajectory under full-batch gradient descent (penalty methodeq (16) with µ = 5 ⇥ 10 4 used).

Original: The convergence pattern for IRM-v1 is highly dependent on theinitialization (shown with different trajectory colors), due to the presence of multiple local minima.
Revised: The convergence pattern for IRM-v1 is highly dependent on theinitialization (shown with different trajectory colors), due to the presence of multiple local minima.

Original: Inthis case, the invariant optimum is especially extremely difﬁcult to converge to as it has higher trainingloss than the nearby non-invariant minima.
Revised: Inthis case, the invariant optimum is especially extremely difﬁcult to converge to as it has higher trainingloss than the nearby non-invariant minima.

Original: In contrast, MRI-v1 always converges to the  invariant
Revised: In contrast, MRI-v1 always converges to the invariantoptimum regardless of initialization since it is a unique minimum. For IRM-relaxed, the invariant

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: optimum regardless of initialization since it is a unique minimum.

Original: For IRM-relaxed, the invariantoptimum is lost at this value of µ , which exists in the limit µ ! 1 .
Revised: optimum is lost at this value of µ , which exists in the limit µ ! 1 .

Original: As a result, all trajectoriesconverge to the non-invariant optimum.
Revised: As a result, all trajectoriesconverge to the non-invariant optimum.

Original: Minimal examples for binary classiﬁcation, based on the shape-texture classiﬁcation and the toyCMNISTa/b datasets, are also analyzed in Supplementary Materials (Figure 3).
Revised: Minimal examples for binary classiﬁcation, based on the shape-texture classiﬁcation and the toyCMNISTa/b datasets, are also analyzed in Supplementary Materials (Figure 3).

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

Original: 6.1 Datasets
Revised: 6.1 Datasets

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

Original: We develop a new versatile dataset designed for domain generalization,called the Shape-Texture dataset, based on the data generating model in Figure 1.
Revised: We develop a new versatile dataset designed for domain generalization,called the Shape-Texture dataset, based on the data generating model in Figure 1.

Original: This dataset isdesigned to test algorithms on linear and non-linear regression and classiﬁcation problems in multipleadjustable settings.
Revised: This dataset isdesigned to test algorithms on linear and non-linear regression and classiﬁcation problems in multipleadjustable settings.

Original: The dataset involves an invariant latent feature ✓ i 2 R , a spurious latent feature ✓ s 2 R , and thelatent feature ✓ y 2 R for the label.
Revised: The dataset involves an invariant latent feature ✓ i 2 R , a spurious latent feature ✓ s 2 R , and thelatent feature ✓ y 2 R for the label.

Original: Each of them represents an angle of orientation, and therefore isnaturally represented on a unit circle in the complex plane as e i ✓ 2 S 1 .
Revised: Each of them represents an angle of orientation, and thereforeis naturally represented on a unit circle in the complex plane as e i ✓ 2 S 1 .

Original: In the linear regression setting,the observed input is simply the concatenated features x = [ e i ✓ i , e i ✓ s ] , and the label is y = e i ✓ y .
Revised: In the linear regressionsetting, the observed input is simply the concatenated features x = [ e i ✓ i , e i ✓ s ] , and the label isy = e i ✓ y .

Original: Whereas in the linear classiﬁcation setting, the input is given by x = [sin( ✓ i ) , sin( ✓ s )] , and the label is y
Revised: Whereas in the linear classiﬁcation setting, the input is given by x = [sin( ✓ i ) , sin( ✓ s )] , andthe label is y

Original: = H (sin( ✓ y ))
Revised: = H (sin( ✓ y ))

Original: where H denotes the Heaviside step function.
Revised: where H is the sign function.

Original: In the non-linear regression andclassiﬁcation setting, the label remains the same as their respective linear versions but the input isan image(Figure 4) composed of two planar waves of differing frequencies; The orientation of thelow frequency wave (a.k.a shape) is given by ✓ i / 2 , and the orientation of the high frequency wave(a.k.a texture) is given by ✓ s / 2 .
Revised: In the non-linear image-based tasks, the input is replaced by an image composed of two planar waves: The orientation of the low frequencywave (a.k.a shape) is given by ✓ i / 2 , and the orientation of the high frequency wave (a.k.a texture)is given by ✓ s / 2 (Figure 4).

Original: For all task settings, the latent angles and the label angle are linearlyrelated as ✓ i ✓ y + ⇠ i , ✓ s ✓ y + ⇠ s , where ⇠ ⇤ ⇠ N (0 ,  ⇤ ) with probability p ⇤ , or ⇠ ⇤ ⇠ U (  ⇡ , ⇡ ) 241 probability p ⇤ for ⇤ 2 { i, s } .
Revised: For all task settings, the latent angles and the label angle are linearlyrelated as ✓ i ✓ y + ⇠ i , ✓ s ✓ y + ⇠ s , where ⇠ ⇤ ⇠ N (0 ,  ⇤ ) with probability p ⇤ , or ⇠ ⇤ ⇠ U (  ⇡ , ⇡ ) 239 probability p ⇤ for ⇤ 2 { i, s } .

Original: p i ,  i are ﬁxed invariant parameters whereas p s ,  s can vary from 242 one environment to another.
Revised: p i ,  i are ﬁxed invariant parameters whereas p s ,  s can vary from 240 one environment to another.

Original: In this work, for all task settings, we use two training environmentsand one testing environment with p i = 0 .
Revised: In this work, for all task settings, we use two training environmentsand one testing environment with p i = 0 .

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

Original: 8 , 1 for the two trainingenvironments, p s = 0 for the testing environment, and  i =
Revised: 8 , 1 for the two trainingenvironments, p s = 0 for the testing environment.  i =

Original:  s = 0  for all environments.
Revised:  s = 0 is used for all environments.

Original: Colored MNIST The Colored MNIST (CMNIST) dataset was introduced by Arjovsky et al.(2019) to analyze IRM-v1’s ability to learn non-linear invariant predictors.
Revised: The Colored MNIST (CMNIST) dataset was introduced by Arjovsky et al.(2019) to analyze IRM-v1’s ability to learn non-linear invariant predictors.

Original: This is a syntheticbinary classiﬁcation task derived from MNIST.
Revised: This is a syntheticbinary classiﬁcation task derived from MNIST.

Original: In this dataset, the color of the image is determinedin correlation with the label whose correlation probability p s changes between environments.
Revised: In this dataset, the color of the image is determinedin correlation with the label whose correlation probability p s changes between environments.

Original: Thecorrelation probability p i between digit and label remains the same across all environments.
Revised: Thecorrelation probability p i between digit and label remains the same across all environments.

Original: Therefore,in this dataset color is the spurious feature and digit is the invariant feature.
Revised: Therefore,in this dataset color is the spurious feature and digit is the invariant feature.

Original: This dataset follows thedata generating model in Figure 1.
Revised: This dataset follows thedata generating model in Figure 1.

Original: We use two versions of CMNIST, CMNISTa and CMNISTb, that are generated with two differentsets of correlation probabilities: CMNISTa uses the same correlation probabilities as Arjovsky et al.(2019) with p i = 0 .
Revised: We use two versions of CMNIST, CMNISTa and CMNISTb, that are generated with two differentsets of correlation probabilities: CMNISTa uses the same correlation probabilities as Arjovsky et al.(2019) with p i = 0 .

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

Original: 9 for the two training environments, andp s = 0 .
Revised: 9 for the two training environments, andp s = 0 .

Original: 1 for the testing environment; CMNISTb uses p i = 0 .
Revised: 1 for the testing environment; CMNISTb uses p i = 0 .

Original: 95 for all environments, p s = 0 .
Revised: 95 for all environments, p s = 0 .

Original: 8 , 1 forthe two training environments, and p s = 0 .
Revised: 8 , 1 forthe two training environments, and p s = 0 .

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

Original: We also use an abstract version of CMNIST introduced by Kamath et al. (2021), called toyCMNISTa/b, in which the image data is replaced by a two-bit data x = [ x d , x c ] , where x d 2 {  1 , 1 }corresponds to the label of the digit (without label noise) and x c 2 {  1 , 1 } corresponds to the color. 
Revised: We also use an abstracted version of CMNIST (Kamath et al., 2021), called toy-CMNISTa/b, inwhich the image data is replaced by a two-bit data x = [ x d , x c ] , where x d 2 {  1 , 1 } corresponds tothe label of the digit (without label noise) and x c 2 {  1 , 1 } corresponds to the color. Note that both Shape-Texture and CMNIST datasets can be equally understood as being generatedfrom Fig 1 with causal directions Y ! Z s or Z s !

Revised: Y .

Original: 6.2 Result
Revised: 6.2 Result

Original: In addition to IRM-v1 and MRI-v1, we report the performance of the ERM algorithm, whichminimizes the average risk across environments without any invariance constraints (Table 2).
Revised: In addition to IRM-v1 and MRI-v1, we report the performance of the ERM algorithm, whichminimizes the average risk across environments without any invariance constraints (Table 2).

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

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

Original: For theshape-texture datasets, we use a small convolutional network (single convolution layer), while for theother image-based datasets, we use the LeNet architecture (LeCun et al., 1998).
Revised: For theshape-texture datasets, we use a small convolutional network (single convolution layer), while for theother image-based datasets, we use the LeNet architecture (LeCun et al., 1998).

Original: Shape-Texture Tasks
Revised: Shape-Texture Tasks

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

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

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

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

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

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

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

Original: IRM-v1’s performance by avoiding the zero-predictor, while MRI-v1 performed similarly irrespectiveof the training method as the the zero-predictor is not a local constrained optimum of MRI-v1 (ALMslightly improves MRI-v1’s performance).
Revised: IRM-v1’s performance by avoiding the zero-predictor, while MRI-v1 performed similarly irrespectiveof the training method as the the zero-predictor is not a local constrained optimum of MRI-v1 (ALMslightly improves MRI-v1’s performance).

Original: Colored MNIST Similar to the shape-texture tasks, IRM-v1 always converges to the null-predictorwhen trained using the penalty method while MRI-v1 achieves a much better test loss and accuracythat IRM-v1 and ERM (Table 2- 3; Figure 6).
Revised: Colored MNIST Similar to the shape-texture tasks, IRM-v1 always converges to the null-predictorwhen trained using the penalty method while MRI-v1 achieves a much better test loss and accuracythat IRM-v1 and ERM (Table 2- 3; Figure 6).

Original: This results for MRI-v1 holds true also when trainedusing ALM, albeit with a slightly improved performance (Table 2- 3; Figure 7).
Revised: This results for MRI-v1 holds true also when trainedusing ALM, albeit with a slightly improved performance (Table 2- 3; Figure 7).

Original: In comparison,similar to the shape-texture tasks, IRM-v1 is able to avoid the null predictor when trained with ALM,but it achieves worse test loss than MRI-v1 on both the datasets (Table 2; Figure 7), especially for the
Revised: In comparison,similar to the shape-texture tasks, IRM-v1 is able to avoid the null predictor when trained with ALM,but it achieves worse test loss than MRI-v1 on both the datasets (Table 2; Figure 7), especially for the

Original: CMNISTb datasets.
Revised: CMNISTb datasets.

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

Original: [Table]
Revised: [Table]

Original:  Discussion
Revised:  DiscussionLimitations

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

Original: The accuracy metric is widely used in the ﬁeld of domaingeneralization (Gulrajani and Lopez-Paz, 2020) to compare the performance of various algorithms.
Revised: The accuracy metric is widely used in the ﬁeld of domain generalization (Gulrajani and Lopez-Paz, 2020) to compare the performance of various algorithms.

Original: While accuracy does provide some insight into the performance of an algorithm, it is not always asensitive enough metric to discriminate between the representations learned by different algorithms.
Revised: While accuracy does provide some insight into the performance of an algorithm, it is not always asensitive enough metric to discriminate between the representations learned by different algorithms.

Original: For example, we ﬁnd that for the toy-CMNISTb dataset, both invariant ( ( w i , w s ) ⇡ (3 , 0) )
Revised: For example, we ﬁnd that for the toy-CMNISTb dataset, both invariant ( ( w i , w s ) ⇡ (3 , 0) )

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

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

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

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

Original: Based on this, we suggest that thereliance over accuracy as the sole metric of comparison between domain generalization algorithmscan be misleading, but loss and accuracy taken together can provide a meaningful comparison.
Revised: Based on this, we suggest that thereliance over accuracy as the sole metric of comparison between domain generalization algorithmscan be misleading, but loss and accuracy taken together can provide a meaningful comparison.

Original: Limitations A limitation of both MRI and IRM (Rosenfeld et al., 2020; Ahuja et al., 2021)
Revised: Limitations A limitation of both MRI and IRM (Rosenfeld et al., 2020; Ahuja et al., 2021)

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

Original: This limits theapplicability of both MRI and IRM on most DomainBed (Gulrajani and Lopez-Paz, 2020) benchmarks.
Revised: This limits theapplicability of both MRI and IRM on most DomainBed (Gulrajani and Lopez-Paz, 2020) benchmarks.

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

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

Original: Similar arguments can be made for other datasets.
Revised: Similar arguments can be made for other datasets.

Original: For example, images with differentstyles across testing and training domains may not satisfy the signiﬁcant support overlap condition.
Revised: For example, images with differentstyles across testing and training domains may not satisfy the signiﬁcant support overlap condition.

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

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

Original: However, note that, while IRM is expected to work for such shifts, IRM-v1 fails for such domainshifts (Ahuja et al., 2020, 2021).This is because IRM-v1 and IRM are not the same, as demonstratedby us and Kamath et al. Checklist The checklist follows the references.
Revised: However, note that, while IRM is expected to work for such shifts, IRM-v1 fails for such domainshifts (Ahuja et al., 2020, 2021).

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

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

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

Original: For example:

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

Original: [Yes] See Section ??.

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

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

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

Original: [N/A]

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

Original: Note that the

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

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

Original: For all authors...

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

Original: contributions and scope?

Original: [Yes]

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

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

Original: Did you discuss any potential negative societal impacts of your work?

Original: [N/A](d)

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

Original: [Yes]2.

Original: If you are including theoretical results...

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

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

Original: Did you include complete proofs of all theoretical results?

Original: [Yes] See Section 3 and 5.3.

Original: If you ran experiments...

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

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

Original: [Yes]

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

Original: (b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?

Original: [Yes] See Supplementary Materials.

Original: [Table]

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

Original: [Table]

Original: ments multiple times)?

Original: [Yes] See Table 2 and 3.

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

Original: [Yes] See Supplementary Materials.

Original: [Table]

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

Original: [Table]

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

Original: [Yes] See Section 6.1 andSupplementary Materials.

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

Original: [Yes] See Supplementary Materials.

Original: [Table]

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

Original: [Yes]

Original: [Table]

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

Original: [Table]

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

Original: [Table]

Original: using/curating?

Original: [N/A](e)

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

Original: information or offensive content?

Original: [N/A]

Original: [Table]

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

Original: [Table]

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

Original: [N/A]

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

Original: Board (IRB) approvals, if applicable?

Original: [N/A]

Original: [Table]

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

Original: [N/A]

Original: A Analytical Solution for Toy Shape-Texture Regression

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

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

Original: e i ✓ y .

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

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

Original: We used p i = 0 .

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

Original: for the testing environment, and  i

Original: =

Original:  s e = 0 for all environments.

Original: Note that

Original: E e [ e i ( ✓ i  ✓ y ) ]

Original: = ⇢ 1 , with prob.

Original: p i 0 , with prob.

Original: 1  p i

Original: E e [ e i ( ✓ s  ✓ y ) ]

Original: = ⇢ 1 , with prob.

Original: p s 0 , with prob.

Original: 1  p s

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

Original: = ⇢ 1 , with prob.

Original: p i p s 0 , with prob.

Original: 1  p i p s which yields

Original: E e [ xx † ] = 

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

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

Original:  =  p

Original: i

Original: p s p i p s 1  (20)

Original: E e [ xy † ] = 

Original: E e [ e i ( ✓ i  ✓ y ) ]

Original: E e [ e i ( ✓ s  ✓

Original: [Equation]

Original: [Equation]

Original: We consider real weights w =

Original: [ w i w s ] .

Original: Since o = wx ,

Original: [Equation]

Original: The risk of environment e is

Original: [Equation]

Original: Therefore, the average risk for training is

Original: [Equation]

Original: MRI constraint

Original: [Equation]

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

Original: = ( p i , 0) as the constrainedoptimum.

Original: [Equation] [Equation] Minimizing ¯ L tr ( f ) subject to the above constraint yields
Revised: This is because IRM-v1 and IRM are not the same, as demonstratedby us and Kamath et al.

Original: [Equation]

Original: as the constrained optima.

Original: IRM constraint

Original: [Equation]

Original: which yields

Original: [Equation]

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

Original: B Training Procedure

Original: PyTorch version 1.10.0 was used.

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

Original: Forall our simulations, we utilized the DomainBed suite (Gulrajani and Lopez-Paz (2020); availableunder MIT license).

Original: For a given task, we use the same training procedure for all the algorithms unlessmentioned otherwise.

Original: The details of the training are described as follows.

Original: B.1 Linear Datasets

Original: For Figure 2, 3, and 8, we used a linear predictor f ( x ) =

Original: w i

Original: z i + w s z s .

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

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

Original: For

Original: Figure 2, we used the SGD optimizer with no momentum.

Original: For Figure 3 and 8, we used the Adamoptimizer with betas of (0.975, 0.999).

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

Original: B.2 Non-Linear Datasets

Original: The results for the non-linear datasets are presented in Table 2-5 and Figure 5-7.

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

Original: , we used a single convolutional layer with12 channels, a kernel size of 5, and stride of 1.

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

Original: For resultsin Table 2 and 3, for optimization using the ALM method we use µ = 10 and  = 10 and foroptimization using PM we use µ = 10000 .

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

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

Original: The learningrate was set to 0.01.

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

Original: Gradient clipping was onlyused when training using PM such that gradient norms were clipped at 2.

Original: [-We used a weight decay of0.
-]
Revised: [-We used a weight decay of0.
-]

Original: for the classiﬁcation setting and 0.005 for the regression setting.

Original: For Figure 5 we use the same training procedure as above for the penalty method except that theseeds, learning rates and algorithm-speciﬁc parameters were chosen at random to generate 50 randomhyperparameter combinations for each algorithm.

Original: The learning rates were chosen randomly fromthe distribution 10 Uniform [  3 ,  1 . 5) .

Original: For IRM-v1 and MRI-v1, µ was chosen randomly from thedistribution 10 Uniform [2 , 5) .

Original: We used DomainBed’s (Gulrajani and Lopez-Paz, 2020)

Original: implementationof IB-IRM, GroupDRO, and MMD.

Original: For IB-IRM,  IRM and  IB were both chosen randomly fromthe distribution 10 Uniform [2 , 5) .

Original: For GroupDRO, ⌘ GroupDRO was chosen from the distribution10 Uniform [  3 , 1) .

Original: For MMD,  MMD was chosen from the distribution 10 Uniform [  1 , 1)

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

Original: We usedmini-batch training with a mini-batch size of 256.

Original: For optimization using the ALM method with =

Original: 10 for both datasets, but µ = 1 for CMNISTa and µ = 10 for CMNISTb datasets for results in

Original: Table 2, 3.

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

Original: We usedthe Adam optimizer with betas of (0.975, 0.999).

Original: The learning rate was set to 0.001, the total numberof gradient steps were set to 10000.

Original: The weight decay was set to 0.001.

Original: Gradient clipping was onlyused when training using PM such that gradient norms were clipped at 2.

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

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

Original: For Figure 6 we used the same training procedure as above for training with PM except µ was chosenrandomly from the distribution 10 Uniform [2 , 5) .

Original: For Figure 7, we used the same training procedure asabove for training with ALM except both  and µ were chosen independently and randomly from thedistribution Uniform [1 , 10) .

Original: For both the ﬁgures, seeds were chosen at random.

Original: The total numberof gradient steps were set to 5000 for Figure 6 and 15000 for Figure 7.

Original: VLCS/Terra-Incognita For both VLCS and Terra-Incognita datasets, we restrict the problem to abinary classiﬁcation task by using only two classes for each of the two datasets, namely opossumand raccoon for the TI dataset and bird and person for VLCS dataset.

Original: The results are presented in

Original: Table 4-5.

Original: For both datasets, we use the same training procedure.

Original: We used mini-batch training with amini-batch size of 256.

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

Original: We used the

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

Original: The learning rate was set to 0.001 and the total numberof gradient steps were set to 2000.

Original: The weight decay was set to 0.001.

Original: No gradient clipping wasused.

Original: All the images were resized to the shape (64,64,3) and were normalized with means (0.485,0.456, 0.406) and standard deviation (0.229, 0.224, 0.225) corresponding to each channel.

Original: We alsoused data augmentation during training.

Original: This was done using Torchvision’s RandomHorizontalFlip

Original: [Table]

Original: transform, RandomGrayscale transform, RandomResizedCrop transform with scale = (0.7, 1.0), and

Original: ColorJitter transform brightness, contrast, saturation, and hue set to 0.3.

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: [Figure]