Invariant Causal Representation Learning
Authors: Chaochao Lu
Yuhuai Wu
José Miguel Hernández-Lobato
Bernhard Schölkopf
Venue:  (2021)

Abstract: Due to spurious correlations, machine learning systems often fail to generalize to environments whose distributions differ from the ones used at training time. Prior work addressing this, either explicitly or implicitly, attempted to find a data representation that has an invariant causal relationship with the outcome. This is done by leveraging a diverse set of training environments to reduce the effect of spurious features, on top of which an invariant classifier is then built. However, these methods have generalization guarantees only when both data representation and classifiers come from a linear model class. As an alternative, we propose Invariant Causal Representation Learning (ICRL), a learning paradigm that enables out-of-distribution generalization in the nonlinear setting (i.e., nonlinear representations and nonlinear classifiers). It builds upon a practical and general assumption: data representations factorize when conditioning on the outcome and the environment. Based on this, we show identifiability up to a permutation and pointwise transformation. We also prove that all direct causes of the outcome can be fully discovered, which further enables us to obtain generalization guarantees in the nonlinear setting. Extensive experiments on both synthetic and real-world datasets show that our approach significantly outperforms a variety of baseline methods.

Under review as a conference paper at ICLR 2021
INVARIANT CAUSAL REPRESENTATION LEARNING
Anonymous authors
Paper under double-blind review
ABSTRACT
Due spurious correlations, machine learning systems often fail to generalize environments whose distributions differ from the ones used at training time. Prior work addressing this, either explicitly or implicitly, attempted to find data repre- sentation that has an invariant causal relationship with the outcome. This is done by leveraging a diverse set of training environments to reduce the effect of spuri- ous features, on top of which an invariant classifier is then built. However, these methods have generalization guarantees only when both data representation and classifiers come from a linear model class. As an alternative, we propose Invariant Causal Representation Learning (ICRL), a learning paradigm that enables out-of- distribution generalization in the nonlinear setting (i.e., nonlinear representations and nonlinear classifiers). It builds upon a practical and general assumption: data representations factorize when conditioning Based on this, we show identifiability up to a permutation and pointwise transfor- mation. We also prove that all direct causes of the outcome can be fully discov- ered, which further enables us to obtain generalization guarantees in the nonlinear setting. Extensive experiments on both synthetic and real-world datasets show that our approach significantly outperforms a variety of baseline methods.
the outcome and the environment.
INTRODUCTION
In recent years, despite various impressive success stories, there is still significant lack of robust- ness in machine learning algorithms. Specifically, machine learning systems often fail to generalize outside of specific training distribution, because they usually learn casier-to-fi spurious correla- tions which are prone to change from training to testing environments. We illustrate this point by considering the widely used example of classifying images of camels and cows (Beery al. 2018). The training dataset has a selection bias, i.e., most pictures of cows are taken in green pastures, while most pictures of camels happen to be deserts. After training a convnet on this dataset, it is found that the model fell into the spurious correlation, i.e., it related green pastures with cows and deserts with camels, and therefore classified green pastures as cows and deserts as camels. The result is that the model failed to classify images of cows when they are taken on sandy beaches.
To address the aforementioned problem, a natural idea to identify which features of the training data present domain-varying spurious correlations with labels and which features describe true cor- relations of interest that are stable across domains. In the example above, the former are the features describing the context (e.g.. pastures and deserts), whilst the latter are the features describing ani- mals (e.g., animal shape). Arjovsky et al. (2019) suggest that one can identify the stable features and build invariant predictors on them by exploiting the varying degrees of spurious correlation naturally present in training data collected from multiple environments. The authors proposed the invariant risk minimization (IRM) approach to find data representations for which the optimal classifier is invariant across all environments. Since this formulation is challenging bi-leveled optimization problem, the authors proved the generalization of IRM across all environments by constraining both data representations and classifiers to be linear (Theorem 9 in Arjovsky et al. (2019)).
Ahuja et al. (2020) studied the problem from the perspective of game theory, with an approach that we call IRMG for short. They showed that the set of Nash equilibria for proposed game are equiv- alent to the set of invariant predictors for any finite number of environments, even with nonlinear data representations and nonlinear classifiers. However, these theoretical results in the nonlinear setting only guarantee that one can learn invariant predictors from training environments, but do not guarantee that the learned invariant predictors can generalize well across all environments including

---

Under review as a conference paper at ICLR 2021
unseen testing environments. In fact, the authors directly borrowed the linear generalization result from et al. (2019) and presented it as Theorem in Ahuja et al. (2020).
In this work propose an alternttive learning paradigm, called Invariant Causal Representation Learning (ICRL), which enables out-of-distribution (OOD) generalization in the nonlinear setting (i.e., nonlinear representations and nonlinear classifier). We first introduce practical and general assumption: the data representation factorizes (i.e., its components are independent of each other) when conditioning on the outcome (e.g., labels) and the environment (represented as index). This assumption builds a bridge between supervised learning and unsupervised learning, leading to a guarantee that the data representation can be identified up permutation and pointwise transfor- mation. We then theoretically show that all the direct causes of the outcome can be fully discovered. Based on this, the challenging bi-leveled optimization problem in IRM and IRMG can be reduced to two simpler independent optimization problems, that is, learning the data representation and learning the optimal classifier can be performed separately. This further enables us to attain generalization guarantees in the nonlinear setting.
Contributions We propose Invariant Causal Representation Learning (ICRL), a novel learning paradigm that enables OOD generalization in the nonlinear setting. (i) We introduce conditional factorization assumption on data representation for the OOD generalization (Assumption D. (ii Base on this assumption, we show that each component of the representation can be identified up to a permutation and pointwise transformation (Theorem I & (iii) We further prove that all the direct causes of the outcome can be fully discovered (Proposition I. (iv) We show that
approach has generalization guarantees in the nonlinear setting (Proposition 2). (v) Empirical results demonstrate that our approach significantly outperforms IRM and IRMG in the nonlinear scenarios.
2 PRELIMINARIES
2.1 IDENTIFIABLE VARIATIONAL AUTOENCODERS
A general issue with variational autoencoder (VABs) (Kingma & Welling 2013 Rezende et al,, 2014) is the lack of identifiability guarantees of the deep latent variable model. In other words, it is generally impossible to approximate the true joint distribution over observed and latent variables, including the true prior and posterior distributions over latent variables. Consider simple latent variable model where O E Rd stands for an observed variable (random vector) and X E IR" for a latent variable. [hemakheem et al. (2020) showed that any model with unconditional latent distri- bution pe( (X) is unidentifiable. That is, we can always find transformations of X which change its value but do not change its distribution. Hence, the primary assumption that they make to obtain an identifiability result is to include a conditionally factorized prior distribution over the latent vari- ables po X|U (U). where U Rm is an additionally observed variable (Hyvarinen al., 2019). More specifically, let 0 = (f. T, A) E θ be the parameters of the conditional generative model:
Po X())) = (rr
(1)
where Pf ( = (O - f())) in which is an independent noise variable with probability density function (c), and the prior probability density function is especifically given by
Pr (X II. (x,))) - exp E *
j=1 (X, (U)],
(2)
where Qi is the hase measure, (U) the normalizing constant, T; = (Tr,,....,,T, the sufficient statistics, (U) (U,, (U)) the corresponding parameters depending on U, and k the dimension of each sufficient statistic that is fixed in advance. It is worth noting that this assumption is not very restrictive as exponential families have universal approximation capabilities Sniperumbudur et al. (017). As in VAES, we maximize the corresponding evidence lower bound:
LIVEE O Φ) :=Ep ( [log (O, - log qe (XO,U))]]
(3)
where we denote by PD the empirical data distribution given by dataset D = {O(,) UO) N This approach is called identifiable VAE (iVAE). Most importantly, it can be proved that iVAE can identify latent variables X up to permutation and pointwise transformation under the conditions stated in Theorem 2 of (Khemakhem et al. 2020).
(=1'
A hrief description of variational autoencoders is given in Appendix
2

---

Under review as a conference paper at ICLR 2021
2.2 INVARIANT RISK MINIMIZATION
Arjovsky et al. 2019) introduced invariant risk minimization (IRM), whose goal is to construct an invariant predictor that performs well across all environments Eaul by exploiting the varying degrees of spurious correlation naturally present in the training data collected from multiple environ- ments Eer, where Eer S Eaul Technically, they consider datasets D. == {(o (if.yi)) j== from multiple training environments
&tr. where of O Rd is the input observation and corresponding labe is yf E _ R The dataset De, collected from environment e, consists of examples iden- tically and independently distributed according to some probability distribution P(O Y The goal of IRM to these multiple datasets learn predictor Y = f(O) that achieves the minimum risk for all the environments. Here we define the risk reached by I in environment e as R² = Eo y [e(f(O"), Y.) Then, the invariant predictor can be formally defined below,
Definition 1 (Invariant Predictor, Arjovsky et al. (2019)). We say that data representation & Ha O - C elicits an invariant predictor w & across environments & if there is a classifier w E Hi C - y simultanoussly optimal for all environments, that is, E arg H.. R (16 b)
for all e E E.
Mathematically, IRM can be phrased as the following constrained optimization problem:
min H.W.
Exee, Rd (w
("
s.t. w E arg min R R2 o 4), Ve Exp. DEH..
(4)
Since this is a generally infeasible bi-leveled optimization problem. Arjovsky et al. (2019) rephrased it tractable penalized optimization problem by transfering the inner optimization routine
penalty term. The main generalization result (Theorem 9 in Arjovsky et al. (2019) states that if both & and come from the class of linear models (i.e., H.g. Ra and H.w 1ax1 ). under certain conditions on the diversity of training environments (Assumption in [Arjvssyy et al. (2019) and the data generation, the invariant predictor w & across Eer obtained by solving Eq. remains invariant in Eau. It is worth noting that Ahuja et al. (2020) reconsidered this IRM problem from the perspective of game theory, called IRMG for short. Although the new formulation they proved that there exist such invariant predictors in &,r when both & and are relaxed to the nonlinear models, their main generalization result in Eall holds only when both & and w are linear models (Theorem 2 in AAuua al. (2020))
3 PROBLEM SETUP
3.1 A MOTIVATING EXAMPLE
In this section, we extend the example which was introduced by Wright (1921) and discussed by Arjovsky al. (2019), and provide further in-depth analysis.
Model 1. Consider the following structural equation model (SEM):
X1 + Gaussian (e)), Y + X1 Gaussian (0, o2(e)): X2 + + Gaussian oa(e)).
where (e 2 0, varying in environment Eatl. and Eall is the of all environments.
To ease exposition. here we consider the simple scenario in which Eqr only contains all modifications varying the noises of X1, X2 and within finite range, i.e., (e(e) E [0, 2- Then, to predict Y from (X1, X2) using a least-square predictor ý= â Xi + 62 for environment e, we can
Case I: regress from Xf, obtain â and å 0.
a (c))+(2(())
Case regress from X2, obtain â₃ 0 and ب a (c) (o))))()))
aale) and d2 ca(e)+ Case regress from (Xf, X.). to obtain di -
02(e)
In general scenarios (i.e., 01 (e) * 0, (2(e) * 0, and a3(e) # 0), the regression using X1 in Case 1 is an invariant correlation: this is the only regression whose coefficients do not vary with the environment e. By contrast, the regressions in both Case 2 and Case 3 have varying coefficients
TTii setup for labels applies to both continuous and categorical data, where the categorical data can be encoded in the one-hot form.
3

---

Under review as a conference paper at ICLR 2021
X
X2
Xyi
Xcp
(b)
(c)
Figure 1: (a) Causal structure of Model (b) A more practical extension Model where X1 and X2 are not directly observed and is their observation. (c) A general version of (b), where we assume there exist multiple unobserved variables. Each of them could be either a parent, a child of Y, or has no connection with Y. Grey nodes denote observed variables and white nodes represent unobserved variables.
depending on the enviromment e. Not surprisingly, only the invariant correlation in Case would generalize well to new test environments.
From practical perspective, let us take closer look at Case 3. Because we do not know in advance that regressing from X, alone will lead an invariant predictor, in practice we may do the regres- sion from all the accessible data (Xf. Xz)) As aforementioned, when (e(e) * 0 for i = 2,2,3, the regression does not work. Actually any empirical risk minimization (ERM) algorithm purely minimizing training error [Vapnik 1992) would not work in this setting, Invariant Causal Prediction (ICP) methods (Peters et al.] 2015) also do not work, since the noise variance Y may change across environments. To this end. Arjovsky et al. (2019) proposed IRM. As aforementioned, how- ever, IRM and IRMG can generalize well to unseen testing environments only in the linear setting. This motivates us develop approach enabling the OOD generalization in the nonlinear setting (i.e., both & and w are from the class of nonlinear models).
A more straightforward way to understand the motivating example is in its corresponding graphical representationf" as shown in Fig. la Following Peters et al. (2015), we treat the environment as random variable E, where E could be any information specific the environment. For simplicity, we let E be the environment index. i.e., E E fl..... N and N is the number of training environ- ments. Note that, here we consider E surrogate variable because itself is not a causal variable [Zhag et al. 2017 Huang et al. 2020). From Fig. [a] it is obvious to see that ICP does not work in this setting, since Y J. E|X1- 1. In fact, a more practical version appearing in real problems present in Fig. [b] where the true variables (X1, X2) are unobserved and we only can observe their transformation O, which is a function of {X X2). In this case, even Y not affected by (i.c., remove the edge E + Y), applying ICP to O still does not work, since each variable (i.e., each dimension) of O is jointly influenced by both X, and X2. By contrast, both IRM and IRMG work when the transformation is linear, but not when it is nonlinear. These analyses also empirically demonstrated in Section 5.1
3.2 THE GENERAL SETTING
Inspired by the two dimensional example (i.e., = (Xx, IR2) described above, we naturally extend it to a more general multi-dimennional setting shown in Fig. Ic Technically, we have
O E O C Rd Y E y C R', X = (Xpr..... ...... p.. X X...... X cs X R=(++k) (lower- dimensional, n(r +) S d), where we assume each X, E R. for simplicity. Is worth emphasising that except X, p... none of (Xp,. (Xp....., Xpr_, has arrow pointed to Y. meaning that all the parents of are absorbed into Xp,. Under this circumstance, more formally, we assume that
Assumption 1. X; X, Y. for any + j-
This assumption is not very restrictive but practically and theoretically reasonable enough to be able to cover various scenarios, due to the following reasons:
'The relation between SEM and its graphical representation formally detined in Appendix B
As an initial work, we do not explicitly consider unobserved confounders in this paper. Precisely, for simplicity. we assume that there are no unobserved confounders between X,, Y O, and E. In fact, in some cases, our approach will not be affected even if there exist unobserved confounders. For example, if there were unobserved confounder between X; and O, would he absorbed into X; and would not affect our approach.
4

---

Under review as a conference paper at ICLR 2021
X
(b)
(c)
(a)
(h)
(k)
(m)
Figure 2: (a) General causal structure over (X,, Y, O, E}. where the arrow from X; to O is must-have connection and the other four might not necessarily. (b) 12 possible causal structures from (a).
Firstly, Assumption alllww us to separately deal with each X, When taking a closer look at each X; in Fig. [cc] it is evident that there exist only five possible connections between X, Y. E, and O, shown in Fig. 2a Among them, only the arrow from X; to O must exist whilst the other four might not necessarily, which leads to 12 possible types structures present in Figs. 2b 2m These structures cover most scenarios in real-world applications, e.g., X; could be either parent or a child Y, either affected or not by E, and even have no connection to Y, etc. It is also worth noting that Assumption [ does not rule out the possibility that there might exist certain correlations between all these latent variables.
Secondly, although it is well known that the key idea behind learning disentangled representations that real-world data generated by few explanatory factors variation, Locatello et al. (2019) show that unsupervised learning of disentangled representations is fundamentally impossible with- out inductive biases, and that well-disentangled models seemingly cannot be identified without su- pervision. They further suggest that disentanglement learning should be explicit about the role of inductive biases and supervision. In this sense, our Assumption reasonably implements this idea.
Thirdly, as described Section 2.1 identifiability of latent variables in iVAE requires condition- which is key
ally factorized prior distribution over the latent variables, as we did in Assumption | condition under which it is guaranteed that latent variables can be identified up to a permutation and pointwise transformation. This condition further inspires us to develop generalization theory in the nonlinear setting that we will formulate in Section
Now, under Assumption I solving Eq. can be reduced to the key question of how to find the subset of X that are the direct causes of Y. from the observational data {O, F,, Y}, which will be discussed in the next section.
4 OUR APPROACH
In this section, we formally introduce our algorithm, namely Invariant Causal Representation Learn- ing (ICRL), which consists of three phases and is summarized in Algorithm The basic idea is that we first identify true latent variables by leveraging iVAE under Assumption 1 (Phase 1), then discover direct causes of Y (Phase 2), and finally learn an invariant predictor based on the identified direct causes (Phase 3).
4.1 PHASE 1: IDENTIFYING TRUE LATENT VARIABLES USING IVAE
Under Assumption I it is straightforward how to identify the true hidden factors X from O with the help of Y and E, by leveraging iVAE. We can directly substitute with (Y, E) in Eqs. and obtain its corresponding generative model:
(5)
(0 (O f(X)).
5

---

Under review as a conference paper at ICLR 2021
Algorithm 1: Invariant Causal Representation Learning
Phase 1: We first learn the iVAE model, including the generative model and its corresponding
inference model, by optimizing the evidence lower bound present in Eq. (8) on the data
{O.Y Y, F). Then, we use the learned iVAE model to infer the corresponding latent variable X from {0, Y, E)). which is guaranteed to be identified up to a permutation and pointwise transformation.
Phase 2: Once obtaining X, according to Theorem 1 discover from them which are the direct causes Pa(Y) of Y by only performing Rule 1.4, Rule 1.8, Rule 2.1, and Rule 3.1 described in Appendix
E.3
Phase 3:
Once obtaining Pa Y), we can separately optimize Eq. and Eq. 110 to learn the invariant data representation & and the invariant classifier w.
Likewise, we also obtain its corresponding prior distribution and lower bound:
x)) PT, 1. b) - II, (x E) exp
(X, (Y, E)],
(7)
Lphasss (0, (f,5)===. :DE [ B)) [log po
PO (O, X Y, F:) - log g (X O, Y, F)]]
(8)
This bound can be further expanded for computational convenience, which is given in Appendix C
More importantly, under this setting we can directly borrow the identifiability result from Khe- makhem al. (2020) and then restate it below by replacing U with (Y,
Theorem 1. Assume that we observe data sampled from a generative model defined according
to Eqs. 7 with parameters 0 == (f,T, and Z Assume the following holds: (i) The set (o E (O) = 0} has measure zero, where Se is the characteristic function of the density Pe defined in Eq. O (ii) The mixing function f in Eq. (6) is injective, and has all second order cross derivatives. (iii) The sufficient statistics Tag in Eq. 17) are twice dif- ferentiable, and (Ti.))gsik are linearly independent on any subset of X of measure greater than zero. (iv) There exist nk + 1 distinct points EY) "..0........"...... Eynh such that the matrix F.)
I, I (X((Y, E)))- {((Y, A(Y,E)))) ....)..... s(Y. - (())) of size rik nk is invertible. Then the parameters O are identifiable up to a permutation and pointwise transformation.
Theorem deals with the general case k 2 2, whose proof is given in Khemakhem et al. 2020 According to Theorem [ we can further have the following result of consistency of estimation.
Theorem 2. Assume the following holds: (i) The family of distributions 4s (x,X Y. E) contains po ( Y, infinite data, iVAE learns the true parameters O" up to permutation and pointwise transformation.
N,YY, E). (ūi We maximize Eyhaser (0, Φ) with respect to both and d. Then in the limit of
An immediate result of Theorem Tand Theorem 2 is as follows,
Theorem 3. Assume the hypotheses of Theorem 7 and Theorem 2 hold, then in the limit of infinite data, iVAE learns the true latent variables X* up to permutation and pointwise transformation.
The proofs of Theorem 2 and Theorem 3 are given in Appendix D Theorem 3 says that we can
leverage iVAE to learn the true conditionally factorized latent variables up to permutation and pointwise transformation, which achieves our goal stated in Assumption I
4.2 PHASE 2: DISCOVERING DIRECT CAUSES
After identifying all the conditionally factorized latent variables X from O, the question that comes
is how to determine which component of X := X .p...... X p... X....... Xx,) is the direct cause of Y. denoted by Pa( (Y)) .. As discussed in Section 3.2 there are totally 12 possible types of structures
over [X, Y, E, O}, as shown in Figs. 2b 2m Apparently, only in the four of them (i.e., Figs. 2c
2f 2i and 2Th does X; serve as a parent of Y Note that, the reason that we also include Figs. 2
and 2 is that we follow Arjovsky et al. (219)) and consider a more general definition of invariance which allows for changes in the noise variance of Y. For convenience, we also put the definition
in Appendix B Surprisingly, given the data (X,, Y, E, O}. we are able to distinguish all the 12
We also provide theorem dealing with the special case l in Appendix D

---

Under review as a conference paper at ICLR 2021
structures by performing conditional independence tests (Spirtes al., 2000 and by leveraging causal discovery algorithms (Jazzing et al. 2013] Peters et al. 2017 Zhang et al,
Zhang al. 2012)
2017 Huang et al. 2020). This is summarized in Proposition and the proof given Appendix
Proposition All the 12 structures shown in Figs. distinguished in parallel.
2h
2m can be independently and completely
Proposition D allows us to efficiently discover all direct causes of Y from X by independently
performing conditional independence
and causal discovery algorithms for each X; in parallel. Besides. for each X; we only need to check if they are one of the four structures (i.e., Figs. 2c 21 and and this check can be also performed in parallel.
4.3 PHASE LEARNING AN INVARIANT PREDICTOR
Once obtaining the invariant causal representation Pa(Y) for Y across the training environments, learning an invariant predictor w & in Eq. (4) is then reduced to two simpler independent optimiza- tion problems: (i) learning the invariant data representation & from O Pa(Y). and (ii) learning the invariant classifier w from PY(Y) to Y. Mathematically, these two optimization problems can be respectively phrased
min E ( min EE
Σ
cEEr Eo (or (a())((( Pa(Y
(9)
DEHo
Σeeen. (")w)= weH
min
min
E.ce.. E (M(u(Pa Y*)
(10)
wEH
Eq. 9 and Eq. 10 guarantee that ICRL can achieve low error across Etr. Also, in Phase 1&2, showed that ICRL enforce invariance across &t,. Now this brings to the question: how enable the OOD generalization? In other words, how does ICRL achieve low error across Exu? As Arjovsky al. (2019) pointed out, low error across Eer and invariance across Eall leads low error across Eoit. because the generalization егrоr of w a 4 respects standard error bounds once the data representation & eliciting an invariant predictor w & across Exll is estimated. Thus, enabling the OOD generalization finally comes to the question: under which conditions does invariance across Err imply invariance across Eatt? Not surprisingly, &tr must contain sufficient diversity to satisfy an underlying invariance across Eatl. Fortunately, the hypotheses of Theorem automatically provides such a guarantee, and we therefore have the following result whose proof Appendix E
Proposition 2. Assume the hypotheses of Theorem and Theorem 2 hold, then in the limit of infinite data, ICRL learns an invariant predictor w & across Eall-
Proposition tells us that under the assumptions given in Theorem T and Theorem if ICRL can learn an invariant predictor w &b across Eer in the limit of infinite data, then such predictor w & is invariant across Eall-
5 EXPERIMENTS
We compare our approach with a variety of methods on both synthetic and real-world datasets. Due space limit, put in the supplement detailed description of the datasets (Appendix (F) and model architectures (Appendix [H). as well as some in-depth analysis on experimental results (Appendix G In all the comparisons, unless stated otherwise, we averaged the performance of the different methods over ten runs.
5.1 SYNTHETIC DATA
As a first experiment, in order to interpret how our approach works, we conduct a series of exper- iments synthetic data generated according to extension of the SEM in Model [T This more practical extension is done by increasing the dimensionality of the two true features X : (X1, X2) to 10 dimensions through a linear nonlinear transformation, cally, the goal is to predict Y from O, where 0 = g(X) and g(-) is called X Transformer. We consider three types of transformations: (a) Identity: o(-) is the identity matrix I E RR2x2 i.e.,
illustrated in Fig. Ib Techni-

---

Under review as a conference paper at ICLR 2021
Table Regression on synthetic data: Comparison of methods in terms MSE (mean + std deviation).
X TRANSFORMER
ALGORITHM
(o) = (0.2,
TRAIN MSE
TEST MSE (e) 5)
MEE (0)
TEST MSE
(oy =
ERM
0,00 0.00
0.00 0.00
0.00 1: 0.00
0.00 0.00
Identity
IRM
0.00 1: 0.00
0.00 2: 0.00
0.00 4: 0.00
0.00 + 0.00
F-RM GAME
0.81 0.37
0.84 1 0.33
1.18 + 0.32
9.58 + 17.
V-IRM GAME
+ 0.44
9.66 10.86
241.01 301.52
ICRL
0.45 0.00
ERM
0.00 100+=0000 0.00
0.00 + 0.00
0.00 + 0.00
0.00 + 0.00
Linear
1.00
+
1.01 0.00
0.00
V-IRM GAME
0.89 0.22
1.89 1.16
14.72 1. 27.57
$80.88 739.91 94
ICRL
0.01 0.03
0.03 1 0.00
0.49 1:9.1000 0.02
1,02 1: 0.03
FRM
6.56 6.01
220.70 220.97
Nonlincar
IRM
F-IRM GAME
64,48 # 12.18
0.08 0.01
964.42 # 634.18
0.20 + 0.04
15006 28 +: 10435.44 44
5.06 + 3.06
360114.72
149.60 104. 86
241228
IRM GAME
0.90 0.201
1.74
24.72 16.34
780.67 729.22999
ICRL
0.003
0.80
30.38
0=g(X)= = X. (b) Linear: g(-) is a random matrix S E 22 i...
x (c)
Nonlinear: 9(-) is implemented by a multilayer perceptron with the 2 dimensional input and the 10 dimensional output, whose parameters are randomly initialized in advance. For simplicity, here we consider the regression task, in which the mean squared error (MSE) is used as metric.
We consider an extremely simple scenario in which we fix σ = 1 and Ø₂ = for all environments and only allow 03 to vary across environments. In this case, d3 controls how much the representation depends on the variable X2, which is responsible for the spurious correlations. Each experiment
draws 1000 samples from each of the five environments d3 {0.2, 2, 5, 20, 100). where the first two are for training and the rest for testing. We compare with several closely related baselines ERM, and two variants of IRMG: F-IRM Game (with & fixed to the identity) and V-IRM Game (with a variable 4).
As shown Table [ in the cases of Identity and Linear, approach is better than IRMG but only comparable with ERM and IRM. This might be because the identifiability result up to a pointwise nonlinear transformation renders the problem more difficult than itself, that is, converting the orig- inal identity or linear problem to a nonlinear problem. In the Nonlinear case, it is clear that the advantage of our approach becomes more obvious as the spurious correlation becomes stronger.
We also perform series of experiments to further analyze our approach, including the analysis on the importance of Assumption [add on the necessity of iVAE in Phase 1, how accurately the direct causes can be recovered in Phase 2, and how well the two optimization problems can be addressed in Phase 3, all of which can be found in Appendix G
5.2 COLORED MNIST AND COLORED FASHION MNIST
In this section, we conduct experiments on two datasets used in IRM and IRMG: Colored MNIST and Colored Fashion MNIST. We follow the same setting of Afuua et al. (020) create these two datasets. The task is to predict a binary label assigned to each image which is originally grayscale but artificially colored in a way that correlated strongly but spuriously with the class label. We com- pare with 1) IRM, 2) two variants of IRMG: F-IRM Game (with & fixed to the identity) and V-IRM Game (with a variable $) 3) three variants of ERM: ERM (on entire training data), ERM e (on each environment e). and ERM GRAYSCALE (on data with spurious correlations), 4) ROBUST MIN MAX (minimizing the maximum loss across the multiple environments). Table 2 shows that our ap- proach significantly outperforms all the others Colored Fashion MNIST. It is worth emphasising that both train and test accuracies of our method closely approach the ones of ERM GRAYSCALE and OPTIMAL, implying that it does approximately learn the true invariant causal representation with nearly no correlation with the color. We can achieve a similar conclusion from the results on Colored MNIST shown in Table 3 However, this dataset seemingly more difficult than the fashion version, because ERM GRAYSCALE tion with the color removed, implying that it might involve some other spurious correlations. In this case, two training environments might be not enough to eliminate all the spurious correlations.
still underperforms even though the spurious correla-
'W also tried ICP, but surprisingly ICP cannot find any parent of Y even in the identity case.
8

---

Under review as a conference paper at ICLR 2021
Table 2: Colored Fashion MNIST: Comparison Table 3: Colored MNIST: Comparison of methods
methods in terms of accuracy (mean + std deviation). in terms of accuracy (mean + std deviation).
ALGORITIM TRAIN ACCURACY TEST ACCURACY
ALGORITIM
TRAIN ACCURACY TEST ACCURACY
ERM
83.17 # 1.01
22.46 0.08
ERM
$6.88 + 0.16
10:45 + 0.06
$1.33 1.35
33.34 8.85
$4.84+
10.86
ERM
$4,30 1.89
13.16 0.82
ERM 2
$6.6 0.20
10:05 + 0.23
ROBUST MIN MAX
81.81
20,22 8.56
ROBUST MIN
$6.25
15.24
F-IRM GAME
02.31 2.335
09:25 5.82
F-IMM GAME
03.37
09.91 2.00
V-IM GAME
08.96
70.
V.IRM GAME
63.87 1.03
48.06
IRM
0.225
06.25 12.42
IRM
50.21 4.33
02.75 9.50
ICRL
71.32 0.43
T3.14 0.56
ICRL
70.34 0.29
06.21 1.42
ERM GRAYSCALE
0.37
74.07 + 0.48
ERM GRAYSCALE
71.81 + 0.47
71.36 # 0.65
OPTIMAL
75
75
OPTIMAL
75
75
6 RELATED WORK
Invariant Prediction Peters et al. (2015) originally introduced the theory of Invariant Causal Pre- diction (ICP), aiming to find the causal feature set (i.e., all direct causes of target variable of interest.) by exploiting the invariance property in causality which has been widely discussed under
the term "autonomy", "modularity", and "stability" (Haavelmo) 1944| Aldrich 1989; Hoover 1990, Pearl, 2009, Dawid et al., 2010: Schölkopf et al., 2012). This invariance property assumed in ICP and its nonlinear extension (Heinze- Deml et al. (2018) is limited, because intervention is allowed on the target variable Y. Besides, ICP methods implicitly assume that variables of interest X are given. Magliacane et al. (2018) and Subbaswamy al. 2019) attempt to find invariant predictors that maximally predictive using conditional independence tests and other graph-theoretic tools, both of which also assume that X are given and further assume that additional information about the structure over X is known. Arjovsky et al. (2019) reformulate this invariance as an optimization- based problem, allowing us learn the invariant data representation from O that required to be linear transformations of X. Ahhja et al. (2020) extend IRM to the nonlinear setting from the perspective of game theory, but their nonlinear theory holds only in training environments.
Domain Generalization Doman generalization emphasizes the ability to transfer acquired knowl- edge to domains unseen during training. A wide range of methods has been proposed for learning Nomain-invarian representations. Khosla et al. (2012) develop max- max-margin classifier that explic- itly exploits the effect of dataset bias and improve generalization ability unseen domains. Fang et al. (2013) propose metric learning approach based on structural SVM such that the neighbors of each training sample consist of examples from both the same and different domains. Muandet et al. (2013) propose a kernel- based optimization algorithm called InmaiinInvariant Component Analysis (DICA), which aims to both minimize the discrepancy among domains and prevent the loss of relationship between input and output features. Ghhifary et al. (2015) train multi-task au- toencoder that recognizes invariances among domains by learning to reconstruct analogs of original inputs from different domains. Motiian et al. (2017) learn embedding subspace where samples from different domains are close if they have the same class labels, and far away if they bear differ- ent class labels. [i et al.) (2018b) minimizes the differences in joint distributions to achieve target domain generalization through the application of a conditional invariant adversarial network. Li al. (2018a) build on the adversarial autoencoders by considering maximum mean discrepancy regularization and aligning the domains' distributions.
CONCLUSION
We developed novel framework to learn invariant predictors from diverse set of training envi- ronments. This framework is based on a practical and general assumption: the data representation can be factorized when conditioning on the outcome and the environment. The assumption leads to a guarantee that the components in the representation can be identified up to permutation and pointwise transformation. This allows us to further discover all the direct causes of the outcome, which enables generalization guarantees in the nonlinear setting. We hope our framework would inspire new ways to address the OOD generalization problem through the causal lens.

---

Under review as a conference paper at ICLR 2021
REFERENCES
Kartik Ahuja, Karthikeyan Shanmugam, Kush Varshney, and Amit Dhurandhar. Invariant risk min- imization games. arXiv preprint arXiv:2002.04692, 2020.
John Aldrich. Autonomy. Oxford Economic Papers. 41(1):15-344 1989.
Martin Arjovsky, Léon Bottou, Ishaan Guirajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint (0XXv::907.0289)) 2019.
Sara Beery, Grant Van Horn, and Pietro Perona. Recognition in terra incognita. In Proceedings of the European Conference on Computer Vision (ECCV), pp. 456-473, 2018.
Philip Dawid, Vanessa Didelez, al. Identifying the consequences of dynamic treatment strate- gies: A decision-theoreti overview. Statistics Surveys, 4:184-231, 2010....0000..
Chen Fang, Ye Xu, and Daniel N. Rockmore. Unbiased metric learning: On the utilization of multiple datasets and web images for softening bias. In ICCV, pp. 1657-1664. IEEE Computer Society, 2013.
Muhammad Ghifary, W. Bastiaan Kleijn, Mengjie Zhang, and David Balduzzi. Domain gener- alization for object recognition with multi-task autoencoders. In ICCV, pp. 2551-2559. IEEE Computer Society, 2015.
Ishaan Guirajani and David Lopez-Paz. In search of lost domain generalization. arXiv preprint arXiv:2007.01434, 2020.
Trygve Haavelmo. The probability approach in econometrics. Econometrica: Journal of the Econo- metric Society. pp. 11i-115, 1944.
Christina Heinze-Demie Jonas Peters, and Nicolai Meinshausen. Invariant causal prediction for nonlinear models. Journal of Causal Inference, 6(2), 2018.
Kevin D Hoover. The logic of causal inference: Econometrics and the conditional analysis of cau- sation. Economics Philosophy, 6(2):207-233 1990.
Patrik O Hoyer, Dominik Janzing, Joris Mooij, Jonas Peters, and Bernhard Schölkopf. Nonlin- causal discovery with additive noise models. In Advances in neural information processing systems, pp. 689-696,14 2009.
Biwei Huang, Kun Zhang, Jiji Zhang, Joseph Ramsey, Ruben Sanchez-Romero, Clark Glymour, and Bernhard Schölkopf. Causal discovery from helerngemeous/nonstttionnayy data. Journal of Machine Learning Research, 21(89):1-53, 2020.
Aapo Hyvarinen, Hiroaki Sasaki, and Richard Tumnr. Nonlinear ica using auxiliary variables and generalized contrastive learning. In The 22nd International Conference Artificial Intelligence and Statistics, pp. 859-868, 2019.
Dominik Janzing, David Balduzzi, Moritz Grosse-Wentrup, Bernhard Schölkopf, et al. Quantifying causal influences. The Annals of Statisties, 41(5):2324-235 2358, 2013.
Ilyes Khemakhem, Diederik Kingma, Ricardo Monti, and Aapo Hyvarinen. Variational autoen- coders and nonlinear ica: A unifying framework. In International Conference on Artificial Intel- ligence and Staristics, pp. 2207-2217, 2020.
Aditya Khosla, Tinghui Zhou, Tomasz Malisiewicz, Alexei Efros, and Antonio Torralba. Undoing the damage of dataset bias. In ECCV(1), volume 7572 of Lecture Notes in Computer Science, pp. 158-171. Springer, 2012.
Diederik Kingma and Max Welling. arXiv:1312.6114 2013.
Auto-encoding variational bayes.
arXiv preprint
Haoliang Li, Sinno Jialin Pan, Shiqi Wang, and Alex C Kot. Domain generalization with adver- sarial feature learning. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 5400-5409, 2018a.
10

---

Under review as a conference paper at ICLR 2021
Ya Li, Xinmei Tian, Mingming Gong, Yajing Liu, Tongliang Liu, Kun Zhang, and Dacheng Tao. Deep domain generalization via conditional invariant adversarial networks. In Proceedings of the European Conference Computer Vision (ECCV), pp. 624 669,, 2018b.
Francesco Locatello, Stefan Bauer, Mario Lucie, Gunnar Raetsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of disentangled representations. In international conference on machine learning, pp. 4114-4124, 2019.
Sara Magliacane, Thijs van Ommen, Tom Claassen, Stephan Bongers, Philip Versteeg, and Joris M Mooij. Domain adaptation by using causal inference to predict invariant conditional distributions. In Advances in Neural Information Processing Systems, pp. 10846-10856, 2018.
Saeid Motiian, Marco Piccirilli, Donald A. Adjeroh, and Gianfranco Doretto. Unified deep super- vised domain adaptation and generalization. In ICCV, pp. 5716-5726. IEEE Computer Society, 2017.
Krikamol Muandet, David Balduzzi, and Bernhard Schölkopf. Domain generalization via invariant feature representation. In ICML (I), volume 28 of JMLR Workshop and Conference Proceedings, pp. 10-18. JMLR.org. 2013.
Judea Pearl. Causality. Cambridge university press, 2009.
Jonas Peters, Peter Bühlmann, and Nicolai Meinshausen. Causal inference using invariant predic- tion: identification and confidence intervals. arXiv preprint arXiv: 1501:0500..1333 2015.
Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Elements of causal inference. The MIT Press. 2017.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference deep generative models. arXiv preprint arXiv:1401, 1401.408, 2014.
Bernhard Schölkopf. Dominik Janzing, Jonas Peters, Eleni Sgouritsa, Kun Zhang, and Joris Mooij. On causal and anticausal learning. arXiv preprint wrXiv:120..6471, 2012.
Peter Spirtes, Clark N Glymour, Richard Scheines, and David Heckerman. Causation, prediction, and search. MIT press, 2000.
Bharath Sriperumbudur, Kenji Fukumizu, Arthur Gretton, Aapo Hyvärinen, and Revant Kumar. Density estimation in infinite dimensional exponential families. The Journal of Machine Learning Research, (8(1):1830-14 2017.
Adarsh Subbaswamy, Bryant Chen, and Suchi Saria. A universal hierarchy of shift-stable distri- butions and the tradeoff between stability and performance. arXiv preprint M.XXv:1905.111. 2019.
Vladimir Vapnik. Principles of risk minimization for learning theory. In Advances in neural infor- mation processing systems, pp. 831-838, 1992.
Hemanth Venkateswara, Jose Eusebio, Shayok Chakraborty, and Sethuraman Panchanathan. Deep hashing network for unsupervised domain adaptation. In Proceedings of the IEEE Conference on Computer Vision and Puttern Recognition, pp. 5018 8-5027,, 2017.
Sewall Wright. Correlation and causation. J. agric. Res., 20:557-580, 1921.
Kun Zhang, Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Kernel-based conditional independence test and application in causal discovery. arXiv preprimt arXiv: 1202.3355 2012.
Kun Zhang, Biwei Huang, Jiji Zhang, Clark Glymour, and Berhaar Schölkopf. Causal discovery data: Skeleton estimation and orientation determination. In from IJCAL: Proceedings of the Conference, volume 2017, pp. 1347. NIH Public Access, 2017.
11

---

Under review as a conference paper at ICLR 2021
A VARIATIONAL AUTOENCODERS
We briefly describe the identifiable variational autoencoders (iVAEs) proposed by Khemakhem et al. (2200) As we know, the framework of variational autoencoders (VAEs) (Kingma & Welling 2013] Rezende et al. 2014) allows us to efficiently learn deep latent-variable models and their correspond- ing inference models. Consider a simple latent variable model where O R stands for observed variable (random vector) and X E R" for a latent variable. The VAE model actually learns a full generative model (o,
pe [ pe (X and an inference model g O) that approxi- mates its posterior pe [ O),, where 0 is vector of parameters of the generative model, Φ a vector of parameters of the inference model, and po X) is a prior distribution over the latent variables. Instead of maximizing the data log-likelihood, we maximize its lower bound LYvE (0, Φ):
log pe (O) LVAE Φ) == (x [log pe ( (Ol KL (4s (X (X)).
where we have used Jensen's inequality, and KL(-II-) denotes the Kullback-Leibler divergence be- tween two distributions.
B DEFINITIONS
For convenience, we restates some definitions here and please refer to the original papers (Arjovsky et al. 2019 Peters et al. 2017) for more details.
Definition 2. A structural equation model (SEM) == (S,N) governing the random vector X (X1, Xa) is a of structural equations:
S (PP((.)),,N),
where Pa(X,) S (X,...., [X,} are called the parents of X,, and the N, are independent noise random variables. We say that "X, causes Xj" if X; Pa(Xx)) We call causal graph of X to the graph obtained by drawing i) one node for each Xi and ii) one edge from X₃ to Xj if X; E Pa(X;). We assume acyclic causal graphs.
Definition 3. Consider SEM C == (S, An intervention C consists of replacing one or with structural
several of its structural equations to obtain an intervened SEM C = equations:
: + (X)))
The variable X is intervened S; # S or N No.
Definition 4. Consider a structural equation model (SEM) S governing the random vector and the learning goal of predicting Y from X. Then, the of all environ- ments (a)(S) indexes all the interventional distributions P(Xe, Y*) obtainable by valid interven- tions An intervention E (S)(S) (8) is valid as long (i) the causal graph remains acyclic, (ii) E EX{((Pa(Y)) = E Pa(Y) (Y)), and (ii)) V [Y (Pa()) remains within a affitte range.
X,,Y, (X1....,
C DERIVATION
In Phase 1, the lower bound is defined by
Cphiee (0. (. :-Ep,
X, E) 20 E)]] [0 + logpt Y E) log q X O, Y. E)]] F)]
Ro (X),,Y
=Epp
[R (X
= Epp
[Rq (X .E)
( + E (
(x [log gø ( (0,Y,,
The first term is the data log-likelihood and the third term has closed-form solution,
E (x log go (X Y, E)] = log(2 + Σ (1 + log of)
where of is simply denote the j-th element of the variational s.d. (α) evaluated at datapoint that is simply function of (O, Y, F)) and the variational parameters Φ.
12

---

Under review as a conference paper at ICLR 2021
Now let us look at the second term,
Po _ log
(X
(x) II ZY E)
=B (x
log
Σ (x)( E))
j=1
(X.)
=B (X1x(X(OYY
Σ log
2Y E)
exp
Σ Ti (X.),
xB (X
(X)),,
j=1
ΣΣΣΣ
Toe (X Az. (Y, EY)
1=1
where we let the base measure (x(x;) 1 and L is the sample size.
D THEOREMS
Theorem 4. Assume that observe data sampled from generative model defined according to Eqs. {O
7 with parameters = (f.T] X) and k = 1. Assume the following holds: (i) The set (P())) -0} has measure zero, where fe the characteristic function of the density Pe defined in Eq.
6 (ii) The mixing function f in Eq. (6) is injective, and all partial derivatives of are continuous. (iti) The sufficient statistics Tı.j in Eq. (7) are differentiable almost everywhere and not monotonic, and (Tij)isiss are linearly independent on any subset of X of measure greater than zero. (iv) There exist nk + distinct points (Y, F.) :.... (Y, Fyyik such that the matrix I. = A((Y A((Y,E)*K) ( A((Y,E)))) of size nk x nk is invertible. Then
(A((Y, k)²) ((YYE))) the parameters = (f. T. X) are identifiable up to a permutation and pointwise transformation.
E PROOFS
E. PROOF OF THEOREM
This proof is similar that of Theorem 4 in Khemakhem al. (2020)
Proof. The loss function in Eq. can be rephrased as follows:
Cphase (0, Φ) = P
KL (go
Y, E)).
If the family of 96 x | E) is flexible enough contain pe Pe X O, O,Y,E), then by optimizing the loss over its parameter Φ we will minimize the KI. term which will eventually reach zero, and the loss will be equal the log-likelihood. Under this circumstance, the iVAE inherits all the properties of maximum likelihood estimation. In this particular case, since our identifiability is guaranteed up to permutation and pointwise transformation, the consistency of MLE means that we converge to the true parameter 0* up to a permutation and pointwise transformation in the limit of infinite data. Because true identifiability is one of the assumptions for MLE consistency, replacing it by identifiability up to permutation and pointwise transformation does not change the proof but only the conclusion.
E.2 PROOF OF THEOREM 3
Proof. Theorem I and Theorem 2 guarantee that in the limit of infinite data. iVAE can learn the true parameters 0* (f". T" X²) up permutation and pointwise transformation. Let (f. T, Å) be the parameters obtained by iVAE, and we therefore have (f. F. Å) ~p (f*,T" X*), where <pp denotes the equivalence up to permutation and pointwise transformation. If there were
13

---

Under review as a conference paper at ICLR 2021
noise, this would mean that the learned f transforms O into X f (f*) (O) up permutation and sighed scaling, If with noise, obtain the posteriors of the latents up to an analogous indeterminacy.
f- (O) that are equal to X
E.3 PROOF PROPOSITION
Proof. The following rules can be independently performed to distinguish all the 12 possible struc- tures shown in Figs. 26 2m For clarity, we divide them into three groups.
Group All the eight structures in this group can be discovered only by performing conditional independence tests.
Rule 1.1 X; IL Y. X; 11 E, and E IL Y, then Fig. 2b is discovered.
Rule 1.2 If X; Y.X; IL E, and E V/_ Y, then Fig. 2h discovered.
Rule 1.3 If Y, Y,X; # E, and E Y, then Fig. 2e is discovered.
Rule 1.4 If PL Y, Y, 1 E, and E V Y, then Fig. discovered.
Rule 1.5 If X, PL Y, X; 1 E. and E 11 Y. then Fig. discovered.
Rule 1.6 If X; H Y, X; # E. E 1 Y. and X; Y E, then Fig. 2k discovered.
Rule 1.7 If X; J. Y,X, 4 E, / Y, and X, IL Y then Fig. 2j is discovered.
Rule 1.8 If X; J. Y, X; y F. E Y, and Y 11 E X, then Fig. 2f is discovered.
Group X, 1 Y. X; 1L E, and E TT Y, then we can discover both Fig. 2c and Fig.
2d
These two structures cannot be further distinguished only by conditional independence tests, because they come from the same Markov equivalence class. Fortunately, we can further distinguish them by running binary causal discovery algorithms (Peters et al., 2017). e.g. ANM (Hoyer al.] 2009) the bivariate fit model that is based on a best-fit criterion relying a Gaussian Process regressor.
Rule 2.1 If X; 1 Y. X; 11 E, and E' TL Y. and chosen binary causal discovery algorithm prefers X; + Y to X; + Y, then Fig. 2c is discovered.
Rule 2.2 If X, LL Y. X, IL E, and E 1l Y. and chosen binary causal discovery algorithm prefers X; < to X; - Y, then Fig. 2d discovered.
U.YYY,,YYYEY YY,, 1 E|Y, and Y 1 E then we can discover both Fig. 21 and Fig. 2m These two structures cannot be further distinguished only by conditional independence tests, because they come from the same Markov equivalence class. They also cannot be distinguished by any binary causal discovery algorithm, since both X; and Y are affected by E. Fortunately, Zhang et al. (2017) provided a heuristic solution to this issue based on the invariance of causal mechanisms, i.e., P(cause) and P(effect]cause change independently. The detailed description of their method is given in Section 4.2 of Zhang et al.) (2017). For convenience, here we directly borrow their final result. Zhang et al. (2017) states that determining the causal direction between X; and Y in Fig. 21 anddd Fig. 2m is finally reduced to calculating the following term:
Group 3 X
P(Y X.) (P X,))
Ax.-Y
log
(11)
where (.) denotes the sample average, P(Y Xi) X() is the empirical estimate of PXY on all data points, and (P(Y) denotes the sample average of P(Y|X,,, which is the estimate in each environment. We take the direction for which A is smaller to be the causal direction.
14

---

Under review as a conference paper at ICLR 2021
Rule 3.1 If X; + Y. X; + E, + Y. X; Y X, H EY Y H E|X,, and
Ax,-Y is smaller than Ay -x,, then Fig. 2 is discovered.
Rule 3.2 X; Y. x H E, E 1
Y.
X y EI X; M E Y J E|X,, and
- is smaller than A - then Fig. 2m discovered.
E.4 PROOF OF PROPOSITION
Proof. Firstly, assumption (iii) and assumption (iv) in Theorem | are the requirements of the set of training environments containing sufficient diversity and satisfying an underlying invariance which holds across all the environments. Interestingly, assumption (iii) elicits Lemma of Khemakhem al. al..), which is closely similar to the linear general position in Assumption 8 offArjovsky et al. (2019). Thus, Lemma can be similarly called the nonlinear general position in our generalization theory. whose proof can be found in Arjovsky environments lic in this nonlinear general position and the other hypotheses of Theorem 1 2 it is guaranteed in Theorem 3 that all the latent factors X can be identified up to permutation and pointwise transformation. Since this identifiability result holds under the assumptions guaranteeing that training environments contain sufficient diversity and satisfy underlying invariance which holds across all the environments, it also holds across all the environments. Thirdly, Proposition I suggests that all the direct causes PY(Y) of Y can be fully discovered, which also holds across all the environments due to the same reason above. Finally, the challenging bi-leveled optimization problem in both IRM and IRMG now can be reduced two simpler independent optimization problems: (i) learning the invariant data representation 4 from O to Pa( Y). and (ii) learning the invariant classifier w from Pa(Y) to Y, as described in Eq. 99) and Eq. (10). For both (i) and (ii), since there exist no spurious correlations between O and Pa(Y) and between Pa(Y) and Y,
al. (2019). Secondly, when the set of training
hold,
learning theory guarantees that in the limit of infinite data, we will converge to the true invariant data representation & and the true invariant classifier
It is worth noting that although assumption (iii) and assumption (iv) in Theorem require compli- cated conditions to satisfy the diversity across training environments for generalization guarantees, it is not the case in practice. As we will observe in our experiments, is often the case that two environments are sufficient to recover invariances.
F DATASETS
For convenience and completeness, we provide the descriptions of Colored MNIST Digits, Colored Fashion MNIST, and Office-Home here. Please refer to the original papers (Arjovsky et al., 2019 Ahuja et al.
2020 Gulrajani & Lopez- Paz 2020: Venkateswara et al. 2017) for more details.
F.1
SYNTHETIC DATA
For the nonlinear transformation, we use the MLPP
Input layer: Input batch (batch size, input dimension)
Layer 1: Fully connected layer, output size = 6, activation = ReLU
Output layer: Fully connected layer, output size = 10
F.2 COLORED MNIST DIGITS
We use the exact same environment as in Arjovsky et al. (2019). Arjovsky et al. (0199) propose to create environment for training to classify digits in MNIST digits data") where the images in MNIST are now colored in such a way that the colors spuriously correlate with the labels. The
www.
load_data
15

---

Under review as a conference paper at ICLR 2021
task is to classify whether the digit is less than 5 (not including 5) environments (two training containing 30,000 points each, one test containing 10,000 points) We 0 if the digit is between
than 5. There are three
add noise to the preliminary label (ū 0 if digit is between 0-4 and 5-9) by flipping with 25 percent probability to construct the final labels. We sample the color id 2 by flipping the final labels with probability Pe. where Pe is 0.2 in the first environment, 0.1 in the second environment, and 0.9 in the third environment. The third environment is the testing green if
environment. We color the digit red if
F.3 COLORED FASHION MNIST
modify the fashion MNIST datase manner similar to the MNIST digits dataset. Fashion MNIST data has images from different categories: "t-shirt", "trouser", "pullover", "dress", "coat", "sandal", "shirt", "sneaker", "bag". "ankle boots". We add colors to the images in such a way that the colors correlate with the labels. The task is classify whether the image is that of foot wear or clothing item. There are three environments (two training. one test) We add noise to the preliminary label (ğ O: "t-shirt", "trouser", "pullover", "dress", "coat", "shirt" and ý 1: "sandal", "sneaker", "ankle boots") by flipping with 25 percent probability to construct the final label. We sample the color id by flipping the noisy label with probability Pe, where pe 0.2 in the first environment, 0.1 in the second environment, and 0.9 in the third environment, which is the test environment. We color the object red if
green if
raw data
data
estimated
estimated
-10
Figure 3: Left: Comparison of the raw data X and the mean of X inferred through the learned inference model in the Nonlinear case. Right: Comparison of the raw data X and the sampled points X using the reparameterization trick in the Nonlinear case. The comparisons clearly show that the inferred X is equal to X up to a permutation and pointwise transformation.
G IN-DEPTH ANALYSIS ON SYNTHETIC DATA
G.1 VERIFYING PHASE
Theorem 3 tells us that we can leverage iVAE to learn the true conditionally factorized latent vari- ables up permutation and pointwise transformation. We empirically verify this point by com- paring the raw data X with the corresponding X inferred through the learned inference model in iVAE. Fig. 3 clearly shows that the inferred X is equal X up to a permutation and pointwise transformation.
mn st. 1add//oadddaaa
16

---

Under review as a conference paper at ICLR 2021
Train MSE
Test MSE
20000
40000
60000 Step
80000 100000
0
20000
40000
60000
80000
100000
Step
(a) Regression results in the Linear case terms of MSE, where the inferred X₂ is the identified cause.
2.5
Train MSE
1.00
MSE
2.0
0.75
8
1.5
0.50
1.0
20000
40000
60000
80000
100000
0
20000
40000
60000
80000 100000
Step
Step
(b) Regression results the Nonlinear case terms of MSE, where the inferred X₂ is the identified cause.
100
100
Train Accuracy
Test Accuracy
X₂
80
0
20000
40000
60000 Step
80000
100000
0
20000
40000 60000 80000 100000
Step
(c) Classification results on Synthetic Data in terms accuracy, where inferred X, the identified cause.
Figure 4: Comparison of the inferred X, and X2 in terms of their final performance.
To show the importance iVAE, we conduct an experiment in which we replace iVAE with the original VAE in Phase 1. As shown in Table the performance of ICRL based iVAE significantly outperforms the one based on VAE. usually occurs in Phase 2 that either all the dimensions or no dimension of X are identified as the parents of Y. This is because all components of are mixed together and will influence one another even when conditioning Y and E.
worth noting that when VAE is instead used in Phase 1,
G.2 VERIFYING PHASE
To show how well our method can identify the direct causes of Y in Phase 2, we compare the final performance when the identified direct cause (i.e., X2) and the identified non-cause (i.e., X2) are respectively used in Phase 3 to learn the predictor. Note that, there might exist a permutation between the inferred { Xx, and the true {X1, X2). J Fig. 4a and Fig. show the results on the regression task, from which we can obviously see that the predictor elicited by the identified cause has much better generalization performance. The classification result (i.e., Y is binarized) in Fig. 4c further demonstrates this point.
G.3 VERIFYING PHASE 3
In this experiment, want to verify how well the data representation & be learned by optimiz- ing the loss in Eq. ( The main idea is to check how well the learned & can purely extract the
17

---

Under review as a conference paper at ICLR 2021
-25
75
-5.0
25
25
0.0
00
-10.0
-500
-122
-50
2
-75
-15
-1500
-10.0
-75
-50
-25
00
25
5.
75
-100
-100
-75
-5.0
-25
0.0
25
5.0
75
Figure 5: Left: Energy plot of X1 fx, (o(X)), Right: Energy plot of X2 fxx(e(X)), Note that, here xx is denoted by fx,-
Table Results on synthetic Data: Comparison iVAE and VAE used Phase terms MSE (mean std deviation).
X TRANSFORMER
ALGORITIMM
TRAIN MSE (7g (0.2,2))
TEST MSE (dy = 5)
TEST MSE (0g = 20)
TEST MSE (0y = 100)
ICRL-VAR VAE ICRL-iVAE
0.26 + 0.09
4.26 + 4.00 0.80 + 0.02
42.03 48.67
1174.96 = 1385.81
Nonlinear
0.31 + 0.03
3.02 + 0.12
30.38 + 3.67
cause X1 from O. For this, first learn Φx, for each X,. Formally, have
X: xx Cgx
where X = (X1,X2)
Then, we observe how X, will change while tuning X1 and X2 respectively. Fig. showss the energy plots when X, is the identified cause of Y and X2 the child of Y. Note that, in the plots, 6x, is denoted by fx,- In theory, we are able to learn an invariant data representation 4x, for X, from O, because there is spurious correlation between X1 and O. By contrast, cannot learn an invariant data representation xx, for X₂, from O, because there exist spurious correlations between X, and O. The results shown in Fig. clearly verify our theory. Specifically, in the left plot. X₁ remains approximately unchanged when changing X2 but it changes when changing X1. However, in the right plot, X₂ changes whether we change X1 or X2.
H MODEL ARCHITECTURES
In this section, we describe the architectures different models used different experiments.
H.1 SYNTHETIC DATA
H.1.1 ERM
Linear ERM
Input layer: Input batch (batch size, input dimension)
Output layer: Fully connected layer, output size
18

---

Under review as a conference paper at ICLR 2021
Nonlinear ERM
Input layer: Input batch (batch size, input dimension)
Layer I: Fully connected layer, output size = 6, activation ReLU
Output layer: Fully connected layer, output size
H.1.2
IRM
Linear Data Representation +
Input layer: Input batch (batch size, inpul dimension)
Output layer: Fully connected layer, output size
Nonlinear Data Representation +
Input layer: Input batch (batch size, input dimension)
Layer Fully connected layer, output size = 6, activation ReLU
Output layer: Fully connected layer, output size
H.1.3 F-IRM GAME
Linear Classifier
Input layer: Input batch (batch size, inpul dimension)
Output layer: Fully connected layer, output size
Nonlinear Classifier
Input layer: Input batch (batch size, input dimension)
Layer I: Fully connected layer, output size = 6, activation ReLU
Output layer: Fully connected layer, output size
H.1.4 V-IRM GAME
Linear Data Representation +
Input layer: Input batch (batch size, input dimension)
Output layer: Fully connected layer, output size
Nonlinear Data Representation +
Input layer: Input batch (batch size, input dimension)
Layer 1: Fully connected layer, output size = 6, activation = ReLU
Output layer: Fully connected layer, output size
I Linear Classifier
Input layer: Input batch (batch size, 2)
Output layer: Fully connected layer, output size
Nonlinear Classitier
Input layer: Input batch (batch size, 2)
Layer Fully connected layer, output size = 6, activation ReL.U
Output layer: Fully connected layer, output size
19

---

Under review as a conference paper at ICLR 2021
H.1.5 ICRL
iVAE Linear Prior
Input layer: Input batch (batch size, input dimension)
Mean Output layer: O, which a vector full of O with the length
Log Variance Output layer: Fully connected layer, output size
iVAE Nonlinear Prior
Input layer: Input batch (batch size, input dimension)
Layer 1: Fully connected layer, output size = 6, activation = ReLU
Mean Output layer: O, which a vector full of 0 with the length
Log Variance Output layer: Fully connected layer, output size 2
iVAE Linear Encoder
Input layer: Input batch (batch size, input dimension)
Mean Output layer: Fully connected layer, output size
Log Variance Output layer: Fully connected layer, output size
iVAE Nonlinear Encoder
Input layer: Input batch (batch size, input dimension)
Layer 1: Fully connected layer, output size 6, activation ReLU
Mean Output layer: Fully connected layer, output size = 2
Log Variance Output layer: Fully connected layer, output size
iVAE Linear Decoder
Input layer: Input batch (batch size, 2)
Mean Output layer: Fully connected layer, output size = output dimension
Variance Output layer: 0.01 x 1, where is a vector full with the length of output dimension
iVAE Nonlinear Decoder
Input layer: Input batch (batch size, 2)
Layer 1: Fully connected layer, output size 6, activation ReLU
Mean Output layer: Fully connected layer, output size output dimension
Variance Output layer: 0.01 x 1, where 1 is a vector full with the length of output dimension
Linear Data Representation 4b
Input layer: Input batch (batch size, input dimension)
Output layer: Fully connected layer, output size
Nonlinear Data Representation +
Input layer: Input batch (batch size, input dimension)
Layer Fully connected layer, output size = 6, activation ReL.U
Output layer: Fully connected layer, output size
20

---

Under review as a conference paper at ICLR 2021
Linear Classifier
Input layer: Input batch (batch size, 1)
Output layer: Fully connected layer, output size =
Nonlinear Classifier w
Input layer: Input batch (batch size, 1)
Layer I: Fully connected layer, output size 6, activation = ReLU
Output layer: Fully connected layer, output size
H.2
COLORED MNIST DIGITS AND COLORED FASHION MNIST
iVAE Prior
Input layer: Input batch (batch size, imput dimension)
Layer 1: Fully connected layer, output size 100, activation ReLU
Mean Output layer: which is vector full of 0 with the length 100
Log Variance Output layer: Fully connected layer, output size 100
iVAE O-Encoder
Input layer: Input batch (batch size, 2, 28, 28)
Layer 1: Convolutional layer, output channels 32, kernel size = 3, stride 2, padding
1, activation ReLU
Layer 2: Convolutional layer, output channels 32, kernel size = 3, stride = 2, padding
1, activation ReLU
Layer 3: Convolutional layer, output channels 32, kernel size 3, stride 2, padding
1, activation = ReLU
Output layer: Flatten
iVAE (Y, F))-Enooder
Input layer: Input batch (batch size, input dimension)
Output layer:
Fully connected layer, output size 100, activation ReLU
MAA (0, Y, E)--errrr/eeeooderr
Input layer: Input batch (batch size, inpul dimension)
Layer Fully connected layer, output size 100, activation ReLU
Mean Output layer: Fully connected layer, output size 100
Log Variance Output layer: Fully connected layer, output size 100
iVAE Decoder
Input layer: Input batch (batch size, 100)
Layer Fully connected layer, output size = 32 x 4, activation ReL.U
Layer Reshape to (batch size, 32. 4,
Layer 3: Deconvolutional layer, output channels 32, kernel size 3, stride 2, padding
= I, outpadding = O, activation = ReLU
Layer 4: Deconvolutional layer, output channels 32, kernel size = 3. stride = 2, padding
1, outpadding 1, activation Rel.U
21

---

Under review as a conference paper at ICLR 2021
Layer 5: Deconvolutional layer, output channels = 2, kernel size 3, stride 2. padding
1, outpadding 1
Mean Output layer: activation Sigmoid
Variance Output layer: 0.01 1, where is matrix full of 1 with the size of 2 28 28.
Data Representation Φ
Input layer: Input batch (batch size, 2, 28, 28)
Layer 1: Convolutional layer, output channels 32, kernel size = 3, stride 2, padding
1, activation = ReLU
Layer 2: Convolutional layer, output channels 32, kernel size = 3, stride = 2, padding
1, activation ReLUU
Layer 3: Convolutional layer, output channels 32, kernel size 3, stride 2, padding
1, activation ReLU
Layer Flatten
Mean Output layer: Fully connected layer, output size 100
Log Variance Output layer: Fully connected layer, output size 100
Classifier w
Input layer: Input batch (batch size, 100)
Layer I: Fully connected layer, output size = 100, activation ReLU
Output layer: Fully connected layer, output size 1, activation Sigmoid
22


Reviews:
[{"review_id": "prtT9EC0R9L", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "This paper proposes an invariant causal representation learning paradigm in the nonlinear setting. Based on a conditional factorial assumption, they proved identifiability up to a linear transform. The ICRL objective, in this case, is able to discover all the direct causes of the outcome, and thus enables OOD generalization.\n\nThe novelty of the paper seems to be in the generalization of the IRM framework to the nonlinear case which is interesting to me. \nThe authors combined iVAE and IRM to solve this problem. Overall the paper is clearly written and easy to follow, but some conceptual issues remain.\n\nHere are my issues with the paper:\n- How to verify assumption 1 in a real case? Although the authors argued this assumption is not very restrictive and similar to the assumption in the iVAE paper, I think there is a difference between these two papers. In the iVAE paper, they assumed the latent variables to be conditionally factorial, while here the authors assume the potential causals (unobserved variables) are independent. \n\n- After discovering direct causes, they still need the IRM phase to learn an invariant predictor. IRM itself can identify spurious causes and learn an invariant predictor, so what is the gain of learning the first two phases? What if some spurious causations are wrongly detected by the second phase, will it affect the predictor?\n\n-The synthetic data experiment is not convincing at all. ICRL outperforms ERM and IRM in a very extreme case, where all the algorithms perform terribly, I don't think I can conclude ICRL is a better algorithm among others from this test case. If the authors can visually show the invariant representation of ICRL is more robust, that would be a good illustration. \n\n- In the colored MNIST experiment, I assume the setting is the same as the IRM paper. While they said their IRM can reach 66.9+-2.5 acc, which is 7 percent higher than this paper and even higher than ICRL. So I wonder what causes this gap.", "extra_scores": {}, "extra_text": {}}, {"review_id": "ePlVFLUOPjG", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "This paper proposes a method for learning invariant (nonlinear) data representations and classifiers, using data from multiple domains. A key step in the method is to discover the direct causes of the outcome of interest from a set of latent variables that are recovered from observed variables via identifiable VAE. The problem being tackled is significant, and the general idea is interesting and sensible. The empirical results also look encouraging. However, there appears to be a major flaw in the theoretical setup. In order to apply identifiable VAE, the method needs to assume that any two latent variables are independent conditional on the outcome variable and the domain index (Assumption 1). But what about latent variables that are direct causes of the outcome variable? If both X_1 and X_2 are direct causes of Y, generically they will not be independent conditional on Y and E, will they? In the motivating example, only one latent variable is a direct cause of the outcome, so this issue does not arise, but the ambition, as I understand it, is to handle any number of direct causes. I don't see how Assumption 1 can be justified as a plausible assumption in the general case where more than one latent variable is a direct cause of the outcome.\n\nMoreover, Theorem 4 is less than fully rigorous and is misleading. For example, the proof of Theorem 4 invokes a method for inferring causal directions in Zhang et al. (2017), but as far as I know, that method does not yet have a rigorous theoretical justification. As it is formulated, Theorem 4 sounds like a theoretical identifiability result, and as such is not rigorously established by the proof given in the paper.", "extra_scores": {}, "extra_text": {}}, {"review_id": "AWqPXUSwB9Y", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "The paper proposes Invariant Causal Representation Learning, which seeks to learn representations for downstream tasks that are based on only causally invariant latent variables so the representation is robust to shifts in the test environment.  \n\nA model is assumed where an environment variable is a cause of all the latent variables and a target. The iVAE algorithm is used to learn the latent variable model. Then, a series of conditional independence tests and bivariate causal discovery methods are used to distinguish which latent variables correspond to causes (effects) of the target. Finally representations are learned from the observed variables to the causal latent variables of the target and then from these variables to the target.\n\nThe approach is evaluated using synthetic data and semi-synthetic data based on MNIST.\n\nThe clarity and organization of the paper could be improved. The algorithm should be moved from the appendix to the main text and the procedure should be described more holistically to give the reader an outline before diving into the details of each component section. The experiments section is also very unclear.\n\nThe main issue that remains unclear to me is how the environment variable E is being used explicitly. It doesn’t seem clear to me that you would generally have access to E, but it is required in the rules necessary to determine the direct causes of Y. What are the E variables being used in each of the experiments?\n\nThe novelty seems somewhat limited. It seems the theoretical results can be divided into (a) results about identifiability of the latent variable model and (b) the method for identifying the direct causes. Some of (a) follows directly from Khemakhem et al. (2020) - it is difficult to determine whether there is sufficient novelty here. (b) follows directly from well known constraint-based and bivariate causal discovery approaches. \n\nThe experiments are unconvincing. The proposed method outperforms existing approaches in a high noise synthetic data example and a kind of adversarial example where grayscale MNIST is colored in a way that is strongly correlated with the class label (the experimental setup and evaluation metric is very confusing here). It would be more convincing to the proposed method evaluated in a more realistic setting.\n\nFurther, there are no (synthetic) experiments which confirm that the proposed method does in fact learn the causally relevant latent variables and its robustness in doing so. Since identifying the correct causal latent variables requires multiple conditional independence tests and bivariate causal discovery methods (on latent variables which may be estimated incorrectly), there is an obvious concern about how robust this procedure is in practice. It would be more convincing to see (e.g.) precision and recall with regard to selecting the correct causal latent variables when the ground truth is known.\n\nIn summary, the paper introducing some interesting ideas, but the clarity could be improved, the novelty may be somewhat limited and the experimental results could be improved.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper aims to provide a framework for learning non-linear feature mappings such that are invariant to environments. The critical concern raised by the reviewers is their assumption: that causal features of the label are conditionally independent given the label. But in any DAG, conditioning on a common child (here, the label) renders the parents dependent. Their assumption thus is not going to hold other than on a measure zero set of parameters.
Decision: Reject