Original: O UT - OF - DISTRIBUTION P REDICTION WITH I NVARIANT R ISK M INIMIZATION : T HE L IMITATION AND A N E FFECTIVE F IX
Revised: O UT - OF - DISTRIBUTION P REDICTION WITH I NVARIANT R ISK M INIMIZATION : T HE L IMITATION AND A N E FFECTIVE F IX

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: This work considers the out-of-distribution (OOD) prediction problem where (1) the training data are from multiple domains and (2) the test domain is unseen in the training.
Revised: This work considers the out-of-distribution (OOD) prediction problem where (1) the training data are from multiple domains and (2) the test domain is unseen in the training.

Original: DNNs fail in OOD prediction because they are prone to pick up spurious correlations.
Revised: DNNs fail in OOD prediction because they are prone to pick up spurious correlations.

Original: Recently, Invariant Risk Minimization (IRM) is proposed to address this issue.
Revised: Recently, Invariant Risk Minimization (IRM) is proposed to address this issue.

Original: Its effectiveness has been demonstrated in the colored
Revised: Its effectiveness has been demonstrated in the colored

Original: MNIST experiment.
Revised: MNIST experiment.

Original: Nevertheless, we ﬁnd that the performance of IRM can be dramatically degraded under strong Λ spuriousness – when the spurious correlation between the spurious features and the class label is strong due to the strong causal inﬂuence of their common cause, the domain label, on both of them (see
Revised: Nevertheless, we ﬁnd that the performance of IRM can be dramatically degraded under strong Λ spuriousness – when the spurious correlation between the spurious features and the class label is strong due to the strong causal inﬂuence of their common cause, the domain label, on both of them (see

Original: Fig. 1).
Revised: Fig. 1).

Original: In this work, we try to answer the questions: why does IRM fail in the aforementioned setting?
Revised: In this work, we try to answer the questions: why does IRM fail in the aforementioned setting?

Original: Why does IRM work for the original colored MNIST dataset?
Revised: Why does IRM work for the original colored MNIST dataset?

Original: Then, we propose a simple and effective approach to ﬁx the problem of
Revised: Then, we propose a simple and effective approach to ﬁx the problem of

Original: IRM.
Revised: IRM.

Original: We combine IRM with conditional distribution matching to avoid a speciﬁc type of spurious correlation under strong Λ spuriousness.
Revised: We combine IRM with conditional distribution matching to avoid a speciﬁc type of spurious correlation under strong Λ spuriousness.

Original: Empirically, we design a series of semi synthetic datasets – the colored MNIST plus, which exposes the problems of IRM and demonstrates the efﬁcacy of the proposed method.
Revised: Empirically, we design a series of semi synthetic datasets – the colored MNIST plus, which exposes the problems of IRM and demonstrates the efﬁcacy of the proposed method.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: E X c Y X s
Revised: [Figure]

Original: Strong empirical results have demonstrated the efﬁcacy of deep neural networks (DNNs) in a variety of areas including computer vision, natural language processing and speech recognition.
Revised: Strong empirical results have demonstrated the efﬁcacy of deep neural networks (DNNs) in a variety of areas including computer vision, natural language processing and speech recognition.

Original: However, such positive results overwhelmingly rely on the assumption that the training, validation and test data consist of independent identical samples of the same underlying distribution.
Revised: However, such positive results overwhelmingly rely on the assumption that the training, validation and test data consist of independent identical samples of the same underlying distribution.

Original: In contrast, in the setting of outof-distribution (OOD) prediction where (1) the training data is from multiple domains and (2) the test set has different distribution from the training, the performance of DNNs can be dramatically degraded.
Revised: In contrast, in the setting of outof-distribution (OOD) prediction where (1) the training data is from multiple domains and (2) the test set has different distribution from the training, the performance of DNNs can be dramatically degraded.

Original: This is because DNNs are prone to pick up spurious correlations which do not hold beyond the training data distribution (Beery et al., 2018;
Revised: This is because DNNs are prone to pick up spurious correlations which do not hold beyond the training data distribution (Beery et al., 2018;

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

Original: For example, when most camel pictures in a training set have a desert in the background, DNNs will pick up the spurious correlation between a desert and the class label, leading to failures when camel pictures come with different backgrounds in a test set.
Revised: For example, when most camel pictures in a training set have a desert in the background, DNNs will pick up the spurious correlation between a desert and the class label, leading to failures when camel pictures come with different backgrounds in a test set.

Original: Therefore, OOD prediction remains an extremely challenging problem for DNNs.
Revised: Therefore, OOD prediction remains an extremely challenging problem for DNNs.

Original: The invariant causal relationships across different domains turns out to be the key to address the challenge of OOD prediction.
Revised: The invariant causal relationships across different domains turns out to be the key to address the challenge of OOD prediction.

Original: The causal graph in Fig.
Revised: The causal graph in Fig.

Original: 1 describes the relationships of the variables in the OOD prediction problem.
Revised: 1 describes the relationships of the variables in the OOD prediction problem.

Original: Although spurious correlations learned in one domain are unreliable in another, invariant causal relationships enable DNNs that capture causal relationships to generalize to unseen domains.
Revised: Although spurious correlations learned in one domain are unreliable in another, invariant causal relationships enable DNNs that capture causal relationships to generalize to unseen domains.

Original: In practice, it is extremely difﬁcult to know whether an input feature is causal or spurious.
Revised: In practice, it is extremely difﬁcult to know whether an input feature is causal or spurious.

Original: Thus, a recipe for training DNNs that capture causal relationships is learning causal feature representations that hold invariant causal relationships with the class label.
Revised: Thus, a recipe for training DNNs that capture causal relationships is learning causal feature representations that hold invariant causal relationships with the class label.

Original: Thus, the main challenge of OOD prediction becomes learning causal feature representations given training data from multiple domains.
Revised: Thus, the main challenge of OOD prediction becomes learning causal feature representations given training data from multiple domains.

Original: Recently, Invariant Risk Minimization (IRM) (Arjovsky et al., 2019) is proposed to learn causal feature representations for OOD prediction.
Revised: Invariant Risk Minimization (IRM) (Arjovsky et al., 2019) is proposed to learn causal feature representations for OOD prediction.

Original: IRM formulates the invariance of causal relationships as a constraint.
Revised: IRM formulates the invariance of causal relationships as a constraint.

Original: It requires that causal feature representations must result in the same optimal classiﬁer across all domains.
Revised: It requires that causal feature representations must result in the same optimal classiﬁer across all domains.

Original: This constraint can be written as the conditional independence, Y ⊥⊥ E | F ( X )  (Chang et al., 2020; Zeng et al., 2019) which can be derived from the causal graph in Fig.
Revised: IRM is written as the conditional independence, Y ⊥⊥ E | F ( X ) in (Chang et al., 2020; Zeng et al., 2019) which can be derived from the causal graph in Fig.

Original: 1 when F ( X ) is a mapping of causal features ( X c in Fig.
Revised: 1 when F ( X ) is a mapping of causal features ( X c in Fig.

Original: 1) with no information loss.
Revised: 1) with no information loss.

Original:  Despite IRM’s success in the colored MNIST dataset (Arjovsky et al., 2019), in this work, we ﬁnd an important issue of IRM that has not been discussed.
Revised: A detailed discussion on IRM and Y ⊥⊥ E | F ( X ) is in Appendix A. Despite IRM’s success in the colored MNIST dataset (Arjovsky et al., 2019), in this work, we ﬁnd an important issue of IRM that has not been discussed.

Original: Speciﬁcally, we consider the situation where strong spurious correlation exists between the spurious features and the class label only in the  training set.
Revised: Speciﬁcally, we consider the situation where strong spurious correlations among the spurious features, the class label and the domain label only hold for training data.

Original: We name this strong Λ spuriousness using the shape of thecausal structure among the three variables in the causal graph (Fig. 1).
Revised: We name this strong Λ spuriousness using the shape of the structure among X s , E and Y in Fig.

Original: Under strong Λ spuriousness, IRM regularized empirical risk can have low values with spurious feature representations that accurately predict the domain label.
Revised: Under strong Λ spuriousness, IRM regularized empirical risk can have low values with spurious feature representations that accurately predict the domain label.

Original: This is because, in this setting, picking up spurious features can achieve high accuracy in predicting both domain and class in the training set, but not in the test.
Revised: This is because, in this setting, picking up spurious features can achieve high accuracy in predicting both domain and class in the training set, but not in the test.

Original: However, the colored MNIST dataset cannot expose this issue because the strong similarity between the two training domains makes it difﬁcult to pick up the weak spurious correlation between the domain label and the spurious features.
Revised: However, the colored MNIST dataset cannot expose this issue because the strong similarity between the two training domains makes it difﬁcult to pick up the weak spurious correlation between the domain label and the spurious features.

Original: To illustrate this problem, we design a new dataset – the colored MNIST plus.
Revised: To illustrate this problem, we design a new dataset – the colored MNIST plus.

Original: As shown in Fig.
Revised: As shown in Fig.

Original: 2, in this dataset, under strong Λ spuriousness, the performance of IRM models is signiﬁcantly degraded.
Revised: 2, in this dataset, under strong Λ spuriousness, the performance of IRM models is signiﬁcantly degraded.

Original: Moreover, to resolve this issue of IRM, we propose an effective solution, which combines IRM with conditional distribution matching (CDM).
Revised: Moreover, to resolve this issue of IRM, we propose an effective solution, which combines IRM with conditional distribution matching (CDM).

Original: The CDM constraint means that the representation distribution of instances from the same class should be invariant across domains.
Revised: The CDM constraint means that the representation distribution of instances from the same class should be invariant across domains.

Original: Theoretically, we show that (1) causal feature representations can satisfy CDM and IRM at the same time and (2) CDM can prevent DNNs from learning spurious feature representations that accurately predict the domain label.
Revised: Theoretically, we show that (1) causal feature representations can satisfy CDM and IRM at the same time and (2) CDM can prevent DNNs from learning spurious feature representations that accurately predict the domain label.

Original: Empirically, on our newly introduced dataset, the proposed method achieves signiﬁcant performance improvement over IRM under strong Λ spuriousness.
Revised: Empirically, on our newly introduced dataset, the proposed method achieves signiﬁcant performance improvement over IRM under strong Λ spuriousness.

Original:  P RELIMINARIES AND IRM
Revised:  P RELIMINARIES AND IRM

Original: Notations.
Revised: Notations.

Original: We use lowercase (e.g., x ), uppercase (e.g., X ) and calligraphic uppercase (e.g., X ) letters for values, random variables and spaces.
Revised: We use lowercase (e.g., x ), uppercase (e.g., X ) and calligraphic uppercase (e.g., X ) letters for values, random variables and spaces.

Original: We let X ∈ X , Y ∈ Y and E ∈ E denote raw input, the class label and the domain label where X , Y and E are the spaces of input, class labels and domains.
Revised: We let X ∈ X , Y ∈ Y and E ∈ E denote raw input, the class label and the domain label where X , Y and E are the spaces of input, class labels and domains.

Original: A DNN model consists of a feature learning function F and a classiﬁer G .
Revised: A DNN model consists of a feature learning function F and a classiﬁer G .

Original: A feature learning function F : X → R d maps raw input X to its d -dimensional representations F ( X ) .
Revised: A feature learning function F : X → R d maps raw input X to its d -dimensional representations F ( X ) .

Original: A classiﬁer G : R d → Y maps a feature representation to a class label.
Revised: A classiﬁer G : R d → Y maps a feature representation to a class label.

Original: We denote their parameters by θ F and θ G , respectively.
Revised: We denote their parameters by θ F and θ G , respectively.

Original: Let θ =
Revised: Let θ =

Original: Concat ( θ F , θ G ) denote the concatenation of them.
Revised: Concat ( θ F , θ G ) denote the concatenation of them.

Original: A domain e of n e instances is denoted by D e = { x e i , y e i } n e i =1 .
Revised: A domain e of n e instances is denoted by D e = { x ei , y ei } n e i =1 .

Original: Let E tr and E ts denote the set of training and test domains, in OOD prediction, we have (1) |E tr | > 1 and (2) E tr ∩ E ts = ∅ .
Revised: Let E tr and E ts denote the set of training and test domains, in OOD prediction, we have (1) |E tr | > 1 and (2) E tr ∩ E ts = ∅ .

Original: Problem Statement.
Revised: Problem Statement.

Original: Given data from multiple training domains { D e } e ∈E tr .
Revised: Given data from multiple training domains { D e } e ∈E tr .

Original: We aim to predict the label y e (cid:48)
Revised: We aim to predict the label y e (cid:48)

Original: i of each instance with features x e (cid:48)
Revised: i of each instance with features x e (cid:48)

Original: i from a test domain { x e (cid:48)
Revised: i from a test domain { x e (cid:48)

Original: i , y e (cid:48) i } n e (cid:48) i
Revised: i , y e (cid:48) i } n e (cid:48) i

Original: =1 , e (cid:48)
Revised: =1 , e (cid:48)

Original: ∈
Revised: ∈

Original: E ts .
Revised: E ts .

Original: Invariant Risk Minimization.
Revised: Invariant Risk Minimization.

Original: IRM (Arjovsky et al., 2019) is a recently proposed method to impose the causal inductive bias: the causal relationships between causal features and the label should be invariant across different domains.
Revised: IRM (Arjovsky et al., 2019) is a recently proposed method to impose the causal inductive bias: the causal relationships between causal features and the label should be invariant across different domains.

Original: It not only aims to address the challenging OOD prediction problem, but also is a pioneer work that guides causal machine learning research towards the development of inductive bias imposing causal constraints.
Revised: It not only aims to address the challenging OOD prediction problem, but also is a pioneer work that guides causal machine learning research towards the development of inductive bias imposing causal constraints.

Original: The effectiveness of IRM and its variants have been demonstrated across various areas including computer vision (Ahuja et al., 2020), natural language processing (Chang et al., 2020),
Revised: The effectiveness of IRM and its variants have been demonstrated across various areas including computer vision (Ahuja et al., 2020), natural language processing (Chang et al., 2020),

Original: CTR prediction (Zeng et al., 2019), reinforcement learning (Zhang et al., 2020a) and ﬁnancial forecasting (Krueger et al., 2020).
Revised: CTR prediction (Zeng et al., 2019), reinforcement learning (Zhang et al., 2020a) and ﬁnancial forecasting (Krueger et al., 2020).

Original: Arjovsky et al.
Revised: Arjovsky et al.

Original: (2019) propose the original formulation of IRM as a two-stage optimization problem:
Revised: (2019) propose the original formulation of IRM as a two-stage optimization problem:

Original: [Equation]
Revised: [Equation]

Original: where R e denotes the loss function of domain e .
Revised: where R e denotes the loss function of domain e .

Original: Then, they show that the IRM constraint can be imposed by adding a regularizer into the loss function.
Revised: Then, they show that the IRM constraint can be imposed by adding a regularizer into the loss function.

Original: It turns out the IRM constraint can be written as the conditional independence Y ⊥⊥ E | F ( X )
Revised: It turns out the IRM constraint can be written as the conditional independence Y ⊥⊥ E | F ( X )

Original: (Chang et al., 2020; Zeng et al., 2019).
Revised: (Chang et al., 2020; Zeng et al., 2019).

Original: This can be derived from the causal graph in Fig.
Revised: This can be derived from the causal graph in Fig.

Original: For brevity, we refer to Y ⊥⊥ E | F ( X ) as the IRM constraint.
Revised: For brevity, we refer to Y ⊥⊥ E | F ( X ) as the IRM constraint.

Original:  T HE L IMITATION OF IRM
Revised:  T HE L IMITATION OF IRM

Original: In this section, we ﬁnd that IRM fails under strong Λ spuriousness – when the spurious correlations among spurious features, the domain label and the class label are strong.
Revised: In this section, we ﬁnd that IRM fails under strong Λ spuriousness – when the spurious correlations among spurious features, the domain label and the class label are strong.

Original: To show it, we design a new dataset – the colored MNIST plus (CMNIST+) to expose this limitation of IRM.
Revised: To show it, we design a new dataset – the colored MNIST plus (CMNIST+) to expose this limitation of IRM.

Original: Then, we try to answer two crucial research questions: (1) why does IRM fail in CMNIST+?
Revised: Then, we try to answer two crucial research questions: (1) why does IRM fail in CMNIST+?

Original: (2) Why does IRM work for the original colored MNIST (CMNIST) dataset?
Revised: (2) Why does IRM work for the original colored MNIST (CMNIST) dataset?

Original: 3. W HY DOES IRM FAIL ?
Revised: 3. W HY DOES IRM FAIL ?

Original: [Figure]
Revised: [Figure]

Original: , in the spurious feature domain cannot

Original: Despite its previous success, results in Fig.
Revised: Despite its previous success, results in Fig.

Original: 2 show that IRM fails in our newly designed colored
Revised: 2 show that IRM fails in our newly designed colored

Original: MNIST plus dataset, under strong Λ spuriousness ( ρ ≥ 0 . 8 ).
Revised: MNIST plus dataset, under strong Λ spuriousness ( ρ ≥ 0 . 8 ).

Original: Here, we show that, under strong Λ spuriousness, at least two types of feature representations, F ( X ) =
Revised: Here, we show that, under strong Λ spuriousness, at least two types of feature representations, F ( X ) =

Original: Y and F ( X )
Revised: Y and F ( X )

Original: =
Revised: =

Original: E , satisfy the
Revised: E , satisfy the

Original: IRM constraint.
Revised: IRM constraint.

Original: When F ( X ) =
Revised: When F ( X ) =

Original: E , the feature representations pick up spurious relationships.
Revised: E , the feature representations pick up spurious relationships.

Original: Satisfying IRM by Perfectly Predicting the Label.
Revised: Satisfying IRM by Perfectly Predicting the Label.

Original: First, we consider a case where there exists a feature representation that perfectly predicts the class label, i.e., F ( X ) =
Revised: First, we consider a case where there exists a feature representation that perfectly predicts the class label, i.e., F ( X ) =

Original: Y .
Revised: Y .

Original: In this case, we know that the feature representation satisﬁes the IRM constraint because Y ⊥⊥ E | F ( X ) =
Revised: In this case, we know that the feature representation satisﬁes the IRM constraint because Y ⊥⊥ E | F ( X ) =

Original: Y always holds.
Revised: Y always holds.

Original: This is reasonable because the causal relationship between Y and itself must be invariant across different domains 1 .
Revised: This is reasonable because the causal relationship between Y and itself must be invariant across different domains.

Original: Satisfying IRM by Perfectly Predicting the Domain Does Not Lead to OOD Generalizable
Revised: Satisfying IRM by Perfectly Predicting the Domain Does Not Lead to OOD Generalizable

Original: Feature Representations.
Revised: Feature Representations.

Original: However, there is another way to satisfy the IRM constraint.
Revised: However, there is another way to satisfy the IRM constraint.

Original: If there exists a feature representation F ( X ) =
Revised: If there exists a feature representation F ( X ) =

Original: E2 , then the IRM constraint can be satisﬁed by F ( X ) .
Revised: E , then the IRM constraint can be satisﬁed by F ( X ) .

Original: This is due to the fact Y ⊥⊥ E | F ( X ) =
Revised: This is due to the fact Y ⊥⊥ E | F ( X ) =

Original: E .
Revised: E .

Original: However, this is not desired behavior of IRM as F ( X ) =
Revised: However, this is not desired behavior of IRM as F ( X ) =

Original: E is a type of spurious feature representation.
Revised: E is a type of spurious feature representation.

Original: This is because P ( Y | E ) can vary across domains.
Revised: This is because P ( Y | E ) can vary across domains.

Original: Using the camel picture example, if domains are the individuals who take the picture, we may ﬁnd more camel pictures taken by person e than from person e (cid:48)
Revised: Using the camel picture example, if domains are the individuals who take the picture, we may ﬁnd more camel pictures taken by person e than from person e (cid:48)

Original: ( P ( Y = camel | E = e ) > P ( Y = camel | E = e (cid:48) ) ).
Revised: ( P ( Y = camel | E = e ) > P ( Y = camel | E = e (cid:48) ) ).

Original: Why is F ( X ) =
Revised: Why is F ( X ) =

Original: E a good solution for the original IRM optimization problem (Eq. 1)?
Revised: E a good solution for the original IRM optimization problem (Eq. 1)?

Original: This is because F ( X ) =
Revised: This is because F ( X ) =

Original: E also leads to low values of the loss function R e since the direct causal inﬂuence of the domain label E on the class label Y is also strong when Λ spuriousness is strong.
Revised: E also leads to low values of the loss function R e since the direct causal inﬂuence of the domain label E on the class label Y is also strong when Λ spuriousness is strong.

Original: So, with a proper classiﬁer G , we can achieve low values of R e with F ( X ) =
Revised: So, with a proper classiﬁer G , we can achieve low values of R e with F ( X ) =

Original: E .
Revised: E .

Original: Then, we describe the CMNIST+ dataset and show experimental and theoretical results to support our claim.
Revised: Then, we describe the CMNIST+ dataset and show experimental and theoretical results to support our claim.

Original: The Colored MNIST Plus (CMNIST+) Dataset.
Revised: The Colored MNIST Plus (CMNIST+) Dataset.

Original: We follow CMNIST (Arjovsky et al., 2019) to create CMNIST+ by resampling and adding colors to instances of MNIST.
Revised: We follow CMNIST (Arjovsky et al., 2019) to create CMNIST+ by resampling and adding colors to instances of MNIST.

Original: The digits 0 − 4 ( 5 − 9 ) are class Y = 

Revised: The digits 0 − 4 ( 5 − 9 ) are class Y =

Original: ( Y = 0 ).
Revised: ( Y = 0 ).

Original: We randomly ﬂip 25% of class labels so that spuriousness can be stronger than the causal relationship between the shape S (causal features) and the class label Y .
Revised: We randomly ﬂip 25% of class labels so that spuriousness can be stronger than the causal relationship between the shape S (causal features) and the class label Y .

Original: Tabledescribes the dataset.
Revised: Table 1 describes the dataset.

Original: The variable C ∈ { G, B, R } denotes the color, which represents the spurious feature X s .
Revised: The variable C ∈ { G, B, R } denotes the color, which represents the spurious feature X s .

Revised: We explain why we use three colors for CMNIST+ in Appendix C.2.

Original: The parameter ρ ∈ (0 . 5 , 1) controls the strength of the spurious correlations between the color (spurious features) C and the class label Y through their common cause, the domain label E .
Revised: The parameter ρ ∈ (0 . 5 , 1) controls the strength of the spurious correlations between the color (spurious features) C and the class label Y through their common cause, the domain label E .

Original: The larger the value of ρ , the stronger the spurious correlations.
Revised: The larger the value of ρ , the stronger the spurious correlations.

Original: Intuitively, in the two training domains of CMNIST+, in addition to the strong spurious correlation between the class label Y and color C , we set the spurious correlation between the domain label E and color C to be strong, too.
Revised: Intuitively, in the two training domains of CMNIST+, in addition to the strong spurious correlation between the class label Y and color C , we set the spurious correlation between the domain label E and color C to be strong, too.

Original: Thus, the CMNIST+ dataset can expose the problem of IRM: it would not penalize the models that ﬁt the domain variable as the feature representation ( F ( X ) = E ).
Revised: Thus, the CMNIST+ dataset can expose the problem of IRM: it would not penalize the models that ﬁt the domain variable as the feature representation ( F ( X ) = E ).

Original: To verify this claim, we show theoretical results to support the experimental results in Fig.
Revised: To verify this claim, we show analysis results to support the experimental results in Fig.

Original: [Table] = [Table]
Revised: The data generating process of CMNIST+ can be found in Append C.2.

Original: [Table]
Revised: [Table]

Original: E =

Original: [Table]
Revised: [Table]

Original: =
Revised: =

Original: [Table]
Revised: [Table]

Original: =
Revised: ρ

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

Original: − [Table] [Table] [Table]
Revised: =

Original: [Table]
Revised: [Table]

Original: ρ

Original:  E =
Revised: [Table] ) /ρ E =0. Y =(

Original: [Table]
Revised: [Table]

Revised: ρ ) /ρ

Original: [Table]
Revised: [Table]

Original: =
Revised: ) / 2 Y =(

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

Revised: ) /ρ

Original:  [Table]
Revised: E = [Table]

Revised: Y =

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

Original:  =
Revised: Y =

Original: [Figure]
Revised: [Figure]

Original: Theoretical Results.
Revised: Analysis Results.

Original: Here, we show the theoretical results on the CMNIST+ to answer two questions: what features will be learned by ERM and IRM under under different Λ spuriousness?
Revised: Here, we show the analysis results on the CMNIST+ to answer two questions: what features will be learned by ERM and IRM under under different Λ spuriousness?

Original: What is the expected test accuracy of these models?
Revised: What is the expected test accuracy of these models?

Original: For simplicity, we assume that the classiﬁer is deterministic, which always predicts the majority class given the feature.
Revised: For simplicity, we assume that the classiﬁer is deterministic, which always predicts the majority class given the feature.

Original: If the number of instances from each class is the same, then it would predict a random label.
Revised: If the number of instances from each class is the same, then it would predict a random label.

Original: We analyze three types of spurious feature representations: (1) those that ﬁt color, (2) those that ﬁt the domain label and (3) those that use a combination of them for prediction.
Revised: We analyze three types of spurious feature representations: (1) those that ﬁt color, (2) those that ﬁt the domain label and (3) those that use a combination of them for prediction.

Original: The ﬁrst case mimics the behavior of ERM that picks up the spurious correlation between color and the class label.
Revised: The ﬁrst case mimics the behavior of ERM that picks up the spurious correlation between color and the class label.

Original: The second one represents a model satisfying IRM by F ( X ) =
Revised: The second one represents a model satisfying IRM by F ( X ) =

Original: E .
Revised: E .

Original: The third one stands for a model satisfying IRM by F ( X ) =
Revised: The third one stands for a model satisfying IRM by F ( X ) =

Original: Concat ( E, C ) .
Revised: Concat ( E, C ) .

Original: One may argue that IRM can be ﬁxed if we simply balance the two classes in each domain.
Revised: One may argue that IRM can be ﬁxed if we simply balance the two classes in each domain.

Original: Theoretical analysis shows that this is an invalid solution.
Revised: Theoretical analysis shows that this is an invalid solution.

Original: We summarizethe theoretical results in Tablewhere ˆ E denotes the domain label predicted by a deterministic classiﬁer using color as the feature.
Revised: We summarize results in Table 2 where ˆ E is the domain label predicted by a deterministic classiﬁer using color as the feature.

Original: This is because in an unseen test domain, we cannot directly use the domain label for prediction.
Revised: This is because in an unseen test domain, we cannot directly use the domain label for prediction.

Original: The highlighted numbers show which type of features would be learned by ERM and IRM.
Revised: The highlighted numbers show which type of features would be learned by ERM and IRM.

Original: We can see as ρ increases, ERM and IRM are more likely to learn spurious feature representations.
Revised: As ρ increases, ERM and IRM are more likely to ﬁt spurious features.

Original: This explains results in Fig.
Revised: This explains results in Fig.

Original: As a DNN may still pick up some causal features (shape in CMNIST+) even when ρ = 0 .
Revised: As a DNN may still pick up some causal features (shape in CMNIST+) even when ρ = 0 .

Original: 9 , the test accuracy of IRM in practice would be greater than 0 .
Revised: 9 , the test accuracy of IRM in practice would be greater than 0 .

Original: 35 without label balancing and 0 .
Revised: 35 without label balancing and 0 .

Original: 2 with label balancing.
Revised: 2 with label balancing.

Original: Similarly, the test accuracy of ERM would be greater than 0 .
Revised: Similarly, the test accuracy of ERM would be greater than 0 .

Original: From Fig.
Revised: From Fig.

Original: 2, we can see the test accuracy of IRM, IRM with label balancing and ERM are slightly better than the aforementioned lower bounds.
Revised: 2, we can see the test accuracy of IRM, IRM with label balancing and ERM are slightly better than the aforementioned lower bounds.

Original: The derivations can be found in Appendix A.1.
Revised: The derivations can be found in Appendix B.

Original: Experimental Setup.
Revised: Experimental Setup.

Original: We use a LeNet-5 (LeCun et al., 1998) instead of a three-layer MLP.
Revised: We use a LeNet-5 (LeCun et al., 1998) instead of a three-layer MLP.

Original: To take input with three colors, we modify the ﬁrst CNN layer to have three channels.
Revised: To take input with three colors, we modify the ﬁrst CNN layer to have three channels.

Original: LeNet-5 has more predictive power such that it can pick up the three types of spurious correlations (color – domain, color – class, and domain – class) or the causal relationship (shape – class).
Revised: LeNet-5 has more predictive power such that it can pick up the three types of spurious correlations (color – domain, color – class, and domain – class) or the causal relationship (shape – class).

Original: We randomly split the instances from the training domains into 80% training and 20% validation.
Revised: We randomly split the instances from the training domains into 80% training and 20% validation.

Original: The model selection is done by picking the one with the lowest validation loss in each run.
Revised: The model selection is done by picking the one with the lowest validation loss in each run.

Original: We report the average test accuracy of the selected models in 10 runs.
Revised: We report the average test accuracy of the selected models in 10 runs.

Original: It is crucial to ensure only data from training domains are used in model selection, since the test domain should be unseen during training and validation.
Revised: It is crucial to ensure only data from training domains are used in model selection, since the test domain should be unseen during training and validation.

Original: During training, we begin applying the IRM penalty at iteration K IRM , set to 200 , 400 , or 600 .
Revised: During training, we begin applying the IRM penalty at iteration K IRM , set to 200 , 400 , or 600 .

Original: By varying K IRM , we aim to examine the following hypothesis: IRM works by pushing the spurious features out of the representations learned by the standard ERM training before the IRM penalty is applied (Krueger et al., 2020).
Revised: By varying K IRM , we aim to examine the following hypothesis: IRM works by pushing the spurious features out of the representations learned by the standard ERM training before the IRM penalty is applied (Krueger et al., 2020).

Original: More details on the setup can be found in Appendix B.1.
Revised: More details on the setup can be found in Appendix C.1.

Original: Experimental Results.
Revised: Experimental Results.

Original: Fig.
Revised: Fig.

Original: 2 shows the performance of IRM and IRM with balanced classes in each domain.
Revised: 2 shows the performance of IRM and IRM with balanced classes in each domain.

Original: We make the following observations: ﬁrst, the test accuracy of IRM drops dramatically under strong Λ spuriousness.
Revised: We make the following observations: ﬁrst, the test accuracy of IRM drops dramatically under strong Λ spuriousness.

Original: Second, we show that a na¨ıve ﬁx for IRM, balancing the two classes in each domain by oversampling the minority class, does not lead to improvement in performance.
Revised: Second, we show that a na¨ıve ﬁx for IRM, balancing the two classes in each domain by oversampling the minority class, does not lead to improvement in performance.

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

Original: [Figure]
Revised: [Figure]

Original: 3. W HY DOES IRM WORK FOR THE O RIGINAL C OLORED MNIST?
Revised: 3. W HY DOES IRM WORK FOR THE O RIGINAL C OLORED MNIST?

Original: [Figure]
Revised: [Figure]

Original: The Colored MNIST (CMNIST) dataset cannot expose the limitation of IRM under strong Λ spuriousness.
Revised: The Colored MNIST (CMNIST) dataset cannot expose the limitation of IRM under strong Λ spuriousness.

Original: This is because its two training domains are quite similar.
Revised: This is because its two training domains are quite similar.

Original: As shown in the large triangle on the left in Fig.
Revised: As shown in the large triangle on the left in Fig.

Original: 3, the values of P ( C | Y, E ) are similar for E = 1 and E = 2 .
Revised: 3, the values of P ( C | Y, E ) are similar for E = 1 and E = 2 .

Original: In addition, the values of P ( Y | E ) are the same for all E .
Revised: In addition, the values of P ( Y | E ) are the same for all E .

Original: This makes it difﬁcult to satisfy the IRM constraint by learning feature representations F ( X ) ≈ E .
Revised: This makes it difﬁcult to satisfy the IRM constraint by learning feature representations F ( X ) ≈ E .

Original: Then, we present experiments to show how IRM gradually goes from working on CMNIST to failing on CMNIST+.
Revised: Then, we present experiments to show how IRM gradually goes from working on CMNIST to failing on CMNIST+.

Original: In Fig.
Revised: In Fig.

Original: 3 Left
Revised: 3 Left

Original: and Right, we observe the differences between between CMNIST and
Revised: and Right, we observe the differences between between CMNIST and

Original: CMNIST+ ( ρ = 0 . 9 ).
Revised: CMNIST+ ( ρ = 0 . 9 ).

Original: In Fig.
Revised: In Fig.

Original: Middle, we create various datasets that interpolate between these two datasets, illustrated by the yellow line in the middle triangle.
Revised: Middle, we create various datasets that interpolate between these two datasets, illustrated by the yellow line in the middle triangle.

Original: We use a parameter w plus ∈ [0 , 1] to control P ( C | Y, E ) of the weights of CMNIST+ ( ρ = 0 . 9 ) in the interpolated dataset:
Revised: We use a parameter w plus ∈ [0 , 1] to control P ( C | Y, E ) of the weights of CMNIST+ ( ρ = 0 . 9 ) in the interpolated dataset:

Original: [Equation]
Revised: [Equation]

Original: where P cmnist+ ( C | Y, E ) and P cmnist ( C | Y, E ) are the values of P ( C | Y, E ) in CMNIST+ and CMNIST.
Revised: where P cmnist+ ( C | Y, E ) and P cmnist ( C | Y, E ) are the values of P ( C | Y, E ) in CMNIST+ and CMNIST.

Original: When w plus = 0 and P ( Y = 1 | E = 1)
Revised: When w plus = 0 and P ( Y = 1 | E = 1)

Original: =
Revised: =

Original: P ( Y = 0 | E = 2) = 0 .
Revised: P ( Y = 0 | E = 2) = 0 .

Original: 5 , the dataset is the same with CMNIST.
Revised: 5 , the dataset is the same with CMNIST.

Original: As w plus and P ( Y = 1 | E = 1) =
Revised: As w plus and P ( Y = 1 | E = 1) =

Original: P ( Y = 0 | E = 2) increase, the dataset becomes more similar to CMNIST+, and becomes the same with CMNIST+ ( ρ = 0 .
Revised: P ( Y = 0 | E = 2) increase, the dataset becomes more similar to CMNIST+, and becomes the same with CMNIST+ ( ρ = 0 .

Original: when w plus = 1 and P ( Y = 1 | E = 1)
Revised: when w plus = 1 and P ( Y = 1 | E = 1)

Original: =
Revised: =

Original: P ( Y = 0 | E = 2) = 0 .
Revised: P ( Y = 0 | E = 2) = 0 .

Original: For the training sets, we set
Revised: For the training sets, we set

Original: P ( Y = 1 | E = 1) =
Revised: P ( Y = 1 | E = 1) =

Original: P ( Y = 0 | E = 2) ∈ { 0 . 5 , 0 . 7 , 0 . 9 } .
Revised: P ( Y = 0 | E = 2) ∈ { 0 . 5 , 0 . 7 , 0 . 9 } .

Original: Fig.
Revised: Fig.

Original: 4 shows results of IRM on the test sets of both CMNIST and CMNIST+.
Revised: 4 shows results of IRM on the test sets of both CMNIST and CMNIST+.

Original: The performance of IRM gradually drops when the values of
Revised: The performance of IRM gradually drops when the values of

Original: P ( C | Y, E ) becomes more similar to CMNIST+ (when the Λ spuriousness becomes stronger).
Revised: P ( C | Y, E ) becomes more similar to CMNIST+ (when the Λ spuriousness becomes stronger).

Original:  A N E FFECTIVE F IX FOR IRM
Revised:  A N E FFECTIVE F IX FOR IRM

Original: In this section, we propose a simple but effective solution to address the limitation of the IRM training.
Revised: In this section, we propose a simple but effective solution to address the limitation of the IRM training.

Original: In a series of experiments, we show that the proposed method can improve the performance of IRM even under strong Λ spuriousness.
Revised: In a series of experiments, we show that the proposed method can improve the performance of IRM even under strong Λ spuriousness.

Original: Since the IRM constraint is too general, we propose an effective solution – combining the conditional distribution matching (CDM) constraint and the IRM constraint.
Revised: Since the IRM constraint is too general, we propose an effective solution – combining the conditional distribution matching (CDM) constraint and the IRM constraint.

Original: The CDM constraint (Li et al., 2018; Long et al., 2018) requires P ( F ( X ) | Y, E )
Revised: The CDM constraint (Li et al., 2018; Long et al., 2018) requires P ( F ( X ) | Y, E )

Original: =
Revised: =

Original: P ( F ( X ) | Y ) , which means the feature representation distribution given the class label should be invariant across different domains.
Revised: P ( F ( X ) | Y ) , which means the feature representation distribution given the class label should be invariant across different domains.

Original: We ﬁrst explain why this can be a reasonable solution.
Revised: We ﬁrst explain why this can be a reasonable solution.

Original: Then, we propose two types of models that combine the two constraints.
Revised: Then, we propose two types of models that combine the two constraints.

Original: Here, we use the two cases from Section 3.1 to explain why combining IRM with CDM can be an effective ﬁx for IRM.
Revised: Here, we use the two cases from Section 3.1 to explain why combining IRM with CDM can be an effective ﬁx for IRM.

Original: Perfectly predicting the label satisﬁes CDM.
Revised: Perfectly predicting the label satisﬁes CDM.

Original: In this case, F ( X ) =
Revised: In this case, F ( X ) =

Original: Y and we already know the
Revised: Y and we already know the

Original: IRM constraint is satisﬁed.
Revised: IRM constraint is satisﬁed.

Original: The CDM condition also holds as P ( F ( X ) | Y, E )
Revised: The CDM condition also holds as P ( F ( X ) | Y, E )

Original: =
Revised: =

Original: P ( Y | Y, E ) = P ( Y | Y ) =
Revised: P ( Y | Y, E ) = P ( Y | Y ) =

Original: P ( F ( X ) | Y ) .
Revised: P ( F ( X ) | Y ) .

Original: Perfectly predicting the domain label violates CDM.
Revised: Perfectly predicting the domain label violates CDM.

Original: In this case, F ( X ) =
Revised: In this case, F ( X ) =

Original: E and the IRM constraint can be satisﬁed, but the CDM constraint does not hold as P ( F ( X ) | Y, E )
Revised: E and the IRM constraint can be satisﬁed, but the CDM constraint does not hold as P ( F ( X ) | Y, E )

Original: =
Revised: =

Original: P ( E | Y, E ) (cid:54) =
Revised: P ( E | Y, E ) (cid:54) =

Original: P ( E | Y ) =
Revised: P ( E | Y ) =

Original: P ( F ( X ) | Y )) .
Revised: P ( F ( X ) | Y )) .

Original: This case implies that we can add CDM to IRM to exclude the undesirable solution of learning spurious features that accurately predict the domain label.
Revised: This case implies that we can add CDM to IRM to exclude the undesirable solution of learning spurious features that accurately predict the domain label.

Original: One way to enforce the CDM constraint is through adversarial training: a discriminator tries to infer the source domain from the feature representation F ( X ) and F ( X ) tries to adjust itself to fool the discriminator.
Revised: One way to enforce the CDM constraint is through adversarial training: a discriminator tries to infer the source domain from the feature representation F ( X ) and F ( X ) tries to adjust itself to fool the discriminator.

Original: From this perspective, we can ﬁnd that the CDM constraint exactly prevents F ( X ) from making use of the domain labels to achieve high training accuracy (i.e., low R e ).
Revised: From this perspective, we can ﬁnd that the CDM constraint exactly prevents F ( X ) from making use of the domain labels to achieve high training accuracy (i.e., low R e ).

Original: With the idea of combining IRM and CDM, we can formulate the optimization problem as:
Revised: With the idea of combining IRM and CDM, we can formulate the optimization problem as:

Original: [Equation]
Revised: [Equation]

Original: where R e denotes the loss function in domain e (e.g., cross entropy loss).
Revised: where R e denotes the loss function in domain e (e.g., cross entropy loss).

Original: To impose the CDM constraint, we aim to minimize certain divergence between feature distributions from different domains, denoted as div ( P ( F ( X ) | Y, E = e ) || P [Equation] F ( X [Equation]
Revised: To impose the CDM constraint, we aim to minimize certain divergence between feature distributions from different domains, denoted as div ( P ( F ( X ) | Y, E = e ) || P ( F ( X ) | Y, E = e (cid:48) ))

Original: | Y

Original: [Equation]

Original: E

Original: [Equation] e [Equation] )) , e (cid:54) =
Revised: , e (cid:54) =

Original: e (cid:48) .
Revised: e (cid:48) .

Original: We propose two choices of divergences, i.e., Maximum Mean Discrepancy (MMD) and KullbackLeibler (KL) divergence (through adversarial training), which result in two algorithms: IRM-MMD and IRM-ACDM (i.e., IRM-Adversarial Conditional Distribution Matching).
Revised: We propose two choices of divergences, i.e., Maximum Mean Discrepancy (MMD) and KullbackLeibler (KL) divergence (through adversarial training), which result in two algorithms: IRM-MMD and IRM-ACDM (i.e., IRM-Adversarial Conditional Distribution Matching).

Original: IRM-MMD.
Revised: IRM-MMD.

Original: In IRM-MMD, we adopt Maximum Mean Discrepancy (MMD) (Long et al., 2015;
Revised: In IRM-MMD, we adopt Maximum Mean Discrepancy (MMD) (Long et al., 2015;

Original: Tolstikhin et al., 2017; Shalit et al., 2017) as the distribution divergence.
Revised: Tolstikhin et al., 2017; Shalit et al., 2017) as the distribution divergence.

Original: The MMD between P and
Revised: The MMD between P and

Original: Q , two d -dimensional distributions of feature representations, can be deﬁned as:
Revised: Q , two d -dimensional distributions of feature representations, can be deﬁned as:

Original: [Equation]
Revised: [Equation]

Original: where Z ∈ R d , f : R d → R maps a feature representation to a real value, k : R d × R d → R denotes the characteristic kernel of f and H is the RHKS of k .
Revised: where Z ∈ R d , f : R d → R maps a feature representation to a real value, k : R d × R d → R denotes the characteristic kernel of f and H is the RHKS of k .

Original: The MMD in Eq.
Revised: The MMD in Eq.

Original: 4 is not directly computable.
Revised: 4 is not directly computable.

Original: So, we use the unbiased estimator of MMD (Gretton et al., 2012a) with N samples z P 1 , ..., z PN from P and M samples z Q 1 , ..., z QM from Q :
Revised: So, we use the unbiased estimator of MMD (Gretton et al., 2012a) with N samples z P 1 , ..., z PN from P and M samples z Q 1 , ..., z QM from Q :

Original: [Equation]
Revised: [Equation]

Original: With MMD deﬁned, we can deﬁne the loss function of IRM-MMD as:
Revised: With MMD deﬁned, we can deﬁne the loss function of IRM-MMD as:

Original: [Equation]
Revised: [Equation]

Original: where L eIRM = 1 n e (cid:80) 
Revised: where L eIRM = 1 n e (cid:80) n

Original: [Equation]
Revised: [Equation]

Original:  =1 R e [Equation] G
Revised: i =1 R e ( G

Original: [Equation]
Revised: [Equation]

Original: F ( x ei )) , y ei )
Revised: F ( x ei )) , y ei )

Original: + α || (cid:53)
Revised: + α || (cid:53)

Original: w | w =1 . 

Revised: w | w =1 .

Original: wR e ( G ( F ( x ei )))
Revised: wR e ( G ( F ( x ei )))

Original: ||
Revised: ||

Original: 2 denotes the IRM regularized loss (Arjovsky et al., 2019) for domain e , hyperparameters α and β control the trade-off between the main loss, the IRM constraint and the CDM constraint.
Revised: 2 denotes the IRM regularized loss (Arjovsky et al., 2019) for domain e , hyperparameters α and β control the trade-off between the main loss, the IRM constraint and the CDM constraint.

Original: IRM-ACDM.
Revised: IRM-ACDM.

Original: In IRM-ACDM, we make div ( P ( F ( X ) | Y, E = e ) || P ( F ( X ) | Y, E = e (cid:48) )) be the
Revised: In IRM-ACDM, we make div ( P ( F ( X ) | Y, E = e ) || P ( F ( X ) | Y, E = e (cid:48) )) be the

Original: Kullback-Leibler (KL) divergence and use adversarial learning to estimate it (Tolstikhin et al., 2017;
Revised: Kullback-Leibler (KL) divergence and use adversarial learning to estimate it (Tolstikhin et al., 2017;

Original: Song et al., 2020).
Revised: Song et al., 2020).

Original: More precisely, we deﬁne the loss function of IRM-ACDM as:
Revised: More precisely, we deﬁne the loss function of IRM-ACDM as:

Original: [Equation]
Revised: [Equation]

Original: where γ ye :=
Revised: where γ ye :=

Original: P ( E = e, Y = y ) .
Revised: P ( E = e, Y = y ) .

Original: In adversarial learning, a conditional discriminator D : R d × Y →
Revised: In adversarial learning, a conditional discriminator D : R d × Y →

Original: E tr with parameters θ D is introduced to predict the domain label of an instance, given its feature representation and class label.
Revised: E tr with parameters θ D is introduced to predict the domain label of an instance, given its feature representation and class label.

Original: As proved in Li et al.
Revised: As proved in Li et al.

Original: ; Song et al.
Revised: ; Song et al.

Original: (2020), the estimation of the
Revised: (2020), the estimation of the

Original: KL divergence above can be reformulated as the minimax game below:
Revised: KL divergence above can be reformulated as the minimax game below:

Original: [Equation]
Revised: [Equation]

Original: where D e ( F ( x ) , y )
Revised: where D e ( F ( x ) , y )

Original: =
Revised: =

Original: ˆ P ( E = e | F ( x ) , y ) is the predicted probability of the instance for domain e by the discriminator D .
Revised: ˆ P ( E = e | F ( x ) , y ) is the predicted probability of the instance for domain e by the discriminator D .

Original: In practice, we solve the minimax game above efﬁciently by the alternative gradient ascent/descent algorithm.
Revised: In practice, we solve the minimax game above efﬁciently by the alternative gradient ascent/descent algorithm.

Original: Experimental Results.
Revised: Experimental Results.

Original: We evaluate IRM-MMD and IRM-ACDM on CMNIST+ to show their efﬁcacy under strong Λ spuriousness.
Revised: We evaluate IRM-MMD and IRM-ACDM on CMNIST+ to show their efﬁcacy under strong Λ spuriousness.

Original: We let the output of the second last layer of LeNet-5 to be the feature representation F ( X ) .
Revised: We let the output of the second last layer of LeNet-5 to be the feature representation F ( X ) .

Original: For IRM-MMD, we use the multiple kernel MMD (MKMMD)
Revised: For IRM-MMD, we use the multiple kernel MMD (MKMMD)

Original: (Gretton et al., 2012b; Long et al., 2015).
Revised: (Gretton et al., 2012b; Long et al., 2015).

Original: For IRM-ACDM, we use a two-layer MLP with |E tr | outputs as the discriminator D .
Revised: For IRM-ACDM, we use a two-layer MLP with |E tr | outputs as the discriminator D .

Original: To show that CDM alone cannot solve OOD prediction under strong Λ spuriousness, we set the weight of the IRM penalty α = 0 in IRM-MMD and IRM-ACDM to obtain the models regularized by MMD and ACDM, respectively. 
Revised: To show that CDM alone cannot solve OOD prediction under strong Λ spuriousness, we set the weight of the IRM penalty α = 0 in IRM-MMD and IRM-ACDM to obtain the models regularized by MMD and ACDM, respectively. We consider EIIL (Creager et al., 2020) that solves a minimax game. In the max step, it learns soft domain labels for instances s.t. the IRMv loss is maximized.

Revised: In the min step, it minimizes the IRMv1 loss.

Original: We also include ERM and oracle.
Revised: We also include ERM and oracle.

Original: The oracle uses the original LeNet-5 architecture with a single-channel CNN as the ﬁrst layer.
Revised: The oracle uses the original LeNet-5 architecture with a single-channel CNN as the ﬁrst layer.

Original: It is trained and tested with instances transformed into grayscale.
Revised: It is trained and tested with instances transformed into grayscale.

Original: Table.
Revised: Table.

Original: 3 shows the performance of IRM-MMD, IRM-ACDM and the baselines.
Revised: 3 shows the performance of IRM-MMD, IRM-ACDM and the baselines.

Original: Under strong Λ spuriousness, compared to IRM (Fig. 2), we can observe the signiﬁcant and consistent performance improvement over IRM resulting from combining CDM with IRM.
Revised: Under strong Λ spuriousness, compared to IRM (Fig. 2), we can observe the signiﬁcant and consistent performance improvement over IRM resulting from combining CDM with IRM.

Original: IRM-MMD and IRM-ACDM also outperform MMD and ACDM.
Revised: IRM-MMD and IRM-ACDM also outperform MMD and ACDM.

Original: This veriﬁes that CDM alone cannot solve the OOD prediction problem.
Revised: This veriﬁes that CDM alone cannot solve the OOD prediction problem.

Revised: EIIL reaches comparable accuracy to IRM-MMD and IRM-ACDM

Revised: when ρ = 0 . 85 , 0 . 9 .

Revised: But EIIL has larger standard deviation ± 10 .

Revised: 32% , ± 13 .

Revised: 03% , ± 13 .

Original: R ELATED W ORK
Revised: 03% with ρ = 0 . 8 , 0 . 85 , 0 . 9 , compared to IRM-ACDM ( ± 4 . 34% , ± 4 . 65% , ± 3 . 09% ) and IRM-MMD ( ± 4 . 56% , ± 2 .

Revised: 50% , ± 6 . 97% ).

Revised: Discussion can be found in Appendix C.3.

Revised:  R ELATED W ORK

Original: OOD Prediction.
Revised: OOD Prediction.

Original: IRM (Arjovsky et al., 2019) formulates causal feature learning as a constraint on the ERM framework (Vapnik, 1992), which imposes the causal inductive bias: causal feature  representations lead to the existence of an optimal classiﬁer for all domains.
Revised: IRM (Arjovsky et al., 2019) formulates causal feature learning as a constraint on the ERM framework (Vapnik, 1992), which imposes the causal inductive bias: causal feature [Table] representations lead to the existence of an optimal classiﬁer for all domains.

Original: Then it is transformed to a regularizer which can be minimized along with empirical risks.
Revised: Then it is transformed to a regularizer which can be minimized along with empirical risks.

Original: Chang et al.
Revised: Chang et al.

Original: (2020) propose a variant of IRM.
Revised: (2020) propose a variant of IRM.

Original: They implement IRM by minimizing the difference between two classiﬁers’ outputs, which take F ( X ) and Concat ( F ( X ) , E ) as the inputs, respectively.
Revised: They implement IRM by minimizing the difference between two classiﬁers’ outputs, which take F ( X ) and Concat ( F ( X ) , E ) as the inputs, respectively.

Original: Ahuja et al.
Revised: Ahuja et al.

Original: [Table] reformulate the optimization problem of IRM from a game theory aspect.
Revised: (2020) reformulate the optimization problem of IRM from a game theory aspect.

Original: IRM is related to robust optimization (Ben-Tal et al., 2009).
Revised: IRM is related to robust optimization (Ben-Tal et al., 2009).

Original: The goal is to optimize the worst domain speciﬁc empirical risk.
Revised: The goal is to optimize the worst domain speciﬁc empirical risk.

Original: Krueger et al.
Revised: Krueger et al.

Original: (2020) extend robust optimization (Ben-Tal et al., 2009) to minimize the empirical risk of the worst domain and maximize that of other domains.
Revised: (2020) extend robust optimization (Ben-Tal et al., 2009) to minimize the empirical risk of the worst domain and maximize that of other domains.

Original: They also propose to minimize the variance of domain speciﬁc empirical risks along with the empirical risk.
Revised: They also propose to minimize the variance of domain speciﬁc empirical risks along with the empirical risk.

Original: Jin et al.
Revised: Jin et al.

Original: (2020) propose to minimize the domain speciﬁc risk between two models, one trained on the same domain, the other trained on the other domains.
Revised: (2020) propose to minimize the domain speciﬁc risk between two models, one trained on the same domain, the other trained on the other domains.

Original: These methods essentially minimize the differences among domain speciﬁc risks.
Revised: These methods essentially minimize the differences among domain speciﬁc risks.

Original: Moreover, data augmentation can also improve the generalizability of DNNs from the data perspective (Ilse et al., 2020; van der Wilk et al., 2018).
Revised: Moreover, data augmentation can also improve the generalizability of DNNs from the data perspective (Ilse et al., 2020; van der Wilk et al., 2018).

Original: However, it requires prior knowledge on the differences between training and test domains, which are not be available in OOD prediction.
Revised: However, it requires prior knowledge on the differences between training and test domains, which are not be available in OOD prediction.

Original: Kuang et al.
Revised: Kuang et al.

Original: (2018) and Qiao et al.
Revised: (2018) and Qiao et al.

Original: (2020) aim to handle the case with only one training domain.
Revised: (2020) aim to handle the case with only one training domain.

Original: Zhang et al.
Revised: Zhang et al.

Original: (2020b) propose to make DNNs more robust against test data generated by unseen interventions.
Revised: (2020b) propose to make DNNs more robust against test data generated by unseen interventions.

Original: They model interventions in training data with a generative model and perform test-time inference to catch unseen interventions.
Revised: They model interventions in training data with a generative model and perform test-time inference to catch unseen interventions.

Original: Compared to the existing work, this work exposes and ﬁxes the issue of IRM in OOD prediction under strong Λ spuriousness.
Revised: Compared to the existing work, this work exposes and ﬁxes the issue of IRM in OOD prediction under strong Λ spuriousness.

Original: Domain Adaptation (DA) assumes that the unlabeled test set can be used during training and validation.
Revised: Domain Adaptation (DA) assumes that the unlabeled test set can be used during training and validation.

Original: From the methodology aspect, distribution matching methods used in DA, such as gradient reversing (Ganin & Lempitsky, 2015) and adversarial CDM (Long et al., 2018) are useful for OOD prediction.
Revised: From the methodology aspect, distribution matching methods used in DA, such as gradient reversing (Ganin & Lempitsky, 2015) and adversarial CDM (Long et al., 2018) are useful for OOD prediction.

Original: Peters et al.
Revised: Peters et al.

Original: (2016) realize the invariance of causal relationships can be used for DA.
Revised: (2016) realize the invariance of causal relationships can be used for DA.

Original: Zhang et al.
Revised: Zhang et al.

Original: (2013) propose reweighting and kernel based distribution matching methods to handle three types of DA problems: target shift, conditional shift and generalized target shift.
Revised: (2013) propose reweighting and kernel based distribution matching methods to handle three types of DA problems: target shift, conditional shift and generalized target shift.

Original: Gong et al.
Revised: Gong et al.

Original: (2016) work on extracting transferable components F ( X ) that ensure P ( F ( X ) | Y ) to be invariant across domains with location-scale transformation.
Revised: (2016) work on extracting transferable components F ( X ) that ensure P ( F ( X ) | Y ) to be invariant across domains with location-scale transformation.

Original: Their method can identify how P ( Y ) changes across domains simultaneously.
Revised: Their method can identify how P ( Y ) changes across domains simultaneously.

Original: Different from DA, we strictly ensure that the test domain is unseen during training and validation to reﬂect the scenario of OOD prediction in real-world applications.
Revised: Different from DA, we strictly ensure that the test domain is unseen during training and validation to reﬂect the scenario of OOD prediction in real-world applications.

Original:  C ONCLUDING R EMARKS
Revised:  C ONCLUDING R EMARKS

Original: This work focuses on the OOD prediction problem under strong Λ spuriousness.
Revised: This work focuses on the OOD prediction problem under strong Λ spuriousness.

Original: Strong Λ spuriousness means the correlations between spurious features and the class label are strong.
Revised: Strong Λ spuriousness means the correlations between spurious features and the class label are strong.

Original: We ﬁnd an important limitation of IRM in OOD prediction under strong Λ spuriousness: it can be satisﬁed by spurious feature representations that are predictive of the domain label.
Revised: We ﬁnd an important limitation of IRM in OOD prediction under strong Λ spuriousness: it can be satisﬁed by spurious feature representations that are predictive of the domain label.

Original: To verify it, we design the
Revised: To verify it, we design the

Original: CMNIST+ dataset which has strong Λ spuriousness between color (spurious features) the class label through their common cause – the domain variable.
Revised: CMNIST+ dataset which has strong Λ spuriousness between color (spurious features) the class label through their common cause – the domain variable.

Original: On CMNIST+, we observe the performance of
Revised: On CMNIST+, we observe the performance of

Original: IRM dramatically drops when the Λ spuriousness becomes stronger.
Revised: IRM dramatically drops when the Λ spuriousness becomes stronger.

Original: Based on this observation, we propose a simple but an effective ﬁx to mitigate this issue of IRM.
Revised: Based on this observation, we propose a simple but an effective ﬁx to mitigate this issue of IRM.

Original: The proposed approach combines CDM and IRM because CDM can also be satisﬁed by causal feature representations.
Revised: The proposed approach combines CDM and IRM because CDM can also be satisﬁed by causal feature representations.

Original: At the same time, CDM can prevent DNNs from picking up the aforementioned spurious feature representations.
Revised: At the same time, CDM can prevent DNNs from picking up the aforementioned spurious feature representations.

Original: Experimental results on CMNIST+ show signiﬁcant performance improvement of the proposed method, demonstrating its effectiveness.
Revised: Experimental results on CMNIST+ show signiﬁcant performance improvement of the proposed method, demonstrating its effectiveness.

Original: Interesting future work includes (1) extension of the proposed method to OOD prediction tasks in complex data (e.g., graphs and time series) and (2) development of general causal inductive bias that can impose various conditional independence. 
Revised: Interesting future work includes (1) extension of the proposed method to OOD prediction tasks in complex data (e.g., graphs and time series) and (2) development of general causal inductive bias that can impose various conditional independence. A IRM AND THE C ONDITIONAL I NDEPENDENCE

Revised: Under general conditions, the conditional independence Y ⊥⊥ E | F ( X ) is a necessary condition for solutions of the original IRM optimization problem.

Revised: In (Arjovsky et al., 2019), IRM is deﬁned as a two-stage optimization problem.

Revised: Any solution F ( X ) to the IRM optimization problem must satisfy Y ⊥⊥ E | F ( X ) .

Revised: However, some F ( X ) satisfying Y ⊥⊥ E | F ( X ) may not be a solution to the original IRM problem.

Revised: For example, generally, F ( X ) =

Revised: E would not minimize the sum of the domain-speciﬁc risk R e .

Revised: However, under strong Λ spuriousness, there exist solutions to the original

Revised: IRM problem that still pick up spurious features.

Revised: Consider the extreme case, in the training data, if

Revised: Y = E , then F ( X ) =

Revised: E is a solution to the problem.

Revised: B A NALYSIS FOR CMNIST+

Revised: Fitting Color for Classiﬁcation.

Revised: In expectation, the ERM model would learn color as the feature representation for classiﬁcation.

Revised: Here, we show the theoretical results for such cases.

Revised: When the model learns color C as the feature representation F ( X ) , we have:

Revised: [Equation]

Revised: The second equality is by Bayes’ rule.

Revised: We know that P ( C ) =

Revised: (cid:80) Y (cid:80) E P

Revised: ( C | Y, E ) P

Revised: ( Y | E ) P ( E ) .

Revised: Then, given P ( C | Y, E ) , P ( Y | E ) , P ( C ) and Eq. 9, we can obtain P ( Y | C ) as shown in Table 4.

Revised: These results imply that the deterministic classiﬁer

Revised: ˆ P ( Y = 1 | C = G )

Revised: = 1 , ˆ P ( Y = 1 | C = B )

Revised: = 1 , ˆ P ( Y = 1 | C = R )

Original: {+= 
+}
Revised: {+= 
+}

Revised: for all ρ ∈ [0 . 55 , 0 . 9] as shown in Table 4.

Revised: So, the accuracy of the deterministic classiﬁer ˆ P ( Y | C ) on the test set is 0 . 2 .

Revised: Its accuracy on the training set is (cid:80) C P ( ˆ Y =

Revised: Y | C ) P ( C )

Revised: where P ( C ) =

Revised: (cid:80) E (cid:80) Y P ( C

Revised: | Y, E ) P ( Y | E ) P ( E ) .

Revised: Fitting Domain Label for Classiﬁcation.

Revised: Since the spurious correlation between the domain label and the class label is strong, and IRM cannot penalize models ﬁtting the domain label, P ( Y | E ) can help us understand the expected behavior of the IRM model.

Revised: Note that the test data is from an unseen domain.

Revised: So, we analyze the model that ﬁrst predicts domain by color, then predicts class label by domain.

Revised: First, we analyze P ( E | C ) as below:

Revised: [Equation]

Revised: where the second and forth qualities are by Bayes’ rule.

Revised: With Eq.

Revised: 10, we can list the values of

Revised: P ( E | C ) in Table 6 for ρ ∈ [0 . 55 , 0 . 9] .

Original: {+Thus, we know the deterministic domain prediction results would be ˆ P ( E = 
+}
Revised: {+Thus, we know the deterministic domain prediction results would be ˆ P ( E = 
+}

Revised: | C = G )

Revised: = 1 , ˆ P ( E = 1 | C = B ) = 0 , and ˆ P ( E = 1 | C = R )

Revised: = 0 .

Revised: Since

Revised: P ( Y | E ) is given in Table 1, we can obtain the deterministic classiﬁer’s predictions as:

Revised: [Equation]

Revised: So, the expected test accuracy of the model would be 0 . 35 .

Revised: Recall that the model ﬁrst predicts domain label by color and then predict class label by the predicted domain.

Revised: By doing this, it would have F ( X ) ≈ E .

Revised: This makes it approximately satisfy the IRM constraint Y ⊥⊥ E | F ( X ) since

Revised: Y ⊥⊥ E | E .

Revised: This implies that the model’s performance can be treated as the expected performance of the IRM model in CMNIST+.

Revised: In terms of the performance of ˆ P ( Y | ˆ E ) on the training set, we can get the results using the same prediction rules as in Eq.

Revised: One may argue that the IRM model can perform well if we balance the two classes in each domain.

Revised: Here, we theoretically show this is not the case.

Revised: By setting P ( Y | E ) = 0 .

Revised: 5 , we can obtain the values of P ( E | C ) in Table 7.

Revised: Note that practically this can be done by oversampling the minority class of each domain in each mini-batch.

Revised: However, since P ( Y | E ) = 0 .

Revised: 5 , the predictions made by the deterministic classiﬁer ˆ P ( Y | E ) would be just random guess, leading to a test accuracy of 0 . 5 .

Revised: Fitting both Domain and Color for Classiﬁcation.

Revised: Here, we consider the model that ﬁrst predicts the domain label by color and then predicts the class label by both the color and the predicted domain label.

Revised: The ﬁrst step is the same as the model ﬁtting the domain label.

Revised: For the second step, we analyze P ( Y | C, E ) as below:

Revised: [Equation]

Revised: With P ( C | E ) =

Revised: (cid:80) Y P ( C | Y, E ) P ( Y | E ) and Eq.

Revised: 12, we obtain values of P ( Y | C, E ) as shown in

Revised: Table 8.

Revised: So, given the predicted domains, ˆ P ( E = 1 | C = G ) = 1 , ˆ P ( E = 1 | C = B ) = 0 , and ˆ P ( E = 1 | C = R ) = 0 , the predictions on the class label are ˆ P ( Y = 1 | C

Revised: =

Revised: G, E = 1) = 1 ,

Original: {+ˆ P ( Y = 
+}
Revised: {+ˆ P ( Y = 
+}

Revised: | C = B, E = 2)

Revised: = 1 , ρ > 0 .

Revised: 8 , ˆ P ( Y = 1 | C

Revised: =

Revised: B, E = 2) = 0 , ρ ≤ 0 .

Revised: 8 and ˆ P

Original: {+( Y = 
+}
Revised: {+( Y = 
+}

Revised: | C

Revised: =

Revised: R, E = 2) = 0 .

Revised: So, the test accuracy is 0 .

Revised: 35 when ρ ≤ 0 .

Revised: 8 and 0 .

Revised: when ρ > 0 . 8 .

Revised: Similarly, when we make P ( Y | E ) = 0 .

Revised: 5 by oversampling the minority class in each domain, we can obtain the predictions shown in Table 9.

Revised: So, the deterministic model would make predictions as ˆ

Revised: P ( Y = 1 | C =

Revised: G, E = 1) = 1 ,

Original: {+ˆ P ( Y = 
+}
Revised: {+ˆ P ( Y = 
+}

Revised: | C = B, E = 2)

Revised: = 1 , and ˆ P ( Y = 1 | C

Revised: = R, E )

Revised: = 0 .

Revised: This would lead to a test accuracy of 0 .

Revised: These results reﬂect the reasons why the IRM model trained with the balanced classes in each domain ( P ( Y | E ) = 0 . 5 ) has worse performance compared to its counterpart trained with the original CMNIST+ data

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: C E XPERIMENTAL S ETUP AND R ESULTS

Revised: Here, we include more details and discussion on experimental setup, datasets and results.

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Table]

Revised: C. E XPERIMENTAL S ETUP

Revised: Here

Revised: [Table]

Revised: we provide more details on experimental setup.

Revised: We perform grid search for hyperparameter tuning.

Revised: For IRM, IRM-MMD and IRM-ACDM, we search the iteration number to plug in the

Revised: IRM penalty term ( K IRM ) in 200 , 400 , 600 and IRM penalty weight α in { 1 , 10 , ..., 10 8 } .

Revised: For

Revised: MMD, ACDM, IRM-MMD, IRM-ACDM, we search CDM penalty weight β in { 1 , 10 , ..., 10 5 } .

Revised: For ACDM and IRM-ACDM, we set the number of steps we train the discriminator D in each iteration to 10 .

Revised: For EIIL, we do the same hyperparameter tuning on the IRM penalty weight ( α ), the iteration to add IRM ( K IRM ) as we do for IRM, IRM-ACDM and IRM-MMD methods.

Revised: We train the soft environment weight q ( E | X, Y ) for { 10 , 100 , 1000 , 10000 } steps.

Revised: C. D ATASETS

Revised: Here, we present more details about the datasets.

Revised: CMNIST.

Revised: CMNIST is introduced by (Arjovsky et al., 2019).

Revised: We can see the two training domains of

Revised: CMNIST are similar to each other in terms of both P ( C | Y, E ) and P ( Y | E ) in Table 10.

Revised: This means

Revised: CMNIST does not cover the case of strong Λ spurious, since the spurious correlations, color–domain and domain–class, are not strong.

Revised: [Table]

Revised: CMNIST+.

Revised: We visualize the CMNIST+ dataset with different values of ρ in Fig.

Revised: In addition, we provide a detailed simulation recipe of CMNIST+ and compare it with that of CMNIST.

Revised: This would also show that CMNIST+ is in accordance with the causal graph in Fig.

Revised: We decide the true label Y ∗ (without noise) of each instance by its original digit label (0-9).

Revised: We randomly split the data into test and training.

Revised: We assign the training instances to the two training domains based on the true label Y ∗ and

Revised: P ( Y | E ) in Table 1.

Revised: This step introduces correlations between Y and E .

Revised: In each training environment, we further randomly split the data into training and validation.

Revised: In each environment, we generate noisy labels Y by randomly ﬂipping them with 25% probability.

Revised: This means Y ∗ → Y .

Revised: Given P ( C | Y, E ) in Table 1, the noisy label and the domain label, we assign color to each instance.

Revised: This step introduces correlations among X s , Y and E .

Revised: Note that there is a difference in what causal relationships mean in traditional causal inference and in OOD prediction.

Revised: In OOD prediction, the deﬁnition of causal relationships is different from a traditional one.

Revised: Traditionally, X c → Y means the generation of Y is (partially) determined by X c .

Revised: It does not necessarily mean P ( Y | X c ) remains the same across domains (Pearl & Bareinboim, 2014).

Revised: However, in OOD prediction, we say there exists a causal relationship X c

Revised: → Y iff P ( Y | X c ) is the same across different domains.

Revised: We also know in the original MNIST dataset, there exists invariant causal relationship

Revised: X c → Y ∗ .

Revised: This implies that, from the data generating process of CMNIST+, we conﬁrm that (1) there exist causal relationships

Revised: X c → Y ∗ → Y , (2) there are correlations among X s , Y and E .

Revised: With these two conclusions, we can claim that the causal graph in Fig.

Revised: 1 is in accordance with CMNIST+ in the OOD prediction problem where the true label Y ∗ is ignored as it is not used in training and evaluation.

Revised: Creating Strong Λ Spuriousness with Two Colors.

Revised: It is possible to setup strong Λ spuriousness with two colors for binary classiﬁcation with two training domains.

Revised: Here, we use the two colors: red (R) and green (G).

Revised: To show it is possible, we use an example with the following setup: P ( Y = 1 | E = 1) = 0 .

Revised: 9 , P ( Y = 1 | E = 2) = 0 .

Revised: Let’s say for E

Revised: = 1 , we set P ( C = G | Y = 1 , E = 1) = 0 .

Revised: 9 , P ( C = G | Y = 0 , E = 1) = 0 .

Revised: Then, we can set P ( C = G | Y = 1 , E = 2)

Revised: = 0 .

Revised: P ( C = G | Y = 0 , E = 2) = 0 .

Revised: 9 for E = 2 .

Revised: This setup makes strong correlations among the color, the class label and the domain label.

Revised: Thus, it is possible to create strong Λ spuriousness with just two colors.

Revised: Our concern with such datasets is that even if strong Λ spuriousness exists in training domains, it is a challenge to create test domains that are diverse enough from the training ones.

Revised: Following the aforementioned setup, for the test domain E = 3 , if we set P ( Y = 1 | E = 3) = 0 .

Revised: P ( C = G | Y = 1 , E = 3) = 0 .

Revised: 5 and P ( C = G | Y = 0 , E = 3) = 0 .

Revised: 5 , it would be right in the middle of the two training domains.

Revised: Unfortunately in this setting, even if IRM fails, it can be difﬁcult to observe it with the test accuracy.

Revised: This is because a model perfectly ﬁts the color features can reach 0 .

Revised: 5 test accuracy.

Revised: This leads to smaller differences between models ﬁtting causal features and those with spurious features.

Revised: Thus, it becomes more challenging to judge whether a model fails in practice, which explains why we use three colors in CMNIST+.

Revised: C. M ORE D ISCUSSION ON R ESULTS OF EIIL

Revised: We evaluate EIIL 1 on CMNIST+.

Revised: EIIL gets test accuracy 43 .

Revised: 40 ± 10 .

Revised: 24 ± 13 .

Revised: 93 ± 13 .

Revised: 03% on CMNIST+ with ρ = 0 . 8 , 0 . 85 , 0 .

Revised: Compared to IRM-ACDM 57 .

Revised: 23 ± 4 .

Revised: 83 ± 4 .

Revised: 85 ± 3 .

Revised: 09% and IRM-MMD 52 .

Revised: 91 ± 4 .

Revised: 83 ± 2 .

Revised: 96 ± 6 .

Revised: 97% , EIIL has comparable mean accuracies when ρ = 0 . 85 , 0 . 9 , but is not stable.

Revised: There are two reasons.

Revised: (1) EIIL relies on IRM and the soft domain weight q ( E | X, Y ) .

Revised: When q ( E | X, Y ) takes values s.t.

Revised: strong Λ spuriousness exists, it makes IRM fail.

Revised: (2) It has issues in model selection.

Revised: EIIL can only select models by validation accuracy as q ( E | X, Y ) is a scalar q i for instance i .

Revised: It is learned for the val/test, so validation loss is not feasible.

Revised: Validation accuracy can be high when the model picks up spurious features.

Revised: While model selection with validation loss considers the regularizers (IRM and CDM), which approximates how well models ﬁt causal features.