Original: ZIN: When and How to Learn Invariance by Environment Inference?
Revised: ZIN: When and How to Learn Invariance by Environment Inference?

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

Original: Affiliation
Revised: Affiliation

Original: Abstract
Revised: Abstract

Original: It is commonplace to encounter heterogeneous data, of which some aspects ofthe data distribution may vary but the underlying causal mechanisms remainconstant.
Revised: It is commonplace to encounter heterogeneous data, of which some aspects ofthe data distribution may vary but the underlying causal mechanisms remainconstant.

Original: When data are divided into distinct environments according to theheterogeneity, recent invariant learning methods have proposed to learn robustand invariant models based on this environment partition.
Revised: When data are divided into distinct environments according to theheterogeneity, recent invariant learning methods have proposed to learn robustand invariant models based on this environment partition.

Original: It is hence tempting toutilize the inherent heterogeneity even when environment partition is not provided.
Revised: It is hence tempting toutilize the inherent heterogeneity even when environment partition is not provided.

Original: Unfortunately, in this work, we show that learning invariant features under thiscircumstance is fundamentally impossible without further inductive biases oradditional information.
Revised: Unfortunately, in this work, we show that learning invariant features under thiscircumstance is fundamentally impossible without further inductive biases oradditional information.

Original: Then, we propose a framework to jointly learn environmentpartition and invariant representation, assisted by additional auxiliary information.
Revised: Then, we propose a framework to jointly learn environmentpartition and invariant representation, assisted by additional auxiliary information.

Original: We derive sufficient and necessary conditions for our framework to provablyidentify invariant features under a fairly general setting.
Revised: We derive sufficient and necessary conditions for our framework to provablyidentify invariant features under a fairly general setting.

Original: Experimental results onboth synthetic and real world datasets validate our analysis and demonstrate animproved performance of the proposed framework.
Revised: Experimental results onboth synthetic and real world datasets validate our analysis and demonstrate animproved performance of the proposed framework.

Original: As a byproduct, our findingsalso raise the need of making the role of inductive biases more explicit in futureworks, when learning invariant models without environment partition.
Revised: As a byproduct, our findingsalso raise the need of making the role of inductive biases more explicit in futureworks, when learning invariant models without environment partition.

Original:  Introduction
Revised:  Introduction

Original: Machine learning algorithms with empirical risk minimization (ERM) generally assume independentand identically distributed (i.i d.) data in training and test sets.
Revised: Machine learning algorithms with empirical risk minimization (ERM) generally assume independentand identically distributed (i.i d.) data in training and test sets.

Original: Due to changing circumstances,selection bias, or time shifts, data distributions are often heterogeneous across different environmentsin practical applications.
Revised: Due to changing circumstances,selection bias, or time shifts, data distributions are often heterogeneous across different environmentsin practical applications.

Original: When the distributions of training and test data are indeed different, modelperformance can severely degrade [44, 15, 35], as ERM based algorithms may exploit the spuriouscorrelations that are particularly useful to ERM in the training data but do not exist in the test data.
Revised: When the distributions of training and test data are indeed different, modelperformance can severely degrade [43, 14, 34], as ERM based algorithms may exploit the spuriouscorrelations that are particularly useful to ERM in the training data but do not exist in the test data.

Original: To tackle the distribution-shift or out-of-distribution generalization problem, a recent line of methodspropose to utilize the causally invariant mechanisms (rather than the spurious correlations in thetraining data) that are supposed to be stable across different environments.
Revised: To tackle the distribution-shift or out-of-distribution generalization problem, a recent line of methodspropose to utilize the causally invariant mechanisms (rather than the spurious correlations in thetraining data) that are supposed to be stable across different environments.

Original: For example, Peterset al.
Revised: For example, Peterset al.

Original: [33] propose to exploit the invariance principle to learn a linear model that merely relies onthe direct causes of the target.
Revised: [32] propose to exploit the invariance principle to learn a linear model that merely relies onthe direct causes of the target.

Original: Such a model is shown to be robust to potential interventions on anyvariable except the target itself.
Revised: Such a model is shown to be robust to potential interventions on anyvariable except the target itself.

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

Original: [5] propose invariant risk minimization (IRM) tocapture invariant correlations by learning representations that elicit an optimal invariant predictoracross multiple training environments.
Revised: [4] propose invariant risk minimization (IRM) tocapture invariant correlations by learning representations that elicit an optimal invariant predictoracross multiple training environments.

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

Original: [3], Jin et al.
Revised: [2], Jin et al.

Original: [19], Krueger et al.
Revised: [18], Krueger et al.

Original: [22], Xie et al.[48] further develop variants of IRM by introducing game theory, regret minimization, variancepenalization, etc., and Xu and Jaakkola [50], Chang et al.
Revised: [21], Xie et al.[47] further develop variants of IRM by introducing game theory, regret minimization, variancepenalization, etc., and Xu and Jaakkola [49], Chang et al.

Original: [8], Lin et al.
Revised: [7], Lin et al.

Original: [24] try to learn invariantfeatures by coupled adversarial neural networks.
Revised: [23] try to learn invariantfeatures by coupled adversarial neural networks.

Original: Noticeably, the aforementioned invariant learning methods require datasets to be explicitly partitionedinto environments (or domains) according to the heterogeneity underlying the data.
Revised: Noticeably, the aforementioned invariant learning methods require datasets to be explicitly partitionedinto environments (or domains) according to the heterogeneity underlying the data.

Original: However, suchan environment partition may be unavailable or hard to obtain in practice [26, 9].
Revised: However, suchan environment partition may be unavailable or hard to obtain in practice [25, 8].

Original: Recently, a lineof works to learn invariance with no environment indexes, i.e., the dataset is assembled by mergingdata from multiple environments.
Revised: Recently, a lineof works to learn invariance with no environment indexes, i.e., the dataset is assembled by mergingdata from multiple environments.

Original: For instance, Creager et al.
Revised: For instance, Creager et al.

Original: [9] propose environmental inferencefor invariant learning (EIIL), a two-stage method by firstly inferring the environments with an ERMbased trained biased model and then performing invariant learning on inferred environments.
Revised: [8] propose environmental inferencefor invariant learning (EIIL), a two-stage method by firstly inferring the environments with an ERMbased trained biased model and then performing invariant learning on inferred environments.

Original: Liuet al.
Revised: Liuet al.

Original: [26] devise an interactive mechanism, called heterogeneous risk minimization (HRM), withenvironment inference and invariant learning on raw feature level, and Liu et al.
Revised: [25] devise an interactive mechanism, called heterogeneous risk minimization (HRM), withenvironment inference and invariant learning on raw feature level, and Liu et al.

Original: [27] extend it torepresentation level with kernelized trick.
Revised: [26] extend it torepresentation level with kernelized trick.

Original: However, as we will show, learning invariant models under this circumstance (with only inputlabel pairs) can be risky , both empirically and theoretically.
Revised: However, as we will show, learning invariant models under this circumstance (with only inputlabel pairs) can be risky , both empirically and theoretically.

Original: As our first contribution, in Section 4,we prove that learning invariant features from heterogeneous data without environment indexesis in theory impossible.
Revised: As our first contribution, in Section 4,we prove that learning invariant features from heterogeneous data without environment indexesis in theory impossible.

Original: Specifically, we provide a counter-example (Example 1) and a generaltheoretic result (Theorem 1), to show that the invariant features are unidentifiable if no environmentinformation is provided.
Revised: Specifically, we provide a counter-example (Example 1) and a generaltheoretic result (Theorem 1), to show that the invariant features are unidentifiable if no environmentinformation is provided.

Original: This impossibility result, similar to the identifiability issue in causal discovery[42, 34], independent component analysis (ICA) [18], and unsupervised learning of disentangledrepresentations [29], motivates us to consider when and how to learn causally invariant features.
Revised: This impossibility result, similar to the identifiability issue in causal discovery[41, 33], independent component analysis (ICA) [17], and unsupervised learning of disentangledrepresentations [28], motivates us to consider when and how to learn causally invariant features.

Original: In this work, we turn to additional auxiliary variables that encode some information about the latentheterogeneity.
Revised: In this work, we turn to additional auxiliary variables that encode some information about the latentheterogeneity.

Original: This strategy can be treated as a case between the two existing directions: 1) it is lessrestrictive than requiring ideal environment partition according to distribution heterogeneity; and2) it also obtains theoretic guarantee which is missing in the case with no environment informationat all.
Revised: This strategy can be treated as a case between the two existing directions: 1) it is lessrestrictive than requiring ideal environment partition according to distribution heterogeneity; and2) it also obtains theoretic guarantee which is missing in the case with no environment informationat all.

Original: Notably, such auxiliary information is often cheaply available for every input in practice[49, 47].
Revised: Notably, such auxiliary information is often cheaply available for every input in practice[48, 46].

Original: Examples include time index of the data in time series forecasting tasks [7, 30], locations(longitude and latitude) of collected satellite data in remote sensing [13, 37], and meta informationassociated with an instance [28].
Revised: Examples include time index of the data in time series forecasting tasks [6, 29], locations(longitude and latitude) of collected satellite data in remote sensing [12, 36], and meta informationassociated with an instance [27].

Original: Based on the additionally observed variables, we proceed to proposea framework to jointly learn environment partition and invariant representation in Section 5.
Revised: Based on the additionally observed variables, we proceed to proposea framework to jointly learn environment partition and invariant representation in Section 5.

Original: Under afairly general setting, we derive both sufficient and necessary conditions for our framework to identifyinvariant features.
Revised: Under afairly general setting, we derive both sufficient and necessary conditions for our framework to identifyinvariant features.

Original: Extensions to other settings are discussed in Section 6 and Appendix B.
Revised: Extensions to other settings are discussed in Section 6 and Appendix B.

Original: Finally, Section 7 presents experimental analysis on both synthetic and real world datasets.
Revised: Finally, Section 7 presents experimental analysis on both synthetic and real world datasets.

Original: Theproposed framework achieves an improved performance over existing methods like EIIL and HRM,and has a comparable performance to IRM with ground-truth environment partition.
Revised: Theproposed framework achieves an improved performance over existing methods like EIIL and HRM,and has a comparable performance to IRM with ground-truth environment partition.

Original:  Related Work
Revised:  Related Work

Original: Invariant learning methods can be interpreted as robust to certain interventions or heterogeneity froma causal perspective [33].
Revised: Invariant learning methods can be interpreted as robust to certain interventions or heterogeneity froma causal perspective [32].

Original: On the other hand, heterogeneous data have also placed challenges andpotential benefits in causal inference tasks.
Revised: On the other hand, heterogeneous data have also placed challenges andpotential benefits in causal inference tasks.

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

Original: [17] consider heterogeneousdatasets to identify variables with changing local mechanism and further the causal directions.
Revised: [16] consider heterogeneousdatasets to identify variables with changing local mechanism and further the causal directions.

Original: Hereconfounders are assumed as functions of the domain index or time.
Revised: Hereconfounders are assumed as functions of the domain index or time.

Original: Tillman and Spirtes [45], Dankset al.
Revised: Tillman and Spirtes [44], Dankset al.

Original: [10], Huang et al.
Revised: [9], Huang et al.

Original: [16] study the multiple dataset setting with non-identical sets of variables forcausal discovery, where the datasets may have different exogenous noise distributions and only partof the variables is present in each dataset.
Revised: [15] study the multiple dataset setting with non-identical sets of variables forcausal discovery, where the datasets may have different exogenous noise distributions and only partof the variables is present in each dataset.

Original: Besides the invariant learning methods described in Section 1, another large class of methods forgeneralizing beyond training data is distributionally robust optimization (DRO) [6, 40, 23, 12, 11, 51].
Revised: Besides the invariant learning methods described in Section 1, another large class of methods forgeneralizing beyond training data is distributionally robust optimization (DRO) [5, 39, 22, 11, 10, 50].

Original: DRO methods propose to optimize the worst-case risk over a set of distributions close to the trainingdistribution.
Revised: DRO methods propose to optimize the worst-case risk over a set of distributions close to the trainingdistribution.

Original: A notable instance of DRO methods is group DRO, which optimizes the worst-case lossover groups in the training data [38].
Revised: A notable instance of DRO methods is group DRO, which optimizes the worst-case lossover groups in the training data [37].

Original: As with most invariant learning methods, group DRO generallyrequires a partition of groups obtained according to group annotation and label of each data sample.
Revised: As with most invariant learning methods, group DRO generallyrequires a partition of groups obtained according to group annotation and label of each data sample.

Original: Obtaining group annotations, however, can be costly or even impossible in practice.
Revised: Obtaining group annotations, however, can be costly or even impossible in practice.

Original: We may not knowa priori the inherent spurious features associated with data.
Revised: We may not knowa priori the inherent spurious features associated with data.

Original: Thus, a number of works have proposedto infer group partition or directly identify the minority group by looking at features produced by biased models, e.g., [41, 2, 9, 25].
Revised: Thus, a number of works have proposedto infer group partition or directly identify the minority group by looking at features produced by biased models, e.g., [40, 1, 8, 24].

Original: Notice that sometimes a small number of labelled annotations arestill needed to find a proper trained model, which is different from the setting considered in this work.
Revised: Notice that sometimes a small number of labelled annotations arestill needed to find a proper trained model, which is different from the setting considered in this work.

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

Original: [25] upweights samples with high losses from the initial ERM model and relieson a small validation set of annotated data to tune parameters.
Revised: [24] upweights samples with high losses from the initial ERM model and relieson a small validation set of annotated data to tune parameters.

Original: Additionally, Nam et al.
Revised: Additionally, Nam et al.

Original: [31], Sanhet al.
Revised: [30], Sanhet al.

Original: [39] try to obtain a more robust model by boosting from the wrongly specified samples, basedon the observation that models with limited capacity tend to learn spurious correlations or shortcuts.
Revised: [38] try to obtain a more robust model by boosting from the wrongly specified samples, basedon the observation that models with limited capacity tend to learn spurious correlations or shortcuts.

Original: Unfortunately, as we will show in Section 4, it is in theory impossible to distinguish spurious orinvariant features from only the observed data consisting of input-label pairs.
Revised: Unfortunately, as we will show in Section 4, it is in theory impossible to distinguish spurious orinvariant features from only the observed data consisting of input-label pairs.

Original:  Preliminaries
Revised:  Preliminaries

Original: Throughout this paper, upper-cased letters such as X ∈ R d and Y ∈ R denote random variables,and lower-cased letters such as x and y denote deterministic instances.
Revised: Throughout this paper, upper-cased letters such as X ∈ R d and Y ∈ R denote random variables,and lower-cased letters such as x and y denote deterministic instances.

Original: Suppose there exist invariantfeatures X v ∈ R d v and spurious (non-invariant) features X s ∈ R d s .
Revised: Suppose there exist invariantfeatures X v ∈ R d v and spurious (non-invariant) features X s ∈ R d s .

Original: We observe a scrambled version
Revised: We observe a scrambled version

Original: X =
Revised: X =

Original: q ( X v , X s
Revised: q ( X v , X s

Original: ) ∈ R d with q being an injective function.
Revised: ) ∈ R d with q being an injective function.

Original: Consider the dataset D :=
Revised: Consider the dataset D :=

Original: { ( x i , y i ) } ni =1where data are collected from multiple environments e ∈ E supp .
Revised: { ( x i , y i ) } ni =1where data are collected from multiple environments e ∈ E supp .

Original: We use superscript e to indicate theenvironment index with a variable, e.g., X e and Y e , when it is useful to make the index explicit.
Revised: We use superscript e to indicate theenvironment index with a variable, e.g., X e and Y e , when it is useful to make the index explicit.

Original: Notice that the true environment index e is not provided in training, unless otherwise stated.
Revised: Notice that the true environment index e is not provided in training, unless otherwise stated.

Original: In this work, we assume an underlying structural causal model (SCM) governing the data generationprocess [32, 34].
Revised: In this work, we assume an underlying structural causal model (SCM) governing the data generationprocess [31, 33].

Original: An SCM considers a set of variables associated with vertices of a directed acyclicgraph, where directed edges represent direct causation.
Revised: An SCM considers a set of variables associated with vertices of a directed acyclicgraph, where directed edges represent direct causation.

Original: Each variable is obtained as a result ofan assignment of a deterministic function depending on the parental variables in the graph and anexogenous random variable.
Revised: Each variable is obtained as a result ofan assignment of a deterministic function depending on the parental variables in the graph and anexogenous random variable.

Original: As an example, we present the following data generation process thatare also considered in [4, 5]:
Revised: As an example, we present the following data generation process thatare also considered in [3, 4]:

Original: [Equation]
Revised: [Equation]

Original: where g v  denote some non-degenerate deterministic functions and ϵ v is independent noise variables.
Revised: where g v , g e denote some non-degenerate deterministic functions and ϵ v , ϵ es are independentnoise variables.

Original: We observe data { X e , Y e } Ee =1 from multiple environments, each with probability P ( X e , Y e ) .
Revised: We observe data { X e , Y e } Ee =1 from multiple environments, each with probability P ( X e , Y e ) .

Original: Thenthe marginal distribution of the mixed data can be represented as P ( X , Y ) =
Revised: Then the marginal distribution of the mixed data can be represented as P ( X , Y ) =

Original: (cid:80) Ee =1 α e P ( X e , Y e )for some α e > 0 and (cid:80) Ee =1 α e = 1 .
Revised: (cid:80) Ee =1 α e P ( X e , Y e ) for some α e > 0 and (cid:80) Ee =1 α e = 1 .

Original: We consider the model as the composition of a feature extractor Φ and a classifier (or predictor)
Revised: We consider the model as the composition of a feature extractor Φ and a classifier (or predictor)

Original: f ωthat is parameterized by ω .
Revised: f ωthat is parameterized by ω .

Original: To seek for out-of-distribution generalization ability, we hope that Φmerely encodes the information of invariant features.
Revised: To seek for out-of-distribution generalization ability, we hope that Φmerely encodes the information of invariant features.

Original: Let ℓ ( · , · ) denote a loss function such as crossentropy loss or squared error.
Revised: Let ℓ ( · , · ) denote a loss function such as crossentropy loss or squared error.

Original: Our goal is to learn a robust and invariant model that minimizes thefollowing objective over all the considered environments e ∈ E supp :
Revised: Our goal is to learn a robust and invariant model that minimizes thefollowing objective over all the considered environments e ∈ E supp :

Original: [Equation]
Revised: [Equation]

Original: While there is no single, common definition of invariant feature in the literature, we treat the directcausal parents of Y as invariant features, as in [33, 5].
Revised: While there is no single, common definition of invariant feature in the literature, we treat the directcausal parents of Y as invariant features, as in [32, 4].

Original: This is because the conditional probability ofthe outcome Y given its parental variables or direct causes remains unchanged with any intervention(not on Y ).
Revised: This is because the conditional probability ofthe outcome Y given its parental variables or direct causes remains unchanged with any intervention(not on Y ).

Original: In contrast, if a non-parental variable is included in the conditioned variables, thenthe conditional probability could change after some intervention.
Revised: In contrast, if a non-parental variable is included in the conditioned variables, thenthe conditional probability could change after some intervention.

Original: As such, the underlying SCMgoverning data generation determines the desired invariance.
Revised: As such, the underlying SCMgoverning data generation determines the desired invariance.

Original:  Impossibility Result
Revised:  Impossibility Result

Original: The first question we ask is whether it is possible to learn invariant models from the heterogeneousdata of multiple environments with unknown environmental indexes.
Revised: The first question we ask is whether it is possible to learn invariant models from the heterogeneousdata of multiple environments with unknown environmental indexes.

Original: The following toy exampleimmediately gives a negative answer.
Revised: The following toy exampleimmediately gives a negative answer.

Original: Example 1.
Revised: Example 1.

Original: Consider that there are two environments e ∈ { 1 , 2 } with α
Revised: Consider that there are two environments e ∈ { 1 , 2 } with α

Original: 1 = α 2 = 0 .
Revised: 1 = α 2 = 0 .

Original: 5 , and weobserve binary-valued features X 1 , X 2 and label Y .
Revised: 5 , and weobserve binary-valued features X 1 , X 2 and label Y .

Original: Suppose the joint distribution of mixed environments is given in Eqn.
Revised: Suppose the joint distribution of mixed environments is given in Eqn.

Original: A learning algorithm that tries to learn invariant feature based on thisdataset would result in a model that (deterministically) depends on either X 1 , X 2 , or both to predict Y .
Revised: A learning algorithm that tries to learn invariant feature based on thisdataset would result in a model that (deterministically) depends on either X 1 , X 2 , or both to predict Y .

Original:  
Revised:  

Original: Y = 0 , w.p. 0 .
Revised: Y = 0 , w.p. 0 .

Original: Y = 1 , w.p. 0 .
Revised: Y = 1 , w.p. 0 .

Original: X 

Revised: X

Original: = X 1 = Y, w.p. 0 .
Revised: = X 1 = Y, w.p. 0 .

Original: X 2 ̸ =
Revised: X 2 ̸ =

Original: [Equation]
Revised: [Equation]

Original: 1 = Y
Revised: 1 = Y

Original: [Equation]
Revised: [Equation]

Original: w.p. 0 .
Revised: w.p. 0 .

Original: [Equation]
Revised: [Equation]

Original: w
Revised: w

Original: [Equation]
Revised: [Equation]

Original: p
Revised: p

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

Original: w
Revised: w

Original: [Equation]
Revised: [Equation]

Original: p
Revised: p

Original: [Equation]
Revised: [Equation]

Original: However, as discussed in last section, whether thelearned features are invariant is determined by theunderlying data generation process.
Revised: However, as discussed in last section, whether thelearned features are invariant is determined by theunderlying data generation process.

Original: We now presenttwo possible data generation processes that generate the same distribution yet have different invariantfeatures:
Revised: We now presenttwo possible data generation processes that generate the same distribution yet have different invariantfeatures:

Original: • X 1 is the invariant feature while X 2 is spurious,
Revised: • X 1 is the invariant feature while X 2 is spurious,

Original: where
Revised: where

Original: p e =1 s
Revised: p e =1 s

Original: = 0 .
Revised: = 0 .

Original: 7 and p e =2 s = 0 .
Revised: 7 and p e =2 s = 0 .

Original: [Equation]
Revised: [Equation]

Original: • X 2 is the invariant feature while X 1 is spurious
Revised: • X 2 is the invariant feature while X 1 is spurious

Original: where p e =1 s = 0 .
Revised: where p e =1 s = 0 .

Original: 25 and p e =2 s = 0 .
Revised: 25 and p e =2 s = 0 .

Original: [Equation]
Revised: [Equation]

Original: ∼ Bernoulli
Revised: ∼ Bernoulli

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

Original: The joint distribution of mixed environments for each case is consistent with Eqn.
Revised: The joint distribution of mixed environments for each case is consistent with Eqn.

Original: Thus, from thejoint distribution P ( X 1 , X 2 , Y ) , a learned model only depends on either X 1 , X 2 or both, and failsto generalize for at least one of the two data generation processes.
Revised: Thus, from thejoint distribution P ( X 1 , X 2 , Y ) , a learned model only depends on either X 1 , X 2 or both, and failsto generalize for at least one of the two data generation processes.

Original: Thus, the true invariant featurescannot be identified.
Revised: Thus, the true invariant featurescannot be identified.

Original: On the other hand, notice that when the partition is given, one can verify that
Revised: On the other hand, notice that when the partition is given, one can verify that

Original: IRM would correctly identify X 1 (resp.
Revised: IRM would correctly identify X 1 (resp.

Original: X 2 ) as the invariant feature for the first (resp.
Revised: X 2 ) as the invariant feature for the first (resp.

Original: second) senario.
Revised: second) senario.

Original: Moreover, such an example is not rare.
Revised: Moreover, such an example is not rare.

Original: Indeed, for any data distribution P ( X , Y ) that can begenerated by some data generation process like (1) (which can be replaced by other forms of SCM),we can find a different data generation process that induces the same distribution yet has differentinvariant/spurious features.
Revised: Indeed, for any data distribution P ( X , Y ) that can begenerated by some data generation process like (1) (which can be replaced by other forms of SCM),we can find a different data generation process that induces the same distribution yet has differentinvariant/spurious features.

Original: Similarly, it is impossible to find the minority group that is determined byspurious features, based on only the data P ( X , Y ) .
Revised: Similarly, it is impossible to find the minority group that is determined byspurious features, based on only the data P ( X , Y ) .

Original: This result is formally summarized in Theorem 1.
Revised: This result is formally summarized in Theorem 1.

Original: Theorem 1.
Revised: Theorem 1.

Original: Let X v and X s be respectively invariant and spurious features, with label Y .
Revised: Let X v and X s be respectively invariant and spurious features, with label Y .

Original: For thejoint distribution P ( X , Y ) consisting of data from multiple environments with P ( X s , Y ) ̸ = 0 , wecan find X ′ v and X ′ s as invariant and spurious features
Revised: For thejoint distribution P ( X , Y ) consisting of data from multiple environments with P ( X s , Y ) ̸ = 0 , wecan find X ′ v and X ′ s as invariant and spurious features

Original: so that: (a) X ′ v ̸ =
Revised: so that: (a) X ′ v ̸ =

Original: X v and X ′ v ∩ X s ̸ = ∅ ;(b) ( X ′ v , X ′ s ) together with some noise variables generate the same distribution P ( X , Y ) .
Revised: X v and X ′ v ∩ X s ̸ = ∅ ;(b) ( X ′ v , X ′ s ) together with some noise variables generate the same distribution P ( X , Y ) .

Original: A proof is provided in Appendix A.1.
Revised: A proof is provided in Appendix A.1.

Original: Under this circumstance, we have to introduce additionalassumptions/conditions or certain “inductive bias” to identify the invariant features.
Revised: Under this circumstance, we have to introduce additionalassumptions/conditions or certain “inductive bias” to identify the invariant features.

Original: The latter maycome implicitly from the proposed algorithm, model, and/or the data.
Revised: The latter maycome implicitly from the proposed algorithm, model, and/or the data.

Original: We believe that this is thereason that makes existing methods achieve improved empirical performance in their consideredscenarios.
Revised: We believe that this is thereason that makes existing methods achieve improved empirical performance in their consideredscenarios.

Original: For example, in [26, 27], the discrepancy of spurious features among clusters are expectedto be larger than that of causal features, and in [9, 31, 39, 25], the ERM model of the first stage shouldheavily or fully rely on the spurious feature.
Revised: For example, in [25, 26], the discrepancy of spurious features among clusters are expectedto be larger than that of causal features, and in [8, 30, 38, 24], the ERM model of the first stage shouldheavily or fully rely on the spurious feature.

Original: However, these inductive biases may not be alwaysguaranteed.
Revised: However, these inductive biases may not be alwaysguaranteed.

Original: Thus, as suggested by Theorem 1, the role of inductive biases shall be discussed moreexplicitly when learning invariant models from heterogeneous data without environment indexes.
Revised: Thus, as suggested by Theorem 1, the role of inductive biases shall be discussed moreexplicitly when learning invariant models from heterogeneous data without environment indexes.

Original: Alternatively, in this work, we consider that there exist additionally observed variable Z with the data,which has been often considered in the literature [17, 18, 20, 49, 47].
Revised: Alternatively, in this work, we consider that there exist additionally observed variable Z with the data,which has been often considered in the literature [16, 17, 19, 48, 46].

Original: The variable Z could be, forexample, the time index in a time series, some kind of class labels, concurrently observed variables,or human annotations on potential unmeasured variable [43].
Revised: The variable Z could be, forexample, the time index in a time series, some kind of class labels, concurrently observed variables,or human annotations on potential unmeasured variable [42].

Original: It is worth noting that Xie et al.
Revised: It is worth noting that Xie et al.

Original: [49]also take advantage of auxiliary information to help improve OOD performance in a semi-supervisedmanner.
Revised: [48]also take advantage of auxiliary information to help improve OOD performance in a semi-supervisedmanner.

Original: However, they consider a different setting with access to unlabeled (out-of-domain) test datatogether with additional unlabeled training data.
Revised: However, they consider a different setting with access to unlabeled (out-of-domain) test datatogether with additional unlabeled training data.

Original: In the next section, we show how to utilize Z tolearn invariant models from the mixed dataset.
Revised: In the next section, we show how to utilize Z tolearn invariant models from the mixed dataset.

Original:  ZIN: Learning Invariance with Additional Auxiliary Information
Revised:  ZIN: Learning Invariance with Additional Auxiliary Information

Original: This section proposes our framework to learn invariant models.
Revised: This section proposes our framework to learn invariant models.

Original: We derive conditions for invarianceidentification and show these conditions are both sufficient and necessary in a fairly general setting.
Revised: We derive conditions for invarianceidentification and show these conditions are both sufficient and necessary in a fairly general setting.

Original: 5.1 Method
Revised: 5.1 Method

Original: We consider that there exists additional auxiliary information Z ∈ R d z in company with the data( X , Y ) .
Revised: We consider that there exists additional auxiliary information Z ∈ R d z in company with the data( X , Y ) .

Original: In the following, we propose ZIN, auxiliary information Z for environmental INference , forinvariant learning from the mixture dataset D without environment indexes.
Revised: In the following, we propose ZIN, auxiliary information Z for environmental INference , forinvariant learning from the mixture dataset D without environment indexes:

Original: We aim to learn a function ρ ( · ) : R d z −→ R K that softly assigns a sample to K environments. Here
Revised: [Equation]

Original:  K  is a pre-specified number (a hyper-parameter) and its choice will be empirically investigatedin our experiments. Let ρ ( k ) ( · ) denote the k -th entry of ρ ( · ) . Then we have ρ ( Z ) ∈ [0 , 1] K with180 (cid:80) k ρ ( k )  ( Z ) = 1 .
Revised: invariance penalty (3) where K ∈ Z is a pre-specified hyper-parameter, ρ ( · ) : R d z −→ R K with ρ ( k ) denoting its k -thentry, R ( ω, Φ) = 1 n (cid:80) ni =1 ℓ ( f w (Φ( x i ))

Original: Denote the ERM loss as R ( ω, Φ) = 1 n (cid:80) ni =1 ℓ ( f w (Φ( x i )) , y i ) and the loss in thek -th inferred environment as R ρ ( k ) ( ω, Φ) = 1 n (cid:80) ni
Revised: y i ) , R ρ ( k ) ( ω, Φ) = 1 n (cid:80) ni

Original: =1 ρ ( k ) ( z i ) ℓ ( f w (Φ( x i )) , y i ) .
Revised: =1 ρ ( k ) ( z i ) ℓ ( f w (Φ( x i )) , y i ) , and

Original: Suppose that the environments have been given according to a fixed ρ ( · )  . 
Revised: R ρ ( k ) ( ω k , Φ) is similarly defined by replacing ω with ω k in R ρ ( k ) ( ω, Φ) . Function ρ ( Z ) ∈

Original: Recall that IRM [5] learnsan invariant representation Φ , upon which there is a classifier ω that is simultaneously optimal in allenvironments. To measure the optimality of ω in environment k , we can fit an environment dependentclassifier ω k on the data from environment k . If ω kachieves a smaller loss than ω , it means ω is notoptimal in this environment.
Revised: [0 , 1] Kwith (cid:80) k

Original: We can further train a set of environment dependent classifiers { ω k } Kk =1(one for each environment), to measure whether ω is simultaneously optimal in all environments.
Revised: ρ ( k ) ( Z ) = 1 is used to softly partition the training data into K environments; the choice of K will be empirically investigated in our experiments.

Original: Thus, when ρ ( · ) is provided, our formulation to learn invariance is then given as follows: [Equation]
Revised: Given a partition ρ ( Z ) , we can then measurehow the model violates the invariance constraint due to the use of spurious information.

Original: [Table] [Equation] If Φ extracts spurious features that are unstable in the inferred environments, R ρ ( k ) ( ω, Φ) will belarger than R ρ ( k ) ( ω k , Φ) and it induces a non-zero penalty. Next, we consider how to learn the partition function ρ ( · ) .
Revised: In particular,we train K independent classifiers { f w k } and compare them with the classifier trained on the wholedataset (cf.

Original: A good partition function should generateenvironments among which the spurious feature exhibit instability.
Revised: the second term in Eqn.

Original: If Φ extracts spurious features,we expect to have a large penalty resulting from the environments. Thus, we seek for an environmentpartition that maximizes the invariance penalty.The overall framework is provided below: [Equation]
Revised: If Φ extracts spurious features that are unstable in theinferred environments, R ρ ( k ) ( ω, Φ) will be larger than R ρ ( k ) ( ω k , Φ) and induce a non-zero penalty.

Original: The idea of our method can be summarized as inferring an environment partition and then learninginvariant models based on the inferred environments.
Revised: The idea of our method can be summarized as inferring an environment partition and then learninginvariant models based on the inferred environments.

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

Original: [9] has used a similar intuitionand adopted a pre-trained biased model to estimate the environments.
Revised: [8] has used a similar intuitionand adopted a pre-trained biased model to estimate the environments.

Original: However, their two-stagemethod cannot be jointly optimized, and the environment partition relies on the given model and lacksa theoretic guarantee.
Revised: However, their two-stagemethod cannot be jointly optimized, and the environment partition relies on the given model and lacksa theoretic guarantee.

Original: In contrast, we will first present sufficient conditions on Z so that the proposedframework can provably identify the invariant features.
Revised: In contrast, we will first present sufficient conditions on Z so that the proposedframework can provably identify the invariant features.

Original: We further show that these conditions arealso necessary, and that violation of these conditions will lead to failure of invariance identification.
Revised: We further show that these conditions arealso necessary, and that violation of these conditions will lead to failure of invariance identification.

Original: Assigning independent weights to each data sample, which is adopted in EIIL [9], is unfortunately aninstance of such violations and may fail to identify invariant features.
Revised: Assigning independent weights to each data sample, which is adopted in EIIL [8], is unfortunately aninstance of such violations and may fail to identify invariant features.

Original: 5.2 Sufficient Conditions for Identifiability
Revised: 5.2 Sufficient Conditions for Identifiability

Original: In this section, we try to understand ZIN from a theoretic perspective.
Revised: In this section, we try to understand ZIN from a theoretic perspective.

Original: We start with a simple yetgeneral setting: X =
Revised: We start with a simple yetgeneral setting: X =

Original: [ X v ; X s ] (i.e., no scramble on the observation), Φ ∈ { 0 , 1 } d is an element-wisefeature selection mask, and f w is general non-linear function; the next section will present theoreticresults with linear feature transformations, the same setting as in [5].
Revised: [ X v ; X s ] (i.e., no scramble on the observation), Φ ∈ { 0 , 1 } d is an element-wisefeature selection mask, and f w is general non-linear function; the next section will present theoreticresults with linear feature transformations, the same setting as in [4].

Original: We also focus on classificationtasks with ℓ ( · , · ) being the cross-entropy loss, as there is an interesting connection to (conditional)
Revised: We also focus on classificationtasks with ℓ ( · , · ) being the cross-entropy loss, as there is an interesting connection to (conditional)

Original: Shannon entropy.
Revised: Shannon entropy.

Original: Extension to other loss functions is left to Appendix B.
Revised: Extension to other loss functions is left to Appendix B.

Original: In this setting, our goal is equivalent to learning the optimal feature mask that merely selects invariantfeatures, i.e., Φ v = [ 1 d v ; 0 d s ] .
Revised: In this setting, our goal is equivalent to learning the optimal feature mask that merely selects invariantfeatures, i.e., Φ v = [ 1 d v ; 0 d s ] .

Original: For a given feature mask Φ , the objective function is equal to
Revised: For a given feature mask Φ , the objective function is equal to

Original: [Equation]
Revised: [Equation]

Original: Then ZIN can correctly identify the invariant features, or solution to Problem 4 is equivalent toΦ v , if and only if ˆ L (Φ v ) <
Revised: Then ZIN can correctly identify the invariant features, or solution to Problem 3 is equivalent toΦ v , if and only if ˆ L (Φ v ) <

Original: ˆ L (Φ) for all Φ ̸ = Φ v .
Revised: ˆ L (Φ) for all Φ ̸ = Φ v .

Original: This observation will be used to establish our main theoretic result.
Revised: This observation will be used to establish ourmain theoretic result.

Original: With a little abuse of notation, we use H ( Y | X ′ ) to denote expected loss ofan optimal classifier over some X ′ and Y , and similarly H ( Y | ρ ( Z ) , X ′ ) to denote the minimumrisk (cid:80)
Revised: With a little abuse of notation, we use H ( Y | X ′ ) to denote expected loss ofan optimal classifier over some X ′ and Y , and similarly H ( Y | ρ ( Z ) , X ′ ) to denote the minimumrisk (cid:80)

Original: Kk =1 R ρ ( k ) ( ω k , Φ) for a given ρ ( Z ) .
Revised: Kk =1 R ρ ( k ) ( ω k , Φ) for a given ρ ( Z ) .

Original: With cross-entropy loss and when ρ ( z i ) gives exactlyone environment, i.e., ρ ( Z ) is one-hot, it can be verified that the optimal loss H ( ·|· ) coincides withconditional entropy.
Revised: With cross-entropy loss and when ρ ( z i ) gives exactlyone environment, i.e., ρ ( Z ) is one-hot, it can be verified that the optimal loss H ( ·|· ) coincides withconditional entropy.

Original: In the following we first state the assumptions for our identifiabitiliy result.
Revised: In the following we first state the assumptions for our identifiabitiliy result.

Original: Assumption 1.
Revised: Assumption 1.

Original: For a given feature mask Φ and any constant ϵ > 0 , there exists f ∈ F such that 221 E [ ℓ ( f (Φ( X )) , Y )]
Revised: For a given feature mask Φ and any constant ϵ > 0 , there exists f ∈ F such that 212 E [ ℓ ( f (Φ( X )) , Y )]

Original: ≤ H ( Y | Φ( X ))
Revised: ≤ H ( Y | Φ( X ))

Original: + ϵ.
Revised: + ϵ.

Original: Assumption 2.
Revised: Assumption 2.

Original: If two features are invariant w.r.t.
Revised: If two features are invariant w.r.t.

Original: an environment partition, then the concatenatedfeatures are also invariant.
Revised: an environment partition, then the concatenatedfeatures are also invariant.

Original: That is, for X 1 , X 2 ⊂ X and ρ ( · ) , if H ( Y | X 1 ) − H
Revised: That is, for X 1 , X 2 ⊂ X and ρ ( · ) , if H ( Y | X 1 ) − H

Original: ( Y | ρ ( Z ) , X 1 ) =
Revised: ( Y | ρ ( Z ) , X 1 ) =

Original: 0and H ( Y | X 2 ) − H
Revised: 0and H ( Y | X 2 ) − H

Original: ( Y | ρ ( Z ) , X 2 )
Revised: ( Y | ρ ( Z ) , X 2 )

Original: = 0 , we have H
Revised: = 0 , we have H

Original: ( Y | X 1 , X 2 ) − H ( Y | ρ ( Z ) , X 1 , X 2 ) = 0 .
Revised: ( Y | X 1 , X 2 ) − H ( Y | ρ ( Z ) , X 1 , X 2 ) = 0 .

Original: Assumption 3.
Revised: Assumption 3.

Original: If a feature representation violates the invariance constraint, adding another featurewould not make the penalty vanish.
Revised: If a feature representation violates the invariance constraint, adding another featurewould not make the penalty diminish.

Original: This implies that there exists a constant δ > 0 so that for spuriousfeature X 1 ⊂ X s and any other feature X 2 ⊂ X , H ( Y | X 1 , X 2 ) − H
Revised: This implies that there exists a constant δ > 0 so that for any X 1 , X 2 ⊂ X , H ( Y | X 1 , X 2 ) − H

Original: ( Y | ρ ( Z ) , X 1 , X 2 ) ≥δ ( H ( Y | X 1 ) − H
Revised: ( Y | ρ ( Z ) , X 1 , X 2 ) ≥ δ ( H ( Y | X 1 ) − H

Original: ( Y | ρ ( Z ) , X 1 )) .
Revised: ( Y | ρ ( Z ) , X 1 )) .

Original: Assumption 4.
Revised: Assumption 4.

Original: For any distinct features X 1 , X 2 , H ( Y | X 1 , X 2 ) ≤ H ( Y | X 1 ) − γ with fixed γ > 0 .
Revised: For any distant features X 1 , X 2 , H ( Y | X 1 , X 2 ) ≤ H ( Y | X 1 ) − γ with fixed γ > 0 .

Original: Assumption 1 is a common assumption that requires the function space F be rich enough such that,given Φ , there is a f ∈ F that can fit P ( Y | Φ( X ))
Revised: Assumption 1 is a common assumption that requires the function space F be rich enough such that,given Φ , there is a f ∈ F that can fit P ( Y | Φ( X ))

Original: well.
Revised: well.

Original: Assumption 4 indicates that any featurecontains some useful information w.r.t.
Revised: Assumption 4 indicates that any featurecontains some useful information w.r.t.

Original: Y , which cannot be explained by other features.
Revised: Y , which cannot be explained by other features.

Original: Otherwise,we can simply remove such a feature (e.g., by variable selection methods [46]), as it does not affectprediction.
Revised: Otherwise,we can simply remove such a feature (e.g., by variable selection methods [45]), as it does not affectprediction.

Original: We next present our sufficient conditions for ZIN to identify invariant features.
Revised: We next present our sufficient conditions for ZIN to identify invariant features.

Original: Condition 1 (Invariance Preserving Condition) .
Revised: Condition 1 (Invariance Preserving Condition) .

Original: Given invariant feature X  v and any function ρ ( · ) , itholds that H ( Y | X v , ρ ( Z ))
Revised: For any feature X iv ∈ X v and any function ρ ( · ) , itholds that H ( Y | X iv , ρ ( Z ))

Original: =
Revised: =

Original: H ( Y | X v ) .
Revised: H ( Y | X iv ) .

Original: Condition 2 (Non-invariance Distinguishing Condition) .
Revised: Condition 2 (Non-invariance Distinguishing Condition) .

Original: For any feature X ks ∈ X s , there exists afunction ρ ( · ) and a constant C > 0 such that H ( Y | X ks )
Revised: For any feature X ks ∈ X s , there exists afunction ρ ( · ) and a constant C > 0 such that H ( Y | X ks )

Original: − H ( Y | X ks , ρ ( Z )) ≥ C .
Revised: − H ( Y | X ks , ρ ( Z )) ≥ C .

Original: We remark that Condition 1 can be met if H ( Y | X v , Z )
Revised: We remark that Condition 1 can be met if H ( Y | X iv , Z )

Original: =
Revised: =

Original: H ( Y | X v )  (a proof is provided in Appendix A.7).
Revised: H ( Y | X iv ) for all X iv ∈ X v (a proofis provided in Appendix A.7).

Original: Condition 1 requires that invariant features should remain invariant w.r.t any environmentpartition induced by ρ ( Z ) .
Revised: Condition 1 requires that invariant features should remain invariantw.r.t any environment partition induced by ρ ( Z ) .

Original: Otherwise, if there exists a partition where an invariant feature becomesnon-invariant, then this feature would induce a positive penalty.
Revised: Otherwise, if there exists a partition where aninvariant feature becomes non-invariant, then this feature would induce a positive penalty.

Original: Condition 2 implies that for eachspurious feature, there exists at least one partition so that this feature is non-invariant in the splitenvironments.
Revised: Condition 2implies that for each spurious feature, there exists at least one partition so that this feature is noninvariant in the split environments.

Original: If a spurious feature does not incur any invariance penalty in all possible environmentpartitions, we can never distinguish it from true invariant features.
Revised: If a spurious feature does not incur any invariance penalty in allpossible environment partitions, we can never distinguish it from true invariant features.

Original: With these conditions, our mainresult follows and a proof is given in Appendix A.2.
Revised: With theseconditions, our main result follows.

Original: Theorem 2 (Identifiability of Invariant Features) .
Revised: Theorem 2 (Identifiability of Invariant Features) .

Original: With Assumptions 1, 3 and 4 and Conditions 1-2, ifϵ <
Revised: With Assumptions 1-4 and Conditions 1-2, ifϵ <

Original: Cγδ 4 γ +2 CδH ( Y ) and λ ∈
Revised: Cγδ 4 γ +2 CδH ( Y ) and λ ∈

Original: [ H ( Y )+1 / 2 δC δC − 4 ϵ − 12 , γ 4 ϵ − 12 ] , then we have ˆ L (Φ v ) <
Revised: [ H ( Y )+1 / 2 δC δC − 4 ϵ − 12 , γ 4 ϵ − 12 ] , then we have ˆ L (Φ v ) <

Original: ˆ L (Φ) for all Φ ̸ =
Revised: ˆ L (Φ) for all Φ ̸ =

Original: Φ v ,where H ( Y ) denotes the entropy of Y .
Revised: Φ v ,where H ( Y ) denotes the entropy of Y .

Original: Thus, the solution to Problem 4 identifies invariant features.
Revised: Thus, the solution to Problem 3 identifies invariant features.

Revised: Proof (Sketch).

Revised: Our proof consists of two steps.

Revised: First, show that the invariance penalty of Φ v willbe upper bounded by ϵ according to Assumption 1, and that a feature mask selecting any spuriousfeature would induce another non-zero penalty depending on the constant C in Condition 2.

Original:  5. Necessary Conditions for Identifiability
Revised: Witha sufficiently large penalty weight λ , this feature mask will induce a larger loss ˆ L (Φ) than ˆ L (Φ v ) . Second, if a feature mask selects only part (not all) of the invariant features, it also results in a losslarger than ˆ L (Φ v ) according to Assumption 4. Please see Appendix A.2 for a complete version. 5. Necessary Conditions for Identifiability

Original: Conditions 1 and 2 may appear rather strong at first glance.
Revised: Conditions 1 and 2 may appear rather strong at first glance.

Original: In this section, we prove that they arealso necessary conditions for ZIN to identify invariant features.
Revised: In this section, we prove that they arealso necessary conditions for ZIN to identify invariant features.

Original: Specifically, Proposition 1 shows thatif Condition 1 is violated, then some invariant features will be excluded in the solution of Problem 4;and in Proposition 2, violation of Condition 2 renders some spurious features included in the solution.
Revised: Specifically, Proposition 1 shows thatif Condition 1 is violated, then some invariant features will be excluded in the solution of Problem 3;and in Proposition 2, violation of Condition 2 renders some spurious features included in the solution.

Original: Proposition 1.
Revised: Proposition 1.

Original: With Assumptions 1, 3 and 4, if Condition 1 is violated, i.e., there exists  ρ ( · ) sothat H ( y | X v ) −H
Revised: If Condition 1 is violated, i.e., there exists X iv ∈ X v and ρ ( · ) so that H ( y | X iv ) −

Original:  ( y | X v , ρ ( Z )) ≥ C ′ > 0 , then there exists a feature mask Φ ′ ̸
Revised: H ( y | X iv , ρ ( Z )) ≥ C ′ > 0 , then there exists a feature mask Φ ′ ̸

Original: = Φ v withˆ L (Φ v ) >
Revised: = Φ v with ˆ L (Φ v ) >

Original: ˆ L (Φ ′ ) .
Revised: ˆ L (Φ ′ ) .

Original: A proof is provided in Appendix A.3, which is similar to the first step of the proof of Theorem 2.
Revised: A proof is provided in Appendix A.3, which is similar to the first step of the proof of Theorem 2.

Original: Aquestion is how Condition 1 could be violated.
Revised: Aquestion is how Condition 1 could be violated.

Original: On the other hand, while we introduce Z as additionalvariables, it is also interesting to know whether we can obtain a valid partition based on only thetraining data ( X , Y ) .
Revised: On the other hand, while we introduce Z as additionalvariables, it is also interesting to know whether we can obtain a valid partition based on only thetraining data ( X , Y ) .

Original: Below we provide examples regarding these questions.
Revised: Below we provide examples regarding these questions.

Original: Corollary 1.
Revised: Corollary 1.

Original: Condition 1 is violated for the following cases: if there is a function ρ ( · ) and an injectivefunction h ( · ) so that (a) ρ ( Z ) =
Revised: Condition 1 is violated for the following cases: if there is a function ρ ( · ) and an injectivefunction h ( · ) so that (a) ρ ( Z ) =

Original: h ( Y ) , (b) ρ ( Z ) =
Revised: h ( Y ) , (b) ρ ( Z ) =

Original: h ( X , Y ) , or (c) ρ ( Z ) =
Revised: h ( X , Y ) , and (c) ρ ( Z ) =

Original: h (Index( X , Y )) .
Revised: h (Index( X , Y )) .

Original: A proof is given in Appendix A.4.
Revised: A proof is given in Appendix A.4.

Original: As h (Index( X , Y )) can be treated as learning independentweights for each sample based on { ( x i , y i ) } i , this case includes EIIL [9] as an instance.
Revised: As h (Index( X , Y )) can be treated as learning independentweights for each sample based on { ( x i , y i ) } i , this case includes EIIL [8] as an instance.

Original: Proposition 2.
Revised: Proposition 2.

Original: With Assumptions 1- 4, if Condition 2 is violated, i.e., there exists a feature X ks ∈ X ssuch that H ( Y | X ks ) 
Revised: If Condition 2 is violated, i.e., there exists a feature X ks ∈ X s such that H ( Y | X ks ) −

Original: − H ( Y | X ks , ρ ( Z ))
Revised: H ( Y | X ks , ρ ( Z ))

Original: = 

Revised: =

Original: for any ρ ( · ) , then there exists a feature mask Φ ′ ̸
Revised: for any ρ ( · ) , then there exists a feature mask Φ ′ ̸

Original: =
Revised: =

Original: Φ vwith ˆ L (Φ v ) >
Revised: Φ v with ˆ L (Φ v ) >

Original: ˆ L (Φ ′ ) .
Revised: ˆ L (Φ ′ ) .

Original: Since there exists a spurious feature X ks that is “invariant” in all possible environment partition,adding this feature to Φ v does not induce any penalty but increases the prediction power.
Revised: Since there exists a spurious feature X ks that is “invariant” in all possible environment partition,adding this feature to Φ v does not induce any penalty but increases the prediction power.

Original: Thus,such a feature mask can achieve a smaller loss than ˆ L (Φ v ) .
Revised: Thus,such a feature mask can achieve a smaller loss than ˆ L (Φ v ) .

Original: For completeness, we provide a proofin Appendix A.5.
Revised: For completeness, we provide a proofin Appendix A.5.

Original: Proposition 2 again indicates in theory that Z should be sufficiently diverse andinformative so that each spurious feature can be recognized.
Revised: Proposition 2 again indicates in theory that Z should be sufficiently diverse andinformative so that each spurious feature can be recognized.

Original:  Beyond Feature Selection: Linear Feature Transformations
Revised:  Beyond Feature Selection: Linear Feature Transformations

Original: In Section 5, both the invariant and spurious features can be directly observed and we focus on theability of ZIN to identify the invariant features.
Revised: In Section 5, both the invariant and spurious features can be directly observed and we focus on theability of ZIN to identify the invariant features.

Original: In this section, we extend feature selection to featurelearning, and show that ZIN is able to learn the invariant features given a scrambled observation in alinear form following [5, 36].
Revised: In this section, we extend feature selection to featurelearning, and show that ZIN is able to learn the invariant features given a scrambled observation in alinear form following [4, 35].

Original: Specifically, we consider the same data generation process as [5]:
Revised: Specifically, we consider the same data generation process as [4]:

Original: [Equation]
Revised: [Equation]

Original: where β ∈ R d v and W ∈ R
Revised: where β ∈ R d v and W ∈ R

Original: d × ( d v + d s ) .
Revised: d × ( d v + d s ) .

Original: We assume that there exists ˜ W ∈ R d v × d so that˜ W ( W [ x v ; x s ]) =
Revised: We assume that there exists ˜ W ∈ R d v × d so that˜ W ( W [ x v ; x s ]) =

Original: x v for all x v and x s .
Revised: x v for all x v and x s .

Original: Both the feature extractor and predictor take a linearform, i.e., Φ takes values in R d × d and ω ∈ R d .
Revised: Both the feature extractor and predictor take a linearform, i.e., Φ takes values in R d × d and ω ∈ R d .

Original: The prediction for X is ω ◦
Revised: The prediction for X is ω ◦

Original: Φ( X ) =
Revised: Φ( X ) =

Original: (Φ X ) T ω .
Revised: (Φ X ) T ω .

Original: A major difficulty in this setting lies in how to characterize the effect of an invariant/spurious featurein a quantitive way: the feature extractor Φ may extract an arbitrarily small portion of spuriousinformation.
Revised: A major difficulty in this setting lies in how to characterize the effect of an invariant/spurious featurein a quantitive way: the feature extractor Φ may extract an arbitrarily small portion of spuriousinformation.

Original: Following [5, 36], we consider a constrained form of Problem 4 for theoretic analysis:
Revised: Following [4, 35], we consider a constrained form of Problem 3 for theoretic analysis:

Original: [Equation]
Revised: [Equation]

Original: As in Conditions 1 and 2, the auxiliary information should also be sufficiently informative so that theinferred environments can be diverse enough but also maintain the underlying invariance.
Revised: As in Conditions 1 and 2, the auxiliary information should also be sufficiently informative so that theinferred environments can be diverse enough but also maintain the underlying invariance.

Original: This is inaccordance to existing conditions for identifiability in the linear case [5, 36].
Revised: This is inaccordance to existing conditions for identifiability in the linear case [4, 35].

Original: In this paper, we utilizesuch a condition, linear general position condition, from [5].
Revised: In this paper, we utilizesuch a condition, linear general position condition, from [4].

Original: For our analysis, we take squared erroras our loss function and consider that ρ ( · ) partitions the environments in a hard manner, i.e., eachdata sample would be assigned to exactly one environment.
Revised: For our analysis, we take squared erroras our loss function and consider that ρ ( · ) partitions the environments in a hard manner, i.e., eachdata sample would be assigned to exactly one environment.

Original: We also assume that the environmentsare non-degenerate, i.e., each inferred environment contains some data samples; otherwise, we cansimply remove such an environment.
Revised: We also assume that the environmentsare non-degenerate, i.e., each inferred environment contains some data samples; otherwise, we cansimply remove such an environment.

Original: Our identifiability result for the linear case then follows.
Revised: Our identifiability result for the linear case then follows.

Original: Theorem 3.
Revised: Theorem 3.

Original: Assume that there exists ρ ( · ) such that the generated environments, denoted as { X k } Kk =1 ,lie in linear general position of degree r , i.e., K > d − r + d/r for some r ∈ N and for all non-zerox ∈ R d : dim (cid:0) span (cid:0)(cid:8)
Revised: Assume that there exists ρ ( · ) such that the generated environments, denoted as { X k } Kk =1 ,lie in linear general position of degree r , i.e., K > d − r + d/r for some r ∈ N and for all non-zerox ∈ R d : dim (cid:0) span (cid:0)(cid:8)

Original: E X k (cid:2)
Revised: E X k (cid:2)

Original: X k X k ⊺ (cid:3) x − E X k ,ϵ v (cid:2)
Revised: X k X k ⊺ (cid:3) x − E X k ,ϵ v (cid:2)

Original: X k ϵ v (cid:3)(cid:9) k
Revised: X k ϵ v (cid:3)(cid:9) k

Original: (cid:1)(cid:1) >
Revised: (cid:1)(cid:1) >

Original: d − r.
Revised: d − r.

Original: If Φ ∈ R d × d has rankr > 0 , then Problem 6 results in the desired invariant predictor.
Revised: If Φ ∈ R d × d has rankr > 0 , then Problem 5 results in the desired invariant predictor.

Original: A proof is given in Appendix A.6.
Revised: A proof is given in Appendix A.6.

Original: We can also replace the above assumption by the non-degenerateconditions in Eqn.
Revised: We can also replace the above assumption by the non-degenerateconditions in Eqn.

Original: (9) and (10) of [36], and a similar identifiability result follows.
Revised: (9) and (10) of [35], and a similar identifiability result follows.

Original:  Experiments
Revised:  Experiments

Original: This section empirically verifies our theoretic analysis and the effectiveness of ZIN on both syntheticand real world datasets.
Revised: This section empirically verifies our theoretic analysis and the effectiveness of ZIN on both syntheticand real world datasets.

Original: We compare ZIN with several existing methods: ERM, IRM, [5], GroupDRO[38], EIIL [9], HRM [26], and LfF [31].
Revised: We compare ZIN with several existing methods: ERM, IRM [4], EIIL [8],and HRM [25], and LfF [30].

Original: We provide the ground-truth partitionto IRM and group
Revised: For IRM, we provide the ground-truth partition and its performanceserves as an upper bound.

Original: DRO. LfF tries to learn a robust model by directly boosting from wrongly specified samples ofshallow neural networks.
Revised: LfF tries to learn a robust model by directly boosting from wronglyspecified samples of shallow neural networks.

Original: Notice that LfF only works for classification tasks.[Table]
Revised: Notice that LfF only works for classification tasks.

Original: We implement both ρ and f in ZIN as a two-layer MLP with 32 hidden units.
Revised: We implement both ρ and f in ZIN as a two-layer MLP with 32 hidden units.

Original: The number of inferredenvironments K is set to be 2 as default.
Revised: The number of inferredenvironments K is set to be 2 as default.

Original: We implement Φ as a two-layer MLP for the syntheticdataset and house price prediction, ResNet-18 [14] for CelebA, and 1D CNN for Landcover.
Revised: We implement Φ as a two-layer MLP for the syntheticdataset and house price prediction, ResNet-18 [13] for CelebA, and 1D CNN for Landcover.

Original: Pleasesee Appendix D for more details.
Revised: Pleasesee Appendix C for more details.

Original: 7.1 Synthetic Dataset
Revised: 7.1 Synthetic Dataset

Original: Temporal Heterogeneity .
Revised: Temporal Heterogeneity .

Original: We consider temporal heterogeneity with distributional shift w.r.t.
Revised: We consider temporal heterogeneity with distributional shift w.r.t.

Original: time.
Revised: time.

Original: Let t ∈ [0 , 1] be time index and
Revised: Let t ∈ [0 , 1] be time index and

Original: X v ( t ) ∈ R the invariant feature.
Revised: X v ( t ) ∈ R the invariant feature.

Original: The data generation process is
Revised: The data generation process is

Original: X v ( t ) ∼ (cid:26) N (1 , σ 2 ) , w.p. 0 .

Original: 5 , N ( − 1 , σ 2 ) , w.p. 0 .

Original: 5 , Y ( t ) ∼ (cid:26) sign( X

Original: [Table]

Original: ) , w.

Original: [Table] [Table] , − [Table]
Revised: [Equation]

Original: ( X

Original: [Table]
Revised: [Table]

Original: ( t )) , w.p. p ′ v , X s ( t ) ∼ (cid:26) N

Original: ( Y ( t ) , σ 2 ) ,

Original: w.p. p s ( t ) , N ( − Y ( t ) , σ 2 ) , w.p. p ′ s ( t ) .

Original: where p v is a constant w.r.t.
Revised: where p v is a constant w.r.t.

Original: t , indicating a stable correlation between Y ( t ) and X v ( t ) , p ′ v = 1 − p v ,and p ′ s ( t )
Revised: t , indicating a stable correlation between Y ( t ) and X v ( t ) , p ′ v = 1 − p v ,and p ′ s ( t )

Original: = 

Revised: =

Original: − p s ( t ) .
Revised: − p s ( t ) .

Original: Notice here p s ( t ) varies with time t .
Revised: Notice here p s ( t ) varies with time t .

Original: A similar setting with two-dimensionalspatial variable is considered in Appendix F.1.
Revised: A similar setting with two-dimensionalspatial variable is considered in Appendix D.1.

Original: Our goal is to learn a model that purely relies on
Revised: Our goal is to learn a model that purely relies on

Original: X v .
Revised: X v .

Original: We simulate two heterogeneous environments along time, namely, { [0 , 0 . 5) , [0 . 5 , 1] } , and p s ( t )will be set differently.
Revised: We simulate two heterogeneous environments along time, namely, { [0 , 0 . 5) , [0 . 5 , 1] } , and p s ( t )will be set differently.

Original: We use tuple of p s ( t ) in the two environments to denote a simulated case.
Revised: We use tuple of p s ( t ) in the two environments to denote a simulated case.

Original: For example, (0 . 999 , 0 . 7) stands for p s ( t )
Revised: For example, (0 . 999 , 0 . 7) stands for p s ( t )

Original: = 0 .
Revised: = 0 .

Original: 999 , t ∈ [0 , 0 . 5) and p s ( t ) = 0 .
Revised: 999 , t ∈ [0 , 0 . 5) and p s ( t ) = 0 .

Original: 7 , t ∈ [0 . 5 , 1] .
Revised: 7 , t ∈ [0 . 5 , 1] .

Original: Weevaluate the performance on four distinct test environments with p s ∈ { 0 . 999 , 0 . 8 , 0 . 2 , 0 . 1 } and p vbeing constant.
Revised: Weevaluate the performance on four distinct test environments with p s ∈ { 0 . 999 , 0 . 8 , 0 . 2 , 0 . 1 } and p vbeing constant.

Original: We use t as the auxiliary information Z .

Original: More implementation details are given in Appendix D
Revised: More implementation details are given in Appendix C

Original: Results .
Revised: Results .

Original: Table 1 reports the test accuracy.
Revised: Table 1 reports the test accuracy.

Original: In all simulated settings, the worst accuracy of ERMis much lower than the mean accuracy, indicating that ERM tends to rely on spurious feature
Revised: In all simulated settings, the worst accuracy of ERMis much lower than the mean accuracy, indicating that ERM tends to rely on spurious feature

Original: X s .
Revised: X s .

Original: EIIL can improve the worst accuracy in some cases, e.g., when p s ( t ) =
Revised: EIIL can improve the worst accuracy in some cases, e.g., when p s ( t ) =

Original: 8) andp s ( t ) =
Revised: 8) andp s ( t ) =

Original: However, its performance is even worse than ERM for some other settings.
Revised: However, its performance is even worse than ERM for some other settings.

Original: This may be attributed to the first stage of EIIL, where the trained biased model is not guaranteed andmay learn both spurious and invariant features.
Revised: This may be attributed to the first stage of EIIL, where the trained biased model is not guaranteed andmay learn both spurious and invariant features.

Original: The proposed method ZIN improves the worst testaccuracy significantly.
Revised: The proposed method ZIN improves the worst testaccuracy significantly.

Original: For instance, when p s ( t ) = (0 . 999 , 0 . 7) and p v = 0 .
Revised: For instance, when p s ( t ) = (0 . 999 , 0 . 7) and p v = 0 .

Original: 9 , ZIN outperforms ERMand EIIL by over 28% and 65%, respectively.
Revised: 9 , ZIN outperforms ERMand EIIL by over 28% and 65%, respectively.

Original: Moreover, ZIN is very close to IRM that knows theground-truth environment partition, showing that ZIN can infer the environments effectively.
Revised: Moreover, ZIN is very close to IRM that knows theground-truth environment partition, showing that ZIN can infer the environments effectively.

Original: Finally,it seems that HRM does not learn a useful model for this classification task.
Revised: Finally,it seems that HRM does not learn a useful model for this classification task.

Original: Ablation Study .
Revised: Ablation Study .

Original: In Appendix F.2, we conduct ablation study to verify our theoretic results in Section 5by choosing different auxiliary variables.
Revised: In Appendix D.2, we conduct ablation study to verify our theoretic results in Section 5 by choosing different auxiliary variables.

Original: We also empirically study the choice of hyper-parameter K .
Revised: We also empirically study the choice of hyperparameter K .

Original: A notable result is that K has a relatively small impact, as shown in Figure 2.
Revised: A notable result is that K has a relatively small impact, as shown in Figure 2.

Original: 7.2 Real World Datasets
Revised: 7.2 Real World Datasets

Original: House Price Prediction .
Revised: House Price Prediction .

Original: This experiment considers a real world regression dataset of house salesprices from Kaggle (https://www.kaggle.com/c/house-prices-advanced-regression-techniques).
Revised: This experiment considers a real world regression dataset of house salesprices from Kaggle (https://www.kaggle.com/c/house-prices-advanced-regression-techniques).

Original: Thetarget variable is house price and each sample contains 17 predictive variables like the built year ofthe house, number of bedrooms, etc.
Revised: Thetarget variable is house price and each sample contains 17 predictive variables like the built year of thehouse, number of bedrooms, etc.

Original: The dataset is split according to the built year, with trainingdataset in period [1900 , 1950] and test dataset in period (1950 , 2000] .
Revised: The dataset is split according to the built year, with training datasetin period [1900 , 1950] and test dataset in period (1950 , 2000] .

Original: We normalize the prices of thehouses with the same built year, and our target is to predict the normalized price.

Original: We choose the builtyear as auxiliary information for ZIN.
Revised: We choose the built year as auxiliaryinformation for ZIN.

Original: As there is no well-defined ground-truth partition for IRM, wemanually split the training dataset equally into 5 segments with 10 -year range in each segment.
Revised: As there is no well-defined ground-truth partition for IRM, we manually splitthe training dataset equally into 5 segments with 10 -year range in each segment.

Original: [Table] [Table] Table 2 reports the mean squared error (MSE), and ZIN again outperforms other methods on testdataset.
Revised: Table 2 reports the mean squared error (MSE), and ZIN again outperforms other methods on testdataset.

Original: We also investigate whether this type of additional information Z , together with X , helps
Revised: We also investigate whether this type of additional information Z , together with X , helps

Original: EIIL and HRM.
Revised: EIIL and HRM.

Original: The corresponding results are marked with (+ Z ) in Table 2.
Revised: The corresponding results are marked with (+ Z ) in Table 2.

Original: We observe that theresults of EIIL and HRM with additional information Z are worse than those without using Z .
Revised: We observe that theresults of EIIL and HRM with additional information Z are worse than those without using Z .

Original: Thisimplies that such information may not be directly utilized by or even can harm the two methods.
Revised: Thisimplies that such information may not be directly utilized by or even can harm the two methods.

Original: CelebA .
Revised: CelebA .

Original: This task is to predict Smiling based on the image from CelebA [28], and by constructionthe target is spuriously correlated with Gender .
Revised: This task is to predict Smiling based on the image from CelebA [27], and by constructionthe target is spuriously correlated with Gender .

Original: We have access to the meta annotations of CeleAimages and they serve as Z in our framework.
Revised: We have access to the meta annotations of CeleAimages and they serve as Z in our framework.

Original: In particular, we pick seven features as potentialauxiliary variables: { Young , Blond Hair , Eyeglasses , High Cheekbones , Big Nose , Bags Under Eyes ,
Revised: In particular, we pick seven features as potentialauxiliary variables: { Young , Blond Hair , Eyeglasses , High Cheekbones , Big Nose , Bags Under Eyes ,

Original: Chubby }.
Revised: Chubby }.

Original: Notice that the annotation of Gender is not provided to ZIN, EIIL, HRM and ERM, but ismade available to IRM to create an oracle environment partition. 
Revised: Notice that the annotation of Gender is not provided to ZIN, EIIL, HRM and ERM, but ismade available to IRM to create an oracle environment partition. [Table]

Original:  Empirical results are given in Table 3, where ZIN(#) represents ZIN using # of the seven features as
Revised: [Table] Empirical results are given in Table [Table] where ZIN(#) represents ZIN using # of the seven features as

Original: Z .
Revised: Z .

Original: ERM achieves the highest training accuracy, while only has 47.58% worst test accuracy.
Revised: ERM achieves the highest training accuracy, while only has 47.58% worst test accuracy.

Original: EIILhas the worst performance on both training and test data, showing that it may not learn the causalfeatures.
Revised: EIILhas the worst performance on both training and test data, showing that it may not learn the causalfeatures.

Original: Finally, we observe that ZIN performs better when more auxiliary information is provided;specifically, ZIN(7) achieves the best test performance based on the mixed dataset.
Revised: Finally, we observe that ZIN performs better when more auxiliary information is provided;specifically, ZIN(7) achieves the best test performance based on the mixed dataset.

Original: [Table]
Revised: [Table]

Original: Landcover .
Revised: Landcover .

Original: Our final task is land cover prediction thatclassifies the land cover type (e.g., grasslands) fromsatellite data [13, 37].
Revised: Our final task is land cover prediction thatclassifies the land cover type (e.g., grasslands) fromsatellite data [12, 36].

Original: We take the same setup from[49]: the time series data input dimension is 46 × 8 ;the target Y is one of six land cover classes; six climate related variables are the auxiliary variables.
Revised: We take the same setup from[48]: the time series data input dimension is 46 × 8 ;the target Y is one of six land cover classes; six climate related variables are the auxiliary variables.

Original: Wealso consider non-African locations as training data and
Revised: Wealso consider non-African locations as training data and

Original: Africa as test data.
Revised: Africa as test data.

Original: In addition, we take location (latitude and longitude) as possible additional information.
Revised: In addition, we take location (latitude and longitude) as possible additional information.

Original: For ZIN, we average climate variables over the timedimension and use those means to predict environments.
Revised: For ZIN, we average climate variables over the timedimension and use those means to predict environments.

Original: In Table 4, we see that ZIN with locationachieves the best OOD performance, while ZIN with climate variables is slightly worse.
Revised: In Table 4, we see that ZIN with locationachieves the best OOD performance, while ZIN with climate variables is slightly worse.

Original: This isbecause climate variables indeed contain some information for predicting the target, violating ourcondition Z .
Revised: This isbecause climate variables indeed contain some information for predicting the target, violating ourcondition Z .

Original: EIIL also has a good performance in this task, and we conjecture that its first stage haslearned the spurious features as desired.
Revised: EIIL also has a good performance in this task, and we conjecture that its first stage haslearned the spurious features as desired.

Original: However, together with previous empirical results, it can berisky as the first stage may not be guaranteed.
Revised: However, together with previous empirical results, it can berisky as the first stage may not be guaranteed.

Revised: 7.3 Discussion on Auxiliary Information for Related Methods

Revised: The proposed framework ZIN relies on additional auxiliary variables.

Revised: It is interesting to ask whetherthis type of information is also useful to related methods such as EIIL and HRM.

Original:   Concluding Remarks
Revised: In many scenarios,the additional variable Z may have no or little information about the target, e.g., 1) in time seriestasks, the time index is rarely used to predict the label; and 2) in the CelebA classification task,features like Young and Eyeglasses are likely to provide little information about predicting Smiling . Then Z may not be useful to ERM and EIIL. In Section 7.2, we also provide Z to EIIL and HRM,and the experimental result shows that it does not help or even harms the test performance.  Concluding Remarks

Original: This paper investigates when and how to learn invariance from heterogeneous data without explicitenvironment indexes.
Revised: This paper investigates when and how to learn invariance from heterogeneous data without explicitenvironment indexes.

Original: We first show that learning invariant models in this case is generally impossible.
Revised: We first show that learning invariant models in this case is generally impossible.

Original: Then we propose ZIN to jointly learn environment partition and invariant representation using someadditional auxiliary variables.
Revised: Then we propose ZIN to jointly learn environment partition and invariant representation using someadditional auxiliary variables.

Original: We provide theoretic guarantees in both feature selection and linearfeature learning scenarios.
Revised: We provide theoretic guarantees in both feature selection and linearfeature learning scenarios.

Original: Experimental results verify our analysis and demonstrate an improvedperformance of ZIN over existing methods.
Revised: Experimental results verify our analysis and demonstrate an improvedperformance of ZIN over existing methods.

Original: Our results also raise the need of future works to makethe role of inductive biases more explicit, when learning invariance from heterogeneous data withoutenvironment indexes.
Revised: Our results also raise the need of future works to makethe role of inductive biases more explicit, when learning invariance from heterogeneous data withoutenvironment indexes.

Original: A limitation of the present work is the lack of theoretic guarantee with generalnonlinear feature extractors, which is also an open problem for IRM even with environment indexes.
Revised: A limitation of the present work is the lack of theoretic guarantee with generalnonlinear feature extractors, which is also an open problem for IRM even with environment indexes.

Original: Checklist
Revised: Checklist

Original: For all authors...
Revised: For all authors...

Original: (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
Revised: (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s

Original: contributions and scope?
Revised: contributions and scope?

Original: [Yes]
Revised: [Yes]

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

Original: [Yes]
Revised: [Yes]

Original: In the conclusion section.
Revised: In the conclusion section.

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

Original: [Yes] InAppendix H.
Revised: [Yes] InAppendix E.

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

Original: [Yes]
Revised: [Yes]

Original: If you are including theoretical results...
Revised: If you are including theoretical results...

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

Original: [Yes]
Revised: [Yes]

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

Original: [Yes]
Revised: [Yes]

Original: If you ran experiments...
Revised: If you ran experiments...

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

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

Original: [Yes]
Revised: [Yes]

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

Original: [Yes]
Revised: [Yes]

Original: (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
Revised: (c) Did you report error bars (e g., with respect to the random seed after running experiments multiple times)?

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

Original: [Yes]
Revised: [Yes]

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

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

Original: [Yes] (b) Did you mention the license of the assets?
Revised: [Yes](b) Did you mention the license of the assets?

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

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

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

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

Original: using/curating?
Revised: using/curating?

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

Original: (e) Did you discuss whether the data you are using curating contains personally identifiable
Revised: (e) Did you discuss whether the data you are using/curating contains personally identifiable

Original: information or offensive content?
Revised: information or offensive content?

Original: N/A
Revised: [N/A

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

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

Original: applicable?
Revised: applicable?

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

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

Original: Board (IRB) approvals, if applicable?
Revised: Board (IRB) approvals, if applicable?

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

Original: (c) Did you include the estimated hourly wage paid to participants and the total amount
Revised: (c) Did you include the estimated hourly wage paid to participants and the total amount

Original: spent on participant compensation?
Revised: spent on participant compensation?

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

Original: A ProofsA.1 Proof of Theorem 1
Revised: A ProofsA.1 Proof of Theorem 1

Original: Proof.
Revised: Proof.

Original: First, we assume Y and X v are univariate variables, i.e., Y, X v ∈ R .
Revised: First, we assume Y and X v are univariate variables, i.e., Y, X v ∈ R .

Original: Let η 1 ∼ Uniform(0 , 1)independent of X s and Y , and set X ′ v =
Revised: Let η 1 ∼ Uniform(0 , 1)independent of X s and Y , and set X ′ v =

Original: X s .
Revised: X s .

Original: Define the conditional cumulative distribution functionand its inverse as:
Revised: Define the conditional cumulative distribution functionand its inverse as:

Original: [Equation]
Revised: [Equation]

Original: By definition, we would have
Revised: By definition, we would have

Original: [Equation]
Revised: [Equation]

Original: Similarly, we can construct X ′ s = f ′ 2 ( X ′ v , Y ′ , η 2 ) =
Revised: Similarly, we can construct X ′ s = f ′ 2 ( X ′ v , Y ′ , η 2 ) =

Original: F − 1 X v | Y, X s ( η 2 ) so that P ( X ′ s | Y ′ , X ′ v ) =
Revised: F − 1 X v | Y, X s ( η 2 ) so that P ( X ′ s | Y ′ , X ′ v ) =

Original: P ( X v | Y, X s ) with η 2 ∼ Uniform(0 , 1) .
Revised: P ( X v | Y, X s ) with η 2 ∼ Uniform(0 , 1) .

Original: Thus, we have
Revised: Thus, we have

Original: [Equation]
Revised: [Equation]

Original: [Equation] Next, we can easily find a function q ′ ( · ) so that X ′ =
Revised: Next, we can easily find a function q ′ ( · ) so that X ′ =

Original: q ′ ( X ′ v , X ′ s )
Revised: q ′ ( X ′ v , X ′ s )

Original: =
Revised: =

Original: q ( X v , X s ) , where we havechosen X ′ v =
Revised: q ( X v , X s ) , where we havechosen X ′ v =

Original: X s and X ′ s
Revised: X s and X ′ s

Original: = X v .
Revised: = X v .

Original: Together with Equation (8), we conclude that P ( X ′ , Y ′ ) =
Revised: Together with Equation (7), we conclude that P ( X ′ , Y ′ ) =

Original: P ( X, Y ) .
Revised: P ( X, Y ) .

Original: Second, we consider that X v has a multi-dimension.
Revised: Second, we consider that X v has a multi-dimension.

Original: We may leave Y to be a univariate variableas the label in many ML problems is scalar-valued.
Revised: We may leave Y to be a univariate variableas the label in many ML problems is scalar-valued.

Original: For this case, we can pick an entry of X v , say,
Revised: For this case, we can pick an entry of X v , say,

Original: X ( j ) v for some j .
Revised: X ( j ) v for some j .

Original: Then we set X ′ v =
Revised: Then we set X ′ v =

Original: ( X ( − j ) v , X s ) where X ( − j ) v denotes the rest entries of X vexcept the j -th.
Revised: ( X ( − j ) v , X s ) where X ( − j ) v denotes the rest entries of X vexcept the j -th.

Original: Then we can similarly use the inverse conditional cumulative distribution functionand an independent uniformly distributed noise variable η 1 to construct Y ′
Revised: Then we can similarly use the inverse conditional cumulative distribution functionand an independent uniformly distributed noise variable η 1 to construct Y ′

Original: =
Revised: =

Original: f ′ ( X ′ v , η 1 ) so that
Revised: f ′ ( X ′ v , η 1 ) so that

Original: P ( Y ′ ≤ y | X ′ v ) =
Revised: P ( Y ′ ≤ y | X ′ v ) =

Original: P ( Y ≤ y | X ( − j ) v , X s ) .
Revised: P ( Y ≤ y | X ( − j ) v , X s ) .

Original: Similarly, we set X ′ s =
Revised: Similarly, we set X ′ s =

Original: X ( j ) v
Revised: X ( j ) v

Original: and we can construct
Revised: and we can construct

Original: X ′ s
Revised: X ′ s

Original: = f ′ 2 ( X ′ v , Y ′ , η 2 ) so that P ( X ′ s | Y ′ ,
Revised: = f ′ 2 ( X ′ v , Y ′ , η 2 ) so that P ( X ′ s | Y ′ ,

Original: X ′ v ) =
Revised: X ′ v ) =

Original: P ( X ( j ) v | Y, X s ) .
Revised: P ( X ( j ) v | Y, X s ) .

Original: We can then get the sameconclusion following the previous proof.
Revised: We can then get the sameconclusion following the previous proof.

Original: A.2 Proof of Theorem 2
Revised: A.2 Proof of Theorem 2

Original: Proof.
Revised: Proof.

Original: Our proof proceeds by two steps.
Revised: Our proof proceeds by two steps.

Original: First, we show that any feature mask that selects at leastone spurious feature would induce a penalty.
Revised: First, we show that any feature mask that selects at leastone spurious feature would induce a penalty.

Original: With sufficiently large λ , the penalty will dominatethe expected risk and then exceed ˆ L (Φ v ) .
Revised: With sufficiently large λ , the penalty will dominatethe expected risk and then exceed ˆ L (Φ v ) .

Original: Second, we show that any proper subset of the invariantfeatures induces a loss larger than ˆ L (Φ v ) .
Revised: Second, we show that any proper subset of the invariantfeatures induces a loss larger than ˆ L (Φ v ) .

Original: Step 1 Suppose that the feature mask contains at least one spurious features.
Revised: Step 1 Suppose that the feature mask contains at least one spurious features.

Original: Denote the selectedfeatures as X + s and the corresponding feature mask as Φ + s .
Revised: Denote the selectedfeatures as X + s and the corresponding feature mask as Φ + s .

Original: We aim to show that
Revised: We aim to show that

Original: [Equation]
Revised: [Equation]

Original: By Assumption 1 with a given ϵ > 0 , we have[Equation]
Revised: By Assumption 1 with a given ϵ > 0 , we have

Original: [Equation]
Revised: [Equation]

Original: One the other hand, we have
Revised: One the other hand, we have

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

Original: where the last inequality is due to Assumption 3 and Condition 2.
Revised: where the last inequality is due to Assumption 3 and Condition 2.

Original: Thus, if we choose ϵ < δC/ 4 andλ > H ( Y )+2 ϵ δC − 4 ϵ , we can get
Revised: Thus, if we choose ϵ < δC/ 4 andλ > H ( Y )+2 ϵ δC − 4 ϵ , we can get

Original: L (Φ [Equation]
Revised: L (Φ v ) <

Original: < L ( [Equation]
Revised: L (Φ + s ) .

Original: Step 2 Denote a proper subset of invariant features as X − v ⊊ X v , and similarly the feature maskas Φ − v .
Revised: Step 2 Denote a proper subset of invariant features as X − v ⊊ X v , and similarly the feature maskas Φ − v .

Original: In Step 1, we have shown that
Revised: In Step 1, we have shown that

Original: [Equation]
Revised: [Equation]

Original: Similar to Equation (10), we have
Revised: Similar to Equation (9), we have

Original: [Equation]
Revised: [Equation]

Original: Then according to Assumption 4, we have
Revised: Then according to Assumption 4, we have

Original: [Equation]
Revised: [Equation]

Original: Thus, if ϵ < γ
Revised: Thus, if ϵ < γ

Original: [Equation]
Revised: [Equation]

Revised: [Equation]

Original: ) , we have
Revised: ) , we have

Original: [Equation]
Revised: [Equation]

Original: In conclusion, with λ ∈ [ H ( Y )+1 / 2 δC δ d C − 4 ϵ − 12 , γ 4 ϵ − 12 ] , we can get[Equation]
Revised: In conclusion, with λ ∈ [ H ( Y )+1 / 2 δC δ d C − 4 ϵ − 12 , γ 4 ϵ − 12 ] , we can get

Original: [Equation]
Revised: [Equation]

Original: Notably, there exists a feasible λ if ϵ <
Revised: Notably, there exists a feasible λ if ϵ <

Original: Cγδ 4 γ +2 CδH ( Y ) .
Revised: Cγδ 4 γ +2 CδH ( Y ) .

Original: The proof is complete by noticing that ϵ can bechosen arbitrarily according to Assumption 1.
Revised: The proof is complete by noticing that ϵ can bechosen arbitrarily according to Assumption 1.

Original: A.3 Proof of Proposition
Revised: A.3 Proof of Proposition

Original: Proof.
Revised: Proof.

Original: Consider the following feature set
Revised: Consider the following feature set

Original: [Equation]
Revised: [Equation]

Original: and the corresponding feature mask is denoted as Φ ¯ v .
Revised: and the corresponding feature mask is denoted as Φ ¯ v .

Original: It corresponds the largest subset of X thatsatisfies the invariance constraint. Note that  Φ ¯ v ̸ = Φ v .
Revised: Note that feature X iv in Proposition 1 is not in X ¯ v , so Φ ¯ v ̸ = Φ v .

Original: By Assumption 1 with a given ϵ , we can have
Revised: By Assumption 1 with a given ϵ ,we can have

Original: [Equation]
Revised: [Equation]

Original: One the other hand, we have
Revised: One the other hand, we have

Original: [Equation]
Revised: [Equation]

Original: which can be shown similarly to Equation (10).
Revised: which can be shown similarly to Equation (9).

Original: Thus, if we choose ϵ < δC ′ / 4 and λ > H ( Y )+2 ϵ δC ′ − 4 ϵ , wewould get 
Revised: Thus, if we choose ϵ < δC ′ / 4 and λ > H ( Y )+2 ϵ δC ′ − 4 ϵ , wewould get L (Φ ¯ v ) < L (Φ v )

Original: [Equation]
Revised: [Equation]

Original: A.4 Proof of Corollary 1
Revised: A.4 Proof of Corollary 1

Original: Proof.
Revised: Proof.

Original: (a) Since h is injective, H ( Y | X v , h ( Y ))
Revised: (a) Since h is injective, H ( Y | X iv , h ( Y ))

Original: =
Revised: =

Original: H ( Y | X v , Y ) =  for any X v .
Revised: H ( Y | X iv , Y ) = 0 for any X iv .

Original: By Assumption 4, we have H ( Y |X v
Revised: By Assumption 4, we have H ( Y |

Original:  ) ≥ H ( Y | X ) + γ ≥ γ .
Revised: X iv ) ≥ H ( Y | X ) + γ ≥ γ .

Original: Then Proposition 1 indicates thatwe cannot identify all the invariant features.
Revised: Then Proposition 1 indicates thatwe cannot identify all the invariant features.

Original: (b) The proof proceeds the same as aboveby noting H ( Y | X v , h
Revised: (b) The proof proceeds the same as aboveby noting H ( Y | X iv , h

Original: ( X , Y ))
Revised: ( X , Y ))

Original: = 0 .
Revised: = 0 .

Original: (c) This case can be shown similarly to (a), because
Revised: (c) This case can be shown similarly to (a), because

Original: H ( Y | X v , h (Index( X , Y )))
Revised: H ( Y | X iv , h (Index( X , Y )))

Original: =
Revised: =

Original: H ( Y | X v , X , Y ) = 0 .
Revised: H ( Y | X iv , X , Y ) = 0 .

Original: A.5 Proof of Proposition
Revised: A.5 Proof of Proposition

Original: Proof.
Revised: Proof.

Original: Denote a feature set X v (+ k ) , which contains the invariant feature set X v as well a spuriousfeature X ks  in Proposition 2.
Revised: Denote a feature set X v (+ k ) , which contains the invariant feature set X v as well a spuriousfeature X ks X s in Proposition 2.

Original: Similar to Eqn. (9), we can show that
Revised: Similar to Equation (8), we can show that

Original: [Equation]
Revised: [Equation]

Original: And similar to Eqn. (10), we also have
Revised: And similar to Equation (9), we also have

Original: [Equation]
Revised: [Equation]

Original: Then it follows that
Revised: Then it follows that

Original: [Equation]
Revised: [Equation]

Original: If ϵ < γ 2(1+2 λ ) , we have
Revised: If ϵ < γ 2(1+2 λ ) , we have

Original: [Equation]
Revised: [Equation]

Original: Thus, we cannot identify all the invariant features from Problem 4.
Revised: Thus, we cannot identify all the invariant features from Problem 3.

Original: A.6 Proof of Theorem
Revised: A.6 Proof of Theorem

Original: Proof.
Revised: Proof.

Original: Step 1: No spurious feature will be learned.

Original: Given a partition { X k } Kk =1 that lies in linear general position of degree r , Arjovsky et al.
Revised: Given a partition { X k } Kk =1 that lies in linear general position of degree r , Arjovskyet al.

Original: [5, Theorem 9] shows that Φ and ω satisfies the normal equations
Revised: [4, Theorem 9] shows that Φ and ω satisfies the normal equations

Original: Φ E X k
Revised: Φ E X k

Original: [ X k X k T ]Φ T ω =
Revised: [ X k X k T ]Φ T ω =

Original: Φ E X k ,Y k
Revised: Φ E X k ,Y k

Original: [ X k Y k ]for all k
Revised: [ X k Y k ] for all k

Original: if and only if Φ elicits the desired invariant predictor Φ T ω =
Revised: if and only if Φ elicits the desired invariant predictor Φ T ω =

Original: ˜ W T β .
Revised: ˜ W T β .

Original: Thus, we only need toshow our solution meets the same normal equations.
Revised: Thus, we only need to show our solution meets the same normal equations.

Original: Let Φ ′ and ω ′ denote a solution to Problem 6.
Revised: Let Φ ′ and ω ′ denotea solution to Problem 5.

Original: According to the constraint and the general linear position condition, we know there exists a partition{ X k } Kk =1 lies in general linear position of degree r , and we have R k ( ω ′ , Φ ′ )
Revised: According to the constraint and the general linear position condition, weknow there exists a partition { X k } Kk =1 lies in general linear position of degree r , and we have R k ( ω ′ , Φ ′ )

Original: =
Revised: =

Original: R k ( ω k , Φ ′ ) .
Revised: R k ( ω k , Φ ′ ) .

Original: Noticethat ω k minimizes the mean squared error on only the k -th environment, hence it must satisfy thenormal equation Φ ′
Revised: Notice that ω k minimizes the mean squared error on only the k -th environment, hence it must satisfy the normal equation Φ ′

Original: E X k
Revised: E X k

Original: [ X k X k T ]Φ ′ T ω
Revised: [ X k X k T ]Φ ′ T ω

Original: k
Revised: k

Original: =
Revised: =

Original: Φ ′ E X k ,Y k
Revised: Φ ′ E X k ,Y k

Original: [ X k Y k ] .
Revised: [ X k Y k ] .

Original: As ω ′ achieves the same minimummean squared error on the k -th environment, ω ′ must satisfy the normal equation, too.
Revised: As ω ′ achieves the same minimum mean squared error on the k -th environment, ω ′ must satisfy thenormal equation, too.

Original: So no spuriousinformation will in included and X Φ

Original: = S X v where S is a invertible function.

Original: Step 2: No invariant information will be discarded.

Original: By Condition 1, we have H ( Y | X v ) − H

Original: ( Y | X v , ρ ( Z )) = 0 .

Original: Then X Φ =

Original: S X v will satisfy theconstrain of Problem 6 for any environment generated by ρ ( Z ) .

Original: Notice that S X v achieves the whensmallest loss only using invariant feature information.

Original: Putting these two steps together completes the proof.
Revised: This completes the proof.

Original: A.7 Proof of Meeting Condition 1
Revised: A.7 Proof of Meeting Condition 1

Original: We show that H ( Y | X v , ρ ( Z ))
Revised: We show that H ( Y | X iv , ρ ( Z ))

Original: =
Revised: =

Original: H ( Y | X v ) for all ρ ( · ) if H ( Y | X v , Z ) =
Revised: H ( Y | X iv ) for all ρ ( · ) if H ( Y | X iv , Z ) =

Original: H ( Y | X v ) holds.
Revised: H ( Y | X iv ) holds.

Original: Proof.
Revised: Proof.

Original: On one hand, because ρ ( Z ) contains less information than Z , so 
Revised: On one hand, because ρ ( Z ) contains less information than Z , so H [Equation] X

Revised: [Equation]

Revised: , ρ

Original: H ( Y | X v , ρ ( Z )) ≥ H( Y | X v , Z )
Revised: [Equation] Z [Equation] ≥ H

Revised: [Equation]

Revised: Y

Original:  =
Revised: [Equation] X [Equation] , Z ) =

Original: H ( Y | X v ) .
Revised: H ( Y | X iv ) .

Original: On the other hand, X v and ρ ( Z ) contain more information than X v , so we have
Revised: On the other hand, X iv and ρ ( Z ) contain more information than X iv , so we have

Original: H ( Y | X v , ρ ( Z )) ≤ H ( Y | X v ) .
Revised: H ( Y | X iv , ρ ( Z )) ≤ H ( Y | X iv ) .

Original: Thus, we conclude H ( Y | X v , ρ ( Z ))
Revised: Thus, we conclude H ( Y | X iv , ρ ( Z ))

Original: =
Revised: =

Original: H ( Y | X v ) .
Revised: H ( Y | X iv ) .

Original: B Extension to Other Loss Functions
Revised: B Extension to Other Loss Functions

Original: Sections 5.2 and 5.3 focus on classification task with cross-entropy loss, to establish sufficientand necessary conditions for invariance identification.
Revised: Sections 5.2 and 5.3 focus on classification task with cross-entropy loss, to establish sufficientand necessary conditions for invariance identification.

Original: Similar results can be shown for other lossfunctions or tasks in a straightforward manner.
Revised: Similar results can be shown for other lossfunctions or tasks in a straightforward manner.

Original: To see this, notice that H ( ·|· ) in Assumptions 1-4 and
Revised: To see this, notice that H ( ·|· ) in Assumptions 1-4 and

Original: Conditions 1-2 is used to represent the optimal expected risk, which coincides with the conditionalentropy with cross-entropy loss and when ρ ( z i ) gives exactly one environment.
Revised: Conditions 1-2 is used to represent the optimal expected risk, which coincides with the conditionalentropy with cross-entropy loss and when ρ ( z i ) gives exactly one environment.

Original: For other lossfunctions l ( · , · ) like squared error, we can use L ( ·|· ) to represent the optimal expected risks, and L ( · )denotes the optimal expected risk with no predictor variables, e.g., variance for the squared error loss.
Revised: For other lossfunctions l ( · , · ) like squared error, we can use L ( ·|· ) to represent the optimal expected risks, and L ( · )denotes the optimal expected risk with no predictor variables, e.g., variance for the squared error loss.

Original: Replacing H ( ·|· ) with L ( ·|· ) , we can follow the same proof procedure to obtain similar sufficient andnecessary conditions.
Revised: Replacing H ( ·|· ) with L ( ·|· ) , we can follow the same proof procedure to obtain similar sufficient andnecessary conditions.

Original: This is summarized in Theorem 4 for completeness.
Revised: This is summarized in Theorem 4 for completeness.

Original: Theorem 4.
Revised: Theorem 4.

Original: Consider Assumptions 1-4 and Conditions 1-2 where H ( ·|· ) is replaced with L ( ·|· )accordingly.
Revised: Consider Assumptions 1-4 and Conditions 1-2 where H ( ·|· ) is replaced with L ( ·|· )accordingly.

Original: If ϵ <
Revised: If ϵ <

Original: [Equation]
Revised: [Equation]

Original: ( Y ) and λ ∈ 
Revised: and λ ∈ [ L

Original: [ L ( Y) +
Revised: [Equation] Y

Original: [Equation]
Revised: [Equation]

Original: − 12 ,
Revised: − 12 ,

Original: [Equation]
Revised: [Equation]

Original: 4 ϵ − 12 ] , we have ˆ L (Φ v ) <
Revised: 4 ϵ − 12 ] , we have ˆ L (Φ v ) <

Original: ˆ L (Φ) , ∀ Φ ̸
Revised: ˆ L (Φ) , ∀ Φ ̸

Original: =Φ v where we assume L ( Y ) < ∞ .
Revised: =Φ v where we assume L ( Y ) < ∞ .

Original: If Condition 1 or Condition 2 is violated, then there exists a featuremask Φ ′ ̸
Revised: If Condition 1 or Condition 2 is violated, then there exists a featuremask Φ ′ ̸

Original: =
Revised: =

Original: Φ v so that ˆ L (Φ v ) >
Revised: Φ v so that ˆ L (Φ v ) >

Original: ˆ L (Φ ′ ) .
Revised: ˆ L (Φ ′ ) .

Original: C Discussion on Auxiliary Information for Related Methods

Original: The proposed framework ZIN relies on additional auxiliary variables.

Original: It is interesting to ask whetherthis type of information is also useful to related methods such as EIIL and HRM.

Original: In many scenarios,the additional variable Z may have no or little information about the target, e.g., 1) in time seriestasks, the time index is rarely used to predict the label; and 2) in the CelebA classification task,features like Young and Eyeglasses are likely to provide little information about predicting Smiling .

Original: Then Z may not be useful to ERM and EIIL.

Original: In Section 7.2, we also provide Z to EIIL and HRM,and the experimental result shows that it does not help or even harms the test performance.

Original: D More Experimental Details
Revised: C More Experimental Details

Original: For the synthetic datasets and the house price prediction dataset, we use the full batch gradient inoptimizing.
Revised: For the synthetic datasets and the house price prediction dataset, we use the full batch gradient inoptimizing.

Original: The ERM method is trained for 4000 epochs.
Revised: The ERM method is trained for 4000 epochs.

Original: The IRM/ZIN method is also trained for4000 epochs, with annealing in the first 2000 epochs.
Revised: The IRM/ZIN method is also trained for4000 epochs, with annealing in the first 2000 epochs.

Original: We train EIIL for 4000 epochs which is dividedequally, i.e., 2000 epochs, in the two stages.
Revised: We train EIIL for 4000 epochs which is dividedequally, i.e., 2000 epochs, in the two stages.

Original: For the CelebA classification task, we use mini-batchtraining with batch size of 128.
Revised: For the CelebA classification task, we use mini-batchtraining with batch size of 128.

Original: All the methods are trained for 50 epochs, with annealing strategy inthe first 25 epochs.
Revised: All the methods are trained for 50 epochs, with annealing strategy inthe first 25 epochs.

Original: For the landcover classification task, we use mini-batch training with batch sizeof 1024.
Revised: For the landcover classification task, we use mini-batch training with batch sizeof 1024.

Original: All the methods are trained for 400 epochs, with anneal strategy in the first 40 epochs.
Revised: All the methods are trained for 400 epochs, with anneal strategy in the first 40 epochs.

Original: Weuse Adam [21] with learning rate 0.001 as our optimizer.
Revised: Weuse Adam [20] with learning rate 0.001 as our optimizer.

Original: Our experiments are run in a Linux workstation with Intel Xeon 3.20GHz CPU, 128GB RAM, and
Revised: Our experiments are run in a Linux workstation with Intel Xeon 3.20GHz CPU, 128GB RAM, and

Original: Nvidia GTX 3090 GPU.
Revised: Nvidia GTX 3090 GPU.

Original: It takes about 5 GPU hours for the CelebA tasks, while the house price taskand Landcover task cost less than 10 minutes for training.
Revised: It takes about 5 GPU hours for the CelebA tasks, while the house price taskand Landcover task cost less than 10 minutes for training.

Original: E Extensions of Section 4E.1 Connection of Impossibility Example to Existing Tasks

Original: In the counter-example in Section 4, we consider binary features and label.

Original: This is inspired by thepopular classification tasks in the literature of IRM, e.g., CMNIST [5], CifarMnist [24], ColoredObject[1] , Waterbirds and CelebA [38].

Original: Take CMNIST for example.

Original: The label takes either 0 or 1.

Original: Theinvariant feature, X 1 , is the semantic feature ‘0’ or ‘1’ of the digit shape.

Original: The spurious feature X 2(color) is also binary: either red or green.

Original: So we will denote the binary invariant and spurious featuresas 0/1 in the following discussion.

Original: Recall the the construction procedure of CMNIST in [5]:

Original: [Equation]

Original: [Equation]

Original: where X 1 (digit shape) is the invariant feature while X 2 (color) is spurious, and p e

Original: =1 s = 0 .

Original: 8 andp e =2 s = 0 .

Original: Notably, in the test domain, p e =3 = 0 . 1 is chosen to simulate further distributional shiftwhere the conditional probability P ( Y | X 2 ) is significantly different from the training domains.

Original: The CMNIST construction process is analogous to the first data generation process in our counterexample in Section 4.

Original: When the two environments are mixed, we observe a joint distribution thatcan also be generated according to a different SCM.

Original: In the next section, we construct a dataset inaccordance to the CMNIST dataset.

Original: E.2 Supportive Experimental Results

Original: We now construct a variant of CMNIST, analogous to the second data generation process in Example1 in Section 4.

Original: In this dataset, color is the invariant feature and the digit shape is the spurious feature.

Original: The correlation between digit shape and label is unstable and it would change dramatically in thetesting domain.

Original: We name this variant of CMNIST as MCOLOR (short for MNIST-COLOR).

Original: We usethe same notations where X 1 stands for the digit shape, X 2 is the color and Y denotes label.

Original: We generate the label in MCOLOR from the color and make the digit shape spuriously correlatedwith the label:

Original: [Equation]

Original: where X 2 (color) is the invariant feature while X 1 (digit shape) is spurious, and p e

Original: =1 s = 0 .

Original: 8 andp e =2 s = 0 .

Original: In the testing domain, we set p e =3 s = 0 .

Original: 1 to simulate the distributional shift

Original: so that

Original: P ( Y | X 1 ) is significantly different from the training domains.

Original: If we mix the training domains of

Original: CMNIST or MCOLOR, we will observe the same distribution given by 

Original: Y = 0 , w.p. 0 .

Original: Y = 1 , w.p. 0 .

Original: [-X 
-]
Revised: [-X 
-]

Original: = X 1 = Y, w.p. 0 .

Original: X 2 ̸ =

Original: X 1 = Y, w.p. 0 .

Original: X 1 ̸ =

Original: X 2 = Y, w.p. 0 .

Original: X 2 =

Original: X 1 ̸ = Y, w.p. 0 .

Original: With this new dataset, we can then compare EIIL, IRM (with environment indexes), and ERM.

Original: Theempirical results are reported in Table 5.

Original: For ERM (oracle), we train the model only on the invariantfeature in MCOLOR, i.e., color.

Original: [Table]

Original: In Table 5, we can see that the EIIL method performs poorly on MCOLOR.

Original: Due to the inductivebias of EIIL, it will deterministically rely on the digit shape as the invariant feature, resulting in poorperformance on this MCOLOR dataset.

Original: However, since we have no prior knowledge of the datageneration process, the TRUE invariant feature could be the color as in the MCOLOR dataset.

Original: LfFwill also rely on either color or digit shape as the invariant feature and will fail on one of the CMNISTand MCOLOR datasets.

Original: On the other hand, if we have the environment indexes, IRM can still learnthe color as the invariant feature on MCOLOR.

Original: We believe that this observation will also hold on

Original: CifarMnist, ColoredObject, Waterbird or CelebA.

Original: [Table]

Original: F More Experiment Results
Revised: D More Experiment ResultsD.1 Spatial Heterogeneity

Original: [Figure] F.1 Spatial Heterogeneity

Original: [Figure]

Original: We next consider spatial heterogeneity that is also commonly encountered in practice, e.g., environments may be divided according to locations of latitude and longitude.
Revised: We next consider spatial heterogeneity that is also commonly encountered in practice, e.g., environments may be divided according to locations of latitude and longitude.

Original: We simulate spatial heterogeneity in the same way the data generation process for time heterogeneity but use a two-dimensional spatial variable r =
Revised: We simulate spatial heterogeneity in the same way the data generation process for time heterogeneity but use a two-dimensional spatial variable r =

Original: [ r 1 , r 2 ] ∈ [0 , 1] 2 .
Revised: [ r 1 , r 2 ] ∈ [0 , 1] 2 .

Original: We simulate four environments for training by equally splitting thespace into four blocks, i
Revised: We simulate four environments for training by equally splitting thespace into four blocks, i.e., { [0 , 0 . 5) ×

Original: [Table] e., { [0 , 0 . 5) × [0 , 0 . 5) , [0 , 0 . 5) × [0 . 5 , 1] , [0 . 5 , 1] × [0 , 0 . 5) , [0 . 5 , 1] × [0 . 5 , 1] } .
Revised: [0 , 0 . 5) , [0 , 0 . 5) × [0 . 5 , 1] , [0 . 5 , 1] × [0 , 0 . 5) , [0 . 5 , 1] × [0 . 5 , 1] } .

Original: Similarly, we denote a simulated case by tuple of p s ( r ) in the four elements.
Revised: Similarly, we denote a simulated case by tuple of p s ( r ) in the four elements.

Original: We also evaluate theperformance on four distinct test environments with p s ∈ { 0 . 999 , 0 . 8 , 0 . 2 , 0 . 1 } and p v being constant.
Revised: We also evaluate theperformance on four distinct test environments with p s ∈ { 0 . 999 , 0 . 8 , 0 . 2 , 0 . 1 } and p v being constant.

Original: More details regarding the implementations are given in Appendix D.
Revised: More details regarding the implementations are given in Appendix C.

Original: The experimental results of spatial heterogeneity in Table 6 show similar performances to thoseof temporal heterogeneity.
Revised: The experimental results of spatial heterogeneity in Table 5 show similar performances to thoseof temporal heterogeneity.

Original: An interesting observation is that ZIN outperforms IRM with groundtruth environments in this simulation, especially when “duplicated” environments exist, e.g., whenp s ( r ) =
Revised: An interesting observation is that ZIN outperforms IRM with groundtruth environments in this simulation, especially when “duplicated” environments exist, e.g., whenp s ( r ) =

Original: We conjecture that in this case the “ground-truth” partition may notbe the most effective for invariant learning due to the heavily overlapped environments.
Revised: We conjecture that in this case the “ground-truth” partition may notbe the most effective for invariant learning due to the heavily overlapped environments.

Original: As shown inthe right panel of Figure 1, ZIN automatically merges duplicated environments.[Figure]
Revised: As shown inthe right panel of Figure 1, ZIN automatically merges duplicated environments.

Original: F.2 Ablation Study on Auxiliary Information and Number of Environments
Revised: D.2 Ablation Study on Auxiliary Information and Number of Environments

Original: We now verify the theoretical results in Section 5 by choosing different inputs as Z .
Revised: We now verify the theoretical results in Section 5 by choosing different inputs as Z .

Original: We adopt asetting of spatial heterogeneity with p s ( r ) =
Revised: We adopt asetting of spatial heterogeneity with p s ( r ) =

Original: 8) and p v = 0 .
Revised: 8) and p v = 0 .

Original: Note that theheterogeneity in this setting is only along the second dimension r 2 .
Revised: Note that theheterogeneity in this setting is only along the second dimension r 2 .

Original: The results are shown in Table 7.
Revised: The results are shown in Table 6.

Original: One can verify that Conditions 1 and 2 are satisfied when r or r 2 is chosen as Z . 
Revised: One can verify that Conditions 1 and 2 are satisfied when r or r 2 is chosen as Z . The mean and worst [Table] [Figure]

Original: The mean and worsttest accuracy implies that ZIN based on r or r 2 can effectively remove the spurious feature.
Revised: [Table] [Figure] [Table] [Figure]

Original:  Sincep s ( r ) does not change along r 1 , choosing r 1 to infer environment partition violates Condition 2,which is reflected by the poor performance of ZIN with r 1 .
Revised: [Table] [Figure] test accuracy implies that ZIN based on r or r 2 can effectively remove the spurious feature. Sincep s ( r ) does not change along r 1 , choosing r 1 to infer environment partition violates Condition 2,which is reflected by the poor performance of ZIN with r 1 .

Original: Using [ X , Y ] as input to ρ ( · ) is alsoa violation by Corollary 1, and the corresponding results in Table 7 confirm our analysis.
Revised: Using [ X , Y ] as input to ρ ( · ) is alsoa violation by Corollary 1, and the corresponding results in Table 6 confirm our analysis.

Original: Lastly,notice that X s ∈ X contains some information of Y and there may exist a function ρ ( · ) so that
Revised: Lastly,notice that X s ∈ X contains some information of Y and there may exist a function ρ ( · ) so that

Original: H ( Y | X v ) >
Revised: H ( Y | X v ) >

Original: H ( Y | X v , ρ ( X )) .
Revised: H ( Y | X v , ρ ( X )) .

Original: This violates Condition 1 and leads to poor results when choosing
Revised: This violates Condition 1 and leads to poor results when choosing

Original: X as input to ρ ( · ) .
Revised: X as input to ρ ( · ) .

Original: We also empirically verify the choice of hyper-parameter K using the same simulated setting.
Revised: We also empirically verify the choice of hyper-parameter K using the same simulated setting.

Original: Figure 2shows that K has a relatively small impact, especially when K ≥ 4 .[Figure] [Table] [Figure]
Revised: Figure 2shows that K has a relatively small impact, especially when K ≥ 4 .

Original: F.3 About the Inferred Environments on CelebA
Revised: E Societal Impact

Original: We visualize the inferred environments of the CelebA dataset in this section. The target Smiling isspuriously correlated with feature Gender , i.e., most females are smiling while most males are not. We set K to be 2 (our framework is insensitive to K as shown in Appendix F.2).

Original: We visualize thespurious correlations in the two inferred environments during training in Fig.

Original: ZIN can generatetwo environments where the spurious correlations differ.

Original: Then we can easily discard the spuriousfeature using IRM methods on the inferred environments.

Original: G How to Choose Z

Original: Conditions 1 and 2 give the theoretical requirements of Z to successfully learn invariant features. Inthis section, we discuss how to find such auxiliary information in applications. G.1 Checking Condition 1 We start with an example in Fig.
Revised: In this work, we propose to utilize the auxiliary information to aid the invariance learning withoutenvironmental indexes.

Original: When we collect a sample (e.g., taking a photo), there often existsa set of meta information (e.g., the time slot, coordinates, temperature).

Original: Consider the label of theimage (e.g., “camel”) is generated by the causal graph on the right panel of Fig.

Original: The invariantfeature X v are X 1 and X 2 in this case.

Original: The meta information is used as Z , which is colored ingreen in the causal graph.

Original: In this example, node Z has correlations with some nodes in the causalsystem. This Z satisfies Condition 1 because H ( Y | X 1 , X 2 ) = H ( Y | X 1 , X 2 , Z ) or equivalently
Revised: This method is also helpful to fairness issues; see, e.g., the discussion aboutout-of-distribution generation and algorithmic fairness in [8].

Original: Y ⊥ Z | X 1 , X 2 .

Original: This example gives us some evidence of choosing Z : Z and Y should be d-separatedby X v .

Original: There are other examples satisfying Condition 1, as shown in Fig.

Original: We also illustrate some failure choices of Z , illustrated in Fig. 6.

Original: Specifically, when Z is a child of Y or a direct cause of Y , Z contains additional information of Y conditional on X 1 and X 2 . Then Condition 1 is violated and this kind of information cannot be used in our framework.
Revised: For the additional information, weshould also avoid using demographic, private, and sensitive information.

Original: G.2 Checking Condition 2

Original: Though Condition 2 is needed to discard all spurious features, it is hard to check in practice becausewe generally have no prior knowledge on all the possible spurious features.

Original: Thus, we can collectas much auxiliary information as possible, which satisfies Condition 1 and also has correlationswith some nodes in the causal system of interest.

Original: Then we can discard as many spurious features aspossible while keeping the invariant features.

Original: A natural question to ask is what would happen if Condition 2 is only satisfied partially , i.e., thecondition holds only for some X

Original: ks ∈ X s .

Original: The experimental results on CelebA in Table 3 already giveus some hints.

Original: As we collect more auxiliary information, we find that the test performance increases.

Original: Even when we have limited (insufficient) auxiliary information, we can also obtain a performancegain over ERM, e.g., ZIN(4) outperforms ERM by a large margin.

Original: In this section, we present anadditional theoretical result: if Condition 2 is satisfied partially, we are able to discard spuriousfeatures partially and still preserve all invariant features.

Original: The main result is given in the following.

Original: Proposition 3.

Original: Assume that Assumptions 1-4 and the following conditions are true:

Original: • Condition 1 holds: given invariant feature X v and any function ρ ( · ) , it holds thatH ( Y | X v , ρ ( Z ))

Original: =

Original: H ( Y | X v ) .

Original: • Condition 2 holds partially : there exists a proper subset of spurious feature X ¯ s ⊊ X sand for any X k ¯ s ∈ X ¯ s , there exists a function ρ ( · ) and a constant C > 0 such thatH ( Y | X k ¯ s )

Original: − H ( Y | X k ¯ s , ρ ( Z)) ≥ C . Further, let X sc :=
Revised: Z

Original: X s \ X ¯ s , and for any function

Original: ρ ( · ) , it holds that H ( Y | X sc , ρ ( Z ))

Original: =

Original: H ( Y | X sc ) .

Original: Let Φ v + sc denote the feature mask that selects both X v and X sc .

Original: If ϵ <

Original: Cγδ 4 γ +2 CδH ( Y ) and λ ∈

Original: [ H ( Y )+1 / 2 δC δC − 4 ϵ − 12 , γ 4 ϵ − 12 ] , then we have ˆ L (Φ v + sc ) <

Original: ˆ L (Φ) for all Φ ̸ = Φ v + sc .

Original: Proposition 3 shows that if Condition 1 is satisfied for X v and Condition 2 is only satisfied for partof the spurious features X ¯ s , then our framework can still successfully discard X ¯ s , while keeping X vand the rest of the spurious features.

Original: Proof. The proof proceeds almost the same as that of Theorem 2, by replacing Φ v with Φ v + sc . Notethat X v + sc behaves as “invariant” under the inferred environments in Proposition 3.
Revised: Condition

Original: Discarding anymember of X v + sc will increase the ERM loss.

Original: Adding any member of X ¯ s also increases the lossdue to a positive penalty (with a sufficient large weight λ ).

Original: G.3 Take-away Message on Choosing Z

Original: To summarize, we make the following suggestions to collect Z in practice:

Original: • Pick Z that satisfies Condition 1 and that is correlated with nodes in the causal graph ofinterest.

Original: Further, the path between Z and Y should be d -separated by X v as shown in

Original: Figures 4 and 5.

Original: Notably, Z cannot be the parent or the child of Y as shown in Figure 6. • Collect as many such Z ’s as possible. H Societal Impact In this work, we propose to utilize the auxiliary information to aid the invariance learning withoutenvironmental indexes.
Revised: Condition

Original: This method is also helpful to fairness issues; see, e.g., the discussion aboutout-of-distribution generation and algorithmic fairness in [9].
Revised: [Table]

Original: For the additional information, weshould also avoid using demographic, private, and sensitive information.
Revised: [Table] [Table] [Table] [Table]

Revised: ✓ 71.

Revised: K (Number of inferred environments)

Revised: [Figure]

Revised: Test mean

Revised: Test worst

Revised: Train