ZIN: When and How to Learn Invariance Without Environment Partition?
Authors: LIN Yong
Shengyu Zhu
Lu Tan
Peng Cui
Venue: NeurIPS 2022 Accept (2022)
Keywords: Invariant Risk Minimization
Out-of-Domain Generalization
Transfer Learning

Abstract: It is commonplace to encounter heterogeneous data, of which some aspects of the data distribution may vary  but the underlying causal mechanisms remain constant.  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. It is hence tempting to utilize the inherent heterogeneity even when environment partition is not provided. Unfortunately, in this work, we show that learning invariant features under this circumstance is fundamentally impossible without further inductive biases or additional information. Then, we propose a framework to jointly learn environment partition and invariant representation, assisted by additional auxiliary information. We derive sufficient and necessary conditions for our framework to provably identify invariant features under a fairly general setting. Experimental results on both synthetic and real world datasets validate our analysis and demonstrate an improved performance of the proposed framework. 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. Codes are available at https://github.com/linyongver/ZIN_official .

ZIN: When and How to Learn Invariance Without Environment Partition?
Yong Lin HKUST ylindf@commen ust.hh.
Shengyu Zhu" Huawei Noah's Ark Lab zhushyu@outlook.cmm com
Lu Tan
Peng Cui
Tsinghua University
Tsinghua University cuip@tsinghua edu, cn
tan121@mails tsinghua. edu..edducco. cn
Abstract
It is commonplace to encounter heterogeneous data, of which some aspects of the data distribution may vary but the underlying causal mechanisms remain constant. 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. hence tempting to utilize the inherent heterogeneity even when enviromment partition not provided. Unfortunately, in this work, we show that learning invariant features under this cireumstance is fundamentally impossible without further inductive biases additional information. Then, we propose framework to jointly learn environment partition and invariant representation, assisted by additional auxiliary information. We derive sufficient and necessary conditions for our framework to provably identify invariant features under fairly general setting. Experimental results on both synthetic and real world datasets validate our analysis and demonstrate an improved performance the proposed framework. 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. Codes are available httpp://gitti htt.aaa......nnnnngg//////// offfiii
Introduction
Machine learning algorithms with empirical risk minimization (ERM) generally assume independent and identically distributed (i.i.d)) data in training and test sets. Due to changing circumstances, selection bias, or time shifts, data distributions are often heterogeneous across different environments in practical applications. When the distributions of training and test data are indeed different, model performance can severely degrade [47
[7 38 as ERM based algorithms may exploit the spurious correlations that are particularly useful ERM the training data but do not exist in the test data.
To tackle the distribution-she or auttof-distribution generalization problem, 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. For example, Peters et al. [36] proposes to exploit the invariance principle learn a linear model that merely relies on the direct causes of the target. Such model is shown to be robust to potential interventions on any variable
*Corresponding author.
36th Conference Neural Information Processing Systems (NeurIPS 2022).

---

except the target itself. Arjovsky et al. 4 propose invariant risk minimization (IRM) to capture invariant correlations by learning representations that elicit an optimal invariant predictor across multiple training environments. Ahuja et al.
21 Jin et al. [21], Krueger et al. (24), Xie et al. [52] Chen et al. 8 further develop variants of IRM by introducing game theory, regret minimization, variance penalization, multi-objective optimization, etc., and Xu and Jaakkola 5 Chang et al. 77. Lin et al.
26 try to learn invariant features by coupled adversarial neural networks. Some recent works report that IRM is less effective when applied to large neural networks [26
[5 Lin al. [27] find that this
can be largely attributed to the overfitting problem, and Zhou et al. [57 56 propose to alleviate this issue through imposing sparsity constraint and sample reweighting, respectively.
Noticeably, the aforementioned invariant learning methods require datasets to be explicitly partitioned into environments (or domains) according to the heterogeneity underlying the data. However, such an environment partition may be unavailable or hard to obtain in practice 22 10 Recently, a line of works try to learn invariance without environment indexes where the dataset is assembled by merging data from multiple environments. For instance, Creager et al. [00] 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. Liu al. [29] devise interactive mechanism, called heterogeneous risk minimization (HRM), with environment inference and invariant learning raw feature level, and Liu et al. 30) extend it to
representation level with kernelized trick.
However, as we will show, learning invariant models under this circumstance (with only input-label pairs) can be risky, both theoretically and empirically. As our first contribution, in Section we prove that learning invariant features from heterogeneous data without environment indexes is in theory impossible. Specifically, we provide a counter-example and a general theoretic result (Theorem []]. to show that the causally invariant features are unidentifiable if no environment information is provided. This impossibility result, similar to the identifiability issue in causal discovery (7], independent component analysis (ICA) 20), and unsupervised learning of disentangled representations [32
motivates us consider when and how to learn invariant features.
In this work, we turn to additional auxiliary variables that encode some information about the latent heterogeneity. This strategy can be treated 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 environment information at all. Notably, such auxiliary information is often cheaply available for every input in practice Examples include time index of the data in time series forecasting tasks 6 333 locations
553 51
(longitude and latitude) of collected satellite data in remote sensing 14 40 and meta information associated with an instance [3I. Based on the additionally observed variables, we proceed to propose framework to jointly learn environment partition and invariant representation in Section [5] Under a fairly general setting, we derive both sufficient and necessary conditions for our framework identify invariant features. Extensions to other settings are discussed in Appendix A
Finally, Section 6 presents experimental analysis on both synthetic and real world datasets. The proposed framework achieves an improved performance over existing methods like EIIL and HRM, and has comparable performance to IRM with ground-truth environment partition.
2 Related Work
Invariant learning methods can be interpreted as robust to certain interventions or heterogeneity from a causal perspective [36]. On the other hand, heterogeneous data have also placed challenges and potential benefits in causal inference tasks. For example, Huang et al. [19] consider heterogeneous datasets to identify variables with changing local mechanism and further the causal directions. Here confounders are assumed as functions of the domain index or time. Tillman and Spirtes [48], Danks al. [.] Huang et al. (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.
2

---

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 (5).
DRO methods propose to optimize the worst-case risk over a set of distributions close to the training distribution. A notable instance of DRO methods is group DRO, which optimizes the worst-case loss over groups the training data As with most invariant learning methods, group DRO generally requires partition of groups obtained according to group annotation and label of each data sample.
Obtaining group annotations, however, can be costly or even impossible in practice. We may not know a priori the inherent spurious features associated with data. 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 TO 28). Notice that sometimes small number labelled annotations are still needed to find proper trained model, which different from the setting considered in this work. For example, Liu et al.
23 upweights samples with high losses from the initial ERM model and relies on small validation set of annotated data to tune parameters. Additionally, Nam et al. [4], Sanh et al. [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 shortcuts. 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.
3 Preliminaries
Throughout this paper, upper-cased letters such as X E Rd and Y E denote random variables, and lower-cased letters such and y denote deterministic instances. Suppose there exist invariant features X,, E Rd and spurious (non-invariant) features X, E Rdd, We observe a scrambled version X g X2) Rd with being injective function. Consider the dataset D {(x), 3/) 31)) ) where data are collected from multiple environments E Suppp- We use superscript to indicate the environment index with variable, e.g., X and Y when it is useful to make the index explicit. Notice that the true environment index e is not provided in training, unless otherwise stated.
In this work, we assume an underlying structural causal model (SCM) governing the data generation process 35 37 An SCM considers a of variables associated with vertices of directed acyclic graph, where directed edges represent direct causation. Each variable is obtained as result of an assignment of deterministic function depending on the parental variables in the graph and an exogenous random variable. As an example, we present the following data generation process that is also considered in B 4
Y gn (x., Ev), X
X -(X; X
(I)
where 9r denote non-degenerate deterministic function and e, is an independent noise variable. We observe data {X y= from multiple environments, each with probability P(X Yo). Then the marginal distribution of the mixed data can represented as P(X, Y) - ar P(Xe, Ye) for
some de > 0 and Σe= 0g = 1.
We consider the model as the composition of a feature extractor t and classifier (or predictor) fo that is parameterized by w. To seek for oft-o--distributiin generalization ability, we hope that merely encodes the information of invariant features. Let f(.--) denote a loss function such as cross-entropy loss and squared error. Our goal is to learn robust and invariant model that minimizes the following objective over all the considered environments
Esuppi
sup eEEvapp
Ex . ((x(")))) Y*)
While there is no single, common definition of invariant feature featurr in the literature, we treat the direct causal parents of Y as invariant features, as in 36 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). In contrast, the conditional probability could change after some intervention. As such, the underlying SCM governing data generation determines the desired invariance.
non-parental variable is included in the conditioned variables, then
3

---

Impossibility Result
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. The following example immediately gives a negative answer.
Example. Consider that there are two environments {1 with 0n = 022 = 0.5, and we observe binary-valued features X1, X2 and label Y. Suppose the joint distribution of mixed environments is given in Eq. 2 A learning algorithm that tries to learn invariant feature based on this dataset would result in a model that (deterministically) depends on either X1, X2, both to predict Y. However, as discussed in last section, whether the learned features are invariant is determined by the underlying data generation process. We now present two possible data generation processes that generate the same distribution yet have different invariant features:
Y = 0. w.p. 0.5,
Y 1, w.p. 0.5,
X2 = X1 = -Y, w.p. 0.6375,
(2)
X2 7 X1 =Y, w.p. 0.1125, X1 X1 X2 =Y, w.p. 0.2125,
X2 X1 Y, w.p. 0.0375.
X₃ is the invariant feature while X2 is spurious, where p. 0.8 and p. 0.9:
w.p. 0.75, Y - 1 w.p. 0.25,
X1,
Y,
w.p. Ps:
X1 (ernoull((0.5),
X = 1 - Y, w.p. 1 - p.
(3)
X₂ is the invariant feature while X2 is spurious where p5= = 0.8 and n== 0.7:
Y= X22, w.p. 0.85, { X2, w.p. 0.15,
Y, w.p. P..
X2 ~ (e))oull((O.5),
Xi 1- w.p. 1 p.
(4)
The joint distribution of the mixed environments for each data generation process is consistent with Eq. (2). Thus, from the joint distribution P(X,, X2, Y), learned model only depends on either X,, X2 or both, and fails to generalize for at least one of the two data generation processes. On the other hand, when the partition is given, one can verify that IRM would correctly identify X, (resp. X2) as the invariant feature for the first (resp. second) scenario.
The above toy example is inspired by commonly-used classification tasks in the IRM literature, e.g., CMNIST 4 and CifarMnist 126). For instance, in CMNIST, invariant feature X1 denotes the semantic feature of the shape of hand-written digits '0" and '1', and spurious feature X₂ represents the color, which is either red or green. Label Y E {0, 1} corresponds to the digit shape and is also binary. Indeed, the construction procedure of CMNIST in 4 can be exactly described by the data generation process in Eq. 3 To further demonstrate the impossibility result, we next conduct an empirical validation on a new dataset in accordance to the data generation process in Eq.
Empirical Validation. We construct a variant of CMNIST according to the second data generation process. In this new dataset, color is the invariant feature and digit shape is spurious. We use the same notations where X₁ stands for digit shape, X2 is the color, and Y denotes label. As described in Eq. (4), the label corresponds to the digit color and the digit shape is spuriously correlated with the label. We name this variant of CMNIST as MCOLOR (short for MNIST-COLOO
In the test domain, we set ps = 0.1 to simulate the distributional shift, same as in CMNIST. We can then compare ERM, EIIL, IRM with environment partition, and LiF (learning from failures) B that tries learn robust model by boosting from wrongly specified samples of shallow neural networks. The empirical results reported in Table For ERM (oracle), train the model only on the invariant feature, i.e., digit shape in CMNIST and color in MCOLOR.
We can see that the EIIL method performs poorly on the MCOLOR dataset. This is due to the inductive bias of EIIL and it indeed relies on the digit shape as the invariant feature. 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. Similarly, ERM and LIF all rely on either color or shape as the invariant feature and would fail on at least one of CMNIST and MCOLOR. On the other hand, if environment partition is available, IRM can still learn the desired invariant feature.

---

Table Experimental results on CMNIST and MCOLOR.
CMNIST
MCOLOR
Method
Fav Partition
Train Ace
Test Ace
Train Ace
Test Acc
ERM (oracle)
75.21...
72.110.1
85.011..0
85.0+1..0
ERM
No
80.0
14.5 #0.1
86.3+0.11
80.1
IRM
Yes
71.4 10.3
66.4 10.3
84.910.1
84.7:10.3
EMT
No
72.5 +0..
67.2 +3.3
74.0++0.
17.8+0.4
LIF
20.3
21.2 40.4
76.620.0
74.210.0
Moreover, such examples are generation process like (I) (which can be replaced by other forms of SCM), we can find different data generation process that induces the same distribution yet has different invariant/spurious features. Similarly, it is impossible to identify the spurious features only based on ((X This result is formally summarized in Theorem
For any data distribution P(X, Y) generated by some data
Theorem 1. Let X and X, be respectively the invariant and spurious features, with label Y. For the joint distribution P(X, Y) consisting of data from multiple environments with P(X,, Y) # we can find X and X, as invariant and spurious features so that:
X + X and X Xx #0;
( X together with some noise variables generate the same distribution P(X,Y).
A proof is provided in Appendix B.I Under this circumstance,, have to introduce additional assumptionscoondiions or certain "inductive bias" to identify the invariant features [5] 37 20 32 50
The latter may come implicitly from the proposed algorithm, model, and/or the data. We believe that this is the reason that makes existing methods achieve improved empirical performance their the discrepancy of spurious features among clusters ERM model of first stage should heavily or fully rely on the spurious feature. However, these inductive biases may the role of inductive biases shall be
considered scenarios. For example, in [29 30..
are expected to larger than that of causal features, and in [TO 34 42 28
be always guaranteed. Thus, suggested by Theorem [ discussed more explicitly when learning invariant models without environment partition.
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 (3] 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 [66] It is worth noting that Xie et al. [53] semi-supervised manner. However, they consider a different setting with access to unlabeled (out-of-domain) test data together with additional unlabeled training data. In the next section, we show how to utilize Z to leam invariant models from the heterogeneous dataset.
also take advantage of auxiliary information to help improve OOD performance
5 ZIN: Learning Invariance with Additional Auxiliary Information
We consider that there exists additional auxiliary information e Rdd in company with the data (X, Y). In this section, we propose ZIN, auxiliary information Z for environmental INference, for invariant learning from the heterogeneous dataset D without environment partition. We will derive conditions for invariance identification and show these conditions are both sufficient and necessary in fairly general setting.
5.1 Method
We aim to learn function p(-) 4. - RR that softly assigns a sample to K environments. Here K' is pre-specified number (a hyper-parameter) and its choice will be empirically investigated in our experiments. Let (-) denote the k-th entry of p(-), with p(Z) E [0, 1] and Σg p (()(()) = 1. Denote the ERM loss as R(w, ②) = + 1 (+(x))) m.) and the loss in the k-th inferred environment R Rg0 0 (w, 4)- Σ" (z.)((f.
5

---

Recall that IRM 4 learns an invariant representation L, upon which there is a classifier fw that is simultaneously optimal in all environments. Suppose that the environments have been given according to a fixed p()) To measure the optimality of fw in the k-th environment, we can fit an environment- dependent classifier Juy on the data from that environment. If fa, achieves a smaller loss, then we know that fa is not optimal in this environment. We can further train a set of environmenn-ddeeeddent classifiers {ww j - one for each environment, measure whether fw is simultaneously optimal in all environments. Thus, when p(-) is provided, our formulation to learn invariance is follows:
K
min max C(0, w1, ,WK, p) == R(w, 2 + w.b {we)
J [Rpees (w, D) Reces (wxx Φ)
(5)
invariance penalty
If & extracts spurious features that are unstable in the inferred environments, Rptes (w, b) will be larger than Ro00 (we, Q), resulting in a non-zero invariance penalty.
Next, we consider how to learn the partition function p(-). A good partition function should generate environments among which the spurious features exhibit instability, so that there is a large penalty Φ extracts spurious features. Thus, we seek for an environment partition that maximizes the invariance penalty. The overall framework is provided below:
min max P. .K)
C( W1,, W1, ,WK, P).
The idea of our method can be summarized as inferring an environment partition and then learning invariant models based on the inferred environments. Creager et al. [10] has used similar intuition and adopted a pre-trained biased model to estimate the environments. However, their two-stage method cannot be jointly optimized, and the environment partition relies on the given model and lacks a theoretical guarantee. In contrast, we will first present sufficient conditions on so that the proposed framework can provably identify the invariant features. We further show that these conditions also necessary, and that violation of these conditions will lead to failure of invariance identification. 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.
5.2 Sufficient Conditions for Identifiability
In this section, we try understand ZIN from a theoretical perspective. We start with simple yet general setting: X [X: X,] (i.e., no scramble on the observation), & E {0, 1]dt dd is an element-wise feature selection mask, and fa is general non-linear function. We also focus on classification tasks with E(::)) being the cross-entropy loss, as there is an interesting connection to (conditional) Shannon entropy. Extensions to other loss functions and linear feature transformations are left to Appendix A. and A.2 respectively.
In this setting, our goal is equivalent learning the optimal feature mask that merely selects invariant features, i..., 4, = [1 [=: 0d. For given feature mask 4, the objective function equal to L()) min,, (wx) L( w. wi, ; Wk: p). Then ZIN can correctly identify the invariant features, or solution to Problem 6 is equivalent to $., if and only ((T))) < L(() for all t of t.. This observation will be used to establish our main theoretic result.
With a little abuse of notation, use H(Y to denote expected loss of an optimal classifier over some X' and Y, and similarly (p))p(( X') to denote the minimum risk >k R, (wk: 2) for a given p( Z). With cross-entropy loss and when p(z:) gives exactly one environment, i.e., pl Z) is one-hot, it can be verified that the optimal loss (.---) coincides with conditional entropy. In the following first state the assumptions for our identifiabitiliy result.
Assumption 1. For a given feature mask &b and any constant e > 0, there exists f E F such that (X)) Y)] < (X)) +
Assumption 2. a feature violates the invariance constraint, adding another feature would not make the penalty vanish, i.e., there exists constant 8 > 0 so that for spurious feature X1 X, and any feature X2 C. X. | X²) H(Y p Z), X1, X2) 2 8 (H)) HY(]p(Z),

---

Assumption 3. For any distinct features X1, X2 HY S (( -y fixe > 0.
Assumption 1 is a common assumption that requires the function space F be rich enough such that, given 2, there exists f & F that can fit P(Y (X))) well. Assumption 2 aims to ensure a sufficient positive penalty if spurious feature included (see Appendix |A.3 for further discussion). Assumption
3 indicates that any feature contains some useful information w.r.t. Y, which cannot be explained by other features. Otherwise, we can simply remove such feature (e.g.. by variable selection methods [49]), it does not affect prediction. We next present our sufficient conditions for ZIN to identify invariant features.
Condition (Invariance Preserving Condition). Given invariant X, and any function p(.), it holds that Xx, px, = ,.))
Condition 2 -on-invariaance Distinguishing Condition). For any feature X s, there exists a function p()) and a constant C > 0 such that [ HYY| k ()) 2 C.
We remark that Condition I can be met if (Y),,, Z) (Y X,) (a proof is provided in Appendix B.6). Condition | requires that invariant features should remain invariant w.r.t any environment partition induced by pl p. Otherwise, if there exists a partition where an invariant feature becomes non-invariant, then this feature would induce a positive penalty. Condition implies that for each spurious feature, there exists at least one partition so that this feature is non-invariant the split environments. If a spurious feature does not incur any invariance penalty in all possible environment partitions, we can never distinguish it from true invariant features. With these conditions, our main result follows and a proof is given in Appendix [8.2]
Theorem 2 (Identifiability of Invariant Features). With Assumptions 73 and Conditions 7 2 if
C- 17+2CSH(Y] and X [ 6C- 2 2 - then we have ((D)) < (4) for all 1 o 4..
H(Y)+1/26C
Ie 1),
where HY) denotes the entropy of Y. Thus, the solution to Problem 6 iidentifies invariant features.
5.3 Necessary Conditions for Identifiability
Conditions D and 2 may appear rather strong. In this section, we prove that they are also necessary conditions for ZIN to identify invariant features. Specifically, Proposition shows that if Condition | is violated, then some invariant features will be excluded in the solution of Problem 6 and in renders some spurious features included in the solution.
Proposition 2
violation of Condition
2
Proposition 1. With Assumptions
3 if Condition I is violated, i.e, there exists p(.) so that 0, then there exists a mask o' + T with L(),,) > ((0)).
Hyy X.) - Hy) p.. p())
2 C'
proof is provided Appendix B. which is similar to the first step of the proof of Theorem question is how Condition D could be violated. On the other hand, while we introduce as obtain valid partition based on only the training data (X, YY. Below we provide examples regarding these questions, with proofs given in Appendix B. As (IIndex(2 X, Y)) can be treated as learning independent weights for each sample based on { bi, this case includes EIIL [10] as an instance.
additional variables, it is also interesting to know whether
Corollary 1. Condition is violated for the following cases: there exists a function p(-) and an injec- tive function lu(-) so that (a) p(Z) = i(Y)), (b) p(Z) = h(X,Y), or (c) p(Z) = h(Index(2 (X, Y)).
We now discuss the necessity of Condition 2 with an additional assumption.
Assumption 4. If two features are invariant w.r.f. environment partition, then the concatenated features are also invariant. That is, for X1, X2 C X and p()), if H(Y)X,) ())))(()). X1) and HY X2) -Y (Ylp((Z). X2) = 0, we have H(Y)X,, X2) - HY)|p(Z),
= 0.
Proposition 2. With Assumptions 7 and 4 if Condition 2 is violated, i.e., there exists a spurious feature X E X, such that H(Y x X p(Z) 0 for any p(-), then there exists
feature mask D' + 2 with (().) > C(2)).
Since there exists spurious feature X that is "invariant" in all possible environment partition, adding this feature to &. does not induce any invariance penalty but increases the prediction power.
7

---

Thus, such a feature mask can achieve a smaller loss than ((..). For completeness, we provide proof in Appendix B.5 Proposition again indicates in theory that Z should be sufficiently diverse and informative so that each spurious feature can be recognized.
agaain
5.4 Choice the Auxiliary Information
Conditions and present now discuss how to find such auxiliary information given certain prior knowledge.
present theoretical requirements of Z to successfully learn invariant features. We
Take Fig. F example. When we collect data point (e.g., taking a photo), there often exists some meta information (e.g., time slot, coordi- nate, and temperature). Suppose that the image is generated according to the causal graph on the right panel of Fig. [ In this case, the in- variant feature X, consists of X1 and X2, and the meta information can be used as Z which colored green in the causal graph. This is because H X X2) H X2, Z) or equivalently Y - Z X2, so that Condi- tion [ holds. This example shows some evidence of choosing the auxiliary information: Z and Y should be d-separated by X,. There are other examples satisfying Condition I which are given Fig.
Auniliary Information:
Time sloe: 2021/0112 17:00 Coordatte: (30 120702) Tempercure: 25.6'C
Figure 1: An example of Z satisfying Condition
S in Appendix E We also illustrate some failure choices in Fig. G Appendix Specifically, if Z is a child or a direct cause of Y. then it contains additional information of conditional on the invariant feature, violating Condition T This kind of information cannot be used in our framework.
To summarize, with some prior knowledge on the path between and Y, we suggest to choose as much auxiliary information possible under the following condition: A feasible choice of Z should satisfy Condition and should be correlated with certain variables in the causal graph of the SCM. The puth belween Z and Y should be d-separated by X (e.g. in Figs. [and 5) or there is path between Z and Y. Notably, Z cannot be the parent or child of Y as shown in Fig. 6
6 Experiments
This section empirically verifies our theoretic analysis and the effectiveness of ZIN on both synthetic and real world datasets. We compare ZIN with several existing methods: ERM. IRM, 4 group DRO [11, EIIL 10I, HRM (29), and LIF (34). We provide the ground-truth partition IRM and group DRO. Notice that LfF only works for classification tasks.
We implement both and f in ZIN as two-layer MLP with 32 hidden units. We adopt the first order approximation of Eq. (5). as described in Appendix C The number inferred environments K is set to as default. We implement as a two-layer MLP for the synthetic dataset and house price prediction, ResNet-18 [6 for CelebA, and CNN for Landcover. More details are provided in Appendix C
6.1 Synthetic Dataset
Temporal Heterogeneity. We consider temporal heterogeneity with distributional shift w.r.t. time. Let [o, 1] be time index and (t) the invariant feature. The data generation process is
(t) [W((, w.p. 0.5, (W(-1,0") w.p. 0.5, Y(t) (X..((X((1)) w.p. Pus X.(t) (N((( o²), w.p. (a(),, (f))(( w.p. pl.
).
w. (N(-((t),')) ), (t).
where Pv is a constant w.r.t. t, indicating a stable correlation between ((t) and X (t). 11
and (4) = (t). Notice here (t) varies with time f. similar setting with two-dimensional spatial variable considered in Appendix D.I Our goal to learn a model that purely relies on
X,, We simulate two heterogeneous environments along time, namely, { 0.5), [0.5, 1]}. and (t) will be set differently. We use tuple of (t) in the two environments denote simulated case. For example, (0.999, 0.7) stands for (t) 0.999, t [0, 0.5) and (t) 0.7, [0.5, 1]. We
8

---

Table 2: Test Mean and Worst accuracy (%) on four temporal heterogeneity synthetic datasets.
Px(t)
0.999, 0.7
0.999,
0.8
0.999. 0.9
Eav Partition
Po
0.9
0.8
0.9
0.8
0.9
0.8
Test Acc
Meun
Worst
Mean
Worst
Mean
Worst
Mean
Worst
Mean
Worst
Meun
Worst
ERM
75.37
57.31
59.65
25.81
68.72
41.97
55.90
15.07
60.61
23.39
52.85
7.57
No
EIIL
38.41
16.80
64.89
49.15
50.77
46.67
68.36
56.35
61.99
53,81
70.10
59.36
HRM
50.00
49.99
49.98
49.93
50.00
49.98
50.01
49.99
50.00
49.98
49.99
49.97
ZIN
87.50
85.36
77.85
75.39
86,35
82.91
76.79
72.77
83.71
75.89
73.55
64.69
Yes
IRM
87.57
85.47
77.99
75.65
86.57
83.25
77.00
73.39
83.99
76.48
73.84
65.33
evaluate the performance on four distinct test environments with ps E {009999, 0.2, 0.1] and pr being constant. We use as the auxiliary information. More details are provided in Appendix C
Results. Table reports the test accuracy. In all simulated settings, the worst accuracy of ERM is much lower than the mean accuracy, indicating that ERM tends rely on spurious feature Xg. EIIL can improve the worst accuracy in some cases, e.g., when (t)
(0.999, 0.8) and (t) (0.999, 0.9). However, its performance is even worse than ERM for some other settings. 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. The proposed method ZIN improves the worst test accuracy significantly. For instance, when (t))) = (0.999, 0.7) and P. = 0.9, ZIN outperforms ERM and EIIL by over 28% and 65%, respectively. Moreover, ZIN is very close to IRM that knows the ground truth environment partition, showing that ZIN can infer the environments effectively. Finally, it seems that HRM does not learn a useful model for this classification task.
Ablation Study. In Appendix D.2 we conduct ablation study to verify our theoretical results in Section 3 by choosing different auxiliary variables. We also empirically study the choice of hyper-parameter K. A notable observation is that K has relatively small impact, as shown in Fig. 3
6.2 Real World Datasets
House Price Prediction. This experiment considers a real world regression dataset of house sales prices from Kaggle https://www.kaggle
The target variable is house price and each sample contains 17 predictive variables like the built year of the house, number of bedrooms, etc. The dataset is split according to the built year, with training dataset in period [1900, 1950] and test dataset in period (1950, 2000). We normalize the prices of the houses with the same built year, and our target is to predict the normalized price. We choose the built year as auxiliary information for ZIN. As there is no well-defined ground-truth partition for IRM, manually split the training dataset equally into 5 segments with 10-year range in each segment.
Table 3 reports the mean squared error (MSE), and ZIN again outperforms other methods on test dataset. We also investigate whether this type of additional information Z. together with X, helps EIIL and HRM. The corresponding results are marked with (+ Z) in Table We observe that the results of EIIL and HRM with additional information Z are worse than those without using Z. This implies that such information may not be directly utilized by or even can harm the two methods.
CelebA. This task is to predict Smiling based on the image from CelebA BI. and by construction the target is spuriously correlated with Gender. We have access to the meta annotations of CeleA images and they serve Z in our framework. In particular, we pick seven features potential auxiliary variables: { Youngg Blond Hair, Eyeglasses, High Cheekbones, Big Nose, Bags Under Eyes, Chubby]. 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.
Empirical results are given in Table 4 where ZIN represents ZIN using # of the seven features as Z. ERM achieves the highest training accuracy, while only has 47.58% worst test accuracy. The worst test performances of ERM, EIIL and L.F indicate that they may not learn the causal features. Finally, we observe that ZIN performs better when more auxiliary information is provided; specifically, ZIN(7) achieves the best test performance based on the mixed dataset. This observation validates discussion on choosing the auxiliary information in Section 5.4

---

Table House price prediction (MSE).
Table 4: Accuracy (%) on CelebA task.
Method
Index
Tmin
Test Mean
Test Worst
Mean
Test
IRM
Yes.
0.1327
0.4456
0.6821
IRM
Yes
81.30 1.3
78.4=-048
75.031: 1.29
group DRO
Yes
0.1213
0.6887
1,0050
group DRO
89.32#0.67
74.28=0.22
58, 08.11=0.61
ERM
No
0.1141
0.4764
0.6703
ERM
No
90,97
70.76=0.26
47.5810,466
EIIL
No
0.6841
0,9625
1,3909
LIF
No
59. 59.89+00.2
52.97-0.56
44.38+22.01
EII((+2)
No
0.6912
0,9701
EIIL
+0.73
71.45-0.29
50,48+
HRM
No
0.3466
0.4621
0.5721
ZIN(U)
No
90.62:00.78
70.79=0.61
47.62 0898
HRM(- Z)
No
0.3190
0,4221
0.5873
83.57
75.20=0.71
63.471:144
ZIN
No
0,3275
0,3389
0,4815
ZIN(7)
No
83.06
76.29-0.60
67.27 1.15
In Appendix [D] Fig. reports the results of different choices of visualizes the inferred environments in this experiment. Again, we find that K has a relatively small impact the proposed algorithm.
K and Fig.
Landcover. Our final task is land cover prediction that classifies the land cover type (e.g., grasslands) We take the same setup
Table 5: Test Accuracy (%) on Landcover task.
Method
ITD Test
OOD Test
from satellite data 14 40 from [53]: the time series data input dimension 46 8; the target Y is one of six land cover classes; six climate related variables are the auxiliary vari-
ERM
52
75.92
58.31
ERM (+climate)
76.58
54.78
I.fF
66.24
61.69
EIlL
72.61
64.79
ables. We also consider non-African locations training data and Africa test data. In addition, we take location (latitude and longitude) as possible ad- ditional information. For ZIN, we average climate variables over the time dimension and use those means to predict environments. In Table5] we see that ZIN with location achieves the best OOD performance, while ZIN with climate variables slightly worse. This is because climate variables indeed contain some information for predicting the target, violating our conditions on Z. ElIl. also has good performance in this task, and we conjecture that its first stage has learned the spurious features desired. However, together with previous empirical results, it can be risky the first stage may not be guaranteed.
ZIN (climate)
72.56
62.50
ZIN (location)
72.18
66.06
6.3 Discussion on Auxiliary Information for Related Methods
The proposed framework ZIN relies on additional auxiliary variables. It is interesting to ask whether this type of information is also useful to related methods such as EIIL and HRM. 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 predict the label; and 2) in the CelebA classification task, features like Young and Eyeglasses are likely to provide little information about predicting Smiling. Then Z may not be useful to ERM and EIIL.. In Section 6.2 we also provide Z to EllL and HRM, and the experimental result shows that it does not help or even harms the test performance.
Concluding Remarks
This paper investigates when and how to learn invariance from heterogeneous data without explicit environment indexes. We first show that learning invariant models this case is generally impossible. Then we propose ZIN to jointly learn environment partition and invariant representation using some additional auxiliary variables. We provide theoretic guarantees in both feature selection and linear feature learning scenarios. Experimental results verify our analysis and demonstrate improved performance of ZIN over existing methods. 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. 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.
10

---

References
[1] Faruk Ahmed, Yoshua Bengio, Harm van Seijen, and Aaron C. Courville. Systematic generali- sation with group invariant predictions. In ICLR, 2021.
[2] Kartik Ahuja, Karthikeyan Shanmugam, Kush Varshney, and Amit Dhurandhar. Invariant risk minimization games. In International Conference on Machine Learning, 2020.
[3] Kartik Ahuja, Jun Wang, Amit Dhurandhar, Karthikeyan Shanmugam, and Kush R. Varshney. Empirical invariant risk minimization? sample complexity perspective. International Conference Learning Representations, 2021.
[4] Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk mini- mization. arXiv preprint 190..0v:997029999 2019.
[5] Aharon Ban, Dick Den Hertog, Anja De Waegenaere, Bertrand Melenberg, and Gijs Rennen. Robust solutions of optimization problems affected by uncertain probabilities. Management Science, 59(2)3341-357, 2013.
[6] Joos-Hendrik Böse, Valentin Flunkert, Jan Gasthaus, Tim Januschowski, Dustin Lange, David Salinas, Sebastian Schelter, Matthias Seeger, and Yuyang Wang. Probabilistic demand forecast- ing scale. Proceedings of the VLDB Endowment, 10(12):1694-1055 2017.
[7] Shiyu Chang, Yang Zhang, Mo Yu, and Tommi Jaakkola. Invariant rationalization. In Interna- tional Conference Machine Learning, 2020.
[8] Yongqiang Chen, Yonggang Zhang, Han Yang, Kaili Ma, Binghui Xie, Tongliang Liu, Bo Han, and James Cheng. Invariance principle meets out-of- generalization on graphs. arXiv preprint arXiv:220..0544/ 2022.
[9] Yongqiang Chen, Kaiwen Zhou, Yatao Bian, Binghui Xie, Kaili Ma. Yonggang Zhang, Han Yang, Bo Han, and James Cheng. Pareto invariant risk minimization. arXiv preprint vrXiv:2206.07766 2022.
[10] Elliot Creager, learning. In International Conference on Machine Learning, 2021.
Jacobsen, and Richard Zemel. Environment inference for invariant
[11] David Danks, Clark Glymour, and Robert Tillman. Integrating locally learned causal structures with overlapping variables. In Advances in Neural Information Processing Systems, 2009.
[12] John C Duchi and Hongseok Namkoong. Learning models with uniform performance via distributionally robust optimization. The Annals of Statistics, 1378)4438 2021.
[13] Rui Gao, Xi Chen, and Anton J Kleywegt. Wasserstein distributionally robust optimization and variation regularization. arXiv preprint arXiv:1712.060550 2020.
[14] Pall Oskar Gislason, Jon Atli Benediktsson, and Johannes R. Sveinsson. Random forests for land cover classification. Pattern Recognition Letters, 27())2294 300, 2006.
[15] Ishaan Gulrajani and David Lopez-Paz. In search lost domain generalization. arXiv preprint arXiv:2007.01434, 2020.
[16] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2016.
[17] Dan Hendrycks and Thomas G. Dietterich. Benchmarking neural network robustness common corruptions and perturbations. In International Conference Learning Representations, 2019.

---

[18] Biwei Huang, Kun Zhang, Mingming Gong, and Clark Glymour. Causal discovery from multiple data sets with non-identical variable sets. In Proceedings of the AAAI Conference on Artificial Intelligence, 2020.
[19] Biwei Huang, Kun Zhang, Jiji Zhang, Joseph Ramsey, Clark Glymour, and Bernhard Schölopf. Causal discovery from beterogenooussoonsaatiooaay data. Journal of Machine Learning Re- search, 21, May 2020.
[20] Aapo Hyvarinen, Hiroaki Sasaki, and Richard Turner. Nonlinear ICA using auxiliary variables and generalized contrastive learning. In International Conference on Artificial Intelligence and Statistics, 2019.
[21] Wengong Jin, Regina Barzilay, and Tommi Jaakkola. Domain extrapolation via regret mini- mization. arXiv preprint arXiv:2006. 2000:20000.0908 03908, 2020.
[22] Ilyes Khemakhem, Diederik Kingma, Ricardo Monti, and Aapo Hyvarinen. Variational autoen- coders and nonlinear ICA: A unifying framework. In International Conference on Artificial Intelligence and Statistics, 2020.
[23] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Interna- tional Conference on Learning Representations, 2014.
[24] David Krueger, Ethan Caballero, Joem-Henrik Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang. Remi Le Priol, and Aaron Courville. Out-of-distribttion generalization via risk extrapo- lation (rex). In International Conference on Machine Learning, 2021.
[25] Jaeho Lee and Maxim Raginsky. Minimax statistical learning with wasserstein distances. In Advances in Neural Information Processing Systems, 2018.
[26] Yong Lin, Lian Qing, and Tong Zhang. An empirical study of invariant risk minimization on deep models. 2021.
[27] Yong Lin, Hanze Dong, Hao Wang, and Tong Zhang. Bayesian invariant risk minimization. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2022.
[28] Evan Zheran Liu, Behzad Haghgoo, Annie S. Chen, Aditi Raghunathan. Pang Wei Koh, Shiori Sagawa, Percy Liang, and Chelsea Finn. Just train twice: Improving group robustness without training group information. In ICML, 2021.
[29] Jiashuo Liu, Zheyuan Hu, Peng Cui, Bo Li, and Zheyan Shen. Heterogeneous risk minimization. In International Conference on Machine Learning, 2021.
[30] Jiashuo Liu, Zheyuan Hu, Peng Cui, Bo Li, and Zheyan Shen. Kernalized heterogeneous risk minimization. In NeurIPS, 2021.
[31] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang, Deep learning face attributes in the wild. In Proceedings of International Conference Computer Vision (ICCV), 2015.
[32] Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Ractsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of disentangled representations. In International Conference on Machine Learning, 2019.
[33] Manfred Mudelsee. Trend analysis of climate time series: A review of methods. Earth-Science Reviews, 190:310-322, 2019.
[34] Jun Hyun Nam, Hyuntak Cha. Sungsoo Ahn, Jaeho Lee, and Jinwoo Shin. Learning from failure: Training debiased classifier from biased classifier. In NeurlPS, 2020.
[35] Judea Pearl. Causality: Models, Reasoning and Inference. Cambridge University Press, USA, 2nd edition, 2009.
12

---

[36] Jonas Peters, Peter Bühlmann, and Nicolai Meinshausen. Causal inference by using invariant prediction: Identification and confidence intervals. Journal of the Royal Statistical Society. Series B (Statistical Methodology), pages 947-1012, 2016.
[37] Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Elements of Causal Inference Foundations and Learning Algorithms. MIT Press, Cambridge, MA, USA, 2017.
[38] Benjamin Recht, Rehecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do ImageNet classifiers generalize to ImageNet? In International Conference Machine Learning, 2019.
[39] Elan Rosenfeld, Pradeep Ravikumar, and Andrej Risteski. The risks of invariant risk minimiza- tion. International Conference on Learning Representations, 2021.
[40] Mare Russwurm, Sherrie Wang, Marco Korner, and David Lobell. Meta-learning for few-shot land cover classification. In Proceedings of the IEEF/CVE Conference on Computer Vision and Pattern Recognition (CVPR) Workshops, June 2020.
[41] Shiori Sagawa, Pang Wei Koh, Tatsunori B Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. In International Conference on Learning Representations, 2020.
[42] Victor Sanh, Thomas Wolf, Yonatan Belinkov, and Alexander M. Rush. Learning from others' mistakes: Avoiding dataset biases without modeling them. In ICML, 2021.
[43] Alexander Shapiro. Distributionally robust stochastic programming. SIAM Journal on Opti-
mization, mization, 27(4)2225-22222) 2017.
[44] Nimit Sharad Sohoni, Jared A. Dunnmon, Geoffrey Angus, Albert Gu, and Christopher Ré. No subclass left behind: Fine-grained robustness in coarse-grained classification problems. In NeurlPS, 2020.
[45] Peter Spirtes, Clark N Glymour, and Richard Scheines. Causation, Prediction, and Search. MIT Press, second edition, 2000.
[46] Megha Srivastava, Tatsunori Hashimoto, and Percy Liang. Robustness to spurious correlations via human annotations. In International Conference on Machine Learning, 2020.
[47] Christian Szegedy, Wojeiech Zaremba, Ilya Sutskever, Joan Bruna, Dumitru Erhan, Ian Good- fellow, and Rob Fergus. Intriguing properties of neural networks. In International Conference on Learning Representations, 2014.
[48] Robert Tillman and Peter Spirtes. Learning equivalence classes of acyclic models with latent and selection variables from multiple datasets with overlapping variables. In Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics, 2011.
[49] Jorge Vergara and Pablo Estevez. A review of feature selection methods based mutual information. Neural Computing and Applications, 24, 01 2014.
[50] Ruoyu Wang, Mingyang Yi, Zhitang Chen, and Shengyu Zhu. Out-of- generalization with causal invariant transformations. In /EEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2022.
[51] Sherrie Wang, William Chen, Sang Michael Xie, George Azzari, and David B. Lobell. Weakly supervised deep learning for segmentation of remote sensing imagery. Remote Sensing, 12(2), 2020.
[52] Chuanlong Xie, Fei Chen, Yue Liu, and Zhenguo Li. Risk variance penalization: From distributional robustness to causality. arXiv preprint arXiv:2006.075.4 2020.
13

---

[53] Sang Michael Xie, Ananya Kumar, Robbie Jones, Fereshte Khani, Tengyu Ma, and Percy Liang. In-N-Out: Pre-training and self-training using auxiliary information for out-of-distribution robustness. In International Conference on Learning Representations, 2021.
[54] Yilun Xu and Tommi Jaakkola. Learning representations that support robust transfer of predic- tors, arXiv preprint urXiv:21100099. 2021.
[55] Mingyang Yi, Lu Hou, Jiacheng Sun, Lifeng Shang. Xin Jiang, Qun Liu, and Zhi-Ming
Ma. Improved OOD generalization via adversarial training and pre-training. In International Conference Machine Learning. 2021.
[56] Xiao Zhou, Yong Lin, Renjie Pi, Weizhong Zhang, Renzhe Xu, Peng Cui, and Tong Zhang Model agnostic sample reweighting for out-of-distribution learning. International Conference on Machine Learning, 2022.
[57] Xiao Zhou, Yong Lin, Weizhong Zhang, and Tong Zhang, Sparse invariant risk minimization. In International Conference on Machine Learning, 2022.
14


Reviews:
[{"review_id": "yjja9FTUBR", "reviewer": "Reviewer_3T6p", "summary": "The authors propose an approach to learning an invariant representation, which simultaneously learns the representation along with an environment partition over samples based on auxiliary measurements, such as time indices or location metadata. They provide sufficient conditions for their proposed method to identify invariant features in both the feature selection and the linear feature learning setting. They also describe necessary conditions for the identifiability of invariant features in the feature selection setting. Finally, experiments demonstrate that the proposed method is effective at learning invariant features, almost matching the performance of IRM which has access to a ground-truth partition of the variables.", "questions": "**Major questions**\n+ What is the intuition behind Assumption 3? In particular if we think of $X_v$ as the parents of $Y$ and $X_s$ as the children of $Y$?\n+ I did not understand Corollary 1. First, I assume it is meant to be (a) _or_ (b) _or_ (c), not _and_. Second, what does “Index” denote? Why is $h$ supposed to be injective - can’t it map the same inputs to the same one-hot vector?\n\n**Suggestions to help reader understanding**\n+ It helps to describe new quantities directly after they are introduced. For example, after Equation (3), mention that $\\rho$ softly partitions environments before introducing the other definitions, mention that $\\mathcal{R}_{\\rho^{(k)}}(\\omega, \\Phi)$ is the loss of $f_w$ in environment $k$, etc.\n+ It may help to separately define the invariance penalty, including the max term. This emphasizes that the max term is only involved in this term, and directly highlights what is being proposed.\n+ It would help to highlight the differences between Theorem 3 and the results for IRM, since they are so similar. It appears simply that you allow the environment partition to be learned from auxiliary information instead of being given, are there any other significant differences?\n\nMinor points\n+ Assumption 3 is unclear: you say adding another feature does not make the penalty diminish, but in the second sentence, it seems that it is allowed to become smaller, just that it is not allowed to vanish to zero.\n+ After Equation (3), you use $w$ to subscript $f$ instead of $\\omega$\n+ Assumption 4: any “distinct” features, not “distant”\n", "limitations": "Limitations and potential negative impacts are adequately addressed.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n+ The paper is well-written in most aspects: assumptions are cleanly stated and intuitively explained after they are introduced. The feature selection example provides an intuitive illustration of the ideas.\n+ The paper is “significant” in an uncommon way. In particular, it seems that it will play an important regularizing role to correct the confusion present in some other works on invariant learning, which have sought to learn an environment partition without metadata. For instance, the authors correctly point out that EIIL can only be successful when the spurious features are *more* correlated with the label than the invariant features are. The presentation of necessary conditions in Propositions 1 and 2 is especially helpful for streamlining future literature on this topic.\n+ The experiments are well-done, with the proposed method performing near the “oracle” baseline of IRM across a variety of different domains (synthetic, house price prediction, a CelebA experiment, and a land cover prediction task).  \n\n**Weaknesses**\n+ Most of the results are not particularly surprising for those familiar with causality. For instance, Example 1 is a standard example of Markov equivalence.\n+ Some of the results seem tautological, e.g., violating Condition 2 seems equivalent with the conclusion of Proposition 2 that a spurious feature will be included in the feature selection.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "YQc6V0k_ff", "reviewer": "Reviewer_Ypmy", "summary": "This paper theoretically proves that it's impossible to infer the environments purely from the heterogeneous data; consequently, additional information is definitely needed. It proposes a framework that jointly learns environment partitions and invariant representations with the additional information. Experimental results demonstrate the effectiveness of this method. ", "questions": "1. What is the auxiliary information $Z$ in the experiment synthetic dataset shown in Section 7.1? There is no $Z$ in the synthetic process. ", "limitations": "I haven't seen any obvious limitations except for the one mentioned in the last part of the paper. ", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: \n\n1. This paper theoretically proves that it's impossible to identify the invariant features purely from heterogeneous data, which is meaningful. The toy example is intuitive and easy to understand. \n\n2. This paper proposes a framework to jointly learn environment partition and invariant representation, which is simple but equipped with a theoretical guarantee, which is also meaningful. \n\n3. This paper is well-written and easy to follow. \n\nWeaknesses:\n\n1. I haven't seen the guarantee to demonstrate how informative the auxiliary information $Z$ should be to identify the invariant features. Looking through Assumptions 1-4 and Conditions 1-2, there seems to be no discussion on $Z$, which should be included in the paper. \n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "RzrgE6Wm5Nu", "reviewer": "Reviewer_TsN9", "summary": "There is a recent line of work that build on methods that learn invariant representations from multiple environments [Peters et al 2016, Arjovsky et al 2019] by attempting to infer the environment label from data. This paper presents a simple counterexample that shows that environment inference is not possible in general---essentially they construct examples where the joint is consistent with two different SCMs, each of which have a different invariant feature; and hence any environment inference procedure that depended only on the joint distribution would be wrong in at least one of the SCMs---and instead argues for environment inference from auxiliary variables. They complement their negative results with theory that shows that auxiliary variables are sufficient for environment inference in the linear and feature selection settings, and give nice experimental results on both simulated and real data.", "questions": "1. Why is assumption 2 necessary? If you weren't making it explicitly, I would have thought it was just a consequence of your definition of invariance (line 116 / 117). What is it ruling out?\n\n2. Related to my weakness comment - many of the selected auxiliary variables seem very similar to an adjustment set that you would control for in order to block paths to the environment. Is there an explicit connection here? Can we interpret the method as an invariant representation learning alternative to adjustment? And if so, does ZIN fail in the presence of unblocked confounding?", "limitations": "Yes", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "This review is going to skew short because I really liked this paper: I've long thought that environment inference was asking too much so I really appreciate a clear counter example that shows where it can fail (of course, that still leaves open the question of sufficient conditions for env inference to succeed without auxiliary info; but that's a separate question), and the auxiliary information perspective is a nice way around the negative examples. I also appreciated that the empirical section focused on more realistic datasets. \n\nWeaknesses:\n* The biggest weakest for me was that it doesn't seem obvious how one should choose an auxiliary variable in practice. There are a nice collection of examples in the experimental section which partially addresses this, but I think it could have been strengthened by including a section where you discuss the process by which you selected the auxiliary variables. For each of the datasets, you could explain what variables are available, why the selected subset make sense as auxiliary variables, and under what conditions you'd expect environment inference to fail.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "-y-YuaiNEPJ", "reviewer": "Reviewer_9XzW", "summary": "This work explores the feasibility of utilizing cheaply available auxiliary features to uncover latent environment partition and help to train an invariant model. The proposed method, namely ZIN, is a min-max game. Specifically, the learned environment partition aims to split samples into groups with different distributions by maximizing the performance gap between environments. The feature extractor and the classifier aim to learn causally invariant features by minimizing such performance gaps. The proposed model is significant when ground environment labels are costly to obtain while additional information that correlates with environments is easy to acquire.", "questions": "1. According to the author, it is theoretically impossible to learn invariant features from heterogeneous data without environment indexes. A counter-example and a general theoretical proof are provided by the author.\na). For counter-example, two examples are symmetric. In other words, the causal graph cannot be identified given the data. Both X1 and X2 can be treated as either invariant features or spurious features. It is not indistinguishable, but more like both are right. Additionally, the data generation process is impractical. It is not a reasonable example, more like a math game.\nb). For the general theoretical result, the proof is under the specific SCM (fully informative invariant features). However, there are other widely used assumptions [1]. Personally, I don't think the timing is ripe to come to such a firm conclusion when a theoretical proof under only one specific assumption is proposed.\n\n2. More information and discussion can be proposed for the methods section. Why the invariance penalty is effective? Why are different classifiers being used for various estimated environments?\n\n3. Only IRM, used as an Oracle baseline, has been compared in the experiments. What about a different class of generalization-enhancing techniques, like group DRO?\n\n4. Figure 1 shows the correlation between inferred environments and true environments in synthetic data, which is important. Also, Figure 2 demonstrates the effect of K - the number of inferred environments, which is crucial for ZIN. Since K, as a hyperparameter, is pre-specified in ZIN. Then how about the situation in real data? Visualization of the environment and the impact of K on real data is essential to demonstrating the efficacy of ZIN. Could you provide, for instance, the distribution of two environments for gender in CelebA?\n\nA small question: It seems that the additional auxiliary variables should have a strong correlation with spurious variables, so that the distribution of spurious variables is distinct in different environments. And the classifier can further remove the dependency. For the regression dataset of house sales price, why the built year is auxiliary information for ZIN? I think it is a cause for predicting the price. \n\n\n\n[1] Invariance Principle Meets Information Bottleneck for Out-of-Distribution Generalization.\n\n\n\nPost rebuttal comment \n---------------------------------------------------\nI thank the authors for providing their feedback and addressing all my concerns.\n\nI was already convinced of the quality of the paper. Reading comments by the authors did nothing other than reinforce my assessment. After the rebuttal, I would like to raise my score from 6 to 8.", "limitations": "There is no negative societal impact of this work.", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "From a personal perspective, the key strength of this paper is its training of an end-to-end model with additional auxiliary information to address environment inference. This idea, for me, is novel and interesting. \n\nHowever, I still have some concerns about the content of this paper. \n\n1. The author makes a very strong conclusion - without inductive biases or additional information, environment inference is impossible, but neither the example nor the theorem convinced me. Does it contest the theoretical viability of resampling-based techniques like LrF?\n\n2. The Methodology section's structure could be enhanced.\n\n3. Some important baselines, such as \"Group DRO\" are missing.\n\n4.Important case studies are missing - including the visualization of environments and the effect of K on real data.\n\nThe following part has more information.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}]

Author rebuttal:

Decision: Accept