Original: ZIN: When and How to Learn Invariance Without Environment Partition?
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 of the data distribution may vary but the underlying causal mechanisms remain constant.
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 the heterogeneity, recent invariant learning methods have proposed to learn robust and invariant models using 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 to utilize 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 this circumstance is fundamentally impossible without further inductive biases or additional 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 environment partition 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 provably identify 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 on both synthetic and real world datasets validate our analysis and demonstrate an improved 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: 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: 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.

Original: Codes are available at https://github.com/linyongver/ZIN_official .
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 independent and 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 environments in 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, 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.
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 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.
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, Peters et al.
Revised: For example, Peters et al.

Original: [36] proposes to exploit the invariance principle to learn a linear model that merely relies on the 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 any variable 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: [4] propose invariant risk minimization (IRM) to capture invariant correlations by learning representations that elicit an optimal invariant predictor across 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: [2], Jin et al.
Revised: [2], Jin et al.

Original: [21], Krueger et al.
Revised: [21], Krueger et al.

Original: [24], Xie et al.
Revised: [24], Xie et al.

Original: [52], Chen et al.
Revised: [52], Chen et al.

Original: [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.
Revised: [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: [7], Lin et al.
Revised: [7], Lin et al.

Original: [26] try to learn invariant features by coupled adversarial neural networks.
Revised: [26] try to learn invariant features by coupled adversarial neural networks.

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

Original: Lin et al.
Revised: Lin et al.

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

Original: [57, 56] propose to alleviate this issue through imposing sparsity constraint and sample reweighting, respectively.
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 partitioned into 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, such an environment partition may be unavailable or hard to obtain in practice [29, 10].
Revised: However, such an environment partition may be unavailable or hard to obtain in practice [29, 10].

Original: Recently, a line of works try to learn invariance without environment indexes where the dataset is assembled by merging data 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: [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.
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: Liu et al.
Revised: Liu et al.

Original: [29] devise an interactive mechanism, called heterogeneous risk minimization (HRM), with environment 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: [30] extend it to representation 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 input-label pairs) can be risky , both theoretically and empirically.
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 indexes is 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 and a general theoretic result (Theorem 1), to show that the causally invariant features are unidentifiable if no environment information 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 [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.
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 latent heterogeneity.
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 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.
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 [53, 51].
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 [6, 33], locations (longitude and latitude) of collected satellite data in remote sensing [14, 40], and meta information associated with an instance [31].
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 propose a 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 a fairly general setting, we derive both sufficient and necessary conditions for our framework to identify invariant 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 in Appendix A.
Revised: Extensions to other settings are discussed in Appendix A.

Original: Finally, Section 6 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: 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.
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 from a causal perspective [36].
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 and potential 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: [19] consider heterogeneous datasets 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: Here confounders 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 [48], Danks et al.
Revised: Tillman and Spirtes [48], Danks et al.

Original: [11], Huang et al.
Revised: [11], Huang et al.

Original: [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.
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 for generalizing beyond training data is distributionally robust optimization (DRO) [5, 43, 25, 13, 12, 55].
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 training distribution.
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 loss over groups in the training data [41].
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 generally requires 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 know a 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 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].
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 are still 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: [28] upweights samples with high losses from the initial ERM model and relies on 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: [34], Sanh et al.
Revised: [34], Sanh et al.

Original: [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.
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 or invariant 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 invariant features 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 = where 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 the environment 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 generation process [35, 37].
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 acyclic graph, 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 of an assignment of a deterministic function depending on the parental variables in the graph and an exogenous 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 that is also considered in [3, 4]:
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 a non-degenerate deterministic function and ϵ v is an independent noise variable.
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: Then the 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) E e =

Revised: (cid:80) E e =

Original: α e P ( X e , Y e ) for some α e > 0 and (cid:80) E e =1 α e = 1 .
Revised: α 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 cross-entropy loss and 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 the following 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 direct causal parents of Y as invariant features, as in [36, 4].
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 of the 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, then the 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 SCM governing 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 heterogeneous data 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 example immediately gives a negative answer.
Revised: The following example immediately gives a negative answer.

Original: [Equation]
Revised: [Equation]

Original: Example.
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 we observe 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 Eq.
Revised: Suppose the joint distribution of mixed environments is given in Eq.

Original: 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 ,
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 2 , or both to predict Y .
Revised: X 2 , or both to predict Y .

Original: However, as discussed in last section, whether the learned features are invariant is determined by the underlying 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 present two possible data generation processes that generate the same distribution yet have different invariant features:
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, where p e =1 s = 0 .
Revised: • X 1 is the invariant feature while X 2 is spurious, where p e =1 s = 0 .

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

Original: [Equation]
Revised: [Equation]

Original: X 1 ∼ Bernoulli(0.5) ,
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: 8 and p e =2 s
Revised: 8 and p e =2 s

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

Original: X
Revised: X

Original: [Equation]
Revised: [Equation]

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

Original: Eq.
Revised: Eq.

Original: Thus, from the joint distribution P ( X 1 , X 2 , Y ) , a learned model only depends on either X 1 ,
Revised: Thus, from the joint distribution P ( X 1 , X 2 , Y ) , a learned model only depends on either X 1 ,

Original: X 2 or both, and fails to generalize for at least one of the two data generation processes.
Revised: X 2 or both, and fails to generalize for at least one of the two data generation processes.

Original: On the other hand, 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) scenario.
Revised: second) scenario.

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

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

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

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

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

Original: Empirical Validation.
Revised: Empirical Validation.

Original: We construct a variant of CMNIST according to the second data generation process.
Revised: We construct a variant of CMNIST according to the second data generation process.

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

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

Original: As described in Eq.
Revised: As described in Eq.

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

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

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

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

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

Original: The empirical results are reported in Table 1.
Revised: The empirical results are reported in Table 1.

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

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

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

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

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

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

Original: [Table]
Revised: [Table]

Original: Moreover, such examples are not rare.
Revised: Moreover, such examples are not rare.

Original: 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.
Revised: 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 identify the spurious features only based on 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 the 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 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:
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: • X ′ v ̸
Revised: • X ′ v ̸

Original: = X v and X ′ v ∩ X s ̸ = ∅ ; • ( 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 B.1.
Revised: A proof is provided in Appendix B.1.

Original: Under this circumstance, we have to introduce additional assumptions/conditions or certain “inductive bias” to identify the invariant features [45, 37, 20, 32, 50].
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 may come 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 the reason that makes existing methods achieve improved empirical performance in their considered scenarios.
Revised: We believe that this is the reason that makes existing methods achieve improved empirical performance in their considered scenarios.

Original: 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.
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 always guaranteed.
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 more explicitly when learning invariant models without environment partition.
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 exists additionally observed variable Z with the data, which has been often considered in the literature [19, 20, 22, 53, 51].
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, 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].
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: [53] also take advantage of auxiliary information to help improve OOD performance in a semi-supervised manner.
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 data together 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 to learn invariant models from the heterogeneous 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: We consider that there exists additional auxiliary information Z ∈ R d z in company with the data ( X , Y ) .
Revised: We consider that there exists additional auxiliary information Z ∈ R d z in company with the data ( X , Y ) .

Original: In this section, we propose ZIN, auxiliary information Z for environmental INference , for invariant learning from the heterogeneous dataset D without environment partition.
Revised: 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 will derive conditions for invariance identification 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 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 investigated in 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 ρ ( · ) , with ρ ( Z ) ∈ [0 , 1] K and (cid:80)
Revised: Let ρ ( k ) ( · ) denote the k -th entry of ρ ( · ) , with ρ ( Z ) ∈ [0 , 1] K and (cid:80)

Original: k ρ ( 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 ω (Φ( x i )) , y i ) and the loss in the k -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 ω (Φ( x i )) , y i ) .
Revised: =1 ρ ( k ) ( z i ) ℓ ( f ω (Φ( x i )) , y i ) .

Original: Recall that IRM [4] learns an invariant representation Φ , upon which there is a classifier f ω that is simultaneously optimal in all environments.
Revised: Recall that IRM [4] learns an invariant representation Φ , upon which there is a classifier f ω that is simultaneously optimal in all environments.

Original: Suppose that the environments have been given according to a fixed ρ ( · ) .
Revised: Suppose that the environments have been given according to a fixed ρ ( · ) .

Original: 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.
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 f ω k achieves a smaller loss, then we know that f ω is not optimal 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 { f ω k } Kk =1 , one for each environment, to measure whether f ω 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 is as follows:
Revised: Thus, when ρ ( · ) is provided, our formulation to learn invariance is as follows:

Original: [Equation]
Revised: [Equation]

Original: 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.
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 generate environments among which the spurious features exhibit instability, so that there is a large penalty if Φ extracts spurious features.
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 environment partition 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 learning invariant 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: [10] has used a similar intuition and 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-stage method cannot be jointly optimized, and the environment partition relies on the given model and lacks a theoretical 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 proposed framework 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 are also 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 [10], is unfortunately an instance 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 theoretical perspective.
Revised: In this section, we try to understand ZIN from a theoretical perspective.

Original: We start with a simple yet general 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-wise feature selection mask, and f ω is a general non-linear function.
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 classification tasks 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: Extensions to other loss functions and linear feature transformations are left to Appendix A. and A.2, respectively.
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 invariant features, i.e., Φ v =
Revised: In this setting, our goal is equivalent to learning the optimal feature mask that merely selects invariant features, i.e., Φ v =

Original: [ 1 d v ; 0 d s ] .
Revised: [ 1 d v ; 0 d s ] .

Original: For a given feature mask Φ , the objective function is equal to ˆ L (Φ) =
Revised: For a given feature mask Φ , the objective function is equal to ˆ L (Φ) =

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

Original: Then ZIN can correctly identify the invariant features, or solution to Problem 6 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 of an optimal classifier over some X ′ and Y , and similarly H ( Y | ρ ( Z ) , X ′ ) to denote the minimum risk (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 exactly one environment, i.e., ρ ( Z ) is one-hot, it can be verified that the optimal loss H ( ·|· ) coincides with conditional 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 that
Revised: For a given feature mask Φ and any constant ϵ > 0 , there exists f ∈ F such that

Original: E [ ℓ ( f (Φ( X )) , Y )] ≤ H ( Y | Φ( X ))
Revised: E [ ℓ ( f (Φ( X )) , Y )] ≤ H ( Y | Φ( X ))

Original: + ϵ.
Revised: + ϵ.

Original: Assumption 2.
Revised: Assumption 2.

Original: If a feature violates the 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: 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
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: ( Y | ρ ( Z ) , X 1 )) .
Revised: ( Y | ρ ( Z ) , X 1 )) .

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

Original: Assumption 2 aims to ensure a sufficient positive penalty if a spurious feature is included (see Appendix A.3 for further discussion).
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 3 indicates that any feature contains 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 [49]), as it does not affect prediction.
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
Revised: We next present our sufficient conditions for

Original: ZIN to identify invariant features.
Revised: 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 ρ ( · ) , it holds 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 a function ρ ( · ) 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: B.6).
Revised: B.6).

Original: Condition 1 requires that invariant features should remain invariant w.r.t any environment partition 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 becomes non-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 each spurious feature, there exists at least one partition so that this feature is non-invariant in the split environments.
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 environment partitions, 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 main result follows and a proof is given in Appendix B.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 Conditions 1-2, if ϵ <
Revised: With Assumptions 1-3 and Conditions 1-2, if ϵ <

Original: [Figure]
Revised: [Figure]

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

Original: In this section, we prove that they are also 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 that if Conditionis violated, then some invariant features will be excluded in the solution of Problem 6; and in
Revised: Specifically, Proposition 1 shows that if Conditionis violated, then some invariant features will be excluded in the solution of Problem 6; and in

Original: Proposition 2, violation of Condition 2 renders some spurious features included in the solution.
Revised: 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, if Condition 1 is violated, i.e., there exists ρ ( · ) so that
Revised: With Assumptions 1-3, if Condition 1 is violated, i.e., there exists ρ ( · ) so that

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

Original: =
Revised: =

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

Original: >
Revised: >

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

Original: A proof is provided in Appendix B.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: A question 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 additional variables, it is also interesting to know whether we can obtain a valid partition based on only the training 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, with proofs given in Appendix B.4.
Revised: Below we provide examples regarding these questions, with proofs given in Appendix B.4.

Original: 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.
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: there exists a function ρ ( · ) and an injective function 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: We now discuss the necessity of Condition 2, with an additional assumption.
Revised: We now discuss the necessity of Condition 2, with an additional assumption.

Original: Assumption 4.
Revised: Assumption 4.

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

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

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

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

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

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

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

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

Original: Proposition 2.
Revised: Proposition 2.

Original: With Assumptions 1, 3 and 4, if Condition 2 is violated, i.e., there exists a spurious feature X ks ∈
Revised: With Assumptions 1, 3 and 4, if Condition 2 is violated, i.e., there exists a spurious feature X ks ∈

Original: X s such that H ( Y | X ks )
Revised: 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: Φ v with ˆ L (Φ v )
Revised: Φ v with ˆ L (Φ v )

Original: >
Revised: >

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 invariance 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 proof in Appendix B.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 and informative 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: 5.4 Choice of the Auxiliary Information
Revised: 5.4 Choice of the Auxiliary Information

Original: Conditions 1 and 2 present theoretical requirements of Z to successfully learn invariant features.
Revised: Conditions 1 and 2 present theoretical requirements of Z to successfully learn invariant features.

Original: We now discuss how to find such auxiliary information given certain prior knowledge.
Revised: We now discuss how to find such auxiliary information given certain prior knowledge.

Original: Take Fig.
Revised: Take Fig.

Original: 1 for example.
Revised: 1 for example.

Original: When we collect a data point (e.g., taking a photo), there often exists some meta information (e.g., time slot, coordinate, and temperature).
Revised: 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: Suppose that the image is generated according to the causal graph on the right panel of Fig.
Revised: Suppose that the image is generated according to the causal graph on the right panel of Fig.

Original: 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.
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: This is because H ( Y | X 1 , X 2 ) =
Revised: This is because H ( Y | X 1 , X 2 ) =

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

Original: This example shows some evidence of choosing the auxiliary information: Z and Y should be d-separated by X v .
Revised: This example shows some evidence of choosing the auxiliary information: Z and Y should be d-separated by X v .

Original: There are other examples satisfying Condition 1, which are given
Revised: There are other examples satisfying Condition 1, which are given

Original: Fig.
Revised: Fig.

Original: 5 in Appendix E. We also illustrate some failure choices in Fig.
Revised: 5 in Appendix E. We also illustrate some failure choices in Fig.

Original: 6 in Appendix E. Specifically, if
Revised: 6 in Appendix E. Specifically, if

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

Original: This kind of information cannot be used in our framework.
Revised: This kind of information cannot be used in our framework.

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

Original: A feasible choice of Z should satisfy Condition 1 and should be correlated with certain variables in the causal graph of the SCM.
Revised: A feasible choice of Z should satisfy Condition 1 and should be correlated with certain variables in the causal graph of the SCM.

Original: The path between Z and Y should be d-separated by X v (e.g., in Figs.
Revised: The path between Z and Y should be d-separated by X v (e.g., in Figs.

Original: 1 and 5) or there is no path between Z and Y .
Revised: 1 and 5) or there is no path between Z and Y .

Original: Notably, Z cannot be the parent or child of Y as shown in Fig. 6.
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 synthetic and 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, [4], group
Revised: We compare ZIN with several existing methods: ERM, IRM, [4], group

Original: DRO [41], EIIL [10], HRM [29], and LfF [34].
Revised: 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: Notice that LfF only works for classification tasks.
Revised: Notice that LfF only works for classification tasks.

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

Original: We adopt the first order approximation of Eq.
Revised: We adopt the first order approximation of Eq.

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

Original: The number of inferred environments 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 synthetic dataset and house price prediction, ResNet-18 [16] 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: More details are provided in
Revised: More details are provided in

Original: Appendix C.
Revised: Appendix C.

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

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-dimensional spatial variable is considered in Appendix D.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] .

Original: We
Revised: We

Original: [Table]
Revised: [Table]

Original: evaluate the performance on four distinct test environments with p s ∈ { 0 . 999 ,
Revised: evaluate the performance on four distinct test environments with p s ∈ { 0 . 999 ,

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

Original: . 2 , 0 . 1 } and p v being constant.
Revised: . 2 , 0 . 1 } and p v being constant.

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

Original: More details are provided in Appendix C
Revised: More details are provided in Appendix C

Original: Results .
Revised: Results .

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

Original: 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
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) and p 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 and may 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 test accuracy 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 ERM and 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 the ground-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 D.2, we conduct ablation study to verify our theoretical results in Section 5 by 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 observation is that K has a relatively small impact, as shown in Fig.
Revised: A notable observation is that K has a relatively small impact, as shown in Fig.

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

Original: Table 3 reports the mean squared error (MSE), and ZIN again outperforms other methods on test dataset.
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 3.
Revised: The corresponding results are marked with (+ Z ) in Table 3.

Original: We observe that the results 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: This implies 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 [31], and by construction the 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 CeleA images 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 potential auxiliary 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 is made available to IRM and group DRO 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 4, 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: The worst test performances of ERM, EIIL and LfF indicate that they may not learn the causal features.
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,
Revised: Finally, we observe that ZIN performs better when more auxiliary information is provided; specifically,

Original: ZIN(7)
Revised: ZIN(7)

Original: achieves the best test performance based on the mixed dataset.
Revised: achieves the best test performance based on the mixed dataset.

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

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

Original: In Appendix D
Revised: In Appendix D

Original: [Table]
Revised: [Table]

Original: Fig.
Revised: Fig.

Original: 3 reports the results of different choices of hyper-parameter K and Fig.visualizes the inferred environments in this experiment.
Revised: 3 reports the results of different choices of hyper-parameter K and Fig.visualizes the inferred environments in this experiment.

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

Original: [Table]
Revised: [Table]

Original: Landcover .
Revised: Landcover .

Original: Our final task is land cover prediction that classifies the land cover type (e.g., grasslands) from satellite data [14, 40].
Revised: 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 [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.
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: We also 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 time dimension 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 5, we see that ZIN with location achieves 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 is because climate variables indeed contain some information for predicting the target, violating our conditions on 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 has learned 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 be risky 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.

Original: 6.3 Discussion on Auxiliary Information for Related Methods
Revised: 6.3 Discussion on Auxiliary Information for Related Methods

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

Original: It is interesting to ask whether this type of information is also useful to related methods such as EIIL and HRM.
Revised: It is interesting to ask whether this 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 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 .
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 .

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

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

Original:  Concluding Remarks
Revised:  Concluding Remarks

Original: This paper investigates when and how to learn invariance from heterogeneous data without explicit environment 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 some additional 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 linear feature 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 improved performance 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 make the role of inductive biases more explicit, when learning invariance from heterogeneous data without environment 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 general nonlinear feature extractors, which is also an open problem for IRM even with environment partition.A More Theoretical Results A.1 Extension to Other Loss Functions
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: Sections 5.2 and 5.3 focus on classification task with cross-entropy loss, to establish sufficient and necessary conditions for invariance identification.

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

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

Original: 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: Replacing H ( ·|· ) with L ( ·|· ) , we can follow the same proof procedure to obtain similar sufficient and necessary conditions.

Original: This is summarized in

Original: Theorem 3 for completeness.

Original: Theorem 3.

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

Original: If ϵ <

Original: [Equation]

Original: and λ

Original: [Equation]

Original: −

Original: [Equation]

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

Original: ˆ L (Φ) for all Φ ̸ = Φ v , where we assume L ( Y ) < ∞ .

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

Original: =

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

Original: ˆ L (Φ ′ ) .

Original: A.2 Beyond Feature Selection: Linear Feature Transformations

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

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

Original: [Equation]

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

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

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

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

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

Original: The prediction for X is ω ◦

Original: Φ( X ) =

Original: (Φ X ) T ω .

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

Original: Following [4, 39], we consider a constrained form of Problem 6 for theoretic analysis:

Original: [Equation]

Original: 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: This is in accordance to existing conditions for identifiability in the linear case [4, 39].

Original: In this paper, we utilize such a condition, linear general position condition, from [4].

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

Original: As in Appendix A.1, we use L ( ·|· ) to represent the optimal expected risks.

Original: 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: Our identifiability result for the linear case then follows.

Original: Proposition 3.

Original: Assume Condition 1 where H ( ·|· ) is replaced with L ( ·|· ) accordingly.

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

Original: E X k (cid:2)

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

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

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

Original: d − r.

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

Original: Problem 8 results in the desired invariant predictor.

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: [4, 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 to show our solution meets the same normal equations.

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

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: Notice that ω k minimizes the mean squared error on only the k -th environment, hence it must satisfy the normal equation Φ ′

Original: E

Original: 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 minimum mean squared error on the k -th environment, ω ′ must satisfy the normal equation, too.

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

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

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

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

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

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

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

Original: Notice that X Φ achieves the smallest loss when only using invariant feature information.

Original: Combining these two steps completes the proof.

Original: A.3 More Discussions on Assumption

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

Original: = 0 .

Original: Then for a spurious feature X 1 , H ( Y | X 1 , X 2 )

Original: [-= 
-]
Revised: [-= 
-]

Original: and

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

Original: =

Original: for any ρ ( · ) , and we cannot identify X 1 as the spurious feature.

Original: Nevertheless, since there are generally exogenous noise variables in the SCM and

Original: H

Original: ( Y | X v ) is positive, we believe that this assumption holds in most cases.

Original: B Proofs B.1 Proof of Theorem

Original: Proof.

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

Original: X s .

Original: Define the conditional cumulative distribution function and its inverse as:

Original: [Equation]

Original: By definition, we would have

Original: [Equation]

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

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

Original: Thus, we have

Original: [Equation]

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

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

Original: =

Original: q ( X

Original: [Equation]

Original: , X

Original: [Equation]

Original: , where we have chosen X ′ v = X s and X ′ s

Original: = X v .

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

Original: P ( X, Y ) .

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

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

Original: X ( j ) v for some j .

Original: Then we set X ′ v =

Original: ( 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 function and an independent uniformly distributed noise variable η 1 to construct Y ′

Original: =

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

Original: P ( Y ′

Original: [Equation]

Original: y

Original: [Equation]

Original: X ′

Original: [Equation]

Original: ) =

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

Original: Similarly, we set X ′ s =

Original: X ( j ) v

Original: and we can construct

Original: X ′ s

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

Original: X ′ v ) =

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

Original: We can then get the same conclusion following the previous proof.

Original: B.2 Proof of Theorem

Original: Proof.

Original: Our proof proceeds by two steps.

Original: 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 dominate the expected risk and then exceed ˆ L (Φ v ) .

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

Original: Denote the selected features as X + s and the corresponding feature mask as Φ + s .

Original: We aim to show that

Original: [Equation]

Original: By Assumption 1 with a given ϵ > 0 , we have

Original: [Equation]

Original: [Equation]

Original: One the other hand, we have

Original: [Equation]

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

Original: [Equation]

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

Original: [Equation]

Original: Similar to Equation (12), we have

Original: [Equation]

Original: Then according to Assumption 3, we have

Original: [Equation]

Original: Thus, if ϵ < γ 2(1+2 λ ) , we have

Original: [Equation]

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

Original: ˆ L (Φ) , ∀ Φ ̸ = Φ v .

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

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

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

Original: B.3 Proof of Proposition

Original: Proof.

Original: Consider the following feature set

Original: [Equation]

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

Original: It corresponds to the largest subset of X that satisfies the invariance constraint.

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

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

Original: [Equation]

Original: One the other hand, we have

Original: [Equation]

Original: +

Original: [Equation]

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

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

Original: L

Original: [Equation]

Original: L

Original: [Equation]

Original: B.4 Proof of Corollary

Original: Proof.

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

Original: =

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

Original: for any X v .

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

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

Original: Then Proposition 1 indicates that we cannot identify all the invariant features.

Original: (b) The proof proceeds the same as above by noting H ( Y | X v , h

Original: ( X , Y ))

Original: = 0 .

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

Original: H ( Y | X v , h (Index

Original: [Equation]

Original: B.5 Proof of Proposition

Original: Proof.

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

Original: (11), we can show that

Original: [Equation]

Original: And similar to Eq.

Original: (12), we also have

Original: [Equation]

Original: Then it follows that

Original: [Equation]

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

Original: [Equation]

Original: Thus, we cannot identify all the invariant features from Problem 6.

Original: B.6 Proof of Meeting Condition

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

Original: =

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

Original: =

Original: H ( Y | X v ) holds.

Original: Proof.

Original: On one hand, because ρ ( Z ) contains less information than Z , we have

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

Original: =

Original: H ( Y | X v ) .

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

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

Original: =

Original: H ( Y | X v ) .

Original: C More Experimental Details

Original: Implementing Approximated ZIN .

Original: Since Eq.

Original: (5) is a challenging minimax formulation, we replace the penalty term given a environment partition with its first order approximation.

Original: Specifically, we consider the following surrogate minimax formulation:

Original: [Equation]

Original: Readers can refer to the Appendix B.3 of [57] for more details about the relationship between Eqs.

Original: (5) and (13).

Original: We further use a two-stage method to approximate the minimax procedure:

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

Original: Kk

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

Original: • Fix ρ and minimize Eq.

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

Original: Training Details .

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

Original: 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 divided equally, i.e., 2000 epochs, in each of the two stages.

Original: 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 in the first 25 epochs.

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

Original: GTX 3090 GPU.

Original: It takes about 5 GPU hours for the CelebA tasks, while the house price task and

Original: Landcover task cost less than 10 minutes for training.

Original: 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: D More Experiment Results D.1 Spatial Heterogeneity

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

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

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

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

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

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

Original: [Table]

Original: p s ( r ) =

Original: 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 in the right panel of Figure 2, ZIN automatically merges duplicated environments.

Original: [Figure]

Original: [Table]

Original: [Figure]

Original: [Table]

Original: [Figure]

Original: [Table]

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

Original: We adopt a setting of spatial heterogeneity with p s ( r ) =

Original: 8) and p v = 0 .

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

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

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

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

Original: H ( Y | X v ) >

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

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

Original: X as input to ρ ( · ) .

Original: [Table]

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

Original: Fig.shows that K has a relatively small impact, especially when K ≥ 4 .

Original: [Figure]

Original: D.3 About the Inferred Environments on CelebA

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

Original: 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 the right panel of Fig. 3).

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

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

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

Original: E Examples of Valid and Invalid Choices of Auxiliary Information

Original: F Societal Impact

Original: 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 about out-of-distribution generation and algorithmic fairness in [10].

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