Invariant Causal Representation Learning for Out-of-Distribution Generalization
Authors: Chaochao Lu
Yuhuai Wu
José Miguel Hernández-Lobato
Bernhard Schölkopf
Venue: ICLR 2022 Poster (2022)

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 relationship with the target. This is done by leveraging a diverse set of training environments to reduce the effect of spurious features and build an invariant predictor. However, these methods have generalization guarantees only when both data representation and classifiers come from a linear model class. We propose invariant Causal Representation Learning (iCaRL), an approach that enables out-of-distribution (OOD) generalization in the nonlinear setting (i.e., nonlinear representations and nonlinear classifiers). It builds upon a practical and general assumption: the prior over the data representation (i.e., a set of latent variables encoding the data) given the target and the environment belongs to general exponential family distributions, i.e., a more flexible conditionally non-factorized prior that can actually capture complicated dependences between the latent variables. Based on this, we show that it is possible to identify the data representation up to simple transformations. We also show that all direct causes of the target can be fully discovered, which further enables us to obtain generalization guarantees in the nonlinear setting. Experiments on both synthetic and real-world datasets demonstrate that our approach outperforms a variety of baseline methods. 

Published as a conference paper at ICLR 2022
INVARIANT CAUSAL REPRESENTATION LEARNING FOR OUT-OF-DISTRIBUTION GENERALIZATION
Chaochao Lu Yuhuai W Jose Miguel Hernández-Lobato" 15- Bernhard Schölkopf" 2+
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 relationship with the target. This is done by leveraging a diverse set of training environments to reduce the effect of spurious features and build an invariant predictor. However, these methods have general- ization guarantees only when both data representation and classifiers come from a linear model class. We propose invariant Causal Representation Learning (iCaRL), an approach that enables out-of -ft--fffrdtiuttiinn (OOD) generalization in the nonlinear setting (i.e., nonlinear representations and nonlinear classifiers). It builds upon practical and general assumption: the prior over the data representation (i.e., a set of latent variables encoding the data) given the target and the environment be- longs to general exponential family distributions, i.e., more flexible conditionally non-factorized prior that can actually capture complicated dependences between the latent variables. Based this, show that it is possible to identify the data representation up to simple transformations. We also show that all direct causes the target can be fully discovered, which further enables us to obtain generalization guarantees in the nonlinear setting. Experiments both synthetic and real-world datasets demonstrate that our approach outperforms a variety of baseline methods.
INTRODUCTION
Modern machine learning algorithms still lack robustness, and may fail to generalize outside of a specific training distribution because they learn easy-to-fit spurious correlations which are prone to change between training and testing environments. We recall the widely used example classifying images of camels and cows (Beery et al., 2018). Here, the training dataset has selection bias, i.e., many pictures of cows are taken on green pastures, while most pictures of camels happen to be in deserts. After training, it is found that the model builds on spurious correlations, i.c., it relates green pastures with
and deserts with camels, and fails to recognize images of cows the beach.
To address this problem, a natural idea is to identify which features of the training data present domain- varying spurious correlations with labels and which features describe true correlations 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 the animals (e.g.. animal shape). By exploiting the varying degrees of spurious correlation naturally present training data collected from multiple environments, one can try to identify stable features and build invariant predictors. Invariant risk minimization (IRM) seeks find data representations (Arjovsky al., 2019) or features (Rojas-Carulle Carulla et al., 2018) for which the optimal predictor is invariant across all environments. The general formulation of IRM is challenging bi-leveled optimization problem, and theoretical guarantees require constraining both data representations and classifiers to be linear (Arjovsky et al., 2019, Theorem 9), or considering the special case of feature selection (Rojas-Carulla al., 2018, Theorem 4). Ahuja al. (2020a) study the problem from the perspective of game theory, with an approach termed invariant risk minimization games (IRMG). They show that the set of Nash equilibria for a proposed game is equivalent to the set of invariant predictors for any finite number of environments, even with nonlinear data representations and nonlinear classifiers. However, these
University of Cambridge, 2PP for Intelligent Systems, Sttanford University. 'Google Research, 'The Alan Turing Institute, "Wok done at University of Toronto, "Equal Supervision, Conresoondence cl 641 @cam. ac. uk.

---

Published as a conference paper at ICLR 2022
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 unseen testing environments.
We propose invariant Causal Representation Learning (iCaRL), a novel approach that enables out- of-distribution (OOD) generalization in the nonlinear setting (i.e., nonlinear representations and nonlinear classifiers)) ). We achieve this by extending and using methods from representation learning and graphical causal discovery. In more detail, we first introduce our main general assumption: when conditioning the target (c.g., labels) and the environment (represented as an index), the prior over the data representation (i.e., set of latent variables encoding the data) belongs to a general exponential family. Unlike the conditionally factorized prior assumed in recent identifiable variational autoencoders (iVAE) (Khemakhem et al., 2020a), this more flexible conditionally non-factorized prior, which can actually capture complicated dependences between the latent variables. We then extend iVAE the case in which the latent variable prior belongs such general exponential family. The combination of this result and the previous general assumption allows us to guarantee that the data representation can be identified up to simple transformations. We then show that the direct causes of the target can be fully discovered by analyzing all possible graphs in a structural equation model setting. Once they are discovered, 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 leads to a practical algorithm and enables us to obtain generalization guarantees in the nonlinear setting.
We propose a general framework for out-of- distribution generalization in the nonlinear setting with the theoretical guarantees both identifiabil- ity and generalizability; (2) We propose general assumption on the underlying causal diagram for prediction (Assumption 1 and Fig. 1c), which covers many real-world scenarios (Section 3.2); (3) We propose a general assumption on the prior over the latent variables (Assumption 2). i.e., a more flexible conditionally non-factorized prior; (4) We prove that an extended iVAE with this conditionally non-factorized prior is also identifiable (Theorems 2&3); (5) We prove that our framework has the theoretical guarantees for OOD generalization in the nonlinear setting (Proposition l).
Overall, we make a number of key contributions: (1)
2
PRELIMINARIES
2.1
IDENTIRIBLE VARIATIONAL AUTOENCODERS
Variational autoencoders (VAEs, see Appendix B) (Kingma & Welling, 2013; Rezende et al., 2014) lack identifiability guarantees. Consider a VAE model where X E IR stands for the observed variables (data) and 2 E R for the latent variables. Khemakhem et al. (2020a) show that VAE with an unconditional prior distribution Po over the latent variables is unidentifiable. However, they also show that it is possible to obtain an identifiable model if one posits a conditionally factorized prior distribution over the latent variables. po (Z)U), where U E R" is an additional observed variable (Hyvärinen et al., 2019). Specifically, let O = (f A) E e conditional generative model
be the parameters of the
(1) where 2 = PX f())) in which e is an independent noise variable with probability density function (c). Importantly, the prior Pr x (Z is assumed to be conditionally factorial, where each element of Z; E Z has a univariate exponential family distribution given U. The conditioning on U is through an arbitrary function X(U) (e.g.. neural net) that outputs the individual exponential family parameters (U) for each Z,. The prior probability density thus takes the form
X = ( P
PT. (Z(U) = II, (((Z)/C,(U) exp Σ j=1 (U))
(2)
where 2, is the base measure, Z, the i-th dimension of Z, CU) the normalizing constant, T: = (Tr.,...., TT,x)) the sufficient statistics, >; = (Aa,1 (U), Ai. (U)) the corresponding natural parameters depending on U, and the dimension of each sufficient statistic that is fixed advance. It is worth noting that this prior is restrictive as it is factorial and therefore cannot capture dependencies.
As in VAEs, the model parameters are estimated by maximizing the corresponding evidence lower bound (ELBO),
LIVAE (0, Φ) :-Epn [Bng(Z|x,0] [logpy [X Z) + log Pr x UZ) - log g ZIX,U)) (3)
'In fact, we not restricted to the classification case and allow the target be either continuous categorical, which will be formally defined in Section 2.2.
2

---

Published as a conference paper at ICLR 2022
where we denote by pp the empirical data distribution given by the dataset D and 96 [ U) denotes an approximate conditional distribution for Z given by recognition network with parameters Φ. This approach is called identifiable VAE (iVAE), Most importantly, it can be proved that under the conditions stated in Theorem 2 of Khemakhem et al. (2020a), iVAE can identify the latent variables Z up to a permutation and a simple componentwise transformation, see Appendix F.
{(x(1,
} / j=1
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 Eall by exploiting data collected from multiple environments Ei,, where &rr S Eail- Technically, they consider datasets De := {a } from multiple training environments E Err. where æf E X C I is the input observation and its corresponding label is yf E E S R The dataset D.. collected from environment e, consists of examples identically and independently distributed according to some probability distribution P(X", )) The goal of IRM is to use these multiple datasets to learn a predictor Y f(X) that performs well for all the environments. Here we define the risk reached by f in environment e as R (f) = E [x)) Y.)], where E(.) is loss function. Then, the invariant predictor can be formally defined follows:
Definition 1 (Invariant Predictor (Arjovsky et al., 2019)). We say that a data representation & e Ho : X + F elicits an invariant predictor w + across environments & if there is classifier w E Hw F - y simultaneously optimal for all environments, that is, w E arg min R (wo+) for all e E E, where o means function composition.
Mathematically, IRM can be phrased as the following constrained optimization problem:
min Exser (w b) s.t. E arg min ( o +), Ve E Exr.
(4)
&E w... 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 to penalty term. The main generalization result (Theorem 9 in Arjovsky et al. (2019)) states that if both & and w come from the class of linear models (i.e., Ho R and Hu = x ) under certain conditions on the diversity of training environments (Assumption Arjovsky et al. (2019)) and the data generation, the invariant predictor w o & obtained by solving Eq. (4) remains invariant in E
@EH.
PROBLEM SETUP
3.1 A MOTIVATING EXAMPLE
In this section, we extend the example which was introduced by Wright (1921) and discussed by Arjovsky et al. (2019), and provide a further in-depth analysis.
Model Consider structural equation model (SEM) with discrete environment variable that modulates the noises in the structural assignments connecting the other variables (cf. Fig. la below): +- Zi + Gaussian 02 o2(E)), Zo + Y + Gaussian((, (0, oa (E)), where Gaussian(O. (0, σ) denotes a Gaussian random variable with zero mean and standard deviation and .I..... 03 are functions of the value E Eall taken by the environment variable E.
Z1 +- Gaussian (0, 0 (E)),
To ease exposition, here we consider the simple scenario in which Eall only contains all modifications varying the noises of Z1, Z2 and Y within a finite range, i.e., o;(e) E [0, of Then, to predict Y from (Z,, Z2) using a least-square predictor ÿz â 0f + 0r2 Zg for environment e, we can
Case 1: regress from Zi, to obtain â = and 2 = 0,
Case 2: regress from Zg. to obtain â1 = and و = の (c))
te)+)
Case regress from (Zf, Zz), to obtain â +
ds(e)
and 0g (2())+
02(e)
In the generic scenario (i.c., a (e) / 0, (e) + 0, and (3(e) of 0), the regression using Z, in Case 1 is an invariant correlation: it the only regression whose coefficients do not vary with By contrast,
TThe setup applies to both continuous and categorical data. If any observation label is categorical, could onc-hot encode it.
3

---

Published as a conference paper at ICLR 2022
Z2
Zpp
Bch
(a)
(b)
(c)
Figure I: (a) Causal structure of Model (b) more practical extension of Model I, where 21 and Z2 are not directly observed and X their observation. (e) 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 direct connection with Y. We allow for arbitrary connections between the latent variables (red dashed lines) as long as the resulting causal diagram including Y is a directed acyclic graph (DAG). Grey nodes denote observed variables and white nodes represent unobserved variables. Dashed lines denote the edges which might vary across environments and even be absent in some scenarios, whilst solid lines indicate that they are invariant across all the environments.
the regressions in both Case 2 and Case have coefficients that depend on e. Therefore, only the invariant correlation in Case will generalize well to new test environments.
Another way to understand Model is through its graphical representation', as shown in Fig. la. We treat the environment as a random variable F,, where E could be any information specific to the environment (Storkey, 2009; Peters et al., 2015; Zhang et al., 2017; Huang et al., 2020). Unless stated otherwise, for simplicity, we let E' be the environment index, i.c., E {L..... NN. where N is the number of training environments. A more realistic version appearing in many settings is shown in Fig. 1b, where the true variables ( Zz2 } are unobserved and we can only observe their transformation X. In this case, Invariant Causal Prediction (ICP) (Peters et al., 2015) will fail when applied to X, even when Y is not affected by E' (i.c., the edge E + Y is removed). The reason is that each variable (i.e., each dimension) of is jointly influenced by both Z1 and Z/ so that ICP is unable to find the variables containing the information only about Zi by searching for subset of variables X. By contrast, both IRM and IRMG work, as long as the transformation is linear. These findings are also empirically illustrated in Appendix K.1. We now go even further and consider a more general causal graph in which Y can have more than one parent child.
3.2 ASSUMPTIONS ON THE CAUSAL GRAPH
We extend the causal graph in Fig. lb more general setting', as encapsulated in Fig. Ic. In particular, we now have X E X' C Rd, Y E y S R Z (Z....., Zp. Ze...... Zez) E Z C R², where n r k and {Z,)ier,- (pi.---e) and {Zy}jer. (c)....)) are multiple scalar causal factors and non-causal factors" of Y. respectively. We denote Z, - (Zp..... 4p.) and Z. = (Z,, Zoz. ) for the ease of clarification. We also assume that Z is of lower dimension than X, that is, S d. We allow for arbitrary connections between the latent variables Z long as the resulting causal diagram including Y is a directed acyclic graph (DAG). We use dashed lines to indicate the causal mechanisms which might vary across environments and even be absent in some scenarios, whilst solid lines indicate that they are invariant across all the environments. To sum up, we assume that the underlying causal graph encapsulated in Fig. Ic satisfies the following assumption :
Assumption 1. (a) Zi depends on one or both of Y and E for any i; (b) The causal graph containing Z and Y is DAG; (c) X IL Y, E Z, implying that p X x is invariant across all the environments; (d) Y 1L Z,. implying that pYY Z,) is invariant across all the environments.
One may be wondering how practical Assumption is in real world applications. Let us explore this in more detail. Assumption la rules out all the useless Z; in the task of predicting Y. This because if Assumption la is violated, meaning that Z, is independent of Y and E and has no influence in predicting Y, then such Z, should be viewed as noise and thus eliminated during learning. Assumption 1b is common assumption in causal discovery (Spirtes et al., 2000; Pearl, 2009; Peters et al., 2015). It also makes sense in Assumption Ic that the generative mechanism p( X Z) is
'The relation between SEM and its graphical representation is formally defined in Appendix D. 'For simplicity, we do not explicitly consider unobserved confounders in this paper. In particular, we assume that there are unobserved confounders between Z. Y. X. and E.
"This means that Z,er. could be either an effect of Y. independent of Y. or spuriously correlated with Y via third set of confounders (i.c., both Zj and Y are affected by a subset of {2, Fizs and E).
For generality, we replace E. with E additionally include the case of multi-dimensional variables.
4

---

Published as a conference paper at ICLR 2022
Algorithm 1: Invariant Causal Representation Learning (iCaRL.)
Phase 1: We first learn NE-iVEE model. including the decoder and its corresponding encoder. by optimizing the objective function in (10) on the data {X, Y. E}. Then, we use the mean the NF-iVAE encoder to infer the latent variables from observations {X; Y. E}. The latent variables are guaranteed to be identified up to permutation and simple transformation.
Phase 2: After inferring Z, we first conduct the PC algorithm to learn a Markov equivalence class of DAGs, and then discover direct causes (parents) of Y among its neighbors by testing all pairs of latent variables with (conditional) independence testing, i.e., finding set of latent variables in which each pair of Z; and Y.
Zj satisfies that the dependency between them increases after additionally conditioning
Phase 3: Having obtained Pa (Y). we solve (11) to learn the invariant classifier When in a new environment, we first infer Pa( from X by solving (12) and then leverage the learned for prediction.
invariant across all the environments. Otherwise, it is impossible to infer Z from X in any unseen environment. Assumption 1d is widely-used default assumption in OOD generalization (Peters ct al., 2015; Arjovsky et al., 2019). In fact, Assumption Id can be further relaxed to the more practical one that E. G 2pp is invariant across all the environments. That is, given Zp. we allow E to only affect the amount of noise in the distribution of Y. because that would not change the expected value of Y since the mean of the noise would be zero. Apparently. Assumption 1, together with the causal graph in Fig. 1c, covers most scenarios (e.g., the ones of Zhang al. (2013); von Kügelgen al. (2020); Sun et al. (2020); Ahuja et al. (2021); von Kügelgen et al. (2021)) and is very flexible model for causal analysis when predicting Y from X.
3.3 ASSUMPTIONS ON THE PRIOR
When the underlying causal graph satisfies Assumption 1, our primary assumption leading to identifi- ability in this general setting is that the conditional prior p( Z Y, EE) belongs to a general exponential family. This is formalized as follows:
Assumption 2. p( Y E) belongs to a general exponential family with parameter vector given by an arbitrary function X(Y, E) and sufficient statistics T ([Z) = [Ty 2 NNN Z)" given by the concatenation of a) the sufficient statistics (Z) = [Ti (Z,)= ,.... (Z.) of a factorized
exponential family, where all the T.(Z)) have dimension larger or equal to 2, and b) the output TNN Z) of a neural network with ReLU activations. The resulting density function is thus given by
Pr. P Y E)- Q(Z)/C(Y, E) exp [T(Z) A(Y, E),,
(5)
where Q is the base measure and C the normalizing constant.
A neural network with ReLU activation has universal approximation power. Therefore, the term TNN(Z) in the above prior distribution will allow us to capture arbitrary dependencies between the latent variables. The distribution in Eq. (5) is more flexible than the conditionally factorized prior assumed by iVAEs. However, surprisingly, the identifiability results of iVAEs also hold when using the more flcxible prior in Eq. (5), as we will show in Section 4.1. This motivates using an extended iVAE model with the above prior to model data generated by the ground truth model in Fig. Ic. However, in this case, the data generating model and the learned model might have different priors. For example, in the ground truth model, the prior for each Ziei, might be pl Zu E), when Z; is only caused by E. By contrast, in the extended iVAE model the prior is p Z Y E). Is this going to affect the identifiability of the latent variables? Well, in practice not because the posterior distribution for Z given X, Y and E would be equivalent in both models (up to identifiability guarantees).
INVARIANT CAUSAL. REPRESENTATION LEARNING
We now introduce our algorithm, invariant Causal Representation Learning (iCaRL), which consists of phases as summarized in Algorithm 1. The idea is that we first identify Z by using an extended iVAE model under Assumptions 1&2 (Phase 1), then discover direct causes of Y among the identified Z (Phase 2), and finally learn an invariant predictor for Y from the discovered causes (Phase 3).
4.1 PHASE 1: IDENTIFYING LATENT VARIABLES USING NF-IVAE
In this section, we describe our identifiable model, namely NF-iVAE, which an extended iVAE with a general non-factorized prior that is able to capture complex dependences between the latent variables. Technically, in the general setting under Assumption 1, it is straightforward to obtain
5

---

Published as a conference paper at ICLR 2022
corresponding generative model by directly substituting U with (Y, EE) in Eq. (I):
(x ZIY, E) P PT > ( E))
(6)
ps x Z = Pe X - f(2)).
(7)
The corresponding ELBO
Lphaee (0, Φ) ==pp ( [r Z) log pT (Z log q0 (Z| Y. E3]
(8)
To obtain an identifiability result, we assume that the prior PT, (Z E) satisfies Assumption 2 (i.e., Eq. (5)). Since the prior a general multivariate exponential family distribution with unknown normalization constant, we cannot learn its parameters (T X) by directly maximizing Eq. (8).
Instead, we use score matching, a well-known method for training unnormalized probabilistic models (Hyvärinen, 2005; Vincent, 201 1), and learn (T, A) by minimizing
(T,)) :=Bpp [ [ [I|V log (2) Y, - Moggr.xxZ/yx [X,]]
plasel
In practice, we can use simple trick of partial integration to simplify the evaluation of Eq. (9).
see Appendix C. Furthermore, we can jointly learn (0, Φ) by combining Eq. (8) and Eq. (9) in the following objective:
Ephase (0.
phisel (f, A.
SSM phase! (f,
(10)
where f,T, & are copies of f. T, X, & that treated constants and whose gradient is not calculated during learning. More details can be found Appendix M.
We now state our main theoretical results:
Theorem 1. Assume that we observe data sampled from a generative model defined according to Eqs. (5-7), with parameters 0 = (f. T. A), where PT P (Z E) satisfies Assumption 2. Furthermore, assume the following holds: (i) The set (X E p (X)
0} has measure zero, where ye is the characteristic function of the density Pe defined in Eq. (7). (ii) Function f in Eq. (7) is injective, and has all second-order cross derivatives. (üii) The sufficient statistics in T3 are all twice differentiable. (iv) There exist l + 1 distinct points (Y, E)" ..... (Y, Ey such that the matrix ()) - E))) of size x k is invertible, where
L= (>((Y, E - -X((, is the dimension T. Then the parameters O are identifiable up to a permutation and a "simple transformation" of the latent variables Z, defined as componentwise nonlinearity making each recovered (Z.) in T (Z) equal to the original up to a linear operation.
Note that, this theorem is inspired by but beyond the main results of iVAEs in that the former is predicated on Assumption 2 which is more flexible than the conditionally factorized prior assumed in i'AEE. It results in several key changes in the proof, clarified in Appendix H. Interestingly, from (iv) we can further see that E is unnecessary when there exist k + 1 distinct points Y Y such that the matrix = (A(Y²) A(Y"), ...., A(Y*) - A(Y of size x k is invertible. Not requiring E would make our approach even more applicable.
We further have the following consistency result for the estimation.
Theorem 2. Assume that the following holds: (i) The family of distributions go ( (Z|X, Y E) con-
tains pe Z Y, E), and gol [Z|X Y E) > everywhere. (ii) We maximize Lphaset Φ) with respect to both O and cb Then in the limit of infinite data, we learn the true parameters 0* up to permutation and simple transformation of the latent variables Z.
As a consequence of Theorems 1&2, we have: Theorem 3. Assume the hypotheses of Theorem I and Theorem 2 hold, then in the limit of infinite data, we identify the true latent variables Z* up to a permutation and simple transformation.
Theorem 3 states that we can use NF-iVAE to infer the true Z" up to permutation and simple transformation. We use the mean of q / X, Y. E) for this task. Note that the noise e may introduce uncertainty in the estimation of Z. However, when X is high dimensional and Z is low dimensional (as common in real world applications), (Z) X, Y. E)) will be highly concentrated and we will still be able to estimate Z with high accuracy. The good results obtained by our method in various experiments seem corroborate this. All three theorems proven in Appendix H.
4.2 PHASE 2: DISCOVERING DIRECT CAUSES
After estimating Z for each data point, the next step is determine which components of Z are direct causes of Y. We denote these components by Pa YY. We first conduct the PC algorithm (Spirtes al., 2000) to learn Markov equivalence class of DAGs, which gives us the direct neighbors of Y.

---

Published as a conference paper at ICLR 2022
denoted by Ne(Y). Then, from Assumption 1, one observation is that in the generic case, for any two latent variables Zi and Zj from Ne(Y), only when both are causes of Y do we have that the dependency between them increases after additionally conditioning on Y. Thus, when there exist at least two causal latent variables in Ne(Y). we can test all pairs of latent variables' with conditional independence testing* (Zhang al., 2012) to discover PY(Y) by comparing p-values from the two tests: IndTest/ t (Zi, Ez) and IndTest (Z E), where IndTest denotes (conditional) independence test. Conversely, if no such a pair is found, it implies that there is at most one causal latent variable. This is highly unlikely case in real world applications, which is left to Appendix G.
4.3 PHASE 3: LEARNING AN INVARIANT PREDICTOR
After having obtained the causal latent variables Pa(Y)) for Y across training environments, we can learn w by solving the following optimization problem:
min WEH
Σ LeEE. [w))== = min
LeEfi E (( Y+)
(11)
wEH.
where (-) could be any loss. Since assume that ((()Pa(( is invariant across Eall (the relaxed version of Assumption Id), the learned w is guaranteed to perform well across Entl.
The remaining question is how to infer Pa( Pa(Y) (Y (i.e.. Zp) from X in a new environment. This can be implemented by leveraging the learned p( (Xx Z). The rationale behind is that p(X 2) is assumed to be invariant across Eall (Assumption lc). In light of this idea, we follow Sun et al. (2020) and infer
Zp from X in any new testing environment by solving the following optimization problem:
max logpy pp, Ze) + A1 112 + 221 1c Zp.Ze
(12)
where the hyperparameters A > and A2 > 0 control the learned Zp and Z. in reasonable scale, and both are selected on first use values of Z sampled from the training set as initial points and then use Adam to optimize for several iterations." Note that the noise e will introduce uncertainty in the estimation of Zp and Z. from X. However, as mentioned before (below Theorem 3), this noise is not going affect the estimation much because the likelihood will be highly concentrated around the ground truth values, as corroborated by our good empirical results.
validation data. For optimization, we follow Schott et al. (2018) to
key question iCaRL performs well across Enul even though it uses only data from E That is, does iCaRL enable OOD generalization, as defined by Arjovsky et al. (2019)? The answer is positive since Theorem A.1 in Arjovsky (2021) indicates that i) any predictor w 4 with optimal OOD generalization uses only Pa(Y) to compute T and ii) the classifier w in this optimal predictor can be estimated using data from any environment e for which the distribution of Pa (YY) has full support, which will always be the case since the conditional prior in Eq. (5) has full support. Finally, Theorem A. in Arjovsky (2021) also indicates that îii) the optimal predictor will be invariant across Eau. Key to these results is that Pa(Y') are available when solving (12). This requires first to identify the latent variables Z from X, Y and E and second, to discover the direct causes of Y. The hypotheses of Theorems 1 and 2 and Assumption 1 provide this guarantee. We therefore have the following result whose proof is in Appendix H.
Proposition 1. Under Assumption and the assumptions of Theorems and the predictor learned by iCaRL across Evr in the limit of infinite data has optimal OOD generalization across Entt.
5 EXPERIMENTS
We compare our approach with a variety of methods on both synthetic and real-world datasets. In all comparisons, unless stated otherwise, we average performance over ten runs. Due space con- straints, we only highlight some key results while pointing to the extensive appendices for more information. The supplement contains all the details of the experiments, e.g., datasets (Appendix J), implementation (Appendix M), hyperparameters and architectures (Appendix N), etc.
5.1 SYNTHETIC DATA
To verify the identifiability of NF-iVAE, we conduct a series of experiments on synthetic data generated according to the causal graph shown in Fig. 2e. Details of the ground truth data generating
We only need to consider those variables in Ne(Y) whose edges connecting Y not oriented by PC. "These conditional independence tests can be performed in parallel to largely accelerate the testing procedure. See more in Appendix G.
Note that, (12) can be optimized either independently for each data point for a number of data points.
7

---

Published as a conference paper at ICLR 2022
(a) Ground truth
(b) VAE
(c) iVAE
(d) NF-iVAE
Synthetic Data
CMNIST
5
554
8 0 5
8
8
P
8
0 5 3 3
8 3
MCC
MCC
0.5.
NF-VVEE
0.0 VAE IVAE NF-IVAE
3 3
3
(e) Causal graph
(f) MCC on synthetic data (g) MCC on CMNIST (h) Intervening CMNIST
Figure 2: (a-d) Visualization the samples (i.e.. 2 (21, 22)) in latent space recovered through different algorithms: (a) Samples from the true distribution; (b-c) Samples from the posterior inferred using VAE and iVAE, respectively. Apparently, our method (d) can recover the original data up to permutation and a simple componentwise transformation. (e) The causal structure with Y having two causes describes the data generating process of the synthetic dataset. (f) Mean correlation coefficient (MCC) scores for VAE, iVAE, and NF-iVAE on synthetic data. (g) MCC scores for VAE, iVAE, and NF-IVAE CMNIST. (h) The effects the CMNIST images of digit (top two rows) and digit (bottom two rows) when intervening causal factor ZiE,, and on non-causal factor (effect) Zyet, respectively.
process are given in Appendix K. The reason we choose this setting is that it is the simplest case satisfying our requirements: a) For ease of visualization, the latent space had better be 2-dimensional; b) To introduce the non-factorized prior given Y and E (i.e., Z; H Z, Y. E), Y has at least two causes. We draw 1000 samples from each of the four environments E (0.2, 2, 3. 5), and thus the whole synthetic dataset consists of 4000 samples. The task is to recover the true latent variable Z (Z,, Z2) using the samples of X, E. and Y. We compare with two widely-used baselines: VAE (Kingma& & Welling, 2013) (without identifiability guarantees) and iVAE (Khemakhem et al., 2020a) (with a conditionally factorized prior for identifiability). Through the aforementioned theoretical analysis. it evident that our method has a more general assumption on the prior leading to identifiability. This is demonstrated empirically in Figs. 2b-2d. Our method NF-iVAE can recover the original data Z up to permutation and simple componentwise transformation, whereas all the other methods fail because they are unable to handle the non- non-factorized case in which Z, 1 Z,) Y. E. We also compute the mean correlation coefficient (MCC) used in Khemakhem et al. (2020a), which can be obtained by calculating the correlation coefficient between all pairs of true and recovered latent factors and then solving a linear sum assignment problem by assigning each recovered latent factor to the true latent factor with which it best correlates. By definition, higher MCC scores indicate stronger identifiability. From Fig. 2f, we can see that the MCC score for NF-iVAE is significantly greater than those of VAE and iVAE, indicating much stronger identifiability. Note that, in Appendix K we additionally compare with more methods, whose differences are further summarized in table.
5.2 COLORED MNIST, COLORED FASHION MNIST, AND VLCS
In this section, we first report experiments on two datasets used in IRM and IRMG: Colored MNIST (CMNIST) and Colored Fashion MNIST (CFMNIST). We follow the same setting of Ahuja et al. (2020a) create these two datasets (see the details in Appendix J). The task to predict binary label assigned to each image which is originally grayscale but artificially colored in a way that the color is correlated strongly but spuriously with the class label. For all the experiments on these two datasets, we set the number of the latent variables to n = 10.
Likewise, we investigate the identifiability of NE-iVAE on CMNIST by computing the MCC score between samples of the true latent variable and of the recovered latent variable. Since the true latent variable on CMNIST is inaccessible to us, we follow Khemakhem et al. (2020b) and compute an average MCC score between samples of latent variables recovered by different models trained with different random initialization. As shown in Fig. 2g. it is evident that the MCC score for NF-iVAE greatly outperforms the others, showing that the latent variables recovered by NF-iVAE have much better identifiability.
Furthermore, we demonstrate the ability of iCaRL to discover the causal latent variables (Phase 2) by visualizing the generated images through performing intervention upon causal latent variable and a non-causal latent variable, respectively. Fig. 2h shows how intervening upon each of them affects the image. Obviously, intervening on a causal latent variable affects the shape of the digit but not its color
8

---

Published as a conference paper at ICLR 2022
Table 1: Colored Fashion MNIST. Comparisons in terms of accuracy (%) (mean + std deviation).
Table 2: VLCS. Comparisons in terms of accuracy (%) (mean 1 std deviation).
METHOD
TRAIN
TEST
METHOD
TEST
ERM
83.17 + 1.01
22,46 + 0.88
ERM
71.4 0.3
ERM 1
81.33 + 1.35
33.34 8.83
IRM
78.1 0.0
ERM 2
84.39 1: 1.89
13.16
0.82
DRO (Sagawa al., 2019)
17.2 1T.2 0.6
ROBUST MIN MAX
82.81 0.11
29.22
8.56
Mixup (Yan et al., 2020)
77.7
F-IRM GAME
62.31 + 2.35
69.25 5.82
CORAL (Sun & Sacnko, 2016)
+
V-IRM GAME
68.96 0.95
70.19
1.47
MMD (Li ct al.. 2018b)
76.7 0.9
IRM
75.01 0.25
55.25
+ 12,42
DANN (Ganin et al., 2016)
78.7 0.3
iCaRL (ours)
74.96 + 0.37
73.61
+ 0.63
C-DANN (Li et al., 2018c)
7..240.4
FRM GRAYSCALE.
74.79 + 0.37
74.67 + 0.48
LaCIM (Sun et al., 2020)
78.4 0.5
OPTIMAL
75
75
iCaRL (ours)
81.8 0.6
(top plots), whilst intervening on a non-causal latent variable, which is an effect in Fig. 2h, affects the color of the digit only (bottom plots). This visually verifies the results of iCaRL in Phase 2.
In terms of the OOD generalization performance, we compare iCaRL 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 (on each environment e), and ERM GRAYSCALE (on data with no spurious correlations), and 4) ROBUST MIN MAZ (minimizing the maximum loss across the multiple environments), Table shows that iCaRL outperforms all other baselines on CFMNIST. It worth emphasising that the train and test accuracies of iCaRL closely approach the ones of ERM GRAYSCALE and OPTIMAL, implying that iCaRL approximately learns the true invariant causal representations with almost no correlation with the spurious color feature. We can draw similar conclusions from the results on CMNIST (Appendix L).
We also report the results on one of the widely used realistic datasets for OOD generalization: VLCCS (Fang et al., 2013). This dataset consists of 10, 729 photographic images of dimension (3, 224, 224) and 5 classes from four domains: Caltech LabelMe, SUN09, and VOC2007. We used the exact experimental setting that is described in Gulrajani & Lopez-Paz (2020). We provide results averaged over all possible train and test environment combination for one of the commonly used hyper-pe preer-parameter tuning procedure: train domain validation. As shown in Table 2, iCaRL achieves state-of- the-art performance when compared to those most popular domain generalization alternatives. We further include experimental evidence another popular dataset: PACS (Li et al., 2017a), all the details of which are placed in Appendix L.
RELATED WORK
Invariant Causal Prediction (ICP), aims to find the causal feature set (i.e., all direct causes of target variable of interest) (Peters et al., 2015) by exploiting the invariance property in causality which has been 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 no intervention is allowed the target variable Y. Besides, ICP methods implicitly assume that the variables of interest Z are given. The works of Magliacane et al. (2018) and Subbaswamy et al. (2019) attempt to find invariant predictors that are maximally predictive using conditional independence tests and other graph-theoretic tools, both of which also assume that the Z are given and further assume that additional information about the structure over Z is known. Mitrovic et al. (2020) analyze data augmentations in self- eeff-supervieed learning from the perspective of invariant causal mechanisms. Arjovsky et al. (2019) reformulate this invariance as an optimization- based problem, allowing us to learn invariant data representation from X constrained to be linear transformation of Z. The risks of this approach have been discussed in Rosenfeld et al. (2020); Kamath et al. (2021); Nagarajan et al. (2020) and its sample complexity is analyzed in Ahuja et al. (2020b).
Another line of related work is in the field of domain generalization, which we discuss in Appendix A.
7 CONCLUSION
We have proposed a novel framework to learn invariant predictors from a diverse set of training environments. is based on practical and general assumption: the prior over the data representation belongs to a general exponential family when conditioning on the target and the environment. This assumption leads to guarantees that the components in the representation can be identified up to a permutation and simple transformation. This allows us to discover all the direct causes of the target, which enables generalization guarantees in the nonlinear setting. We hope our framework will inspire new ways to address the OOD generalization problem through causal lens.

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGEMENTS
We are thankful to Wenlin Chen for contributing Step III the proof of Theorem 4 and for generalizing the proof in Theorem 5 to the case in which each (2;)) in Ty ()) contains arbitrary sufficient statistics instead of just Z; and z," We also thank Ilyes Khemakhem, Aapo Hyvärinen, and Arthur Gretton for their helpful discussions, and the anonymous reviewers for their constructive comments on an earlier version of this paper.
REFERENCES
Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg. Dandelion Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Mike Schuster, Jonathon Shlens, Benoit Steiner, llya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg. Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. TensorFlow: Large-scale machine learning on heterogeneous systems, 2015. URL https: http://wwwww/t/o/// www.otensorfoow.. org/. Software available from tensorflow.org.
Olah,
Kartik Ahuja, Karthikeyan Shanmugam, Kush Varshney, and Amit Dhurandhar. Invariant risk minimization games. arXiv preprint (.Kiv:2002.046992 2020a.
Kartik Ahuja, Jun Wang. Amit Dhurandhar, Karthikeyan Shanmugam, and Kush R Varshney. Empirical invariant risk minimization? sample complexity perspective. arXiv preprint arXiv:2010. 16412, 2020b.
Kartik Ahuja, Ethan Caballero, Dinghuai Zhang, Yoshua Bengio, Ioannis Mitliagkas, and Irina Rish. Invariance principle meets information bottleneck for out-of-distribution generalization. arXiv preprint (06.0//0666667,, 2021.
John Aldrich. Autonomy. Oxford Economic Pupers, 41(1):15-34, 1989.
Martín Arjovsky. Out of distribution generalization machine learning. CoRR, abs/2103.02667, 2021.
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Pazz Invariant risk minimization. arXiv preprint arXiv:1907.02893 2019.
Yogesh Balaji, Swami Sankaranarayanan, and Rama Chellappa. Metareg: Towards domain gen- cralization using meta-regulaarzzatoo.. Advances in Neural Information Processing Systems, 31: 998-1008, 2018.
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.
Shai Ben-David, John Blitzer, Koby Crammer, Fernando Pereira, et al. Analysis of representations for domain adaptation. Advances in neural information processing systems, 19:137, 2007.
Shai Ben- Devid. John Blitzer, Koby Crammer, Alex Kulesza, Fernando Pereira, and Jennifer Wortman Vaughan. A theory of learning from different domains. Machine learning, 79(1):151-175, 2010.
Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspectives. IEEE transactions on pattern analysis and machine intelligence, 1788-1888888888 2013.
Pierre Comon. Independent component analysis, new concept? Signal processing, 36(3):287-314, 1994.
Koby Crammer, Michael Kearns, and Jennifer Wortman. Learning from multiple sources. Journal of Machine Learning Research, 98)), 2008
Povilas Daniusis, Dominik Janzing, Joris Mooij, Jakob Zscheischler, Bastian Steudel, Kun Zhang. and Bernhard Schölkopf. Inferring deterministic causal relations. arXiv preprint arXiv: 240.34455 2012.
10

---

Published as a conference paper at ICLR 2022
A Philip Dawid, Vanessa Didelez, et al. Identifying the consequences of dynamic treatment strategies: decision-theoretic overview. Statistics Surveys, 4: 4:14--23,, 2010.
Zhengming Ding and Yun Fu. Deep domain generalization with structured low-rank constraint, IEEE Transactions Image Processing, 27(1):304-313, 2017.
Qi Dou, Daniel C Castro, Konstantinos Kamnitsas, and Ben Glocker. Domain generalization via -node--ggoottcc learning of semantic features. arXiv preprint 1rXiiv:9100135000 2019.
Sarah Erfani, Mahsa Baktashmotlagh, Masud Moshtaghi, Xuan Nguyen, Christopher Leckie, James Bailey, and Rao Kotagiri. Robust domain generalisation by enforcing distribution invariance. In Proceedings of the Twenty-Fifth International Joint Conference on Artificial Intelligence (IJCAI- 16), pp. 1455-1461. AAAI Press, 2016.
Chen Fang, Ye Xu, and Daniel N Rockmore. Unbiased metrie learning: On the utilization of multiple datasets and web images for softening bias. Proceedings of the IEEE International Conference Computer Vision, pp. 1657-1664, 2013.
José AR Fonollosa. Conditional distribution variability measures for causality detection. In Cause Effect Pairs in Machine Learning, pp. 339-347. Springer, 2019.
Yaroslav Ganin, Bygeniya Ustinova, Hana Ajakan, Paseal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky, Domain-adversariie training of neural networks. The Journal of Machine Learning Research, 17(1):2096-203) 2016.
Muhammad Ghifary, W. Bastiaan Kleijn, Mengjie Zhang, and David Balduzzi. Domain generalization for object recognition with multi-task autoencoders. In ICCV, pp. 2551-2550 IEEE Computer Society, 2015.
Ishaan Gulrajani 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 Econometric Society, pp. iii-115, 1944.
Christina Heinze-Deml, 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 causation. Economics & Philosophy, 6(2):207 -234, 1990.
Patrik O Hoyer, Dominik Janzing, Joris M Mooij, Jonas Peters, and Bernhard Schölkopf. Nonlinear causal discovery with additive noise models. In Advances in neural information processing systems, pp. 689-696, 2009.
Biwei Huang, Kun Zhang, Jiji Zhang, Joseph Ramsey, Ruben Samchez-Romero, Clark Glymour, and Bernhard Schölkopf. Causal discovery from hetenngeneous/nonstttoonay data. Journal of Machine Learning Research, (198))11-53 2020.
Aapo Hyvärinen. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
Aapo Hyvärinen and Erkki Oja. Independent component analysis: algorithms and applications. Neural networks, 13(4-5):41 1-430, 2000.
Aapo Hyvärinen, Hiroaki Sasaki, and Richard Turner. Nonlinear ICA using auxiliary variables and generalized contrastive learning. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 859-868 868, 2019.
Dominik Janzing, Joris Mooij. Kun Zhang, Jan Lemeire, Jakob Zscheischler, Povilas Daniušis, Bastian Steudel, and Berhhard Schölkopf. information-geommtrie approach inferring causal directions. Artificial Intelligence, 182:1-31,200 2012.
Pritish Kamath, Akilesh Tangella, Danica Sutherland, and Nathan Srebro. Does invariant risk minimization capture invariance? arXiv preprint arXiv:2101.01 2021.
11

---

Published as a conference paper at ICLR 2022
llye Khemakhem, Diederik Kingma, Ricardo Monti, and Aapo Hyvärinen. Variational autoencoders and nonlinear ICA: unifying framework. International Conference on Artificial Intelligence and Statistics, pp. 2207-2217, 2020a.
llyes Khemakhem, Ricardo Monti, Diederik Kingma, and Aapo Hyvarinen. Ice-beem: Identifiable conditional energy-based deep models based on nonlinear ica. Advances in Neural Information Processing Systems, 33, 2020b.
Diederik P Kingma and Max Welling. arXiv:1312.61 14, 2013.
Auto-encoding variational bayes.
arXiv preprint
Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M Hospedales. Deeper, broader and artier domain generalization. In Proceedings of the IEEE international conference on computer vision, pp. 5542-5550, 2017a.
Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy Hospedales. Learning to generalize: Meta-learning for domain generalization. Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018a
Da Li, Jianshu Zhang, Yongxin Yang, Cong Liu, Yi-Zhe Song, and Timothy M Hospedales. Episodic training for domain generalization. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 1446-1455, 2019a.
Haoliang Li, Sinno Jialin Pan, Shiqi Wang, and Alex C Kot. Domain generalization with adver- sarial feature learning. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 540055400000 000-5400,222 2018b.
Wen Li, Zheng Xu, Dong Xu, Dengxin Dai, and Lue Van Gool. Domain generalization and adaptation using low rank exemplar svms. IEEE transactions on pattern analysis and machine intelligence, 40(5):1114-) 111114--22 1127, 2017b.
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 on Computer Vision (ECCV), pp. 624-639, 2018c.
Yiying Li, Yongxin Yang. Wei Zhou, and Timothy Hospedales. Feature-critic networks for heteroge- neous domain generalization. In International Conference on Machine Learning, pp. 3915-3924. PMLR, 2019b.
Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Ractsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of dis- entangled representations. In international conference on machine learning, pp. 4114-4124 2019.
Sara Magliacane, Thijs 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.
Massimiliano Mancini, Samuel Rota Bulò, Barbara Caputo, and Elisa Ricci. Best sources forward: domain generalization through source source-specific nets. In 2018 25th IEEE international conference image processing (ICIP), pp.
IEEE, 2018.
Jovana Mitrovic, Brian McWilliams, Jacob Walker, Lars Buesing, and Charles Blundell. Representa- tion learning via invariant causal mechanisms. arXiv preprint arXiv:2010.07922. 2020.
Saeid Motian, Marco Piccirilli, Donald A. Adjeroh, and Gianfranco Doretto. Unified deep supervised domain adaptation and generalization. In /CCV, pp. 5716-5726 IEEE Computer Society, 2017.
Krikamol Muandet, David Balduzzi, and Bernhard Schölkopf. Domain generalization via invariant feature representation. In ICML, volume 28 of JMLR Workshop and Conference Proceedings, pp. 10-18, 2013.
Vaishnavh Nagarajan, Anders Andreassen, and Behnam Neyshabur. Understanding the failure modes of out-of-distribution generalization. arXiv preprint arXiv:2010. 2020.
12

---

Published as a conference paper at ICLR 2022
Li Niu, Wen Li, and Dong Xu. Multi-view domain generalization for visual recognition. In Proceedings of the IEEE international conference on computer vision, pp. 4193-4201, 2015.
Sinno Jialin Pan and Qiang Yang. A survey on transfer learning. IEEE Transactions on knowledge and data engineering, 22(10):1345-135, 2009.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Vito,o Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' d''AlcééBar Buue E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp. 8024 8035. Curran Associates, Inc., 2019.
Vishal M Patel, Raghuraman Gopalan, Ruonan Li, and Rama Chellappa. Visual domain adaptation: survey of recent advances. IEEE signal processing magazine, 32(3):53 69, 2015.
Judea Pearl. Causality. Cambridge university press, 2009.
Jonas Peters, Peter Bühlmann, and Nicolai Meinshausen. Causal inference using invariant prediction: identification and confidence intervals. arXiv preprint arXiv: 1.XX::1501.013333 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.4082, 2014.
M. Carull,C B. Schölkopf, R.
Turner, Peters. Invariant models for causal transfer learning. Journal of Machine Learning Research, 19(36):1-34 2018.
Elan Rosenfeld, Pradeep Ravikumar, and Andrej Risteski. The risks of invariant risk minimization. arXiv preprint arXiv:2010.0576,, 2020.
Shiori Sagawa, Pang Wei Koh, Tatsunori Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. arXiv preprint arXiv: 191119919991 0.07311,, 2019.
B. Schölkopf, D. Janzing, J. Peters, E. Sgouritsa, K. Zhang, and J. M. Mooij. On causal and anticausal learning. In Proceedings of the 29th International Conference Machine Learning (ICML), pp. 1255-1262, 2012.
Lukas Schott, Jonas Rauber, Matthias Bethge, and Wieland Brendel. Towards the first adversarially robust neural network model on mnist. arXiv preprint urXiv:1805.09190 2018.
Peter Spirtes, Clark N Glymour, Richard Scheines, and David Heckerman. Causation, prediction, and search. MIT press, 2000.
Amos Storkey. When Training and Test Sets Are Different: Characterizing Learning Transfer. 2009.
Adarsh Subbaswamy, Bryant Chen, and Suchi Saria. A universal hierarchy of shift-stable distributions and the tradeoff between stability and performance. arXiv preprint arXiv: /905.. 1374, 2019.
Baochen Sun and Kate Saenko. Deep coral: Correlation alignment for deep domain adaptation. In European conference computer vision, pp. 443-450 Springer, 2016.
Xinwei Sun, Botong Wu, Chang Liu, Xiangyu Zheng, Wei Chen, Tao Qin, and Tie-yan Liu. Latent causal invariant model. arXiv preprint arXiv:2011.02200 2020.
Hemanth Venkateswara, Jose Fusebio, Shayok Chakraborty, and Sethuraman Panchanathan. Deep hashing network for unsupervised domain adaptation. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition, pp. 5018-5027, 2017.
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural compula- fion, 23(7):1661-1674 201
13

---

Published as a conference paper at ICLR 2022
Julius Kügelgen, Yash Sharma, Luigi Gresele, Wieland Brendel, Bernhard Schölkopf. Michel Besserve, and Francesco Locatello. Self-supervised learning with data augmentations provably isolates content from style. arXiv preprint (exXix:2206.046111 2021.
J. von Kügelgen, A. Mey, M. Loog, and B. Schölkopf. Semi-supervised learning, causality and the conditional cluster assumption. Conference on Uncertainty in Artificial Intelligence (UAI), 2020.
Changzhang Wang, You Zhou, Qiang Zhao, and Zhi Geng. Discovering and orienting the edges connected to a target variable in a dag via a sequential local learning approach. Computational Statistics & Data Analysis, 77:252-266, 2014.
Haohan Wang, Aaksha Meghawat, Louis-Philippe Morency, and Eric P Xing. Select-additive learning: Improving generalization in multimodal sentiment analysis. In 20/7 IEEE International Conference Multimedia and Expo (ICME), pp. 949-9544 IEEE, 2017.
Haohan Wang, Zexue He, Zachary C Lipton, and Eric P Xing. Learning robust representations by projecting superficial statistics out, arXiv preprint wrXiv:1903:062555 2019.
Garrett Wilson and Diane Cook. A survey of unsupervised deep domain adaptation. ACM Transactions Intelligent Systems and Technology (TIST), 11(5):1-46, 2020.
Sewall Wright, Correlation and causation. J. agric. Res., 20:557-580, 1921.
Shen Yan, Huan Song, Nanxiang Li, Lincan Zou, and Liu Ren. Improve unsupervised domain adaptation with mixup training. arXiv preprint arXiv:2001.0066 2020.
Kun Zhang, Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Kernel-based conditional independence test and application in causal discovery. arXiv preprint arXiv: arXiv: 1202.3775, 2012.
Kun Zhang, Bernhard Schölkopf, Krikamol Muandet, and Zhikun Wang. Domain adaptation under target and conditional shift. In International Conference on Machine Learning, pp. 819-827. PMLR, 2013.
Kun Zhang, Mingming Gong, and Bernhard Schölkopf. Multi-source domain adaptation: A causal view. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 29, 2015.
Kun Zhang, Biwei Huang, Jiji Zhang, Clark Glymour, and Bernhard Schölkopf. Causal discovery data: Skeleton estimation and orientation determination. In
from nonstationaryyheerrggenen IJCAI: Proceedings of the Conference, volume 2017, pp. 1347. NIH Public Access, 2017.
Han Zhao, Remi Tachet Des Combes, Kun Zhang, and Geoffrey Gordon. On learning invariant representations for domain adaptation. In International Conference on Machine Learning, pp. 7523-7532. PMLR, 2019.
14

---

Published as a conference paper at ICLR 2022
A DOMAIN GENERALIZATION
The goal of domain generalization (DG) (Muandet et al., 2013) is OOD generalization: learning predictor that performs well at unseen test domains. Unlike domain adaptation (Pan & Yang. 2009; Ben-David et al., 2007: 2010; Crammer et al., 2008; Patel et al., 2015; Zhao et al., 2019; Wilson & Cook, 2020; Zhang et al., 2015), DG assumes that the test domain data are not available during training. One thread of DG is to explore techniques from kernel methods (Muandet et al., 2013; Niu al., 2015; Erfani et al., 2016; Li et al., 2017b). Muandet et al. (2013) propose kernel-based optimization algorithm that learns an invariant transformation by minimizing the discrepancy among domains and preventing the loss of relationship between input and output features. Another line of DG work is using end-to-end methods from deep learning: (a) reducing the differences of representations across domains through adversarial or similar techniques (Ghifary al., 2015; Wang et al., 2017; Motiian et al., 2017; Li et al., 20:8b;;); (b) projecting out superficial domain-specific statistics to reduce sensitivity to the domain ( (Wang et al., 2019); (c) fusing representations from an ensemble of models across domains (Ding & Fu, 2017: Mancini et al., 2018). Meta-learning can also be applied to domain generalization, by dividing source domains into meta-training and meta-test sets, and aiming for low generalization error on meta-test sets after training on meta-training sets (Balaji et al., 2018; Dou et al., 2019; Li et al., 201 Sa; 2019a;b). Recently, an extensive empirical survey of many DG algorithm (Gulrajani & Lopez-Paz, 2020) suggested that with current models and data augmentation techniques, plain ERM may be competitive with the state-of-the-art. It is worth noting that Sun et al. (2020) also propose an approach to learning latent causal factors for prediction. However, their assumptions over the underlying causal graph are restricted due to two reasons: 1) they only consider the scenarios when Z and Y are generated concurrently, which excludes the cases in which some part of Z could also be affected by Y in some manner; 2) They assume that the causal latent factors Zp and the non-causal latent factors Z, are independent when conditioning on E, that is, Zp 11. Zel E. In practice, during model learning, they actually further assume that Z4 11 Zj] E for any * j so that VAE could be leveraged to learn the model. In this sense, their approach has the same issue as the one in IVEE i.c., unable to deal with the montfactori cases. This point is also verified in Table 2, where our approach greatly outperforms theirs.
B VARIATIONAL AUTOENCODERS
We briefly describe the framework of variational autocncoders (VAEs), which allows us to efficiently learn deep aatent-variabl models and their corresponding inference models (Kingma Welling, 2013; Rezende et al., 2014). Consider simple latent variable model where X E IR stands for an observed variable and Z IR" for a latent variable. A VAE method learns a full generative model Po(X, Z) Pa X Z )o Z) and an inference model go X),, typically factorized Gaussian distribution whose mean and variance parameters are given by the output of neural network with input X. This inference model approximates the posterior pe Z X). where is a vector of parameters of the generative model, & a vector of parameters of the inference model, and po Z)) a prior distribution over the latent variables. Instead of maximizing the data log-likelihood we maximize its lower bound LVAR Φ):
log logpo( PopPo (X) 2 (ONE(O,0)) Φ) == (xx()) [logro X [] - KI. (1s 2 X | ()),
where we have used Jensen's 's inequality, and KL(-II-) denotes the Kullback-Leible divergence between two distributions.
C DERIVATION
C.I IVAE
In Section 2.1, the evidence lower bound of iVAE is defined by
LIVAE (0, Φ) :=Epp [T [Z X,U) [logps(X] X Z)] KL (gs X, U U) > (2(
=Epp [F x,U) [log X Z) log UU) log q [X,UU))
15

---

Published as a conference paper at ICLR 2022
C.2 SCORE MATCHING
In Section 4.1, we follow Hyvärinen (2005) and simple trick of partial integration simplify the evaluation of the score matching objective LSM in Eq. (10) of the main text:
phasel
phisel (T) Epp (aa(2ix.Y
[|| log 4 X, E)
Vz lo pT (Z)Y, EX]"
Ep
ExgiZiX.YEE)
Σ
8 YZ, E
1
θPT.A E)
+ const.
azz
2
az,
where the last equality is due to the Theorem Hyvärinen (2005).
D DEFINITIONS FOR SEM AND IRM
Definition 2. A structural equation model (SEM) M == (S, governing the random vector Z (Zi...., Za) is a of structural equations:
S; Z4 & (P(P((.,) No),
where Pa(Zi) S [Z)..... (Z.) are called the parents of Zi, and the N; are independent noise random variables. We say that "Z; causes Zj" " if Zi E Pa(Z)), We call causal graph of Z to the graph obtained by drawing i) one node for each Zi, and ii) one edge from Zi to Z, if Z, E Zz)) We assume acyclic causal graphs.
Definition 3. Consider a SEM : (S. An intervention e on M consists of replacing one or several of its structural equations to obtain an intervened SEM M == (S (8), with structural equations:
S Z (P (ZF),
The variable Ze is intervened S; + St N₂ + No.
Definition 4. Consider structural equation model (SEM) S governing the random vector (Zy..... Ln. Y) and the learning goal of predicting Y from Z. Then, the set of all environ- ments (S)(S) indexes all the interventional distributions PIZ., Y") obtainable by valid interven- tions e. intervention e E (S)(S) valid as long as (i) the causal graph remains acyclic, (ii) E ((Y)))((Y) = E [ (Y),, and (iii) V (Yo(Pa(Y)) remains within a finite range.
E DEFINITIONS AND LEMMAS FOR THE EXPONENTIAL FAMILIES
Definition 5. (Exponential family) A multivariate exponential famity is set of distributions whose probability density function can be written as
((Z)
p( 2
(13)
C(A)
where - is the base measure, C(0) is the normalizing constant, T 2 R is the sufficient statistics, and f E R&& is the natural parameter. The size k 2 is the dimension of the sufficient statistics T and depends the latent space dimension Note that k is fixed given n. Definition 6. (Strongly exponential distributions) multivariate exponential family distribution
Q.Z))
p(Z)
C(0) (xp((T(Z),0))
(14)
is strongly exponential, if
(30 E R s.t. (T(Z). 0) = const, VZ 3) (I(Z) 0 0r where is the Lebesgue measure.
AZ C R"
(15)
The density of strongly exponential distribution has almost surely the exponential component and can only be reduced to the base measure on a set of measure zero. Note that all common multivariate exponential family distributions (e.g. multivariate Gaussian) strongly exponential.
F DEFINITIONS FOR IDENTIFIABILITY
Definition Let the domain of the parameters = (f. Let be an equivalence relation on θ. deep generative model is said to be ~ identifiable if
po (x Pa (x
~~6
(16)
16

---

Published as a conference paper at ICLR 2022
(b)
(d)
E
E
E
(a)
(h)
(k)
Figure 3: (a) General causal structure over (X,, X, E, where the arrow from Z, to X is must-have connection and the other four might not be necessary. (b) Ten possible causal structures from (a) under Assumptions 1&2.
The elements in the quotient space θ are called the identifiability classes.
Definition 8. Lel ~ be an equivalence relation on defined by: (X)) = A (x)) t AX E EO, (17)
(f, ~ (f, = 3A, s.t. s.t. T(f where A E Rxk an invertible matrix and E R is vector.
Definition 9, Let ~p be an equivalence relation on defined by: (f. T.T.A X) rup (f.T,A) > 3P, s.t. T(f (X)) =
+ vX E E0, (18)
where P Bkxk is block permutation matrix and c R is vector.
G PHASE 2: DISCOVERING SINGLE CAUSE
G.1 SOME SPECIAL CASES IN MULTI-CAUSE SETTINGS
These conditional independence tests can be performed in parallel to largely accelerate the testing procedure. Note that, practice, it might occur that there exist some Z, which is independent of any other Z, when conditioning on Y and E. It is probably because such Z, is deterministic transformation of Y. In this special case, 2012) to determine whether or not Z; a cause Y. Also note that, in some scenarios in which the dependence signals between pair of causal latent variables might be weak due to the data issue, we can test the conditional independence of such a latent variable with all the other causal latent variables. If it is conditionally dependent on more than half of them or its average p-value is larger than the pre-specified threshold, will select it as one cause of Y.
IGCI (Daniusis et al., 2012; Janzing al.,
G.2 DISCOVERING SINGLE CAUSE
In the single-cause case, by following Wang al. (2014), we leverage the MB-by-MB algorithm first construct a local structure around Y and then discover the single parent Y according to the constructed local graph. One obvious advantage of this approach is in efficiency, because there is no need construct the whole causal graph containing all the latent variables and Y
We could have an even more efficient solution to the single-cause case in some special scenarios where we assume that Z, IL Z, IY, E for any i o j. In fact, this assumption covers more scenarios than the common assumption that Z, JL Z; for i + in latent variable models, e.g., disentanglement (Bengio al., 2013; Locatello al., 2019), autoencoders (Kingma & Welling, 2013; Rezende et al., 2014), ICA (Comon, 1994; Hyvärinen & Oja, 2000), etc. Y is caused by at most one Z, and Zj. and no matter whether Z, and Z, are caused by E or not, then Z; JL Z, Y, E holds, but may well have Z, J.. Z, (e.g., if Y causes both Z, and Z,, or there chain Z, - Y + Zz)) Y causes or is caused by at most one of Z, and Zj. and at most one of Z; and Zj is caused by E. then both Z; 11 Z, and Z, IL Z, Y. E hold.
Y.. Z, Z,} form collider Z, Y Zy, and no matter whether 2, and Z, are caused by or not, then Z; 2,1 E hold, but we may have Z, J Z, (e.g., when both Z, and Z, are caused by E).
17

---

Published as a conference paper at ICLR 2022
Under this assumption, are able to separately look into each Z, given Y and E, without consider- ing any other Zj. Fig. 5a shows that there exist only five possible connections between Z; Y, E, and X. Among them, only the arrow from Z; to X must exist, because X is generated from Z. The other four arrows might not be present, with the exception that there must be at least one connection between Z₁ and Y or E (Assumption la). This leaves ten possible structures, shown in Figs. 3b-3k.
Given data (2,, Y. E, } in which the 2, are obtained using lø 2 X, Y. E) (for example, as given by the mean of this distribution), we are able to distinguish all structures in Figs. 3b-3k by using causal discovery algorithms (Peters et al., 2017; Zhang et al., 2017; Huang et al., 2020) and performing conditional independence tests (Spirtes et al., 2000; Zhang et al., 2012). This is summarized in Proposition below. Its proof can be found in Appendix H. which also describes the specific assumptions made. In practice, we can assess in parallel whether or not each Z, is direct cause of Y, which accelerates this phase significantly.
Proposition 2. Under the assumptions stated in Appendix H, the ten structures shown in Figs. 3b-3k can be identified using causal discovery methods consistent in the infinite sample limit.
Note that there are only four cases in which Z; is a parent of Y (i.e., Figs, 3b, 3e, 3g, and 3j). We can identify these by applying rules 1.2, 1.6, 2.1, and 3.1 from Appendix H.5.
H PROOFS
H.1 PROOF OF THEOREM 1
The proof of this theorem consists of three parts.
In Part I, we prove that the parameters are ~ identifiable (Definition 8) by using assumption (i), the first half of assumption (ii), and assumption (iv) of Theorem 1.
In Part II, based on the result in Part I, we further prove that the parameters are ~>p identifiable (Definition 9) by additionally using Assumption 2, the second half of assumption (ii) and assumption (iii) of Theorem 1.
In Part III, we combine the results (Theorems 4&5) in both Part and Part into one theorem (Theorem 1), which completes the proof.
is worth noting that, compared the proof iVAE, the main changes our proof consist of
Part I, In step III.
has been updated to account for vectors of sufficient statistics whose entries can be arbitrary functions of all entries in the random variable vector, while in the previous proof the sufficient statistics contained entries that are functions of individual entries in the random variable vector.
The assumption of "The sufficient statistics in T are all linearly independent." is not required in proof, but it is in the proof of iVAE.
Part II.
has been updated to account for the part of the sufficient statistics which is the output of a deep neural network with ReLU activation functions.
H.1.1 PART I
For notational simplicity, in the proof of this part we denote (Y, EY by U. Hence, our generative model defined according to Eqs. (6-8) in the main text now becomes:
po X, Z UU) py
(19)
Pf (X Z p.(X f()))
(20)
P Z ( Q(Z)/((U) exp [T())
(21)
Theorem 4. Suppose that observe data sampled from deep generative model defined according Eqs. (19-21) with parameters (f. J, A). Assume that
(i) The set (X E Ol (X) = 0} has measure zero, where Pe is the characteristic function of the density Pe defined in Eq. (20);
18

---

Published as a conference paper at ICLR 2022
(ii) The mixing function f in Eq. (20) is injective;
(iii) There exist l points U ...... ,U E U such that the matrix
L [-((U')) - - ER
(22)
is invertible.
Then the parameters (S.T. }} are ~A identifiable.
Proof. Define vol( det((B" B) for any full column rank matrix B. Suppose that we have two
sets of parameters 0 = (f. X) and 0 (f,T,
T. X) such that po X (U) Pá X Pa U), v(x,
O x U. We want show 0 ~ 0.
Step The proof of this step is similar to Step in the proof of Theorem in Khemakhem et al. (2020a). We transform the equality of the marginal distributions over observed data into the equality of noise-free distributions. For all pairs (X,U) E O U. we have
Po (X Po (XU)
(23)
Pj (X 2))7,AZZ
(24)
- / palx ((())
Px -(Z))PTX
(25)
= pc(( (1 vol(J, (X)))X / Pc(( X) vol j
(x
(26)
/ Pe X = / Pe(X {)P)
(27)
(PITAU (X) = (0j (X)
(28)
F'rrmuullw F10g,....((()))(())
(29)
FliSTAUl(w) = -F0G (X) = -
(30)
(31)
where
in Eq. (26), denotes the Jacobian, and we made the change of variable = f(Z) on the left hand side, and X = f(Z) on the right hand side.
in Eq. (27), we introduced
Py.T,..(() é (f (X) vol(Js (x))ll((X)
(32)
on the left hand side, and similarly on the right hand side.
in Eq. (28), we used for the convolution operator.
in Eq. (29), we used F[.] to designate the Fourier transform, and where Pe Flye] (by definition of the characteristic function).
in Eq. (30), dropped (e(w) from both sides it is non-zero almost everywhere (by assumption (i)).
Step II. In this step, remove all terms that either function of X or U. First, by replacing both sides of Eq. (31) by their corresponding expressions from Eq. (32), we have
(f vol(J, (X)) = Pr. U) vol(J
(33)
Then, by taking logarithm on both sides of Eq. (33) and replacing PT. A by its expression from Eq. (21), we obtain
log vol vol(J--((X)) + log Q(f- log (()) ( (T(
log vol(J (X)) log (1) (1)00000000000000000000000011010
(34)
Let U U1 U'.... U E l4 be the + points defined in assumption (iii). We evaluate the above
equation these points obtain is + equations, and subtract the first equation from the remaining
19

---

Published as a conference paper at ICLR 2022
k equations obtain:
2(09)
X(U*) X(U))) log 2(U7)
2(U0)
= - (x)), X(U)) + log
1= 1,
k. (35)
7(U7))
Let L be defined as in assumption (iii) and L defined similarly for X. Note that L is invertible by assumption, but L is not necessarily invertible. Letting b E IK* in which by log 2 )U0))
) z(U')
we
have
L' = L' +b Left multiplying both sides of the above equation by L-T gives T(f TX) = AT)) +c where A = L- E R and L R
(36)
(37)
Step III. To complete the proof, we need to show that is invertible. Let Zi E 2, X = fZz), (= = 0, ... k. We evaluate Eq. (37) at these k + 1 points to obtain k + equations and subtract the first equation from the remaining equations obtain
( T())), T(z)) T())
RER
= -TF (.)) -Tf (38)
RRER*
We need to show that for given Zo Z. there exist & points Z.. Zx E Z such that the columns of are linearly independent. Suppose, for contradiction, that the columns R would never be linearly independent for any choice of I..... Z5 E Z. Then the function g(Z) ZZ) & T - T())) would live in I
lower dimensional subspace, and thus could find non-zero vector .. A) = O and thus E R orthogonal to that subspace. This would imply that (T - T(o)). -(Zo) (T A) = (T(Zn), A) const, VZ E 2, which contradicts the assumption that the prior strongly exponential (Definition 6). Therefore, we have shown that there exist
points
2 such that invertible. Since AR and A is not function of Z, A must be invertible. This completes the proof.
Zo Z..... Z
H.1.2 PART II
Theorem 5. Suppose that all assumptions in Theorem 4 hold. Let the sufficient statistics
= = ()) TZ given by the concatenation of a) the sufficient statistics (Z) = (T( T, (Z,) T of a factorized exponential family, where all the Tx(Z,) have dimension larger equal to 2, and b) the output N (Z) of a neural network with ReLU activations. (note that a neural network with ReL. activation has universal approximation power and should able to capture any dependencies of interest). Let ke be the dimension of T, and thus that k' 2 2n. Assume that
(i) the sufficient statistics T have all second-order own derivatives;
(ii) the mixing function f has all second-order cross derivatives.
Then the parameters (f.T.X A} are rup identifiable.
Proof. Let v o Z. Since assumptions in Theorem 4 hold, we have T(Z) = AT (z()) + c,
(39)
where A BRak is invertible. We want to show that A is block permutation matrix.
20

---

Published as a conference paper at ICLR 2022
Step I. In this step, we show that e is componentwise function. First we differentiate both sides of Eq. (39) with respect to Zg and Ze (8 + 1) to obtain
0T(Z) = =A pol Z az.
a (T Z )
(Z)
(40)
OZ.
OZ
8T
(2y(2) (z)
(T((( ()) Po)(Z) (41)
A
02,04. ==
ou, Z Joy (Z)
oz,
oz, ==1 Du, (Z)
oz,
By construction, the second-order cross derivatives of T and T are all zero. Therefore, we have
(T))
(z:
(v;(Z)
(v ())
(z)
0 =
(42)
du; 2)² az.
az.
(2) 8z,07.
1=1
The above equation can be written in the matrix vector form:
0 AT (Z)'', (Z) + AT (Z)o",((),
(43)
where
O (v) O'T(2( (z)
(T ())
(Z)
E R
(44)
dm (2)²
dun (Z)²
VZ) = 001((2)
(Z)
(Z)
(v)
T
E ER",
(45)
82.
00.
oz.
OZ.
and
(TT(((
(T()(
(2)
E ERax
(46)
O0 (Z)
(Z)
T821(2 (Z)
(2) T
(Z) = 82,0Z.
E ER"
(47)
02,0Z.
Now by concatenating 7' = [T" (T) E 2x22 and v" (Z) =
"2, E R
(48)
we obtain
AT Z (Z)
(49)
Finally, we take the rows of T (Z) that corresponds to the factorized strongly exponential family distribution part and denote them by T (Z) E RR x2n, By Lemma 5 the iVAE paper (Khemakhem
et al., 2020a) and the assumption that k' 2n, we have that the rank of Tyr(() is 2n. Since k 2
k' 2 2n, the rank of T"" (Z) is also 2n. Since the rank of A is k the rank of ATT (Z) E R is 2n. This implies that (Z) must be a zero vector. In particular, we have that v (Z) 0, Vs + t. Therefore, we have shown that v is a componentwise function.
Step II. To complete the proof, we need show that A is a block permutation matrix. With- out loss of generality, we assume that the permutation in v is the identity. That is v(Z) = [vi (Zi),... ".,..,,,,,,,) (Zn)) for some nonlinear univariate scalar functions Vi, "., Gn. Since f and if are bijective, we have that v is also bijective and v (Z) = [o (Z)) ..... e, (Z.)).. We denote (Z)) = T(v( 2)) + A and plug it into Eq. (39) to obtain T(Z) = AT)) Applying - to the variables Z both sides gives
T(e T()= Let be the index of an entry in the sufficient statistics T that corresponds to the the factorized strongly exponential family distribution part Ty. For all s + t, we have
AT
(50)
0 (T (Z)): a7,
Σ OT(Z), 1=1
axj 07,
(51)
Since the entries of T are linearly independent (if they were not linearly independent, then T can be compressed into a smaller vector by removing the redundant entries), we have that aij is zero for any
j such that 8T(Z), 82," + O. This includes the entries j in the sufficient statistics T that corresponds to 1) the factorized strongly exponential family distribution part which do not depend on Z,; and 2) the neural network part.
21

---

Published as a conference paper at ICLR 2022
Therefore, when & is the index of an entry in the sufficient statistics T that corresponds to factor in the factorized strongly exponential family distribution part Ty, the only non-zero úij are the ones that map between (z,) and (v (Z.)). where T; are the factors in Ty that only depend on Z and T; defined similarly. Therefore, we can construct an invertible submatrix A3 with all non-zero elements azj for all that corresponds to factor such that
(52) where T. are the factors in Ty, that only depends on Zi and e; are the corresponding elements of c. This means that the matrix A is block permutation matrix. For each i = 1, the block A; of A affinely transforms T, (2,) into T. (ex()). There is also an additional block A' which affinely transforms (Z) into ()). This completes the proof.
(2,)) = (1;
=
+ c.. 1,
H.1.3 PART III
Now we combine Theorem 4 in Part I and Theorem 5 in Part II into one theorem, which completes the proof of Theorem 1.
H.2 PROOF OF THEOREM 2
We recall that the loss function in Phase is as follows: Cpbase (0, Φ) <pisel (f. T. A.A. Φ)
LSM phuse!
X, Φ).
(53)
where
CBLBO (0. (b) :Bp [ [ (oggy (X Z) + log PT, ( E)
log d0 (Z Y. E
(54)
SM
(T,)) :Bp, pp
[ V log gol Z Y E
Vz log Z| E) [1]
(55)
Proof. If the family of (Z E) is flexible enough contain pe (Z Y, E). then by optimizing the loss over its parameter Φ, we will minimize the score matching term LSM phase1 in Eq. (55), which will eventually reach zero. If we assume that the model is not degenerate and phusel = implies that Vz log 40 Z| Y, E) and Zz log Pr. (Z)Y, E) are equal. This implies that log d [ X Y. E) = log
that 1s > O everywhere, then having that LSM
Pr x Z Y E) + for some constant c. But e is necessarily 0 because both go 7 X Y, E) and Pr x Z Y E)) are pdf's. Therefore, the ELBO term phusel in Eq. (54) will be equal to the log- likkelihoood, meaning that the loss Lphusel in Eq. (53) will be equal to the log-likelihood. Under this circumstance, the estimation in Eq. (53) inherits all the properties of maximum likelihood estimation (MLE). In this particular case, since our identifiability is guaranteed up to permutation and componentwise transformation, the consistency of MLE means that converge to the true parameter 0" up permutation and componentwise 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 componentwise transformation does not change the proof but only the conclusion.
H.3 PROOF OF THEOREM 3
Proof. Theorem and Theorem 2 guarantee that in the limit of infinite data, NE-iVAE can learn the true parameters 0* := (f+, T X*) up to a permutation and componentwise transformation of the latent variables. Let (f: P) be the parameters obtained by NF-iVAE. We, therefore, have (f.T.A) Å) ~P (f",T", A*), where
denotes the equivalence up to permutation and componentwise transformation. If there were no noise, this would mean that the learned f transforms X into 2- (X) that are equal to Z (f+)-1 (X) up to permutation and componentwise transformation (Definition 9). If with noise, we obtain the posterior distribution of the latent variables up to an analogous indeterminacy.
H.4 PROOF OF PROPOSITION 1
Proof. Theorem A.I in (Arjovsky, 2021) has showed that i) any predictor w with optimal OOD generalization uses only Pa(Y) to compute t; ii) the classifier w in this optimal predictor can be estimated using data from any environment e for which the distribution of Pa(Y has full support;
22

---

Published as a conference paper at ICLR 2022
iii) the optimal predictor will be invariant across Eaul In iCaRL, the hypotheses of Theorems 1 and and Assumption guarantee that Pa(Y) can be recovered by first identifying the latent variables Z from X, Y and E and then discovering the direct causes of Y through solving Eq. (12). Furthermore, since the conditional prior in Eq. (5) of the main text has full support, the distribution of Pa(Y always has full support. Also, under Assumption I we have that p(Y (Pa(Y) )) is invariant across Eall- Hence, the classifier in this optimal predictor can be estimated using data from any environment e. We therefore have that the resulting optimal predictor will be invariant across Eall- This completes the proof.
H.S PROOF OF PROPOSITION 2
Proof. The following rules can be independently performed to distinguish all the 10 structures shown in Figs. 3b-3k. For clarity, we divide them into three groups. Note that, since these rules rely on different algorithms of causal discovery and conditional independence tests, unless stated otherwise, we assume that the assumptions of these algorithms are satisfied during the proof process.
Group All the six structures in this group can be discovered only by performing conditional independence tests.
Rule 1.1 If Z; IL Y. Z 1 E. and Y, then Fig. 3d is discovered.
Rule 1.2 If Z3 J Y, Z E, and V Y. then Fig. 3g is discovered.
Rule 1.3 If Z; Jí. Y, Z V/. E, and Y. then Fig. discovered.
Rule 1.4 If If jL Y. 2; 1 E, E 1 Y. and Z; JL Y|E, then Fig. 3i is discovered.
Rule 1.5 If 1 Y. 1 E. E 1 Y. and Z, 1. E|Y, then Fig. 3h is discovered.
Rule 1.6 If Z; H Y, Z 4 E, E 1 Y, and Y IL ElZ, then Fig. 3e discovered.
Group 2 1f J Y. Z; IL E. and IL Y, then we can discover both Fig. 3b and Fig. 3c. 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 et al., 2009) for continuous data and CDS (Fonollosa, 2019) for continuous and discrete data.
Rule 2.1 If Li J Y. 24 IL E, and E IL Y, and a chosen binary causal discovery algorithm prefers Z; + Y to Z - Y. then Fig. 3b is discovered.
Rule 2.2 If 2 ₁ Y, Z, 11 E, and E IL Y, and a chosen binary causal discovery algorithm prefers Z, <- Y 7, Y. then Fig. 3c is discovered.
Z, 1 Y Z, 1 EYY, and Y IL EIZ,, then we can discover both Fig. 3j and Fig. 3k. These two structures cannot 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 Z, and Y are affected by E. Fortunately, Zhang et al. (2017) provided heuristic solution to this case based on the invariance of causal mechanisms, i.e., P(cause) and P(effect|cause detailed deseription 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 Z; and Y in Fig. 3j and Fig. 3k is finally reduced to calculating the following term:
Group 3
If ZY
change independently. The
P(Y
AZ log (P(Y (2)))
(56)
where (-) denotes the sample average, P(Y [2))) is the empirical estimate of P(Y on all data
points, and (P Z denotes the sample average of P which is the estimate of P(Y 22)
in each environment. We take the direction for which A is smaller to be the causal direction.
23

---

Published as a conference paper at ICLR 2022
Rule 3.1 If Z Tf Y. Z. K E, E #L Y. Z, 1 Y Z; t EY. Y H E Z, and Az,-Y is smaller than Ay +7,, then Fig. 3j is discovered.
Rule 3.2 If Zi Y Y. Z y E, E rf Y. 2, 1 Y, Zi 1 E Y y E and Ay is smaller than A 7,,Y, then Fig. 3k is discovered.
I ILLUSTRATIONS FOR MODEL LEARNING
As described in Section 3.3, in the ground truth model, the prior for each Zie,, is either p( p(Z(E) E) or p( (. depending on whether Z; caused by or not. By contrast, in the NF-iVAE model the prior is p( YY. E). Is this going to affect the identifiability of the latent variables? Well, in practice not because the posterior distribution for Z given X, Y and E would be equivalent both models (up to identifiability guarantees).
J DATASETS
For convenience and completeness, we provide descriptions of Colored MNIST Digits and Colored Fashion MNIST here. Please refer to the original papers (Arjovsky et al., 2019; Ahuja et al., 2020a; Gulrajani & Lopez-Paz, 2020; Venkateswara et al., 2017) for more details.
J.1 SYNTHETIC DATA
For the nonlinear transformation, we use the MLP:
Input layer: Input batch (batch size, input dimension)
Layer Fully connected layer, output size = 6, activation ReLU
Output layer: Fully connected layer, output size 10
J.2 COLOREED MNIST DIGITS
We use the exact same environment as in Arjovsky et al. (2019). Arjovsky et al. (2019) propose to create environment for training to classify digits in MNIST data where the images in MNIST are now colored in such a way that the colors spuriously correlate with the labels. The task is to classify whether the digit is less than 5 (not including 5). There are three environments (two training containing 30,000 points each, one test containing 10,000 points) We add noise to preliminary label ( 0 if the digit is between 0-4 and ğ if the digit is between 5-9) by flipping it with 25 percent probability to construct the final labels. We sample the color id by flipping the final labels with probability Pe, where is 0.2 in the first environment, 0. in the second environment, and 0.9 in the third environment. The third environment is the testing environment. We color the digit red if 2 1 or green if 2 = O.
J.3 COLORED FASHION MNIST
We modify the fashion MNIST " in manner similar 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 the images in such way that the colors correlate with the labels. The task is to classify whether the image that of foot wear or clothing item. There are three environments (two training, one test) We add noise to the preliminary label (ÿ 0: "t-shirt", "trouser", "pullover", "dress". "coat", "shirt" and ý = 1: "sandal". "sneaker". "ankle boots") by flipping it with 25 percent probability to construct the final label. We sample the color id z by flipping the noisy label with probability pe, where pe 0.2 in the first environment, 0.1
10 hhtpps: http::/://wwwwteaarr..oom
tensor htwp://www.ttnoooooooooooo/pp/p/ http:///wwaaaaaaoo/////////////a
Load data 11 "https::// mni st/ooad_ddaa
24

---

Published as a conference paper at ICLR 2022
Table 3: Results on synthetic data. Comparisons are in terms of MSE (mean - std deviation).
x
METHOD
TRAIN
TEST
(dg (0.2,2)
(Ca 100)
ERM
0.00
0.00
0.00 0.00
Identity
IRM
0.00
0.00
0.00 + 0.00
F-IMM GAME
0.98 0.23
1.03 0.04
V-IRM GAME
0.99 2.74
1.07 2.26
iCaRL (ours)
0.01 0.03
1.00 +.001
ERM
0.00
0.00
0.00 + 0.00
Linear
IRM
0.00 0.00
0.00 + 0.00
F-IRM GAME
0.99 0.01
1.08 0.06
V-IRM GAME
1.00 5.98
1.05 + 0.04
iCaRL (ours)
0.01 0.03
0.04
FRM IRM F-IRM GAME
0.06 0.01
220.79 + 229.97
Nonlinear
0.08 0.01
149.60
104.85
1.06 0.09
196.59
150.71
V-IRM GAME
1.00 0.01
170 46 125.62
iCaRL (ours)
0.29 0.04
28.16 + 2.54
in the second environment, and 0.9 in the third environment, which is the test environment. We color the object red if or green if = 0.
K IN-DEPTH ANALYSIS ON SYNTHETIC DATA
K.1 COMPARISONS WITH SFATE-OF-THE
We first conduct series of experiments on synthetic data generated according to an extension of the SEM in Model 1. The extension is to map the variables Z := (Z1, Z2) into 10 dimensional observation X through linear or nonlinear transformation. Our goal is to predict Y from X, where X = g( (2) We consider three transformations: (a) Identity: g(.) is the identity matrix I E R
Z. (b) Linear: 9()) is random matrix S 2R i.e., X = (Z Z S. (c) Nonlinear: g(-) is given by a neural network with 2-dimensional input and 10-dimensional output, whose parameters are randomly set in advance. Since this is regression task, we use the mean squared error (MSE) as a metric of performance. Note that, in this problem, there is only one causal latent variable Z1. meaning that the conditional prior in Eq. (5) will not exhibit dependencies. Because of this, in this case we do not include N (Z) term in our NF-iVAE conditional prior. In the following section we do consider settings with many potential causal latent variables and, in that case, do include N (Z) in the NF-iVAE prior.
Le., X == (Z
We consider a simple scenario in which we fix 01 = and 02 = for all environments and only allow 03 to vary across environments. In this case, d3 controls the spurious correlations between Ze and Y. Each experiment draws 1000 samples from each of the three environments 03 = {0.2,, 100), where the first two are for training and the third for testing. We compare with the following baselines: ERM, and two variants of IRMG: F-IRM Game (with + fixed to the identity) and V-IRM Game (with variable t).
As shown in Table 3, in the cases of Identity and Linear, our approach is better than IRMG but only comparable with ERM and IRM. This might be because the identifiability result up to simple nonlinear transformation renders the problem more difficult by converting the original identity or linear problem into nonlinear problem. In the Nonlinear case, the gains of iCaRL are very clear.
Table 4: Comparisons of assumptions on the prior leading to identifiability in different algorithms.
METHOD
Assumption on the Prior for Identifiability
VAE (Kingma & Welling, 2013)
Non-idennifiabilityy with pr (Z) II. p(z)) -2 I)
iVAE (Khemakhem et al., 2020a)
Pr.x Y, E) = II, (Z.)/C, E) expl_ c Tr, E)]
ICE- BeEBBeMM (Khemakhem et al., 2020b)
(Z E) = ((2) IL, exp D =1 Tz (Y,)) E)
NF-IVAB
A E) ((Z)/((Y. E) exp T XYY
AYY E)
We also tried ICP, but ICP was unable find any parent of Y even in the identity case.
25

---

Published as a conference paper at ICLR 2022
E u (0,2,2,3,53
(57)
21 N(ZI(E,))
(58)
21
22
22 ~ (222/2222
(59)
Y ~ + Z2:1)
(60)
X g(1 Z2)
(61)
(b)
Figure 4: (a) Causal structure with Y having two causes. (b) Data generating process corresponding (a), where 4d-3 denotes the discrete uniform distribution, .V(-) the Gaussian distribution, and g(-)
is given by a neural network with 2-dimensional input and 10-dimensional output, whose parameters randomly set in advance.
(b) VAE
(c)
(a) Data generating model
(d) ICE-BeeM
(e) ICaRL-NF-IVAE
Figure 5: Visualization the samples (i.e., 2 - (21, 22)) in latent space generated through different algorithms. (a) Samples from the true distribution. (b-e) Samples from the posterior of different algorithms. Apparently, our method (e) can recover the original data up to a permutation and simple componentwise transformation.
K.2 VISUALIZATION OF IDENTIFIABILITY OF NE-IVAE
To further verify identifiability of NE-iVAE, conduct series of experiments on synthetic data generated according to the data generating process (Fig. 4b) corresponding to the causal graph shown in Fig. 4a. The reason we choose this setting is that it is the simplest case satisfying requirements: a) For ease of visualization, the latent space had better be 2-dimensional; b) To introduce the non- non-fcctor prior given Y and E (i.e., Z, V Z, Y E), Y has at least two causes.
We draw 1000 samples from each of the four environments E 20 2, 5), and thus the whole synthetic dataset consists of 4000 samples. We compare with the following baselines: VAE Kingma & Welling (2013) (without identifiability guarantees). iVAE Khemakhem et al. (2020a) (with conditionally factorized prior for identifiability), and ICE BeeM Khemakhem et al. (2020b). is worth noting that in ICE-BeeM the primary assumption leading to identifiability is similar to that in iVAE, where the base measure Q Z)) could be arbitrary to capture the dependences between latent variables but the exponential term still has to factorize across components (dimensions). All these are summarized in Table Clearly, from the table we can see that our method has more general assumption on the prior leading to identifiability. This is also demonstrated empirically in Fig. 5. Our method iCaRL can recover the original data Z up a permutation and simple componentwise
26

---

Published as a conference paper at ICLR 2022
Table 5: Colored MNIST. Comparisons in terms of
accuracy (%) (mean std deviation).
METHOD
TRAIN
TEST
ERM
84,88 1 0.16
10.45 0.60
ERM
84.84 0.21
10.86
0.52
ERM
84.95
10.05
ROBUST MAX
84.25 0.43
15.24 2.45
F-IMM GAME
63.37 + 1.14
59.91 1 2.60
V-RRM GAME
63.97 + 1.03
49.06 + 3.43
IRM
59.27 + 4.39
62.75
+ 9.59
iCaRL (ours)
70.56 + 0.81
68.75
1.45
ERM GRAYSCALE
71.81 + 0.47
71.36 + 0.65
OPTIMAL
75
75
Table 6: PACS, Comparisons in terms of accuracy (%) (mean + deviation).
METHOD
TEST
ERM
85.7 0.5
IRM
84.4
DRO (Sagawa et al., 2019)
84.1 = U.4
Mixup (an al., 2020)
84.3 1 0.5
CORAL (Sun Sacnko, 2016)
86.0
MMD (Li et al., 2018h)
85.0 100+-00 0.2
DANN (Ganin et al., 2016)
$4.6 = 1.1
C-DANN
al., 2018c)
82.8
LaCIM (Sun et al., 2020)
83.5 1.2
iCaRL (ours)
88.7 + 0.6
transformation, whereas all the other methods fail because they are unable handle the case in which Z; LL Y, E.
L IN-DEPTH ANALYSIS ON MORE REALISTIC DATA
L.1 COLORED MNIST
We compare iCaRL with 1) IRM, 2) two variants of IRMG: F-IRM Game (with & fixed to the identity) and V-IRM Game (with variable +). 3) three variants of ERM: ERM (on entire training data), ERM e (on each environment e), and ERM GRAYSCALE (on data with no spurious correlations). and 4) ROBUST MIN MAZ (minimizing the maximum loss across the multiple environments). Table shows that iCaRL outperforms all other baselines on Colored MNIST. However, this dataset seems more difficult because even ERM GRAYSCALE, where the spurious correlation with color is removed, falls well short of the optimum.
L.2 PACS
We report the results on another one of the widely used realistic datasets for OOD generalization:
PACS (Li et al., 2017a). This dataset consists of 9, 991 images of dimension (3, 224, 224) and 7 classes from four domains: art, cartoons, photos, and sketches. We used exact experimental setting that is described in Gulrajani & Lopez-Paz (2020). We provide results averaged over all possible train and test environment combination for one the commonly used hyper-parameter tuning procedure: train domain validation. As shown in Table 6, iCaRI. achieves state-of-the-aar. performance when compared to those most popular domain generalization alternatives.
27

---

Published as a conference paper at ICLR 2022
M IMPLEMENTATION DETAILS
M.1 JOINT TRAINING
As described in Section 4.1, we can jointly learn (0, Φ) by optimizing the following objective:
Lehasel (0, ()) -LLLOO (f. S,,,,A X, Φ)
CSM
(f: T, X, Φ)
(62)
-Epo (2
log (X b oo Y E) log (o B3]
BpD [B 5 8)
| log 44 Z X Y, E) V log Z E)|
(63)
-Ep
Bu (2) X R.)
loggy (og 2) log Pr.
E)
log go ( Y, B)
- Ep
B (
Σ
8 P (Z
E
OpT A E)
07,
2
07g
+ const.
(64)
where the last equality is due to the equation in Appendix C.2, and f. S Φ are copies of f.T, X, Φ that are treated constants and whose gradient is not calculated during learning. In practice, f. T, Å, & can be easily implemented through either "detach" in PyTorch Paszke et al. (2019) or
in TensorFlow Abadi et al. (2015).
M.2 THE GENERAL NON-FACTORIZED PRIOR
In the experiments, the general non-factorized prior in Assumption 2 is implemented as follows: IZ E) = (NN(Z:araml) NN(Y, E: + concat(2, NN(Y, E: )
Ty)(2)
(2)
Ay
where (:.-) is the dot product of two vectors, and concat(-.) means the concatenation of two vectors. Now let us explain each term in details.
Firstly, concat(Z, 2²) is vector of the latent variables and their squared values, and NN(Y, E.ppaam)) is deep neural network parameterized by param3 that com- putes vector of natural parameters as a function of Y and E. Hence, the term (concat(2, 2²), NN(Y, E; p:ra)))) is equivalent to the factorized exponential family, which also satisfies that each (2,) has dimension larger or equal to 2.
Secondly, NN : parami) is neural network that receives input vector of latent variables and outputs another vector representing complicated nonlinear transformations of those variables. NN(Y, E; param2) is another neural network that generates corresponding vector of natural parameters. Hence, the term (NN Z; paraml), NN(Y, E; param))) will allow this prior to capture the dependencies between the latent variables Z.
N HYPERPARAMETERS AND ARCHITECTURES
In this section, describe the hyperparameters and architectures of different models used in different experiments. Unless stated otherwise, we have A and À = 1, both of which are selected on training/validation data.
N.1 SYNTHETIC DATA
We used Adam optimizer for training with learning rate set to 1e-3 and batch size set 128.
N.1.1
ERM
Linear ERM
Input layer: Input batch (batch size, input dimension)
Output layer: Fully connected layer, output size
28

---

Published as a conference paper at ICLR 2022
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
N.1.2 IRM
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, inpul dimension)
Layer Fully connected layer, output size = 6, activation ReLU
Output layer: Fully connected layer, output size
N.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
N.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 Fully connected layer, output size = 6, activation ReLU
Output layer: Fully connected layer, output size
Linear Classifier w
Input layer: Input batch (batch size, 2)
Output layer: Fully connected layer, output size
Nonlinear Classifier w
Input layer: Input batch (batch size, 2)
Layer Fully connected layer, output size = 6, activation ReL.U
Output layer: Fully connected layer, output size
29

---

Published as a conference paper at ICLR 2022
N.1.5 ICARL
NF-IVAE -L-nner Prior
Input layer: Input batch (batch size, input dimension)
Output layer: Fully connected layer, output size 4
NF-iVAE -yNonline Prior
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
NF-iVAE T Prior
Input layer: Input batch (batch size, inpur dimension)
Layer I: Fully connected layer, output size 6. activation = ReLU
Output layer: Fully connected layer, output size
NF-iVAE Awx-Nonlinear Prior
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
NF-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 2
NF-iVAE Nonlinear Encoder
Input layer: Input batch (batch size, inpul dimension)
Layer Fully connected layer, output size = 6, activation ReLU
Mean Output layer: Fully connected layer, output size =
Log Variance Output layer: Fully connected layer, output size
NF-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 1 is vector full of with the length of output dimension
NF-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 is vector full with the length of output dimension
30

---

Published as a conference paper at ICLR 2022
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 Fully connected layer, output size = 6, activation ReLU
Output layer: Fully connected layer, output size
N.2 CMNIST AND CEMNIST
Considering that the results of most baselines come from IRMG (Ahuja al., 2020a), for fair comparison, we follow the same setting of IRMG terms of hyper-parameters and validation considerations. For example, the batch size is set to 256, and the learning rate is 10-4. 4. We also did the environment data for validation. Please find more details Ahuja al. (2020a).
NF-iVAE -Prirr
Input layer: Input batch (batch size, input dimension)
Layer I: Fully connected layer, output size = 50, activation ReLU
Output layer: Fully connected layer, output size = 45
NF-iVAE ANx--Prior
Input layer: Input batch (batch size, inpul dimension)
Layer I: Fully connected layer, output size 50, activation ReLU
Output layer: Fully connected layer, output size 45
NF-iVAE -Prior
Input layer: Input batch (batch size, impul dimension)
Layer Fully connected layer, output size = 50, activation ReLU
Output layer: Fully connected layer, output size 20
NF-iVAE X-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 Convolutional layer, output channels = 32, kernel size = 3, stride = 2, padding = 1, activation = ReLU
Layer Convolutional layer, output channels 32, kernel size 3, stride 2, padding 1, activation = ReLU
Output layer: Flatten
NF-iVAE (Y, En--nooder
Input layer: Input batch (butch size, input dimension)
Output layer: Fully connected layer, output size 100, activation ReL.U
31

---

Published as a conference paper at ICLR 2022
NF-iVAE (X, Y, B)-Merrer//Eecoddr
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 10
Log Variance Output layer: Fully connected layer, output size 10
NF-iVAE Decoder
Input layer: Input batch (batch size, 10)
Layer I: Fully connected layer, output size = 32 x 4 x 4, activation = ReLU
Layer Reshape to (batch size, 32, 4, 4)
Layer 3: Deconvolutional layer, output channels 32, kernel size 3, stride padding
1, outpadding = 0, activation ReL.U
Layer Deconvolutional layer. output channels 32, kernel size 3, stride padding
1, outpadding = 1, activation = ReLU
Layer 5: Deconvolutional layer, output channels 2, kernel size 3, stride = 2, padding = 1,
outpadding =
Mean Output layer: activation Sigmoid
Variance Output layer: 0.01 1, where a matrix full of with the size of 2 28 x 28.
Classifier
Input layer: Input batch (batch size, 50)
Layer Fully connected layer, output size 100, activation ReLU
Output layer: Fully connected layer, output size 1, activation Sigmoid
N.3 VLCS AND PACS
used the exact experimental setting that is described Gulrajani & Lopez-Paz (2020). Specifically, we trained our model over all possible train and test environment combination for one of the commonly used
tuning procedure: train domain validation. We use ResNet-50 as an encoder and reverse the architecture of ResNet-50 decoder. We set the number of the latent variables 50. We do the hyperparameter search by exactly following the guides given in Gulrajani & Lopez-Paz (2020).
32


Reviews:
[{"review_id": "pFb7-XEaP17", "reviewer": "Reviewer_Ciwc", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "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": "strengths\n- a general underlying causal model that covers many real-world scenarios is peoposed\n- assumptions of the general causal model are clearly presented and explained\n- theoretical guarantee on the identifiability of the latent varibles is provided\n- the results of synthetic data experiments are carefully analyzed and can well corroborate the theoretical claims\n\n\nweakness\n- the algorithm seems a bit cumbersome\n- there appears to be a technical flaw in Phase 2 of the algorithm\n- some popular benchmark data sets are missing", "summary_of_the_paper": "This paper proposes Invariant Causal Representation Learning (iCaRL), which learns causal representation for OOD generalization and can be viewd as a generalization of IRM to the settings with both nonlinear representations and nonlinear classidiers. The problem is important but part of the algorithm seems not technically correct.", "summary_of_the_review": "This paper introduces an interesting idea about a general causal model, and presents related assumptions and theoretical results clearly. There seems to be some technical flaws in the proposed algorithm and the experimental results could be improved by adding some popular benchmark  real-world data set.\n\n\ndetailed comments\n- the general casual structure (Fig. 1(c)) lasks an edge from E to Y compared with Fig. 1(b). Will this difference affect its expressiveness such that some data-generating processes of Fig. 1(b) can not be characterized by Fig. 1(c)\n- when determining Pa(Y) in phase 2, the algorithm compares the p-values of independent testing IndTest(X_i, X_j) and conditional independence testing IndTest(X_i, X_j | Y). Since these two methods derive test statistic from differernt rationale, I was wondering whether the p-values of these two test can be directly compared to conclude the evidence of independence.\n- I was wondering whether phase 2 is able to distinguish Y's parents between the its ancestors. it seems an ancestor of a direct cause of Y will also be included in Pa(Y) after phase 2.\n- this paper only provedes the results on VLCS data set while existing works usually evaluate the prerformance on at least two real-world data set. It would be much better to include one more real-world data set (e.g. PACS)", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "ljfW4ODiUN", "reviewer": "Reviewer_rp1x", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "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": "Section 4.1 presents the theoretical foundation, with a complete and clear proof. This is a valuable contribution.\n\nSection 4.2 is less than rigorous: a) “one observation is that for any two latent variables Xi and Xj , only when both are causes of Y do we have that the dependency between them increases after conditioning on Y .” Note that Xi and Xj are actually correlated due to the confounder E, how to justify this claim? Or do you mean conditional on E? b) ”when there exist at least\ntwo causal latent variables, it is trivial to test all pairs of latent variables with independence testing (Gretton et al., 2007) and conditional independence testing (Zhang et al., 2012) to discover all of them by comparing p-values from these two tests.” I am not sure whether the p-values can be compared. Also, p-value itself is not a statistical measure. Please justify the methods in this section. \n\nSection 4.3 gives the algorithm wrt. both training datasets and new test data. Here is my main concern: I wonder whether this method is for DG and OOD generalization, or indeed a DA method. Note that “ Phase 3:  … When in a new environment, we first infer Pa(Y) from O by solving Eq. (12) and then leverage the learned w for prediction.” This means, the method requires knowledge of test domain data, although in a not very complicated way. To me, if test domain data are needed in adjusting the algorithm, then it should be a DA method. Note that this is different from simple normalization, as this operation can be included in the trained model. Further, what if we have a test dataset as a mixture of several domains, then is this method affected? Please clarify.\n\nMinor concerns and questions:\n1.\tThe causal graph is explained well wrt practical considerations. However, it is not clear how practical the assumptions in Theorem 1 and 2 are. For instance, is it true to have O=f(x) injective in VLCS datasets? Why is it needed to assume n <= d?\n2.\t“It also makes sense in Assumption 1c that the generative mechanism p(O| X) is invariant across all the environments. Otherwise, it is impossible to infer X from O in any unseen environment” Is it possible to assume only $P(O|X_pa)$ remains invariant? \n3.\tFootnote 7 is missing. Also maybe move some footnotes to the main text in later version (e.g., camera ready version if accepted)? Footnote kind of affecs reading flow.\n4.\tIn practical algorithms, how to select the dimension of X? \n5.\t(very minor concern) maybe to include one or two more real datasets like PACS?\n", "summary_of_the_paper": "This paper proposes invariant Causal Representation Learning (iCaRL) for OOD generalization in the nonlinear setting. The work extends iVAE to a somewhat more general setting and shows the direct cause of the target can be discovered. iCARL is then developed based on the direct cause. . Extensive experiments verify the effectiveness of the proposed method.", "summary_of_the_review": "In summary, the paper has adequate theoretical contributions in terms of identifiability. The algorithm part seems less justified. Also, the method appears to be a DA method, not for OOD generalization, thus authors should compare their algorithm to recent DA methods. I look forward to authors' response.\n\n== after reading response ==\nBased on the current response, I decide to increase my evaluation.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Blyz1IperJ5", "reviewer": "Reviewer_it5B", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "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": "**Strengths**\n\nThis paper has significant technical and empirical novelty. Moreover, it is very well written and communicates the central contributions clearly.\n\n**Weaknesses**\n\nOverall, I only have 2 issues with this paper. **First**, a very minor concern regarding readability. Since your solution builds off of the VAE literature in part, I think that it would be more readable to use $X$ for observed random variables and $Z$ for latent random variables, rather than $O$ and $X$ respectively. While $O$ makes sense since it is the first letter in observed, I think it is burdensome to ask computer science readers to disassociate it from its usage in \"Big O notation\".\n\n**Second**, and more significantly. I'm not convinced that \"Out-of-Distribution Generalization\" is a sensible term, and I'm also concerned that it is misleading. I think you clearly present one graphical model that defines one distribution. All environments are a part of that distribution, and I would argue that unseen environments are out-of-sample rather than out-of-distribution.\n\nFor me, out-of-distribution would correspond to a fundamental change in the relationship between variables. For example, let's say $Y$ is a binary variable indicating whether or not a sequence of characters $O$ signifies addition. Under one distribution the plus sign (+) (represented by some $X_p$) would have a strong causal association with $Y$. This would be an invariant causal association across whether the characters were printed, hand written, displayed on a screen, blue, purple, etc. If we never observed an orange \"+\" during training, under the assumption of the same distribution (graphical model) it makes sense that we would want to generalize to this *out-of-sample* environment. But what if we see samples from a different graphical model, say one in which $Y$ now signifies the logical or operation and under this distribution all plus signs are instead associated with that $Y$? If our original model saw samples from this distribution it would confidently and incorrectly classify such examples as signifying addition, thus failing to generalize to the new distribution. This is just one example of why I think it is incorrect to characterize what this method does as out-of-distribution generalization. I am open to further discussion about this point, but I think it would be more precise to use the term out-of-sample in place of out-of-distribution.", "summary_of_the_paper": "This paper proposes a general framework for so-called \"out-of-distribution generalization\" that can handle non-linear associations and causal links between variables. It provides proofs for identifiability and generalization. It proposes a general causal model that is sensible for many prediction problems. It proposes a conditionally non-factorized prior. It provides a tractable method that can be used in practice.", "summary_of_the_review": "This paper has significant technical and empirical novelty. Therefore I recommend acceptance. However, I think we need to have a discussion about whether Out-of-Distribution Generalization is a sensible term.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper sets up a complex algorithm for out-of-distribution generalization. The algorithm requires first, a generalization of identification results for variational autoencoders, the followed by second, a causal discovery subroutine, and third, learning an invariant predictor using the discovered causes. The procedure reads sound, and the results on common benchmarks look good, though I do not know how practical the approach would be in general.
Decision: Accept (Poster)