Original: Repeated Environment Inference for Invariant Learning
Revised: Repeated Environment Inference for Invariant Learning

Original: Anonymous Authors

Original: Abstract
Revised: Abstract

Original: We study the problem of invariant learning when the environment labels are unknown.
Revised: We study the problem of invariant learning when the environment labels are unknown.

Original: We focus on the Invariant representation notion when the
Revised: We focus on the Invariant representation notion when the

Original: Bayes optimal conditional label distribution is the same across different environments.
Revised: Bayes optimal conditional label distribution is the same across different environments.

Original: Previous work conducts the Environment Inference (EI) by maximizing the penalty term in the Invariant Risk
Revised: Previous work conducts the Environment Inference (EI) by maximizing the penalty term in the Invariant Risk

Original: Minimization (IRM) framework.
Revised: Minimization (IRM) framework.

Original: The EI step uses a reference model which focuses on spurious correlations to efficiently reach a good environment partition.
Revised: The EI step uses a reference model which focuses on spurious correlations to efficiently reach a good environment partition.

Original: However, it is not clear how to find such a reference model.
Revised: However, it is not clear how to find such a reference model.

Original: In this work, we propose to repeat the EI process and retrain an ERM model on the majority environment inferred by the EI step in the previous step.
Revised: In this work, we propose to repeat the EI process and retrain an ERM model on the majority environment inferred by the EI step in the previous step.

Original: Under mild assumptions, we find that this iterative process helps learn a representation capturing the spurious correlation better than the single step.
Revised: Under mild assumptions, we find that this iterative process helps learn a representation capturing the spurious correlation better than the single step.

Original: This results in better Environment Inference and better Invariant Learning.
Revised: This results in better Environment Inference and better Invariant Learning.

Original: We show that this method outperforms baselines on both synthetic and real-world datasets.
Revised: We show that this method outperforms baselines on both synthetic and real-world datasets.

Original: 1. Introduction
Revised: 1. Introduction

Original: In conventional machine learning,the training data is assumed to  independently and identically distributed ( iid ) as the test data.
Revised: In conventional machine learning, training data is assumed to be independently and identically distributed ( iid ) as the test data.

Original: This  is usually violated in the real world.
Revised: This assumption is usually violated in the real world.

Original: Samples used in training might not be representative of the whole data distribution in many applications.
Revised: Samples used in training might not be representative of the whole data distribution in many applications.

Original: Therefore, performance suffers when machine learning algorithms are deployed in new domains or environments where spurious correlations learnt from the training data do not hold.
Revised: Therefore, performance suffers when machine learning algorithms are deployed in new domains or environments where spurious correlations learned from the training data do not hold.

Original: To solve this problem, recent literature advocates focusing on extracting causal relationships (Peters et al., 2017) from data.
Revised: To solve this problem, recent literature advocates focusing on extracting causal relationships (Peters et al., 2017; Pearl, 2019) from data.

Original: However, causal discovery and inference in the highdimensional data remains a big challenge in practice.{+practice (Sch¨olkopf et al., 2021). Published at the ICML 
+}
Revised: However, causal discovery and inference in high-dimensional data remain a big challenge in {+practice (Sch¨olkopf et al., 2021). Published at the ICML 
+}

Original:  Recent work (Arjovsky et al., 2019; Ilse et al., 2020;Chuang et al.) have proposed to learn stable correlations that are in-
Revised: Workshop on Spurious Correlations, Invariance, and Stability . Baltimore, Maryland, Copyright Recent works (Arjovsky et al., 2019; Ilse et al., 2020;

Original: variant across different domains because they have been shown to strongly relates with Causality (Heinze-Deml et al., 2018; Peters et al., 2016).
Revised: Chuang et al.) have proposed to learn stable correlations that are invariant across different domains because they have been shown to strongly relates with Causality (HeinzeDeml et al., 2018; Peters et al., 2016).

Original: However, Invariant Learning methods usually require predefined environment labels which are not always available.
Revised: However, Invariant Learning methods usually require predefined environment labels, which are not always available.

Original: For example, fair machine learning can be regarded as an Invariant Learning problem where different sub-populations correspond to different environments and the sensitive features are usually not available or even hard to define (Corbett-Davies & Goel, 2018).
Revised: For example, fair machine learning can be regarded as an Invariant Learning problem where different sub-populations correspond to different environments, and the sensitive features are usually unavailable or even hard to define (Corbett-Davies & Goel, 2018).

Original: On the other hand, environment labels play an important role in invariant learning.
Revised: On the other hand, environment labels play an important role in invariant learning.

Original: Figure 1 shows how the performance of IRM  increases when training samples are grouped differently (Experiment B) than the ground-truth group labels (Experiment A).
Revised: Figure 1 shows how the performance of IRM (Arjovsky et al., 2019) increases when training samples are grouped differently (Experiment B) from the ground-truth group labels (Experiment A).

Original: We randomly shuffle the background for digits with different probabilities to create two training environments and fix the background for others.
Revised: To create two training environments, we randomly shuffle the background for digits with different probabilities.

Original: IRM benefits from the regrouping as it reflects a better differentiation between the two environments in terms of the domain-specific “optimal” representation, for predicting the labels.
Revised: IRM benefits from the regrouping as the new grouping reflects a better differentiation between the two environments. See Figure 1.

Original: Creager et al.
Revised: Creager et al.

Original: (2021) proposed EIIL which enables invariant learning when environment labels are not available.
Revised: (2021) proposed EIIL, enabling invariant learning when environment labels are unavailable.

Original: It conducts Environment Inference (EI) using the representation learned by a reference model, by maximizing the penalty term in the Invariant Risk Minimization (IRM) framework.
Revised: It conducts Environment Inference (EI) using the representation learned by a reference model by maximizing the penalty term in the Invariant Risk Minimization (IRM) framework.

Original: Moreover, it is shown that EIIL usually works when the reference model captures “spurious” correlations.
Revised: Moreover, it is shown that EIIL usually works when the reference model captures “spurious” correlations.

Original: However, it is not clear how to find such a model besides learning an
Revised: However, it is unclear how to find such a model besides learning an

Original: ERM model from all the training data.
Revised: ERM model from all the training data.

Original: In this paper, we find that environment partition obtained in a single EI step after ERM is  sub-optimal for invariant learning and propose a simple repeated EI strategy to improve the quality of Environment Inference and consequently the performance of Invariant Learning.
Revised: In this paper, we find that environment partition obtained in a single EI step after ERM is often sub-optimal for invariant learning. We propose a simple repeated EI strategy to improve the quality of Environment Inference and, consequently, the performance of Invariant Learning.

Original: Our contribution can be summarized as follows:
Revised: Our contribution can be summarized as follows:

Original: We propose a novel method that generates environment labels by repeating the EI step and training the ERM model on the majority environment.
Revised: We propose a novel method that generates environment labels by repeating the EI step and training the ERM model on the majority environment.

Original: We call it Repeated Environment Inference for Invariant Learning (REIIL).
Revised: We call it Repeated Environment Inference for Invariant Learning (REIIL).

Original: Under mild assumptions, we show that our strategy better facilitates downstream invariant learning.
Revised: Under mild assumptions, we show that our strategy better facilitates downstream invariant learning.

Original: Our approach outperforms previous methods including EIIL on previously used and a newly introduced dataset.
Revised: Our approach outperforms previous methods, including EIIL on previously used and a newly introduced dataset.

Original: 2. Related Work
Revised: 2. Related Work

Original: IRM and Environment Inference Arjovsky et al.
Revised: IRM and Environment Inference Arjovsky et al.

Original: (2019) proposes a training objective for learning the invariant representations, under the assumption that the  optimal conditional P ( y | ϕ ( x )) remains invariant across domains.
Revised: (2019) proposes a training objective for learning the invariant representations, under the assumption that the Bayes optimal conditional P ( y | ϕ ( x )) remains invariant across domains.

Original: However, recent works (Rosenfeld et al., 2020; Kamath et al., 2021) have shown that IRM fails when target domains are not sufficiently similar to training domains.
Revised: However, recent works (Rosenfeld et al., 2020; Kamath et al., 2021) have shown that IRM fails when target domains are not sufficiently similar to training domains.

Original: Besides, it also requires good environment labels beforehand.
Revised: Besides, it also requires good environment labels beforehand.

Original: Several methods have been proposed to infer the environment labels before or along with the invariant learning.
Revised: Several methods have been proposed to infer the environment labels before or along with the invariant learning.

Original: EIIL (Creager et al., 2021) takes a two-step strategy and first partitions the training data into environments using a biased reference model and then train a IRM model.
Revised: EIIL (Creager et al., 2021) takes a two-step strategy and first partitions the training data into environments using a biased reference model before using environment-based invariant learning.

Original: However, it is unclear how to get a perfectly biased reference model to obtain the optimal environment partition.
Revised: However, it is unclear how to get a perfectly biased reference model to obtain a good environment partition.

Original: Recently, Adversarial
Revised: Adversarial

Original: Invariant Learning (Ye et al., 2021) uses a minimax game to robustly train the predictors while inferring the worstcase environment partition.
Revised: Invariant Learning (Ye et al., 2021) uses a minimax game to robustly train the predictors while inferring the worstcase environment partition.

Original: In this paper, we follow the
Revised: In this paper, we follow the

Original: EIIL setup and study how tobetter conduct environment inference before downstream invariant learning.
Revised: EIIL setup and study how to conduct environment inference better for downstream invariant learning.

Original: Other Types of Invariance in Domain Generalization
Revised: Other Types of Invariance in Domain Generalization

Original: Besides IRM, other definitions of Invariance exist in the domain generalization literature.
Revised: Besides IRM, other definitions of invariance exist in the domain generalization literature.

Original: Dated back to early works in domain adaptation (Gretton et al., 2009; Sun & Saenko, 2016), marginal feature alignment motivates the learning of the domain-invariant features, that is P ( ϕ ( x )) being invariant across domains (Ganin et al., 2016).
Revised: Dated back to early works in domain adaptation (Gretton et al., 2009; Sun & Saenko, 2016), marginal feature alignment motivates the learning of the domain-invariant features, that is P ( ϕ ( x )) being invariant across domains (Ganin et al., 2016).

Original: However, the shortcoming of these methods is also obvious that it ignores the relation between the features and the labels sothat it is sub-optimal especially under label shift (Zhao et al., 2019; 2020).
Revised: However, the shortcoming of these methods is also obvious: it ignores the relationship between the features and the labels, so it still suffers under label shift (Zhao et al., 2019; 2020).

Original: Other work utilize tools like Information Bottleneck and propose learning objectives based on the mutual information (Li et al., 2021).
Revised: Other works utilize tools like Information Bottleneck and propose learning objectives based on the mutual information (Li et al., 2021).

Original: Li et al.
Revised: Li et al.

Original: (2018) focus on conditional domain invariant representation and make sure that the P ( ϕ ( x ) | y ) remains invariant across domains.
Revised: (2018) focus on conditional domain invariant representation and make sure that the P ( ϕ ( x ) | y ) remains invariant across domains.

Original: Focusing on only one type of invariant learning has been shown to be sub-optimal (Shui et al., 2022), but choosing the right type of invariant learning algorithm can be a tricky question.
Revised: Focusing on only one type of invariance is inadequate (Shui et al., 2022), and choosing the right type of invariant learning algorithm can be tricky. Some works also propose using domain-specific features along with domain-invariant features to improve performance in domain generalization tasks (Ding & Fu, 2017; Chattopadhyay et al., 2020; Bui et al., 2021). However, it is hard to say whether exploiting domain-specific features will always help in unseen domains without theoretical guarantees.

Original: In our paper, we aim to learn the invariant P ( y | ϕ
Revised: Our paper aims to learn the invariant P ( y | ϕ

Original: ( x )) distribution when environment labels are not available.
Revised: ( x )) distribution when environment labels are unavailable.

Original: Representation Learning under Domain Shift
Revised: Representation Learning under Domain Shift

Original: Other representation learning techniques havebeen also  found helpful under domain shift.
Revised: Other representation learning techniques have also been found helpful under domain shift.

Original: For example, Liu et al.
Revised: For example, Liu et al.

Original: (2021) proposes that upweighting underperforming samples in a second iteration of training is sufficient to improve worst-case group performance.
Revised: (2021) proposes that upweighting underperforming samples in a second iteration of training is sufficient to improve worst-case group performance.

Original: Kirichenko et al.
Revised: Kirichenko et al.

Original: (2022) suggest that last layer retraining of Neural Networks with interesting samples might be sufficient for robustness against spurious correlations.
Revised: (2022) suggest that last layer retraining of Neural Networks with interesting samples might be sufficient for robustness against spurious correlations.

Original: Contrastive Learning approaches (Wang et al., 2022; Kim et al., 2021) have also shown promise for domain adaptation and generalization.
Revised: Contrastive Learning approaches (Wang et al., 2022; Kim et al., 2021) have also shown promise for domain adaptation and generalization.

Original: However, these methods do not explicitly learn any invariance and are usually hard to analyze.
Revised: However, these methods do not explicitly learn any invariance and are usually hard to analyze.

Original: In our paper, we focus on the more principled
Revised: In our paper, we focus on the more principled

Original: IRM setting and overcome the difficulty of obtaining biased models for environment inference.
Revised: IRM setting and overcome the difficulty of obtaining biased models for environment inference.

Original: 3. Method 3.1. IRM method and Environment Inference
Revised: 3. Method 3.1. IRM method and Environment Inference

Original: The practical version of IRM (called IRMv1) has the following objective.
Revised: The practical version of IRM (called IRMv1) has the following objective,

Original: [Equation]
Revised: [Equation]

Revised: where samples are taken from multiple environments and this information is used to regularize the training.

Original: In EIIL, an Environment Inference step produces the required environment labels.
Revised: As multiple environments are not always available, EIIL proposes an Environment Inference step to produce the required environment labels.

Original: It uses a fixed reference classifier Φ .
Revised: It uses a fixed reference classifier Φ and the IRM penalty term to do this.

Original: The EI step maximizes the following with respect to  q ,
Revised: The EI step maximizes the following with respect to soft environment assignment q ,

Original: [Equation]
Revised: [Equation]

Original: where R e (Φ , q ) =
Revised: where R e (Φ , q ) =

Original: 1 N (cid:80) i q i ( e ) l (Φ( x i ) , y i ) .
Revised: 1 N (cid:80) i q i ( e ) l (Φ( x i ) , y i ) .

Revised: This soft assignment is used to produce hard environment labels which are then used for invariant learning.

Original: In summary, EIIL(Creager et al., 2021) takes the following steps:
Revised: In summary, EIIL takes the following steps:

Original: Takes an input reference model ˜Φ . 2.
Revised: Takes an input reference model ˜Φ . 2.

Original: Maximizes C EI to get environment partition.
Revised: Maximizes C EI to get environment partition.

Original: q ∗ = arg max q C EI (˜Φ , q ) .
Revised: q ∗ = arg max q C EI (˜Φ , q ) .

Original: Minimizes C IL (any Invariant Learning objective of choice) to get the final model.
Revised: Minimizes C IL (any Invariant Learning objective of choice) to get the final model.

Original: [Equation]
Revised: [Equation]

Original: Note that here q produces a binary partition of environments, and we assume this binary partitioning throughout this paper.
Revised: Note that here q is used to produce a binary partition , and we assume this binary partitioning throughout this paper.

Original: EIIL works under the assumption that the reference model ( ˜Φ ) focuses only on the spurious features.
Revised: EIIL works under the assumption that the reference model ( ˜Φ ) focuses only on the spurious features.

Original: It is only when ˜Φ =
Revised: It is only when ˜Φ =

Original: Φ Spurious , does the environment inference (EI) process maximally violates the Invariance Principle (EIC):
Revised: Φ Spurious , does the environment inference (EI) pro- cess maximally violates the Invariance Principle (EIC):

Original: E [ y | Φ( x ) = h, e 1 ] = E [ y | Φ( x ) = h, e 2 ] , ∀ h ∈ H , e 1 , e 2 ∈ E obs
Revised: [Equation]

Original: The authors use Φ ERM as an approximation for Φ Spurious , but this does not always work.
Revised: The authors use Φ ERM as an approximation for Φ Spurious , but this does not always work.

Original: As underlined in the paper explicitly via various experiments, it remains unexplored to find better reference models (which would be worse than
Revised: As underlined in the paper explicitly via various experiments, it remains unexplored to find better reference models (which would be worse than

Original: ERM in terms of generalization performance) to fully exploit the EI step in EIIL.
Revised: ERM in terms of generalization performance) to fully exploit the EI step in EIIL.

Original: 3.2. Proposed method
Revised: 3.2. Proposed method

Original: We present an assumption about the training distribution and a conjecture about the resulting EI step using a model trained from such a training distribution.
Revised: We present an assumption about the training distribution and a conjecture about the resulting EI step using a model trained from such a training distribution.

Original: We then reason how finding a model that satisfies the conjecture under the assumption would help the downstream invariant learning.
Revised: We then reason how finding a model that satisfies the conjecture under the assumption would help the downstream invariant learning.

Original: We finally propose a strategy to approximate such a model.
Revised: We finally propose a strategy to approximate such a model.

Original: Assumption 3.1.
Revised: Assumption 3.1.

Original: Assume a data generating graph x → y → z , where x and z are observed features and y is the target label.
Revised: Assume a data generating graph X → Y → Z , where X and Z are observed features, Y is the target feature and the anti-causal mechanism is unstable.

Original: A training distribution is sampled following the graph with noise.
Revised: A training dataset is sampled following this graph.

Original: Assume that in the training distribution, anti-causal features are more informative than the causal features, i.e.,
Revised: Assume that in the training dataset, anti-causal features are more informative than the causal features, i.e.,

Original: [Equation]
Revised: [Equation]

Original: where I measuresthe overall mutual information between features and the labels across samples.
Revised: where I measures mutual information between features.

Original: Conjecture 3.2.
Revised: Conjecture 3.2.

Original: There exists a model trained on such training datasets  such that when it is used as a reference model for the EI step for two environments, the resulting majority environment  satisfies the following:
Revised: There exists a model trained on such training datasets ( tr ) such that when it is used as a reference model for the EI step (2), the resulting majority environment ( e maj ) satisfies the following:

Original: [Equation]
Revised: [Equation]

Original: tr [Equation]
Revised: Remarks : By induction, if we find such a model after one

Original: Y [Equation] [Equation] ) where e maj and tr denote the majority environment after the EI step and all the training data, respectively.
Revised: EI step, there exists another such model in the next step, when the majority environment is the training dataset.

Original: Remarks : By induction, we conclude that if a model satisfies the conjecture after one step of EI, there exists such a model in the next step when we regard the majority environment as the training distribution, until a point of diminishing returns.
Revised: This process repeats until a point of diminishing returns.

Original: Therefore, we propose a simple repetition of reference model training on a biased subset of samples to get a more biased reference model better approximating Φ Spurious .
Revised: Therefore, we propose a simple repetition of reference model training on a biased subset of samples to get a more biased reference model better approximating Φ Spurious .

Original: We use ERM to approximate such a model in each step.
Revised: We use ERM to approximate such a model in each step.

Original: The strategy can be defined in the following steps:
Revised: The strategy can be defined in the following steps:

Original: Perform EI step from EIIL using the ERM reference model trained on the whole dataset. 2.
Revised: Perform EI step from EIIL using the ERM reference model trained on the whole dataset. 2.

Original: Instead of training an Invariant Learning algorithm directly on the obtained partition, retrain an ERM model on the majority environment obtained from the previous step.
Revised: Instead of training an Invariant Learning algorithm directly on the obtained partition, retrain an ERM model on the majority environment obtained from the previous step.

Original: Repeat this step for n iterations.
Revised: Repeat this step for n iterations.

Original: Use the finally obtained partition for downstream Invariant Learning methods.
Revised: Use the finally obtained partition for downstream Invariant Learning methods.

Original: The repetition of the EI step is based on the proposed conjecture that majority subsets obtained would show even stronger correlation between spurious and target labels and, the ERM model trained on it would iteratively approximate Φ Spurious .
Revised: The repetition of the EI step is based on the proposed conjecture that majority subsets obtained would show even stronger correlation between spurious and target labels and, the ERM model trained on it would iteratively approximate Φ Spurious .

Original: 4. Experimental Results
Revised: 4. Experimental Results

Original: We show the effectiveness of our proposed method on various datasets under different challenging conditions and find that it recovers good environment splits and provides boosted performance in invariant learning.
Revised: We show the effectiveness of our proposed method on various datasets under different challenging conditions and find that it recovers good environment splits and provides boosted performance in invariant learning.

Original: We also provide empirical evidence to support our conjecture.
Revised: We also provide empirical evidence to support our conjecture.

Original: 4.1. Datasets
Revised: 4.1. Datasets

Original: Colored MNIST (CMNIST) was originally introduced in the IRM paper (Arjovsky et al., 2019).
Revised: Colored MNIST (CMNIST) was originally introduced in the IRM paper (Arjovsky et al., 2019).

Original: It has a synthetic binary classification task where color is introduced as an anti-causal spurious correlation.
Revised: It has a synthetic binary classification task where color is introduced as an anti-causal spurious correlation.

Original: CIFAR-Background MNIST (CBMNIST)
Revised: CIFAR-Background MNIST (CBMNIST)

Original: We introduce this new dataset, which is based on the simple concept of putting MNIST (Deng, 2012) digits on CIFAR (Krizhevsky et al., 2009) backgrounds.
Revised: We introduce this new dataset, which is based on the simple concept of putting MNIST (LeCun et al., 1998) digits on CIFAR (Krizhevsky et al., 2009) backgrounds which would act as spurious correlations.

Original: The classification task for the resultant images remains same as the MNIST digit classification task.
Revised: The classification task for the resultant images remains same as the MNIST digit classification task.

Original: This process results ina dataset which is easily accessible, modifiable and has more complex “real-world” features than simple colors. The resulting images have the same size as CIFAR images ( 32 × 32 × 3 ).
Revised: CBMNIST [I] was created in the following manner

Original: In our experiments, we select a randommapping of CIFAR and MNIST classes and put the CIFAR images on
Revised: : • One random CIFAR image from each class is chosen to act as the background for MNIST images.

Original:  MNIST digit backgrounds.
Revised: This results in 10 different backgrounds with different visual properties. • Each MNIST class is mapped with a randomly selected CIFAR class.

Original: This mapping  acts as thespurious anti-causal correlation.
Revised: This mapping between classes acts as the anti-causal spurious correlation in the training set.

Original: We only randomly shuffle this background for a small number of samples in the dataset to introduce noise.
Revised: For the test set, a completely different random mapping of classes is chosen such that there is no overlap of digitbackground combinations with the training set.

Original: For one variant, CBMNIST [I], we split the training MNIST training  data in two equally sized environments (Figure 1) and put CIFAR backgrounds on them using the  training mapping.
Revised: • Training data is split into two environments with samples each (from 60000 MNIST training samples). Test data has the remaining 10000 samples. • All training samples are filled with the mapped CIFAR backgrounds.

Original: With probability ∼ 1% and ∼ 2% respectively in each environment, we change the background.
Revised: In Environment 1, backgrounds of ∼ 1% of samples are shuffled .

Original: Figure 2 shows example images from one CBMNIST [I] training environment.
Revised: Shuffled here means a random background which is not equal to the one in the train or test mapping of that class (one of the remaining 8), is applied to this sample. Similarly, in Environment 2, ∼ 2% samples have their background shuffled . The resulting dataset has ∼ 97% samples exhibiting a strong anti-causal relationship between the target class label and the background. The rest ∼ 3% samples have shuffled backgrounds so the only invariant features in them are the actual digit shapes.

Original: In the test set, we use completely different backgrounds.
Revised: This ensures that the dataset satisfies the assumption (3.1) for our conjecture (3.2) to hold. CBMNIST [II] has a slight variation of the concept of shuffled . Here shuffled samples have their background switched randomly to exactly one other pre-decided background rather than 8 in CBMNIST [I].

Original: Design details and other variants [II,
Revised: In this way, even the ∼ 3% shuffled samples exhibit a spurious correlation with the background and a particular digit is only ever seen with at most 2 different backgrounds in the training set. The test set remains the same. CBMNIST [III] has the percentage of samples having shuffled backgrounds change from CBMNIST [I]. We changed the numbers 1% and 2% to 10% and 20% respectively, simulating a case where even simple ERM can focus on invariant causal features because of lots of background variations.

Original: III] of this dataset can be seen in Appendix A.
Revised: It is unclear whether our assumption (3.1) holds in this case.

Original: 4.2. Implementation details
Revised: 4.2. Implementation details

Original: We reuse the code [anonymous link], model architectures and hyperparameters from EIIL, which followsthe IRM for CMNIST experiments.
Revised: We reuse the code [github link], model architectures and hyperparameters from EIIL, which follows IRM for CMNIST experiments.

Original: For CMNIST, we increase the number of training steps to 900.
Revised: For CMNIST, we increase the number of training steps to 900.

Original: For CBMNIST, we use an MLP with a similar structure but instead of down sampling as in CMNIST, we use all the 32 × 32 × 3 input dimensions to help the model learn any complex background features.
Revised: For CBMNIST, we use an MLP with a similar structure but instead of down sampling as in CMNIST, we use all the 32 × 32 × 3 input dimensions to help the model learn any complex background features.

Original: For REIIL, we found that in n = 9 iterations, the performance usually saturates.
Revised: For REIIL, we found that in n = 9 iterations, the performance usually saturates.

Original: Note that for EIIL and REIIL, the environment labels are not used in the training.
Revised: Note that for EIIL and REIIL, the environment labels are not used in the training.

Original: We also used L2 regularization on the parameters for all methods.
Revised: We also used L regularization on the parameters for all methods.

Original: 4.3. Result Analysis 
Revised: 4.3. Result Analysis 4.3.1. B ETTER I NVARIANT L EARNING

Original: Better Invariant Learning Results on CMNIST and CBMNIST [I] can be seen in Figure 3.
Revised: Test performance on CMNIST and CBMNIST variants can be seen in Figure 4.

Original: In boththe datasets, REIIL outperforms other methods.
Revised: In both datasets, REIIL outperforms other methods.

Revised: Now we analyze each experiment as follows.

Original: For CMNIST, the performance gains are not big becausethe Φ ERM approximates Φ Spurious quite well in the first EIIL step.
Revised: In CMNIST, the performance gains are insignificant because Φ ERM approximates Φ Spurious quite well in the first EIIL step itself.

Original: However, REIIL’s performance is more stable than EIIL across 5 runs.
Revised: However, REIIL’s performance is more stable than EIIL’s across five runs.

Original: For CBMNIST [I], the performance gains are quite substantial.
Revised: In CBMNIST [I], the performance gains are quite substantial.

Original: Because the shuffled samples arevery scarce and  IRM fails to beat even ERM.
Revised: Because the shuffled samples are scarce and poorly split, IRM fails to beat even ERM.

Original: But REIIL consistently performs better than other methods.
Revised: But REIIL consistently performs better than other methods.

Original: Results on other CBMNIST variants and synthetic data are deferred to the Appendix B.
Revised: For CBMNIST [II], the results were quite similar. Interestingly, in both these cases,

Original: [Figure]
Revised: ERM outperforms the native IRM, and using the EI step proves critical for performance gains. ERM emerges as the best-performing model in CBMNIST [III], probably because the higher percentage of shuffled background samples can regularize the ERM model’s training. The initial reference model is also very far from Φ Spurious , making the performance of EIIL and even REIIL worse than IRM. REIIL still beats EIIL, but it cannot reach the performance of even IRM due to the poor starting point.

Original: Dynamics of Repeated EI We study the  environment inference dynamics in the REIIL.
Revised: It highlights the importance of a good reference model for the success of Environment Inference. One thing to note is the superior performance of the reweighting method to Invariant Learning methods. Inspired from (Liu et al., 2021), we trained a Weighted ERM (WERM) model on the EI obtained environment splits (REIWERM). WERM weighs samples within each environment equally rather than weighing each sample equally in the whole dataset as in ERM.

Original: Repeated EI finds more biased reference models which produce better environment splits.
Revised: This model surprisingly matched the performance of IRM. We defer investigating this further for future work.

Original: Recall that a good environment split regroups data such that the spurious features are more informative about the labels in the majority environment (for example, corresponding to the non-shuffled images in CBMNIST [I]).
Revised: [Figure] 4.3.2. D YNAMICS OF R EPEATED EI If datasets follow our assumption (3.1) and we find a model satisfying Conjecture (3.2) using ERM, we claim that repeating the EI step produces better environment splits such that the spurious features are more informative about the labels in the majority environment (for example, corresponding to the non-shuffled images in CBMNIST [I]).

Original: The minority environment’s test accuracy would be low under the reference model trained in the majority environment.
Revised: This would imply that the minority environment’s test accuracy would be low under the reference model trained using the majority environment.

Revised: It would also imply that most of the causally informative samples (in the case of CBMNIST, these would be shuffled samples) would fall in the minority environment.

Original: Figure 4 shows the accuracies of the reference models in each step for the samples in the minority environment, and the percentage of shuffled samples assigned in the minority environment with each EI step.
Revised: Figure 5 shows the accuracies of the reference models in each step for the samples in the minority environment and the percentage of shuffled samples assigned in the minority environment with each EI step.

Original: The accuracy in the minority domain keeps decreasing while the proportion of shuffled samples is getting higher.
Revised: We see that accuracies keep decreasing while the proportion of shuffled samples increases in the minority environment.

Revised: Note that the training accuracies in the majority environment always reach close to 1, implying that the ERM model successfully trains to convergence in each step.

Original: For CMNIST,this effect is not as pronounced as CBMNIST [I], and shuffled samples  correspond to those which do not exhibit the color-based spurious correlation. 
Revised: For CMNIST, shuffled samples would correspond to those which do not exhibit the colorbased spurious correlation. This effect is not as pronounced as in CBMNIST [I], but follows a similar pattern.

Original: This provides empirical evidence that the ERM model trained on the majority environment obtained from the previous EI step  captures spurious features (background or color)increasingly and produces a better environment partition.
Revised: Therefore, we can conclude that an ERM model trained on the majority environment obtained from the previous EI step better captures spurious features (background or color) and produces a better environment partition.

Original: [Figure]
Revised: [Figure]

Original: 5.Conclusion and Discussion
Revised: 5. Discussion

Original: In this paper, we propose a novel strategy to find a more biased reference model for environment inference.
Revised: What if our assumption is violated?

Original: By repeating the EI step iteratively and training the reference model on the majority environment, our method helps find a better environment partition for downstream Invariant
Revised: It was shown by (Gulrajani & Lopez-Paz, 2020) that carefully trained ERM models to work quite well and often outperform Domain

Original: Learning tasks.
Revised: Generalization methods.

Original:  Our assumption often holds in practice where most training samples exhibit strong spurious correlations. We conduct experiments on both CMNIST and our newly designed CBMNIST and demonstrate that our method achieves better predictive performance in invariant learning.
Revised: We also see that ERM outperforms other methods in CBMNIST [III]. Our assumption is supposed to hold in more challenging situations where ERM usually fails.

Original: What if our assumption is violated? In practice, it may not be easy to test whether our assumption is valid.
Revised: However, in practice, it may not be easy to test whether our assumption is valid, especially when the variables in the causal graph are hard to extract directly from data as in the image data.

Original: Therefore, we conduct experiments on the synthetic data (Appendix B) where there is no clear “optimal” data split.
Revised: We then conduct experiments on the synthetic data (Appendix A) where there is no clear optimal data split due to the significant noise level.

Original: We find thateven though REIIL usually achieves the best test error, it may not recover the models close to the true ones.
Revised: We find that REIIL usually achieves the best test error. However, it may not recover the models close to the true ones.

Revised: There also exists a model selection problem as the parameter tuning according to the validation data may make the model farther from the ground truth one.

Original: We  leave further explorationof this setting toour future work.
Revised: We refer more discussion to Appendix A and leave further exploration to future work.

Original: What if there are more than two environments?
Revised: What if there are more than two environments?

Revised: For simplicity, we only investigate environment inference for two environments in the paper.

Original: With multiple types of spurious correlations (or multiple training environments), a good environment split is hard to define.
Revised: We are aware that a good environment split is hard to define with multiple types of spurious correlations.

Original: It is not clear whether a binary partition is sufficient.
Revised: It is unclear whether a binary partition is sufficient.

Original: Therefore, we want to investigate these more complicated scenarios in invariant learning in future.A. CBMNIST
Revised: We defer further investigation of these more complicated scenarios in invariant learning for the future.

Original: CBMNIST [I] This dataset was created in the following manner : • One random CIFAR image from each class is chosen to act as background for MNIST images. This results in 10 different backgrounds with different visual properties. • Each MNIST class is mapped with a randomly selected CIFAR class.
Revised: 6. Conclusion

Original: This mapping between classes acts as the anti-causal spurious correlation in the training set. For the test set, a completely different random mapping of classes is chosen such that there is no overlap. • Training data is split in two environments with 25000 samples each (from 60000 MNIST training samples). Test data has the remaining 10000 samples.
Revised: In this paper, we propose a novel strategy to find a more biased reference model for environment inference.

Original: • All trainingsamples are filled with the mapped CIFAR background from the corresponding class. In Environment 1, backgrounds of ∼ 1% of samples are shuffled . Shuffled here means a random background which is not equal to the one in train or test mapping of that class (one of the remaining 8), is assigned to this sample. Similarly, in Environment 2 ∼ 2% samples have their background shuffled.
Revised: By repeating the EI step iteratively and training the reference model on the majority environment, our method helps find a better environment partition for downstream Invariant Learning tasks.

Original: The resulting dataset has ∼ 97% samples exhibiting a strong anti-causal relationship with the target class label and the background.

Original: The rest ∼ 3% samples have shuffled backgrounds so the only invariant features in them are the actual digit shapes. The causal graph can be seen in Figure 5. This ensures that the dataset satisfies the assumption for our conjecture to hold. to give the actual observed background X back .
Revised: Our assumption often holds in practice when most training samples exhibit strong spurious correlations and simple ERM methods fail to generalize.

Original: CBMNIST [II] only had a slight variation of the concept of shuffled .

Original: The shuffled samples don’t have their background shuffled randomly to one of the remaining 8, but to exactly one other background (which is already decided).

Original: In this way, even the ∼ 3% shuffled samples exhibit a spurious correlation with the background and a particular digit is only ever seen with at most 2 different backgrounds in the training set.

Original: The test set remains the same with a completely different mapping of classes and backgrounds.

Original: This correlates more to the CMNIST setting where only 2 colors are used and the shuffled samples do not exhibit the anti-causal spurious correlation with color.

Original: CBMNIST [III] only had the percentage of samples having shuffled backgrounds change from CBMNIST

Original: [I].

Original: In this case, we changed the numbers ∼ 1% and ∼ 2% to ∼ 10% and ∼ 20% respectively.

Original: This simulates a case where even ERM has a lot of variation in the dataset to learn invariant causal features from. B. Other Experiments CBMNIST Results on the variations of CBMNIST can be seen in Figure 6. For CBMNIST [II] the results were similar, but a lot changed with CBMNIST [III] where ERM emerged as the best performing model.
Revised: We conduct experiments on both CMNIST and our newly designed CBMNIST to demonstrate the effectiveness of our method.

Original: This is because a higher percentage of shuffled background samples can regularize the ERM model’s training to learn the invariant feature which is the digit shape.

Original: This also changes the fact that for the first EI step, a reference model which is very far from Φ Spurious is used.

Original: This in turn produces a bad environment split and hampers the progress of REIIL.

Original: REIIL still beats EIIL but because of the poor starting point, it was not able to reach the performance of even IRM.

Original: Interestingly, this could be an issue with the downstream Invariant Learning method. Inspired from (Liu et al., 2021), we trained a Weighted ERM (WERM) model on the EI obtained environment splits (REIWERM). Weighted ERM simply weighs samples from each environment equally rather than weighing each individual sample equally as in ERM. This model surprisingly matched the performance of IRM.
Revised: A. Synthetic Data Experiments

Original: We defer investigating this further for future work.

Original: Synthetic Data We used all variants of the synthetic data defined in IRM (Arjovsky et al., 2019) as well as the synthetic data defined in (Creager et al., 2021).
Revised: We use all variants of the synthetic data defined in IRM (Arjovsky et al., 2019) as well as the synthetic data defined in (Creager et al., 2021) to test REIIL.

Original: The data was originally split in 3 environments with noises 0 . 2 , 2  and 5 respectively.
Revised: The data is initially split into 3 environments with noises 0 . 2 , 2 , and 5 respectively.

Original: We found an interesting detail in the previous implementations where ERM was trained  on all three environments and IRM/EIIL were trained  only on the first two and  the lastwas used as a validation set to pick the IRM regularizer.
Revised: In the previous implementations, they trained ERM on all three environments, trained IRM/EIIL only on the first two, and used the last as a validation set to pick the IRM regularizer.

Original: We changed this detail by introducing a validation set explicitly using noise = 3 . 5 and treating the last environment with noise = 5 as the test set.
Revised: We change this setup by introducing a validation set explicitly using noise = 3 . 5 and treating the last environment with noise = 5 as the test set.

Original: The IRM data experiment results can be summarized in Table 1.
Revised: The IRM data experiment results can be summarized in Table 1.

Original: In almost all experiments, IRM found the best solutions.
Revised: In almost all experiments, IRM finds the best solutions.

Original: Interestingly, REIIL almost always found models with  lowest validation and test errors (always lower than EIIL) but that did not translate to the corresponding solutions having lower Causal and Non-Causal Errors  with the ground truth solutions.
Revised: Interestingly, REIIL almost always finds models with the lowest validation and test errors (always lower than EIIL), but that does not translate to the corresponding solutions having lower Causal and Non-Causal Errors compared with the ground truth solutions.

Original: We found that the validation set in both previous and current settings had high noise ( 5 and 3 . 5 ), which is not representative of the real causal graph (which would have no noise).
Revised: We find that the validation set in both previous and current settings had high noise ( 5 and 3 . 5 ), which is not representative of the real causal graph (which would have no noise).

Original: Therefore, we infer that selecting a model based on  lower validation error does not necessarily yield the best solutions.
Revised: Therefore, we infer that selecting a model based on a lower validation error does not necessarily yield the best solutions.

Original: But in a learning setting with  fixed number of samples, finding a model with lower validation error is usually the best we can do.
Revised: But in a learning setting with a fixed number of samples, finding a model with lower validation error is usually the best we can do.

Original: REIIL often achieves the lowest validation error, which is always lower than EIIL, suggesting that repeating the EI step improves performance.
Revised: REIIL often achieves the lowest validation error, which is always lower than EIIL, suggesting that repeating the EI step improves performance.

Original: We found similar results with the synthetic data defined in EIIL, where REIIL usually found models with  lowest validation and test errors.
Revised: We found similar results with the synthetic data defined in EIIL, where REIIL usually found models with the lowest validation and test errors.

Original: [Table]
Revised: Discussion This questions the traditional methods of model and hyperparameter selection based on validation sets. In this case, the validation set was biased, and it is clear that the model which performs best for this set does not yield the causal solution. Even if the validation set was unbiased, it might not generalize to biased sets because the solution does not account for the noise explicitly.

Original: [Figure]
Revised: [Table]