Original: ZIN: When and How to Learn Invariance by Environment Inference? Anonymous Author(s) Affiliation
Revised: ZIN: When and How to Learn Invariance Without Environment Partition?

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 of the data distribution may vary but the underlying causal mechanisms remain constant.

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 the heterogeneity, recent invariant learning methods have proposed to learn robust and invariant models using this environment partition.

Original: It is hence tempting toutilize the inherent heterogeneity even when environment partition is not provided.
Revised: It is hence tempting to utilize 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 this circumstance is fundamentally impossible without further inductive biases or additional 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 environment partition 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 provably identify 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 on both synthetic and real world datasets validate our analysis and demonstrate an improved performance of the proposed framework.

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

Revised: Codes are available at https://github.com/linyongver/ZIN_official .

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 independent and 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 environments in 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, model performance can severely degrade [47, 17, 38], as ERM based algorithms may exploit the spurious correlations 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 methods propose to utilize the causally invariant mechanisms (rather than the spurious correlations in the training data) that are supposed to be stable across different environments.

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

Original: [33] propose to exploit the invariance principle to learn a linear model that merely relies onthe direct causes of the target.
Revised: [36] proposes to exploit the invariance principle to learn a linear model that merely relies on the 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 any variable 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) to capture invariant correlations by learning representations that elicit an optimal invariant predictor across 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: [21], 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: [24], Xie et al. [52], Chen et al. [8, 9] further develop variants of IRM by introducing game theory, regret minimization, variance penalization, multi-objective optimization, etc., and Xu and Jaakkola [54], 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: [26] try to learn invariant features by coupled adversarial neural networks.

Revised: Some recent works report that IRM is less effective when applied to large neural networks [26, 15].

Revised: Lin et al.

Revised: [27] find that this can be largely attributed to the overfitting problem, and Zhou et al.

Revised: [57, 56] propose to alleviate this issue through imposing sparsity constraint and sample reweighting, respectively.

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 partitioned into 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, such an environment partition may be unavailable or hard to obtain in practice [29, 10].

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 line of works try to learn invariance without environment indexes where the dataset is assembled by merging data 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: [10] propose environmental inference for invariant learning (EIIL), a two-stage method by firstly inferring the environments with an ERM based trained biased model and then performing invariant learning on inferred environments.

Original: Liuet al.
Revised: Liu et 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: [29] devise an interactive mechanism, called heterogeneous risk minimization (HRM), with environment inference and invariant learning on raw feature level, and Liu et al.

Original: [27] extend it torepresentation level with kernelized trick.
Revised: [30] extend it to representation 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 input-label pairs) can be risky , both theoretically and empirically.

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 indexes is 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 and a general theoretic result (Theorem 1), to show that the causally invariant features are unidentifiable if no environment information 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 learncausally invariant features.
Revised: This impossibility result, similar to the identifiability issue in causal discovery [45, 37], independent component analysis (ICA) [20], and unsupervised learning of disentangled representations [32], motivates us to consider when and how to learn 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 latent heterogeneity.

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 less restrictive than requiring ideal environment partition according to distribution heterogeneity; and 2) it also obtains theoretic guarantee which is missing in the case with no environment information at 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 [53, 51].

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, 33], locations (longitude and latitude) of collected satellite data in remote sensing [14, 40], and meta information associated with an instance [31].

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 propose a 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 a fairly general setting, we derive both sufficient and necessary conditions for our framework to identify invariant features.

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

Original: Finally, Section 7 presents experimental analysis on both synthetic and real world datasets.
Revised: Finally, Section 6 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: The proposed 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 from a causal perspective [36].

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 and potential 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: [19] consider heterogeneous datasets 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: Here confounders are assumed as functions of the domain index or time.

Original: Tillman and Spirtes [45], Dankset al.
Revised: Tillman and Spirtes [48], Danks et al.

Original: [10], Huang et al.
Revised: [11], 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: [18] study the multiple dataset setting with non-identical sets of variables for causal discovery, where the datasets may have different exogenous noise distributions and only part of 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 for generalizing beyond training data is distributionally robust optimization (DRO) [5, 43, 25, 13, 12, 55].

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 training distribution.

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 loss over groups in the training data [41].

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 generally requires 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 know a 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 proposed to infer group partition or directly identify the minority group by looking at features produced by biased models, e.g., [44, 1, 10, 28].

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 are still 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: [28] upweights samples with high losses from the initial ERM model and relies on 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: [34], Sanh et 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: [42] try to obtain a more robust model by boosting from the wrongly specified samples, based on 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 or invariant 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 invariant features 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 = where 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 the environment 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 generation process [35, 37].

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 acyclic graph, 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 of an assignment of a deterministic function depending on the parental variables in the graph and an exogenous 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 that is 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 denote a non-degenerate deterministic function and ϵ v is an independent noise variable.

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 e =
+} α e P ( X e , Y e )for some α e > 0 and (cid:80) Ee =1 α e = 1 .
Revised: (cid:80) {+E e =
+} α e P ( X e , Y e ) for some α e > 0 and (cid:80) E e =1 α e = 1 .

Original: We consider the model as the composition of a feature extractor Φ and a classifier (or predictor) f ωthat is parameterized by ω .
Revised: We consider the model as the composition of a feature extractor Φ and a classifier (or predictor) 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 cross-entropy loss and 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 the following 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 direct causal parents of Y as invariant features, as in [36, 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 of the 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, then the conditional probability could change after some intervention.

Original: As such, the underlying SCMgoverning data generation determines the desired invariance.
Revised: As such, the underlying SCM governing 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 heterogeneous data of multiple environments with unknown environmental indexes.

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

Revised: [Equation]

Original: Example 1.
Revised: Example.

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 we observe 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 Eq.

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 .   Y = 0 , w.p. 0 . Y = 1 , w.p. 0 .
Revised: A learning algorithm that tries to learn invariant feature based on this dataset would result in a model that (deterministically) depends on either X 1 ,

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

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

Original: X 2 ̸ =

Original: [Equation]

Original: 1 = Y[Equation] w.p. 0 .[Equation]
Revised: X 2 , or both to predict Y .

Original: w

Original: [Equation]

Original: p

Original: [Equation]

Original: [Equation]

Original: w

Original: [Equation] p [Equation] 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 the learned features are invariant is determined by the underlying data generation process.

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

Original: • X 1 is the invariant feature while X 2 is spurious, 
Revised: • X 1 is the invariant feature while X 2 is spurious, where p e =1 s = 0 .

Original: where p e =1 s
Revised: 8 and p e =2 s

Original: = 0 .
Revised: [Equation]

Original: 7 and p e =2 s = 0 .
Revised: X 1 ∼ Bernoulli(0.5) ,

Original: [Equation]
Revised: [Equation]

Original: • X 2 is the invariant feature while X 1 is spurious where p e =1 s = 0 .
Revised: • X 2 is the invariant feature while X 1 is spurious where p e =1 s = 0 .

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

Original: [Equation]
Revised: [Equation]

Original:  ∼ Bernoulli
Revised: X 2 ∼ Bernoulli(0.5) ,

Original: [Equation]
Revised: [Equation]

Revised: X

Original: [Equation]
Revised: [Equation]

Original: The joint distribution of  mixed environments for each case is consistent with Eqn.
Revised: The joint distribution of the mixed environments for each data generation process is consistent with Eq.

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.Thus, the true invariant featurescannot be identified.
Revised: Thus, from the joint distribution P ( X 1 , X 2 , Y ) , a learned model only depends on either X 1 , X 2 or both, and fails to generalize for at least one of the two data generation processes.

Original: On the other hand,notice that when the partition is given, one can verify that IRM would correctly identify X 1 (resp.
Revised: On the other hand, when the partition is given, one can verify that 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) scenario.

Revised: The above toy example is inspired by commonly-used classification tasks in the IRM literature, e.g., CMNIST [4] and CifarMnist [26].

Revised: For instance, in CMNIST, invariant feature X 1 denotes the semantic feature of the shape of hand-written digits ‘0’ and ‘1’, and spurious feature X 2 represents the color, which is either red or green.

Revised: Label Y ∈ { 0 , 1 } corresponds to the digit shape and is also binary.

Revised: Indeed, the construction procedure of CMNIST in [4] can be exactly described by the data generation process in Eq.

Original: Moreover, such an example is not rare.
Revised: To further demonstrate the impossibility result, we next conduct an empirical validation on a new dataset in accordance to the data generation process in Eq. Empirical Validation. We construct a variant of CMNIST according to the second data generation process. In this new dataset, color is the invariant feature and digit shape is spurious.

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

Revised: As described in Eq.

Revised: (4), the label corresponds to the digit color and the digit shape is spuriously correlated with the label.

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

Revised: In the test domain, we set p e =3 s = 0 .

Revised: 1 to simulate the distributional shift, same as in CMNIST.

Revised: We can then compare ERM, EIIL, IRM with environment partition, and LfF (learning from failures) [34] that tries to learn a robust model by boosting from wrongly specified samples of shallow neural networks.

Revised: The empirical results are reported in Table 1.

Revised: For ERM (oracle), we train the model only on the invariant feature, i.e., digit shape in CMNIST and color in MCOLOR.

Revised: We can see that the EIIL method performs poorly on the MCOLOR dataset.

Revised: This is due to the inductive bias of EIIL and it indeed relies on the digit shape as the invariant feature.

Revised: Since we have no prior knowledge of the data generation process, the true invariant feature can be the color, e.g., in the MCOLOR dataset.

Revised: Similarly, ERM and LfF all rely on either color or shape as the invariant feature and would fail on at least one of CMNIST and MCOLOR.

Revised: On the other hand, if environment partition is available, IRM can still learn the desired invariant feature.

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: [Table] Moreover, such examples are not rare. For any data distribution P ( X , Y ) generated 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 different invariant/spurious features.

Original: Similarly, it is impossible to find the minority group that is determined byspurious features, based ononly the data P ( X , Y ) .
Revised: Similarly, it is impossible to identify the spurious features only based on 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 the 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 the joint distribution P ( X , Y ) consisting of data from multiple environments with P ( X s , Y ) ̸ = 0 , we can find X ′ v and X ′ s as invariant and spurious features so that:

Original: so that: (a) X ′ v ̸=
Revised: • 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 ̸ = ∅ ; • ( 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 B.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 additional assumptions/conditions or certain “inductive bias” to identify the invariant features [45, 37, 20, 32, 50].

Original: The latter maycome implicitly from the proposed algorithm, model, and/or the data.
Revised: The latter may come 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 the reason that makes existing methods achieve improved empirical performance in their considered scenarios.

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 [29, 30], the discrepancy of spurious features among clusters are expected to be larger than that of causal features, and in [10, 34, 42, 28], the ERM model of the first stage should heavily or fully rely on the spurious feature.

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

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

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 exists additionally observed variable Z with the data, which has been often considered in the literature [19, 20, 22, 53, 51].

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, for example, the time index in a time series, some kind of class labels, concurrently observed variables, or human annotations on potential unmeasured variable [46].

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: [53] also take advantage of auxiliary information to help improve OOD performance in a semi-supervised manner.

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 data together 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 to learn invariant models from the heterogeneous 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: We consider that there exists additional auxiliary information Z ∈ R d z in company with the data ( X , Y ) . In this section, we propose ZIN, auxiliary information Z for environmental INference , for invariant learning from the heterogeneous dataset D without environment partition.

Original: We  derive conditions for invarianceidentification and show these conditions are both sufficient and necessary in a fairly general setting.
Revised: We will derive conditions for invariance identification 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 ) .

Original: 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.
Revised: We aim to learn a function ρ ( · ) : R d z −→ R K that softly assigns a sample to K environments.

Original: Here
Revised: Here

Original: K is a pre-specified number (a hyper-parameter) and its choice will be empirically investigatedin our experiments.
Revised: K is a pre-specified number (a hyper-parameter) and its choice will be empirically investigated in our experiments.

Original: Let ρ ( k ) ( · ) denote the k -th entry of ρ ( · ) . Then we have ρ ( Z ) ∈ [0 , 1] K with 180 (cid:80)k
Revised: Let ρ ( k ) ( · ) denote the k -th entry of ρ ( · ) , with ρ ( Z ) ∈ [0 , 1] K and (cid:80)

Original:  ρ ( k ) ( Z ) = 1 .
Revised: k ρ ( k ) ( Z ) = 1 .

Original: Denote the ERM loss as R ( ω, Φ) =
Revised: Denote the ERM loss as R ( ω, Φ) =

Original: 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: 1 n (cid:80) ni =1 ℓ ( f ω (Φ( x i )) , y i ) and the loss in the k -th inferred environment as R ρ ( k ) ( ω, Φ) = 1 n (cid:80) ni

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

Original: Suppose that the environments have been given according to a fixed ρ ( · ) . Recall that IRM [5] learnsan invariant representation Φ , upon which there is a classifier  ω that is simultaneously optimal in allenvironments.
Revised: Recall that IRM [4] learns an invariant representation Φ , upon which there is a classifier f ω that is simultaneously optimal in all environments. Suppose that the environments have been given according to a fixed ρ ( · ) .

Original: To measure the optimality of  ω in environment k , we can fit an environment dependentclassifier ω k on the data from environment k .
Revised: To measure the optimality of f ω in the k -th environment, we can fit an environmentdependent classifier f ω k on the data from that environment.

Original: If  ω k achieves a smaller loss than ω , it means ω is notoptimal in this environment.
Revised: If f ω k achieves a smaller loss, then we know that f ω is not optimal in this environment.

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: We can further train a set of environment-dependent classifiers { f ω k } Kk =1 , one for each environment, to measure whether f ω is simultaneously optimal in all environments.

Original: Thus, when ρ ( · ) is provided, our formulation to learn invariance isthen given as follows:
Revised: Thus, when ρ ( · ) is provided, our formulation to learn invariance is as follows:

Original: [Equation]
Revised: [Equation]

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.
Revised: If Φ extracts spurious features that are unstable in the inferred environments, R ρ ( k ) ( ω, Φ) will be larger than R ρ ( k ) ( ω k , Φ) , resulting in a non-zero invariance penalty.

Original: Next, we consider how to learn the partition function ρ ( · ) .
Revised: Next, we consider how to learn the partition function ρ ( · ) .

Original: A good partition function should generateenvironments among which the spurious feature exhibit instability. If Φ extracts spurious features,we expect to have a large penalty resulting from the environments.
Revised: A good partition function should generate environments among which the spurious features exhibit instability, so that there is a large penalty if Φ extracts spurious features.

Original: Thus, we seek for an environmentpartition that maximizes the invariance penalty.
Revised: Thus, we seek for an environment partition that maximizes the invariance penalty.

Original: The overall framework is provided below:
Revised: The overall framework is provided below:

Original: [Equation]
Revised: [Equation]

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 learning invariant 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: [10] has used a similar intuition and 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-stage method cannot be jointly optimized, and the environment partition relies on the given model and lacks a theoretical 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 proposed framework 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 are also 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 [10], is unfortunately an instance 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 theoretical perspective.

Original: We start with a simple yetgeneral setting: X =
Revised: We start with a simple yet general 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-wise feature selection mask, and f ω is a general non-linear function.

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

Original: Extension to other loss functions is left to Appendix B.
Revised: Extensions to other loss functions and linear feature transformations are left to Appendix A. and A.2, respectively.

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 invariant features, 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 ˆ L (Φ) =

Original: [Equation]
Revised: min ω max ρ, { ω k } L (Φ , ω, ω 1 , · · · , ω k , ρ ) .

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 6 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 our main 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 of an optimal classifier over some X ′ and Y , and similarly H ( Y | ρ ( Z ) , X ′ ) to denote the minimum risk (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 exactly one environment, i.e., ρ ( Z ) is one-hot, it can be verified that the optimal loss H ( ·|· ) coincides with conditional 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 that221 E [ ℓ ( f (Φ( X )) , Y )] ≤ H ( Y | Φ( X ))
Revised: For a given feature mask Φ and any constant ϵ > 0 , there exists f ∈ F such that E [ ℓ ( f (Φ( X )) , Y )] ≤ H ( Y | Φ( X ))

Original: + ϵ.
Revised: + ϵ.

Original: Assumption 2.
Revised: Assumption 2.

Original: If two features are invariant w.r.t. an environment partition, then the concatenatedfeatures are also invariant. That is, for X 1 , X 2 ⊂ X and ρ ( · ) , if H ( Y | X 1 ) − H ( Y | ρ ( Z ) , X 1 ) ={+invariance constraint, adding another feature would not make the penalty vanish, i.e., there exists a constant δ > 
+}
Revised: If a feature violates the {+invariance constraint, adding another feature would not make the penalty vanish, i.e., there exists a constant δ > 
+}

Original: 0and H ( Y | X 2 ) − H ( Y | ρ ( Z ) , X 2 ) = 0 ,we have H ( Y | X 1 , X 2 ) − H ( Y | ρ ( Z ) , X 1 , X 2 ) = 0 .
Revised: so that for spurious feature X 1 ⊂ X s and any feature X 2 ⊂ X , H ( Y | X 1 , X 2 ) − H ( Y | ρ ( Z ) , X 1 , X 2 ) ≥ δ ( H ( Y | X 1 ) − H

Original: Assumption 3.

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

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

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

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 distinct 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 exists f ∈ F that can fit P ( Y | Φ( X ))

Original: well.
Revised: well.

Revised: Assumption 2 aims to ensure a sufficient positive penalty if a spurious feature is included (see Appendix A.3 for further discussion).

Original: Assumption 4 indicates that any featurecontains some useful information w.r.t.
Revised: Assumption 3 indicates that any feature contains 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 [49]), as it does not affect prediction.

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: Given invariant feature X v and any function ρ ( · ) , it holds that H ( Y | X v , ρ ( Z ))

Original: =
Revised: =

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

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 a function ρ ( · ) 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 v , Z )

Original: =
Revised: =

Original: H ( Y | X v ) (a proof is provided in Appendix
Revised: H ( Y | X v ) (a proof is provided in Appendix

Original: A.7).
Revised: B.6).

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 invariant w.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 an invariant 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 2 implies that for each spurious feature, there exists at least one partition so that this feature is non-invariant 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 all possible 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 these conditions, our main result follows and a proof is given in Appendix B.2.

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-3 and Conditions 1-2, if ϵ <

Original: Cγδ 4 γ +2 CδH ( Y ) and λ ∈ [ H ( Y )+1 / 2 δC δC − 4 ϵ − 12 , γ 4 ϵ − 12 ] , then we have ˆ L (Φ v ) <
Revised: [Figure] Y ) and λ ∈ [ 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 6 identifies invariant features.

Original: 5. Necessary Conditions for Identifiability
Revised: 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.

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 are also 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 that if Conditionis violated, then some invariant features will be excluded in the solution of Problem 6; 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: With Assumptions 1-3, if Condition 1 is violated, i.e., there exists ρ ( · ) so that

Original:  ( y | X v , ρ ( Z )) ≥ C ′ > 0 , then there exists afeature mask Φ ′ ̸
Revised: H ( y | X v ) − H ( y | X v , ρ ( Z )) ≥ C ′ > 0 , then there exists a 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 B.3, which is similar to the first step of the proof of Theorem 2.

Original: Aquestion is how Condition 1 could be violated.
Revised: A question 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 additional variables, it is also interesting to know whether we can obtain a valid partition based on only the training data ( X , Y ) .

Original: Below we provide examples regarding these questions.
Revised: Below we provide examples regarding these questions, with proofs given in Appendix B.4.

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

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: there exists a function ρ ( · ) and an injective function 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 ) , or (c) ρ ( Z ) =

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

Original: A proof is given in Appendix A.4.
Revised: We now discuss the necessity of Condition 2, with an additional assumption. Assumption 4. If two features are invariant w.r.t. an environment partition, then the concatenated features are also invariant.

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: That is, for X 1 , X 2 ⊂ X and ρ ( · ) , if H ( Y | X 1 ) − H ( Y | ρ ( Z ) , X 1 ) = and H ( Y | X 2 ) − H ( Y | ρ ( Z ) , X 2 )

Revised: = 0 , we have H

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

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: With Assumptions 1, 3 and 4, if Condition 2 is violated, i.e., there exists a spurious 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 invariance 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 proof in Appendix B.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 and informative so that each spurious feature can be recognized.

Original: Beyond Feature Selection: Linear Feature Transformations
Revised: 5.4 Choice of the Auxiliary Information

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: Conditions 1 and 2 present theoretical requirements of Z to successfully learn 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]. Specifically, we consider the same data generation process as [5]:
Revised: We now discuss how to find such auxiliary information given certain prior knowledge.

Original: [Equation]
Revised: Take Fig.

Original: where β ∈ R d v and W ∈ R d × ( d v + d s ) .
Revised: 1 for example. When we collect a data point (e.g., taking a photo), there often exists some meta information (e.g., time slot, coordinate, and temperature).

Original: We assume that there exists ˜ W ∈ R d v × d so that˜ W ( W [ x v ; x s ]) =
Revised: Suppose that the image is generated according to the causal graph on the right panel of Fig.

Original: x v for all x v and x s . Both the  feature extractor and predictor take a linearform, i.e., Φ takes values in R d × d and ω ∈ R d .
Revised: In this case, the invariant feature X v consists of X 1 and X 2 , and the meta information can be used as Z which is colored in green in the causal graph.

Original: The prediction for X is ω ◦ Φ( X  ) =
Revised: This is because H ( Y | X 1 , X 2 ) =

Original: (Φ X  ) T ω .
Revised: H ( Y | X 1 , X 2 , Z ) or equivalently Y ⊥ Z | X 1 , X 2 , so that Condition 1 holds.

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: This example shows some evidence of choosing the auxiliary information: Z and Y should be d-separated by X v .

Original: Following [5, 36], we consider a constrained form of Problem 4 for theoretic analysis:
Revised: There are other examples satisfying Condition 1, which are given

Original: [Equation]
Revised: Fig.

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: 5 in Appendix E. We also illustrate some failure choices in Fig.

Original: This is inaccordance to existing conditions for identifiability in the linear case [5, 36]. In this paper, we utilizesuch a condition, linear general position condition, from [5]. 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: 6 in Appendix E. Specifically, if

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. Our identifiability result for the linear case then follows.
Revised: Z is a child or a direct cause of Y , then it contains additional information of Y conditional on the invariant feature, violating Condition 1. This kind of information cannot be used in our framework.

Original: Theorem 3. 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) E X k (cid:2)
Revised: To summarize, with some prior knowledge on the path between Z and Y , we suggest to choose as much auxiliary information as possible under the following condition: A feasible choice of Z should satisfy Condition 1 and should be correlated with certain variables in the causal graph of the SCM.

Original:  Xk X k ⊺ (cid:3) x − E X k ,ϵ v (cid:2) X k ϵ v (cid:3)(cid:9) k
Revised: The path between Z and Y should be d-separated by X v (e.g., in Figs.

Original: (cid:1)(cid:1) > d − r.
Revised: 1 and 5) or there is no path between Z and Y .

Original: If Φ ∈ R d × d has rankr > 0 , then Problem 6 results in the desired invariant predictor. A proof is given in Appendix A.6. We can also replace the above assumption by the non-degenerateconditions in Eqn. (9) and (10) of [36], and a similar identifiability result follows.
Revised: Notably, Z cannot be the parent or child of Y as shown in Fig. 6.

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 synthetic and 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], group DRO [41], EIIL [10], HRM [29], and LfF [34].

Original: We provide the ground-truth partition to IRM and group DRO.
Revised: We provide the ground-truth partition to IRM and group DRO.

Original: LfF tries to learn a robust model by directly boosting from wrongly specified samples ofshallow 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.

Revised: We adopt the first order approximation of Eq.

Revised: (5), as described in Appendix C.

Original: The number of inferredenvironments K is set to be 2 as default.
Revised: The number of inferred environments 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 synthetic dataset and house price prediction, ResNet-18 [16] for CelebA, and 1D CNN for Landcover.

Original: Pleasesee Appendix D for more details.
Revised: More details are provided in Appendix C.

Original: 7.1 Synthetic Dataset
Revised: 6.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]
Revised: [Equation]

Original: , −

Original: [Table]

Original: ( X

Original: [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-dimensional spatial 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] .

Revised: We

Revised: [Table]

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: evaluate the performance on four distinct test environments with p s ∈ { 0 . 999 , [Table] [Table] . 2 , 0 . 1 } and p v being constant.

Original: We use t as the auxiliary information Z .
Revised: We use t as the auxiliary information.

Original: Moreimplementation details are given in Appendix D
Revised: More details are provided in Appendix C

Original: Results .
Revised: Results .

Original: Table 1 reports the test accuracy.
Revised: Table 2 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 ERM is 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) and p 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 and may learn both spurious and invariant features.

Original: The proposed method ZIN improves the worst testaccuracy significantly.
Revised: The proposed method ZIN improves the worst test accuracy 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 ERM and 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 the ground-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 theoretical 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 hyper-parameter K .

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

Original: 7.2 Real World Datasets
Revised: 6.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 sales prices 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: The target variable is house price and each sample contains 17 predictive variables like the built year of the house, 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 dataset in 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.
Revised: We normalize the prices of the houses 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 auxiliary information 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.[Table]
Revised: As there is no well-defined ground-truth partition for IRM, we manually split the training dataset equally into 5 segments with 10 -year range in each segment.

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

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

Original: We observe that theresults of EIIL and HRM with additional information Z are worse than those without using Z .
Revised: We observe that the results 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: This implies 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 [31], and by construction the 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 CeleA images 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 potential auxiliary 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 is made available to IRM and group DRO to create an oracle environment partition.

Original: Empirical results are given in Table 3, where ZIN(#) represents ZIN using # of the seven features as Z .
Revised: Empirical results are given in Table 4, where ZIN(#) represents ZIN using # of the seven features as 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 worstperformance on both training and test data, showing that it may not learn the causalfeatures.
Revised: The worst test performances of ERM, EIIL and LfF indicate that they may not learn the causal features.

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.

Revised: This observation validates our discussion on choosing the auxiliary information in Section 5.4.

Original: [Table]
Revised: [Table]

Revised: [Table]

Original: Landcover .
Revised: In Appendix D [Table] Fig. 3 reports the results of different choices of hyper-parameter K and Fig.visualizes the inferred environments in this experiment.

Revised: Again, we find that K has a relatively small impact on the proposed algorithm.

Original:  Our final task is land cover prediction thatclassifies the land cover type (e.g., grasslands) fromsatellite data [13, 37].
Revised: [Table] Landcover . Our final task is land cover prediction that classifies the land cover type (e.g., grasslands) from satellite data [14, 40].

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 [53]: 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 Africa as test data.
Revised: We also consider non-African locations as training data and 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 time dimension 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 5, we see that ZIN with location achieves 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 is because climate variables indeed contain some information for predicting the target, violating our conditions on 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 has learned 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 be risky as the first stage may not be guaranteed.

Revised: 6.3 Discussion on Auxiliary Information for Related Methods

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

Revised: It is interesting to ask whether this 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 series tasks, 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 6.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 explicit environment 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 some additional 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 linear feature 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 improved performance 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 make the role of inductive biases more explicit, when learning invariance from heterogeneous data without environment 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 general nonlinear feature extractors, which is also an open problem for IRM even with environment partition.

Original: Checklist
Revised: A More Theoretical Results A.1 Extension to Other Loss Functions

Original: For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
Revised: Sections 5.2 and 5.3 focus on classification task with cross-entropy loss, to establish sufficient and necessary conditions for invariance identification.

Original: [Yes]
Revised: Similar results can be shown for other loss functions or tasks in a straightforward manner. To see this, notice that H ( ·|· ) in Assumptions 1-3 and Conditions 1-2 is used to represent the optimal expected risk that coincides with the conditional entropy, with cross-entropy loss and when ρ ( z i ) gives exactly one environment. For other loss functions 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: (b) Did you describe the limitations of your work? [Yes]
Revised: Replacing H ( ·|· ) with L ( ·|· ) , we can follow the same proof procedure to obtain similar sufficient and necessary conditions. This is summarized in

Original: In the conclusion section.
Revised: Theorem 3 for completeness.

Original: (c) Did you discuss any potential negative societal impacts of your work?
Revised: Theorem 3. Consider Assumptions 1-3 and Conditions 1-2 where H ( ·|· ) is replaced with L ( ·|· ) accordingly.

Original: [Yes] InAppendix H.
Revised: If ϵ < [Equation]

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

Original: [Yes]
Revised: [Equation]

Original: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
Revised: −

Original: [Yes]
Revised: [Equation]

Original: (b) Did you include complete proofs of all theoretical results? [Yes] If you ran experiments...
Revised: ϵ − 12 ] , we have ˆ L (Φ v ) < ˆ L (Φ) for all Φ ̸ = Φ v , where we assume L ( Y ) < ∞ .

Original: (a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)?
Revised: If Condition 1 or Condition 2 is violated, then there exists a feature mask Φ ′ ̸

Original: [Yes]
Revised: =

Original: (b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
Revised: Φ v so that ˆ L (Φ v ) > ˆ L (Φ ′ ) . A.2 Beyond Feature Selection: Linear Feature Transformations

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

Original: (c) Did you report error bars (e.g., with respect to  the random seed after running experiments multiple times)?
Revised: In this section, we extend feature selection to feature learning, and show that ZIN is able to learn the invariant features given a scrambled observation in a linear form following [4, 39]. Specifically, we consider the same data generation process as [4]:

Original: [Yes]
Revised: [Equation]

Original: (d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)? [Yes]
Revised: where β ∈ R d v and W ∈ R d × ( d v + d s ) . We assume that there exists ˜ W ∈ R d v × d so that ˜ W ( W [ x v ; x s ]) =

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
Revised: x v for all x v and x s . Both the feature extractor and predictor take a linear form, i.e., Φ takes values in R d × d and ω ∈ R d . The prediction for X is ω ◦

Original: [Yes]
Revised: Φ( X ) =

Original: (b) Did you mention the license of  the assets?
Revised: (Φ X ) T ω . A major difficulty in this setting lies in how to characterize the effect of an invariant/spurious feature in a quantitive way: the feature extractor Φ may extract an arbitrarily small portion of spurious information. Following [4, 39], we consider a constrained form of Problem 6 for theoretic analysis:

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

Original: (c) Did you include any new assets either in  the supplemental material or as a URL?
Revised: As in Conditions 1 and 2, the auxiliary information should also be sufficiently informative so that the inferred environments can be diverse enough but also maintain the underlying invariance.

Original: [N/A] (d) Did you discuss whether and how consent was obtained from people whose data you’re
Revised: This is in accordance to existing conditions for identifiability in the linear case [4, 39].

Original: using/curating?
Revised: In this paper, we utilize such a condition, linear general position condition, from [4].

Original: [N/A]
Revised: For our analysis, we take squared error as our loss function and consider that ρ ( · ) partitions the environments in a hard manner, i.e., each data sample would be assigned to exactly one environment. As in Appendix A.1, we use L ( ·|· ) to represent the optimal expected risks.

Original: (e) Did you discuss whether the data you are using curating contains personally identifiable information or offensive content?
Revised: We also assume that the environments are non-degenerate, i.e., each inferred environment contains some data samples; otherwise, we can simply remove such an environment.

Original: N/A
Revised: Our identifiability result for the linear case then follows. Proposition 3. Assume Condition 1 where H ( ·|· ) is replaced with L ( ·|· ) accordingly. Suppose 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-zero x ∈ R d : dim (cid:0) span (cid:0)(cid:8)

Revised: E X k (cid:2)

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

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

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

Revised: d − r.

Revised: If Φ ∈ R d × d has rank r > 0 , then

Revised: Problem 8 results in the desired invariant predictor.

Revised: Proof.

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

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

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

Revised: Φ E X k

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

Revised: Φ E X k ,Y k

Revised: [ X k Y k ] for all k

Original: If you used crowdsourcing or conducted research with human subjects.
Revised: if and only if Φ elicits the desired invariant predictor Φ T ω = ˜ W T β . Thus, we only need to show our solution meets the same normal equations. Let Φ ′ and ω ′ denote a solution to Problem 8.

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

Original: (a) Did you include the full text of instructions given to participants and screenshots, if
Revised: R k ( ω k , Φ ′ ) . Notice that ω k minimizes the mean squared error on only the k -th environment, hence it must satisfy the normal equation Φ ′ E X k

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

Original: applicable?
Revised: k = Φ ′ E X k ,Y k [ X k Y k ] .

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

Revised: Thus, no spurious information will be included and X Φ =

Revised: [ S X v ; 0 d s ] where S ∈ R d v × d v is an invertible matrix and 0 d s denotes a d s -dim vector of all zeros.

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

Revised: Under Condition 1 where H ( ·|· ) is replaced with

Revised: L ( ·|· ) , we have L ( Y | X v ) − L

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

Revised: Then X Φ will satisfy the constraint of Problem 8.

Original: [N/A]
Revised: Notice that X Φ achieves the smallest loss when only using invariant feature information. Combining these two steps completes the proof. A.3 More Discussions on Assumption This assumption aims to ensure that the invariance penalty cannot be arbitrarily small if a spurious feature, together with other features, is selected by a feature mask.

Original: (b) Did you describe any potential participant risks, with links to Institutional Review
Revised: For example, in the extreme case where X 2 is the only invariant feature (i.e., X v consists of only X 2 ) and can perfectly predict Y , we would have H ( Y | X 2 ) = 0 .

Original: Board (IRB) approvals, if applicable?
Revised: Then for a spurious feature X 1 , H ( Y | X 1 , X 2 )

Original: [N/A]{+= 
 and+}
Revised: {+= 
 and+}

Original: (c) Did you include the estimated hourly wage paid to participants and  the total amount
Revised: H ( Y | X 1 , X 2 , ρ ( Z )) = for any ρ ( · ) , and we cannot identify X 1 as the spurious feature. Nevertheless, since there are generally exogenous noise variables in the SCM and

Original: spent on participant compensation?
Revised: H

Original: [N/A]
Revised: ( Y | X v ) is positive, we believe that this assumption holds in most cases.

Original: A ProofsA.1 Proof of Theorem1
Revised: B Proofs B.1 Proof of Theorem

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 function and 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 [Equation] , X [Equation]

Original:  X s and X ′ s
Revised: , where we have chosen X ′ v = 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 (10), 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 variable as 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 v except 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 function and 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 ′ [Equation] y [Equation]

Original:  P ( Y ≤ y | X ( − j ) v , X s ) .
Revised: X ′ [Equation] ) = 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 same conclusion following the previous proof.

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

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 least one 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 dominate the 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 invariant features 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 selected features 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
Revised: By Assumption 1 with a given ϵ > 0 , we have

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

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

Original: [Equation]

Original: [Equation]
Revised: [Equation]

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

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

Original: [Equation]
Revised: [Equation]

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 mask as Φ − 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 (12), we have

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

Original: Thus, if ϵ < γ [Equation] ) , we have
Revised: Thus, if ϵ < γ 2(1+2 λ ) , 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 ϵ − 1 2 , γ 4 ϵ − 12 ] , we can get ˆ L (Φ v ) <

Original: [Equation]
Revised: ˆ L (Φ) , ∀ Φ ̸ = Φ v .

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 be chosen arbitrarily according to Assumption 1.

Original: A.3 Proof of Proposition
Revised: B.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.
Revised: It corresponds to the largest subset of X that satisfies the invariance constraint.

Original: Note that Φ ¯ v ̸ = Φ v .
Revised: Note that Φ ¯ v ̸ = Φ v .

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

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

Revised: +

Revised: [Equation]

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

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 ϵ , we would get L [Equation] L

Original: [Equation]
Revised: [Equation]

Original: A.4 Proof of Corollary1
Revised: B.4 Proof of Corollary

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 v , h ( Y ))

Original: =
Revised: =

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

Original: for any X v .
Revised: for any X v .

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

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

Original: Then Proposition 1 indicates thatwe cannot identify all the invariant features.
Revised: Then Proposition 1 indicates that we 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 above by noting H ( Y | X v , 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 v , h (Index

Original: = H ( Y | X v , X , Y ) = 0 .
Revised: [Equation]

Original: A.5 Proof of Proposition
Revised: B.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 spurious feature X ks in Proposition 2.

Original: Similar to Eqn.
Revised: Similar to Eq.

Original: (9), we can show that
Revised: (11), we can show that

Original: [Equation]
Revised: [Equation]

Original: And similar to Eqn.
Revised: And similar to Eq.

Original: (10), we also have
Revised: (12), 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. A.6 Proof of Theorem
Revised: Thus, we cannot identify all the invariant features from Problem 6.

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

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

Original: Φ E X k

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

Original: Φ E X k ,Y k

Original: [ X k Y k ]for all k

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

Original: ˜ W T β .

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

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

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 ( ω ′ , Φ ′ )

Original: =

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

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

Original: E X k

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

Original: k

Original: =

Original: Φ ′ E X k ,Y k

Original: [ X k Y k ] .

Original: As ω ′ achieves the same minimummean squared error on the k -th environment, ω ′ must satisfy the normal 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.

Original: A.7 Proof of Meeting Condition1
Revised: B.6 Proof of Meeting Condition

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

Original: =
Revised: =

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

Original: =
Revised: =

Original: H ( Y | X v ) holds.
Revised: H ( Y | X v ) 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 , we have

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

Original: =
Revised: =

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

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

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

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

Original: =
Revised: =

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

Original: B Extension to Other Loss Functions
Revised: C More Experimental Details

Original: Sections 5.2 and 5.3 focus on classification task with cross-entropy loss, to establish sufficientand necessary conditions for invariance identification. Similar results can be shown for other lossfunctions or tasks in a straightforward manner. To see this, notice that H ( ·|· ) in Assumptions 1-4 and 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: Implementing Approximated ZIN .

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. Replacing H ( ·|· ) with L ( ·|· ) , we can follow the same proof procedure to obtain similar sufficient andnecessary conditions. This is summarized in Theorem 4 for completeness. Theorem 4.
Revised: Since Eq.

Original: Consider Assumptions 1-4 and Conditions 1-2 where H ( ·|· ) is replaced with L ( ·|· )accordingly. If ϵ < [Equation] ( Y ) and λ ∈
Revised: (5) is a challenging minimax formulation, we replace the penalty term given a environment partition with its first order approximation.

Original: [ L ( Y ) + [Equation] − 12 ,
Revised: Specifically, we consider the following surrogate minimax formulation:

Original: [Equation]
Revised: [Equation]

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

Original: ˆ L (Φ) , ∀ Φ ̸

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

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

Original: =

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

Original: ˆ 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: For the synthetic datasets and the house price prediction dataset, we use  the full batch gradient inoptimizing.
Revised: Readers can refer to the Appendix B.3 of [57] for more details about the relationship between Eqs. (5) and (13). We further use a two-stage method to approximate the minimax procedure:

Revised: • Minimize R ( ω, Φ) over ( ω, Φ) and simultaneously maximize (cid:80)

Revised: Kk

Revised: =1 ∥∇ ω R ρ ( k ) ( ω, Φ) ∥over ρ .

Revised: • Fix ρ and minimize Eq.

Revised: (13) over ( ω, Φ) .

Revised: Training Details .

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

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 for 4000 epochs, with additional 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 divided equally, i.e., 2000 epochs, in each of 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-batch training 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 in the first 25 epochs.

Original: For the landcover classification task, we use mini-batch training with batch sizeof 1024. All the methods are trained for 400 epochs, with anneal strategy in the first 40 epochs.
Revised: For the Landcover task, we use mini-batch training with batch size of 1024, and all the methods are trained for 400 epochs with annealing strategy in the first 40 epochs.

Original: Weuse Adam [21] with learning rate 0.001 as our optimizer.
Revised: We use Adam [23] 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 Nvidia

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

Original: It takes about 5 GPU hours for the CelebA tasks, while the house pricetaskand Landcover taskcost less than 10 minutes for training. E Extensions of Section 4E.1 Connection of Impossibility Example to Existing Tasks In the counter-example in Section 4, we consider binary features andlabel. This is inspired by thepopular classification tasks in the literature of IRM, e.g., CMNIST [5], CifarMnist [24], ColoredObject[1] , Waterbirds and CelebA [38].
Revised: It takes about 5 GPU hours for the CelebA tasks, while the house price task and

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 ̸ = X 2 = Y, w.p. 0 . X 2 = X 1 ̸ = Y, w.p. 0 .
Revised: Landcover task cost less than 10 minutes for training.

Original: With this new dataset, we can then compare EIIL, IRM (with environment indexes), and ERM. Theempirical results are reported in Table 5. For ERM (oracle), we train the model only on the invariantfeature in MCOLOR, i.e., color. [Table]
Revised: Notably, since IRM suffers from overfitting problem when applied to large models like ResNet-18 [27, 57], we fix the feature extraction backbone (that is, use a pre-trained model) in the CelebA experiment to allieviate this issue.

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] F More Experiment Results [Figure] F.1 Spatial Heterogeneity
Revised: D More Experiment Results D.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 also 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 the space 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 the performance 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 6 show similar performances to those of 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., when [Table] p 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 not be 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.
Revised: As shown in the right panel of Figure 2, ZIN automatically merges duplicated environments.

Original: [Figure]
Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

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 a setting 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 the heterogeneity 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 7.

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 .

Original: The mean and worsttest accuracy implies that ZIN based on r or r 2 can effectively remove the spurious feature.
Revised: The mean and worst test accuracy implies that ZIN based on r or r 2 can effectively remove the spurious feature.

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: Since p 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 also a violation by Corollary 1, and the corresponding results in Table 7 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 ρ ( · ) .

Revised: [Table]

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]
Revised: Fig.shows that K has a relatively small impact, especially when K ≥ 4 .

Original: [Figure]
Revised: [Figure]

Original: F.3 About the Inferred Environments on CelebA
Revised: D.3 About the Inferred Environments on CelebA

Original: We visualize the inferred environments of the CelebA dataset in this section.
Revised: We visualize the inferred environments of the CelebA dataset in this section.

Original: The target Smiling isspuriously correlated with feature Gender , i.e., most females are smiling while most males are not.
Revised: The target Smiling is spuriously correlated with feature Gender , i e., most females are smiling while most males are not.

Original: We set K to be 2 (our framework is insensitive to K as shown in Appendix F.2).
Revised: We set K to be 2 (our framework is insensitive to K as shown in the right panel of Fig. 3).

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

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

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

Original: G How to Choose Z
Revised: E Examples of Valid and Invalid Choices of Auxiliary Information

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

Original: We start with an example in Fig.

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.

Original: This Z satisfies Condition 1 because H ( Y | X 1 , X 2 )

Original: =

Original: H ( Y | X 1 , X 2 , Z ) or equivalently

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

Original: Y or a direct cause of Y , Z contains additional information of Y conditional on X 1 and X 2 .

Original: Then

Original: Condition 1 is violated and this kind of information cannot be used in our framework.

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

Original: ( Y |

Original: X k ¯ s , ρ ( Z )) ≥ C .

Original: Further, let X sc :=

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.

Original: The proof proceeds almost the same as that of Theorem 2, by replacing Φ v with Φ v + sc .

Original: Notethat X

Original: v + sc behaves as “invariant” under the inferred environments in Proposition 3.

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.

Original: • Collect as many such Z ’s as possible.

Original: H Societal Impact
Revised: F Societal Impact

Original: In this work, we propose to utilize the auxiliary information to aid the invariance learning withoutenvironmental indexes.
Revised: In this work, we propose to utilize the auxiliary information to aid the invariance learning without environmental indexes.

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

Original: For the additional information, weshould also avoid using demographic, private, and sensitive information.
Revised: For the additional information, we should also avoid using demographic, private, and sensitive information.