Year: 2021

Abstract:
Due to spurious correlations, machine learning systems often fail to generalize to environments whose distributions differ from the ones used at training time. Prior work addressing this, either explicitly or implicitly, attempted to find a data representation that has an invariant 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 prove that all direct causes of the target can be fully discovered, which further enables us to obtain generalization guarantees in the nonlinear setting. Extensive experiments on both synthetic and real-world datasets show that our approach outperforms a variety of baseline methods.

heading: 1 INTRODUCTION
text: 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 of classifying images of camels and cows (Beery et al., 2018). Here, the training dataset has a 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.e., it relates green pastures with cows and deserts with camels, and fails to recognize images of cows on 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 in training data collected from multiple environments, one can try to identify stable features and build invariant predictors. Invariant risk minimization (IRM) seeks to find data representations (Arjovsky et al., 2019) or features (Rojas-Carulla et al., 2018) for which the optimal predictor is invariant across all environments. The general formulation of IRM is a 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 et al., 2018, Theorem 4). Ahuja et 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 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 outof-distribution (OOD) generalization in the nonlinear setting (i.e., nonlinear representations and nonlinear classifiers1). 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 on the target (e.g., labels) and the environment (represented as an index), the prior over the data representation (i.e., a 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 is a more flexible conditionally non-factorized prior, which can actually capture complicated dependences between the latent variables. We then extend iVAE to the case in which the latent variable prior belongs to such a general exponential family. The combination of this result and the previous general assumption allows us to a guarantee that the data representation can be identified up to simple transformations. We then theoretically 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.
Overall, we make a number of key contributions: (1) We propose a general framework for out-ofdistribution generalization in the nonlinear setting with the theoretical guarantees on both identifiability and generalizability; (2) We propose a 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 1, 2&3); (5) We prove that our framework has the theoretical guarantees for OOD generalization in the nonlinear setting (Proposition 1).

heading: 2 PRELIMINARIES

heading: 2.1 IDENTIFIABLE VARIATIONAL AUTOENCODERS
text: Variational autoencoders (VAEs, see Appendix B) (Kingma & Welling, 2013; Rezende et al., 2014) lack identifiability guarantees. Consider a VAE model where O ∈ Rd stands for the observed variables (data) and X ∈ Rn for the latent variables. Khemakhem et al. (2020a) show that a VAE with an unconditional prior distribution pθ(X) 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, pθ(X|U), where U ∈ Rm is an additional observed variable (Hyvärinen et al., 2019). Specifically, let θ = (f ,T ,λ) ∈ Θ be the parameters of the conditional generative model pθ(O,X|U) = pf (O|X)pT ,λ(X|U), (1) where pf (O|X) = p (O − f(X)) in which is an independent noise variable with probability density function p ( ). Importantly, the prior pT ,λ(X|U) is assumed to be conditionally factorial, where each element of Xi ∈ X has a univariate exponential family distribution given U . The conditioning onU is through an arbitrary function λ(U) (e.g., a neural net) that outputs the individual exponential family parameters λi(U) for each Xi. The prior probability density thus takes the form
pT ,λ(X|U) = ∏
i Qi(Xi)/Zi(U) exp [∑k j=1 Ti,j(Xi)λi,j(U) ] , (2)
where Qi is the base measure, Xi the i-th dimension of X , Zi(U) the normalizing constant, Ti = (Ti,1, . . . , Ti,k) the sufficient statistics, λi(U) = (λi,1(U), . . . , λi,k(U)) the corresponding natural parameters depending on U , and k the dimension of each sufficient statistic that is fixed in 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, LiVAE(θ,φ) :=EpD [ Eqφ(X|O,U) [log pf (O|X) + log pT ,λ(X|U)− log qφ(X|O,U)] ] , (3)
where we denote by pD the empirical data distribution given by the dataset D = {( O(i),U (i) )}N i=1 and qφ(X|O,U) denotes an approximate conditional distribution for X given by a recognition
1In fact, we are not restricted to the classification case and allow the target to be either continuous or categorical, which will be formally defined in Section 2.2.
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 variablesX up to a permutation and a simple componentwise transformation, see Appendix F.

heading: 2.2 INVARIANT RISK MINIMIZATION
text: Arjovsky et al. (2019) introduced invariant risk minimization (IRM), whose goal is to construct an invariant predictor f that performs well across all environments Eall by exploiting data collected from multiple environments Etr, where Etr ⊆ Eall. Technically, they consider datasets De := {(oei ,yei )} ne i=1 from multiple training environments e ∈ Etr, where oei ∈ O ⊆ Rd is the input observation and its corresponding label is yei ∈ Y ⊆ Rs.2 The datasetDe, collected from environment e, consists of examples identically and independently distributed according to some probability distribution P (Oe,Y e). The goal of IRM is to use these multiple datasets to learn a predictor Y = f(O) that performs well for all the environments. Here we define the risk reached by f in environment e as Re(f) = EOe,Y e [`(f(Oe),Y e)], where `(·) is a loss function. Then, the invariant predictor can be formally defined as follows: Definition 1 (Invariant Predictor (Arjovsky et al., 2019)). We say that a data representation Φ ∈ HΦ : O → C elicits an invariant predictor w ◦ Φ across environments E if there is a classifier w ∈ Hw : C → Y simultaneously optimal for all environments, that is, w ∈ arg minw̄∈Hw Re(w̄◦Φ) for all e ∈ E , where ◦ means function composition.
Mathematically, IRM can be phrased as the following constrained optimization problem: min
Φ∈HΦ,w∈Hw ∑ e∈Etr Re(w ◦ Φ) s.t. w ∈ arg min w̄∈Hw
Re(w̄ ◦ Φ),∀e ∈ Etr. (4) Since this is a generally infeasible bi-leveled optimization problem, Arjovsky et al. (2019) rephrased it as a tractable penalized optimization problem by transfering the inner optimization routine to a 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., HΦ = Rn×n and Hw = Rn×1), under certain conditions on the diversity of training environments (Assumption 8 in Arjovsky et al. (2019)) and the data generation, the invariant predictor w ◦ Φ obtained by solving Eq. (4) remains invariant in Eall.

heading: 3 PROBLEM SETUP

heading: 3.1 A MOTIVATING EXAMPLE
text: 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 1. Consider a structural equation model (SEM) with a discrete environment variable E that modulates the noises in the structural assignments connecting the other variables (cf. Fig. 1a below): X1 ← Gaussian(0, σ1(E)), Y ← X1 + Gaussian(0, σ2(E)), X2 ← Y + Gaussian(0, σ3(E)), where Gaussian(0, σ) denotes a Gaussian random variable with zero mean and standard deviation σ, and σ1, . . . , σ3 are functions of the value e ∈ Eall taken by the environment variable E.
To ease exposition, here we consider the simple scenario in which Eall only contains all modifications varying the noises of X1, X2 and Y within a finite range, i.e., σi(e) ∈ [0, σ2max]. Then, to predict Y from (X1, X2) using a least-square predictor Ŷ e = α̂1Xe1 + α̂2X e 2 for environment e, we can
• Case 1: regress from Xe1 , to obtain α̂1 = 1 and α̂2 = 0, • Case 2: regress from Xe2 , to obtain α̂1 = 0 and α̂2 = σ1(e)+σ2(e) σ1(e)+σ2(e)+σ3(e) ,
• Case 3: regress from (Xe1 , X e 2), to obtain α̂1 = σ3(e) σ2(e)+σ3(e) and α̂2 = σ2(e) σ2(e)+σ3(e) .
In the generic scenario (i.e., σ1(e) 6= 0, σ2(e) 6= 0, and σ3(e) 6= 0), the regression using X1 in Case 1 is an invariant correlation: it is the only regression whose coefficients do not vary with e. By contrast, the regressions in both Case 2 and Case 3 have coefficients that depend on e. Therefore, only the invariant correlation in Case 1 will generalize well to new test environments.
2The setup applies to both continuous and categorical data. If any observation or label is categorical, we one-hot encode it.
Another way to understand Model 1 is through its graphical representation3, as shown in Fig. 1a. We treat the environment as a random variable E, where E could be any information specific to the environment (Storkey, 2009; Peters et al., 2015; Zhang et al., 2017; Huang et al., 2020). For simplicity, we let E be the environment index, i.e., E ∈ {1, . . . , N}, 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 {X1, X2} are unobserved and we can only observe their transformationO. In this case, Invariant Causal Prediction (ICP) (Peters et al., 2015) will fail when applied toO, even when Y is not affected by E (i.e., the edge E → Y is removed). The reason is that each variable (i.e., each dimension) ofO is jointly influenced by both X1 and X2 so that ICP is unable to find the variables containing the information only about X1 by searching for a subset of variablesO. By contrast, both IRM and IRMG work, as long as the transformation is linear. These findings are also empirically illustrated in Section 5.1. We now go even further and consider a more general causal graph in which Y can have more than one parent or child.

heading: 3.2 ASSUMPTIONS ON THE CAUSAL GRAPH
text: We extend the causal graph in Fig. 1b to a more general setting4, as encapsulated in Fig. 1c. In particular, we now have O ∈ O ⊆ Rd, Y ∈ Y ⊆ Rs, X = (Xp1 , . . . , Xpr , Xc1 , . . . , Xck) ∈ X ⊆ Rn, where n = r + k, and {Xi}i∈Ip={p1,...,pr} and {Xj}j∈Ic={c1,...,ck} are multiple scalar causal factors and non-causal factors5 of Y , respectively. We denote Xp . = (Xp1 , . . . , Xpr ) and Xc . = (Xc1 , . . . , Xck) for the ease of clarification. We also assume that X is of lower dimension thanO, that is, n ≤ d. We allow for arbitrary connections between the latent variablesX as 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. 1c satisfies the following assumption6: Assumption 1. (a) Xi depends on one or both of Y and E for any i; (b) The causal graph containingX and Y is a DAG; (c)O ⊥ Y ,E|X , implying that p(O|X) is invariant across all the environments; (d) Y ⊥ E|Xp, implying that p(Y |Xp) is invariant across all the environments.
One may be wondering how practical Assumption 1 is in real world applications. Let us explore this in more detail. Assumption 1a rules out all the useless Xi in the task of predicting Y . This is because if Assumption 1a is violated, meaning that Xi is independent of Y and E and has no influence in predicting Y , then such Xi should be viewed as noise and thus eliminated during learning. Assumption 1b is a common assumption in causal discovery (Spirtes et al., 2000; Pearl, 2009; Peters et al., 2015). 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 inferX fromO in any unseen environment. Assumption 1d is a widely-used default assumption in OOD generalization (Peters et al., 2015; Arjovsky et al., 2019). In fact, Assumption 1d can be further relaxed to the more
3The relation between SEM and its graphical representation is formally defined in Appendix D. 4For simplicity, we do not explicitly consider unobserved confounders in this paper. In particular, we assume
that there are no unobserved confounders betweenX,Y ,O, and E. 5This means that Xj∈Ic could be either an effect of Y , independent of Y , or spuriously correlated with Y via a third set of confounders (i.e., both Xj and Y are affected by a subset of {Xi}i 6=j and E). 6For generality, we replace Y and E with Y and E to allow for multi-dimensional variables.
Algorithm 1: Invariant Causal Representation Learning (iCaRL) Phase 1: We first learn a NF-iVAE model, including the decoder and its corresponding encoder, by
optimizing the objective function in Eq. (10) on the data {O,Y ,E}. Then, we use the mean of the NF-iVAE encoder to infer the latent variablesX from observations {O,Y ,E}. The latent variables are guaranteed to be identified up to a permutation and simple transformation. Phase 2: After inferringX , 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 a set of latent variables in which each pair of Xi and Xj satisfies that the dependency between them increases after additionally conditioning on Y . Phase 3: Having obtained Pa(Y ), we can solve the optimization problems in Eq. (11) to learn the invariant classifier w. When in a new environment, we first infer Pa(Y ) fromO by solving Eq. (12) and then leverage the learned w for prediction.
practical one that E[Y |Xp] is invariant across all the environments. That is, given Xp, 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 et al. (2013); von Kügelgen et al. (2020); Sun et al. (2020); Ahuja et al. (2021); von Kügelgen et al. (2021)) and is a very flexible model for causal analysis when predicting Y fromO.

heading: 3.3 ASSUMPTIONS ON THE PRIOR
text: When the underlying causal graph satisfies Assumption 1, our primary assumption leading to identifiability in this general setting is that the conditional prior p(X|Y ,E) belongs to a general exponential family. This is formalized as follows:
Assumption 2. p(X|Y ,E) belongs to a general exponential family with parameter vector given by an arbitrary function λ(Y ,E) and sufficient statistics T (X) = [Tf (X)T,TNN (X)T]T given by the concatenation of a) the sufficient statistics Tf (X) = [T1(X1)T , · · · ,Tn(Xn)T ]T of a factorized exponential family, where all the Ti(Xi) have dimension larger or equal to 2, and b) the output TNN (X) of a neural network with ReLU activations. The resulting density function is thus given by
pT ,λ(X|Y ,E) = Q(X)/Z(Y ,E) exp [ T (X)Tλ(Y ,E)], (5)
where Q is the base measure and Z the normalizing constant.
A neural network with ReLU activation has universal approximation power. Therefore, the term TNN (X) 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 flexible 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. 1c. 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 Xi∈Ip might be p(Xi|E), when Xi is only caused by E. By contrast, in the extended iVAE model the prior is p(X|Y ,E). Is this going to affect the identifiability of the latent variables? Well, in practice not because the posterior distribution forX givenO, Y and E would be equivalent in both models (up to identifiability guarantees).

heading: 4 INVARIANT CAUSAL REPRESENTATION LEARNING
text: We now introduce our algorithm, invariant Causal Representation Learning (iCaRL), which consists of 3 phases as summarized in Algorithm 1. The basic idea is that we first identify X by using an extended iVAE model under Assumptions 1&2 (Phase 4.1), then discover direct causes of Y among the identifiedX (Phase 4.2), and finally learn an invariant predictor for Y from the discovered causes (Phase 4.3).

heading: 4.1 PHASE 1: IDENTIFYING LATENT VARIABLES USING NF-IVAE
text: In this section, we describe our identifiable model, namely NF-iVAE, which is 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 a corresponding generative model by directly substituting U with (Y ,E) in Eq. (1):
pθ(O,X|Y ,E) = pf (O|X)pT ,λ(X|Y ,E), (6) pf (O|X) = p (O − f(X)). (7)
The corresponding lower bound is LVAEphase1(θ,φ) :=EpD [ Eqφ(X|O,Y ,E) [log pf (O|X) + log pT ,λ(X|Y ,E)− log qφ(X|O,Y ,E)] ] . (8)
To obtain an identifiability result, we assume that the prior pT ,λ(X|Y ,E) satisfies Assumption 2 (i.e., Eq. (5)). Since the prior is a general multivariate exponential family distribution with unknown normalization constant, we cannot learn its parameters (T ,λ) by directly maximizing Eq. (8). Instead, we use score matching, a well-known method for training unnormalized probabilistic models (Hyvärinen, 2005; Vincent, 2011), and learn (T ,λ) by minimizing LSMphase1(T ,λ) :=EpD [ Eqφ(X|O,Y ,E) [ ||∇X log qφ(X|O,Y ,E)−∇X log pT ,λ(X|Y ,E)||2 ]] . (9)
In practice, we can use a simple trick of partial integration to simplify the evaluation of Eq. (9), see Appendix C. Furthermore, we can jointly learn (θ,φ) by combining Eq. (8) and Eq. (9) in the following objective:
Lphase1(θ,φ) = LVAEphase1(f , T̂ , λ̂,φ)− LSMphase1(f̂ ,T ,λ, φ̂), (10) where f̂ , T̂ , λ̂, φ̂ are copies of f ,T ,λ,φ that are treated as constants and whose gradient is not calculated during learning. More details can be found in 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 θ := (f ,T ,λ), where pT ,λ(X|Y ,E) satisfies Assumption 2. Furthermore, assume the following holds: (i) The set {O ∈ O|ϕ (O) = 0} has measure zero, where ϕ is the characteristic function of the density p defined in Eq. (7). (ii) Function f in Eq. (7) is injective, and has all second-order cross derivatives. (iii) The sufficient statistics in Tf are all twice differentiable. (iv) There exist k + 1 distinct points (Y ,E)0, . . . , (Y ,E)k such that the matrix L = ( λ((Y ,E)1)− λ((Y ,E)0), . . . ,λ((Y ,E)k)− λ((Y ,E)0) ) of size k×k is invertible, where k is the dimension of T . Then the parameters θ are identifiable up to a permutation and a “simple transformation" of the latent variablesX , defined as a componentwise nonlinearity making each recovered Ti(Xi) in Tf (X) 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 iVAEs. It results in several key changes in the proof, clarified in Appendix H. Interestingly, from (iv) we can further see thatE is unnecessary when there exist k+ 1 distinct points Y 0, . . . ,Y k such that the matrix L = ( λ(Y 1)− λ(Y 0), . . . ,λ(Y k)− λ(Y 0) ) of size k × 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 qφ(X|O,Y ,E) contains pθ(X|O,Y ,E), and qφ(X|O,Y ,E) > 0 everywhere. (ii) We maximize Lphase1(θ,φ) with respect to both θ and φ. Then in the limit of infinite data, we learn the true parameters θ∗ up to a permutation and simple transformation of the latent variablesX .
As a consequence of Theorem 1&2, we have: Theorem 3. Assume the hypotheses of Theorem 1 and Theorem 2 hold, then in the limit of infinite data, we identify the true latent variablesX∗ up to a permutation and simple transformation.
Theorem 3 states that we can use NF-iVAE to infer the true X up to a permutation and simple transformation. We use the mean of qφ(X|O,Y ,E) for this task. Note that the noise may introduce uncertainty in the estimation ofX . However, whenO is high dimensional andX is low dimensional (as common in real world applications), qφ(X|O,Y ,E) will be highly concentrated and we will still be able to estimateX with high accuracy. The good results obtained by our method in various experiments seem to corroborate this. All three theorems are proven in Appendix H.

heading: 4.2 PHASE 2: DISCOVERING DIRECT CAUSES
text: After estimating X for each data point, the next step is to determine which components of X are a direct cause of Y . We denote these components by Pa(Y ). We first conduct the PC algorithm
(Spirtes et al., 2000) to learn a Markov equivalence class of DAGs, which gives us the direct neighbors of Y , denoted by Ne(Y ). Then, from Assumption 1, one observation is that for any two latent variables Xi and Xj 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 variables7 with conditional independence testing8 (Zhang et al., 2012) to discover Pa(Y ) by comparing p-values from the two tests: IndTest(Xi, Xj |E) and IndTest(Xi, Xj |Y ,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 a highly unlikely case in real world applications, which is left to Appendix G.

heading: 4.3 PHASE 3: LEARNING AN INVARIANT PREDICTOR
text: 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 w∈Hw ∑ e∈Etr Re(w) = min w∈Hw ∑ e∈Etr EPa(Y e),Y e [`1(w(Pa(Y e)),Y e)] , (11)
where `1(·) could be any loss. Since we assume that E(Y |Pa(Y )) is invariant across Eall (the relaxed version of Assumption 1d), the learned w is guaranteed to perform well across Eall. The remaining question is how to infer Pa(Y ) (i.e.,Xp) fromO in a new environment. This can be implemented by leveraging the learned p(O|X). The rationale behind is that p(O|X) is assumed to be invariant across Eall (Assumption 1c). In light of this idea, we follow Sun et al. (2020) and infer Xp fromO in any new testing environment by solving the following optimization problem:
max Xp,Xc
log pf (O|Xp,Xc) + λ1||Xp||22 + λ2||Xc||22, (12)
where the hyperparameters λ1 > 0 and λ2 > 0 control the learned Xp and Xc have a reasonable scale, both of which are selected on training/validation data. For optimization, we follow Schott et al. (2018) to first use values ofX sampled from the training set as initial points and then use Adam to optimize for several iterations. Note that the noise will introduce uncertainty in the estimation of Xp andXc fromO. However, as we mentioned before (below Theorem 3), this noise is not going to affect the estimation much because the likelihood will be highly concentrated around the ground truth values, as corroborated by our good empirical results.
A key question is if iCaRL performs well across Eall even though it uses only data from Etr. 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 ◦ Φ with optimal OOD generalization uses only Pa(Y ) to compute Φ and ii) the classifier w in this optimal predictor can be estimated using data from any environment e for which the distribution of Pa(Y e) has full support, which will always be the case since the conditional prior in Eq. (5) has full support. Finally, Theorem A.1 in (Arjovsky, 2021) also indicates that iii) the optimal predictor will be invariant across Eall. Key to these results is that Pa(Y e) are available when solving (12). This requires first to identify the latent variables X from O, 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 1 and the assumptions of Theorems 1 and 2, the predictor learned by iCaRL across Etr in the limit of infinite data has optimal OOD generalization across Eall.

heading: 5 EXPERIMENTS
text: 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 to space constraints, 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.

heading: 5.1 SYNTHETIC DATA
text: 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
7We only need to consider those variables in Ne(Y ) whose edges connecting to Y are not oriented by PC. 8These conditional independence tests can be performed in parallel to largely accelerate the testing procedure.
See more in Appendix G.
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., Xi 6⊥ Xj |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 X = (X1, X2) using the samples ofO,E, andY . 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 is 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 dataX up to a permutation and a simple componentwise transformation, whereas all the other methods fail because they are unable to handle the non-factorized case in which Xi 6⊥ Xj |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 a table.

heading: 5.2 COLORED MNIST, COLORED FASHION MNIST, AND VLCS
text: 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) to create these two datasets (see the details in Appendix J). The task is to predict a 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 NF-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 a 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
(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 e (on each environment e), and ERM GRAYSCALE (on data with no spurious correlations), and 4) ROBUST MIN MAX (minimizing the maximum loss across the multiple environments). Table 1 shows that iCaRL outperforms all other baselines on CFMNIST. It is 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: VLCS (Fang et al., 2013). This dataset consists of 10, 729 photographic images of dimension (3, 224, 224) and 5 classes from four domains: Caltech101, 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-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 on another popular dataset: PACS (Li et al., 2017a), all the details of which are placed in Appendix L.

heading: 6 RELATED WORK
text: Invariant Causal Prediction (ICP), aims to find the causal feature set (i.e., all direct causes of a 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 on the target variable Y . Besides, ICP methods implicitly assume that the variables of interestX are given. The works 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 theX are given and further assume that additional information about the structure over X is known. Mitrovic et al. (2020) analyze data augmentations in self-supervised learning from the perspective of invariant causal mechanisms. Arjovsky et al. (2019) reformulate this invariance as an optimization-based problem, allowing us to learn an invariant data representation from O constrained to be a linear transformation of X . 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.

heading: 7 CONCLUSION
text: We have proposed a novel framework to learn invariant predictors from a diverse set of training environments. It is based on a 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 a causal lens.

heading: REPRODUCIBILITY STATEMENT
text: Proofs of all our theoretical results are given in Appendix H with disclosure of all assumptions. All datasets used are publicly available or instructions are provided on how to generate them in Appendix J. Implementation details are given in Appendix M. A description of the hyperparameters and network architectures used is included in Appendix N. Source code for all our experiments will be made available on GitHub after deanonymisation, providing also a complete description of our experimental environment, configuration files and instructions on the reproduction of our experiments.

heading: A DOMAIN GENERALIZATION
text: The goal of domain generalization (DG) (Muandet et al., 2013) is OOD generalization: learning a 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 et al., 2015; Erfani et al., 2016; Li et al., 2017b). Muandet et al. (2013) propose a 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 et al., 2015; Wang et al., 2017; Motiian et al., 2017; Li et al., 2018b;c); (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 a low generalization error on meta-test sets after training on meta-training sets (Balaji et al., 2018; Dou et al., 2019; Li et al., 2018a; 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 whenX and Y are generated concurrently, which excludes the cases in which some part ofX could also be affected by Y in some manner; 2) They assume that the causal latent factors Xp and the non-causal latent factors Xc are independent when conditioning on E, that is, Xp ⊥ Xc|E. In practice, during model learning, they actually further assume that Xi ⊥ Xj |E for any i 6= j so that VAE could be leveraged to learn the model. In this sense, their approach has the same issue as the one in iVAE, i.e., unable to deal with the non-factorized 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 autoencoders (VAEs), which allows us to efficiently learn deep latent-variable models and their corresponding inference models (Kingma & Welling, 2013; Rezende et al., 2014). Consider a simple latent variable model where O ∈ Rd stands for an observed variable and X ∈ Rn for a latent variable. A VAE method learns a full generative model pθ(O,X) = pθ(O|X)pθ(X) and an inference model qφ(X|O), typically a factorized Gaussian distribution whose mean and variance parameters are given by the output of a neural network with input O. This inference model approximates the posterior pθ(X|O), where θ is a vector of parameters of the generative model, φ a vector of parameters of the inference model, and pθ(X) is a prior distribution over the latent variables. Instead of maximizing the data log-likelihood, we maximize its lower bound LVAE(θ,φ): log pθ(O) ≥ LVAE(θ,φ) := Eqφ(X|O) [log pθ(O|X)]− KL (qφ(X|O)||pθ(X)) , where we have used Jensen’s inequality, and KL(·||·) denotes the Kullback-Leibler divergence between two distributions.

heading: C DERIVATION
text: C.1 IVAE
In Section 2.1, the evidence lower bound of iVAE is defined by LiVAE(θ,φ) :=EpD [ Eqφ(X|O,U) [log pf (O|X)]− KL (qφ(X|O,U)||pT ,λ(X|U)) ] =EpD [ Eqφ(X|O,U) [log pf (O|X) + log pT ,λ(X|U)− log qφ(X|O,U)] ] .

heading: C.2 SCORE MATCHING
text: In Section 4.1, we follow Hyvärinen (2005) and use a simple trick of partial integration to simplify the evaluation of the score matching objective LSMphase1 in Eq. (10) of the main text: LSMphase1(T ,λ) :=EpD [ Eqφ(X|O,Y ,E) [ ||∇X log qφ(X|O,Y ,E)−∇X log pT ,λ(X|Y ,E)||2
]] =EpD [ Eqφ(X|O,Y ,E) [ n∑
j=1
[ ∂2pT ,λ(X|Y ,E)
∂X2j +
1
2
( ∂pT ,λ(X|Y ,E)
∂Xj
)2]]] + const.
where the last equality is due to the Theorem 1 in Hyvärinen (2005).

heading: D DEFINITIONS FOR SEM AND IRM
text: Definition 2. A structural equation model (SEM) C := (S, N) governing the random vectorX = (X1, . . . , Xd) is a set of structural equations: Si : Xi ← fi(Pa(Xi), Ni), where Pa(Xi) ⊆ {X1, . . . , Xd} \ {Xi} are called the parents of Xi, and the Ni are independent noise random variables. We say that “Xi causes Xj” if Xi ∈ Pa(Xj). We call causal graph of
X to the graph obtained by drawing i) one node for each Xi, and ii) one edge from Xi to Xj if Xi ∈ Pa(Xj). We assume acyclic causal graphs. Definition 3. Consider a SEM C := (S, N). An intervention e on C consists of replacing one or several of its structural equations to obtain an intervened SEM Ce := (Se, Ne), with structural equations: Sei : Xei ← fei (Pae(Xei ), Nei ), The variableXe is intervened if Si 6= Sei or Ni 6= Nei . Definition 4. Consider a structural equation model (SEM) S governing the random vector (X1, . . . , Xn,Y ), and the learning goal of predicting Y from X . Then, the set of all environments Eall(S) indexes all the interventional distributions P (Xe,Y e) obtainable by valid interventions e. An intervention e ∈ Eall(S) is valid as long as (i) the causal graph remains acyclic, (ii) E [Y e|Pa(Y )] = E [Y |Pa(Y )], and (iii) V [Y e|Pa(Y )] remains within a finite range.

heading: E DEFINITIONS AND LEMMAS FOR THE EXPONENTIAL FAMILIES
text: Definition 5. (Exponential family) A multivariate exponential family is a set of distributions whose probability density function can be written as
p(X) = Q(X) Z(θ) exp(〈T (X),θ〉), (13)
where Q : X → R is the base measure, Z(θ) is the normalizing constant, T : X → Rk is the sufficient statistics, and θ ∈ Rk is the natural parameter. The size k ≥ n is the dimension of the sufficient statistics T and depends on the latent space dimension n. Note that k is fixed given n. Definition 6. (Strongly exponential distributions) A multivariate exponential family distribution
p(X) = Q(X) Z(θ) exp(〈T (X),θ〉) (14)
is strongly exponential, if (∃θ ∈ Rk s.t. 〈T (X),θ〉 = const, ∀X ∈ X ) =⇒ (l(X ) = 0 or θ = 0), ∀X ⊂ Rn, (15)
where l is the Lebesgue measure.
The density of a 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) are strongly exponential.

heading: F DEFINITIONS FOR IDENTIFIABILITY
text: Definition 7. Let Θ be the domain of the parameters θ = {f ,T ,λ}. Let ∼ be an equivalence relation on Θ. A deep generative model is said to be ∼–identifiable if
pθ(O) = pθ̃(O) =⇒ θ ∼ θ̃. (16) The elements in the quotient space Θ\ ∼ are called the identifiability classes. Definition 8. Let ∼A be an equivalence relation on Θ defined by:
(f ,T ,λ) ∼A (f̃ , T̃ , λ̃) ⇐⇒ ∃A, c s.t. T (f−1(O)) = AT̃ (f̃−1(O)) + c, ∀O ∈ O, (17) where A ∈ Rk×k is an invertible matrix and c ∈ Rk is a vector. Definition 9. Let ∼P be an equivalence relation on Θ defined by:
(f ,T ,λ) ∼P (f̃ , T̃ , λ̃) ⇐⇒ ∃P, c s.t. T (f−1(O)) = P T̃ (f̃−1(O)) + c, ∀O ∈ O, (18) where P ∈ Rk×k is a block permutation matrix and c ∈ Rk is a vector.

heading: G PHASE 2: DISCOVERING SINGLE CAUSE

heading: G.1 SOME SPECIAL CASES IN MULTI-CAUSE SETTINGS
text: These conditional independence tests can be performed in parallel to largely accelerate the testing procedure. Note that, in practice, it might occur that there exist some Xi which is independent of any other Xj when conditioning on Y and E. It is probably because such Xi is a deterministic
transformation of Y . In this special case, we can use IGCI (Daniusis et al., 2012; Janzing et al., 2012) to determine whether or not Xi is a cause of Y . Also note that, in some scenarios in which the dependence signals between a 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, we will select it as one cause of Y .

heading: G.2 DISCOVERING SINGLE CAUSE
text: In the single-cause case, by following Wang et al. (2014), we leverage the MB-by-MB algorithm to first construct a local structure around Y and then discover the single parent of Y according to the constructed local graph. One obvious advantage of this approach is in efficiency, because there is no need to 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 Xi ⊥ Xj |Y ,E for any i 6= j. In fact, this assumption covers more scenarios than the common assumption that Xi ⊥ Xj for i 6= j in latent variable models, e.g., disentanglement (Bengio et al., 2013; Locatello et al., 2019), autoencoders (Kingma & Welling, 2013; Rezende et al., 2014), ICA (Comon, 1994; Hyvärinen & Oja, 2000), etc. If Y is caused by at most one of Xi and Xj , and no matter whether Xi and Xj are caused by E or not, then Xi ⊥ Xj |Y ,E holds, but we may well have Xi 6⊥ Xj (e.g., if Y causes both Xi and Xj , or if there is a chain Xi → Y → Xj). If Y causes or is caused by at most one of Xi and Xj , and at most one of Xi and Xj is caused by E, then both Xi ⊥ Xj and Xi ⊥ Xj |Y ,E hold. If {Y , Xi, Xj} form a collider Xi → Y ← Xj , and no matter whether Xi and Xj are caused by E or not, then Xi ⊥ Xj |E hold, but we may have Xi 6⊥ Xj (e.g., when both Xi and Xj are caused by E). Under this assumption, we are able to separately look into each Xi given Y and E, without considering any other Xj . Fig. 7a shows that there exist only five possible connections between Xi, Y , E, andO. Among them, only the arrow from Xi toO must exist, becauseO is generated from X . The other four arrows might not be present, with the exception that there must be at least one connection between Xi and Y or E (Assumption 1a). This leaves ten possible structures, shown in Figs. 3b-3k.
Given data {X̂i,Y ,E,O} in which the X̂i are obtained using qφ(X|O,Y ,E) (for example, as given by the mean of this distribution), we are able to distinguish all ten 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 2 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 Xi is a 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 Xi 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.

heading: H PROOFS

heading: H.1 PROOF OF THEOREM 1
text: The proof of this theorem consists of three parts.
In Part I, we prove that the parameters θ are ∼A 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 I and Part II into one theorem (Theorem 1), which completes the proof.
It is worth noting that, compared to the proof in iVAE, the main changes in our proof consist of
• Part I, In step III.
– It 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 our proof, but it is in the proof of iVAE.
• Part II.
– It 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.

heading: H.1.1 PART I
text: For notational simplicity, in the proof of this part we denote (Y ,E) by U . Hence, our generative model defined according to Eqs. (6-8) in the main text now becomes:
pθ(O,X|U) = pf (O|X)pT ,λ(X|U), (19) pf (O|X) = p (O − f(X)), (20)
pT ,λ(X|U) = Q(X)/Z(U) exp [ T (X)Tλ(U)]. (21)
Theorem 4. Suppose that we observe data sampled from a deep generative model defined according to Eqs. (19-21) with parameters (f ,T ,λ). Assume that
(i) The set {O ∈ O|ϕε(O) = 0} has measure zero, where ϕε is the characteristic function of the density pε defined in Eq. (20);
(ii) The mixing function f in Eq. (20) is injective;
(iii) There exist k + 1 points U0,U1, · · · ,Uk ∈ U such that the matrix L = [λ(U1)− λ(U0), · · · ,λ(Uk)− λ(U0)] ∈ Rk×k (22)
is invertible.

heading: Then the parameters {f ,T ,λ} are ∼A identifiable.
text: Proof. Define vol(B) = √
det(BTB) for any full column rank matrix B. Suppose that we have two sets of parameters θ = (f ,T ,λ) and θ̃ = (f̃ , T̃ , λ̃) such that pθ(O|U) = pθ̃(O|U), ∀(O,U) ∈ O × U . We want to show θ ∼A θ̃.
Step I. The proof of this step is similar to Step I in the proof of Theorem 1 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 (O,U) ∈ O × U , we have
pθ(O|U) = pθ̃(O|U) (23) =⇒ ∫ X pf (O|X)pT ,λ(X|U)dX = ∫ X pf̃ (O|X)pT̃ ,λ̃(X|U)dX (24)
=⇒ ∫ X pε(O − f(X))pT ,λ(X|U)dX = ∫ X pε(O − f̃(X))pT̃ ,λ̃(X|U)dX (25)
=⇒ ∫ O pε(O − Ō)pT ,λ(f−1(Ō)|U) vol(Jf−1(Ō))dŌ = ∫ O pε(O − Ō)pT̃ ,λ̃(f̃ −1(Ō)|U) vol(Jf̃−1(Ō))dŌ
(26) =⇒ ∫ Rd pε(O − Ō)p̃f ,T ,λ,U (Ō)dŌ = ∫ Rd pε(O − Ō)p̃f̃ ,T̃ ,λ̃,Ũ (Ō)dŌ (27)
=⇒ (p̃f ,T ,λ,U ∗ pε)(O) = (p̃f̃ ,T̃ ,λ̃,Ũ ∗ pε)(O) (28) =⇒ F [p̃f ,T ,λ,U ](ω)ϕε(ω) = F [p̃f̃ ,T̃ ,λ̃,Ũ ](ω)ϕε(ω) (29)
=⇒ F [p̃f ,T ,λ,U ](ω) = F [p̃f̃ ,T̃ ,λ̃,Ũ ](ω) (30) =⇒ p̃f ,T ,λ,U (O) = p̃f̃ ,T̃ ,λ̃,Ũ (O) (31)
where
• in Eq. (26), J denotes the Jacobian, and we made the change of variable Ō = f(X) on the left hand side, and Ō = f̃(X) on the right hand side.
• in Eq. (27), we introduced p̃f ,T ,λ,U (O) , pT ,λ(f
−1(O)|U) vol(Jf−1(O))IO(O) (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 ϕε = F [pε] (by definition of the characteristic function).
• in Eq. (30), we dropped ϕε(ω) from both sides as it is non-zero almost everywhere (by assumption (i)).
Step II. In this step, we remove all terms that are either a function ofO or U . First, by replacing both sides of Eq. (31) by their corresponding expressions from Eq. (32), we have
pT ,λ(f −1(O)|U) vol(Jf−1(O)) = pT̃ ,λ̃(f̃ −1(O)|U) vol(Jf̃−1(O)). (33)
Then, by taking logarithm on both sides of Eq. (33) and replacing pT ,λ by its expression from Eq. (21), we obtain
log vol(Jf−1(O)) + logQ(f−1(O))− logZ(U) + 〈 T (f−1(O)),λ(U) 〉 = log vol(Jf̃−1(O)) + log Q̃(f̃ −1(O))− log Z̃(U) + 〈 T̃ (f̃−1(O)), λ̃(U) 〉 . (34)
Let U0,U1, · · · ,Uk ∈ U be the k + 1 points defined in assumption (iii). We evaluate the above equation at these points to obtain k + 1 equations, and subtract the first equation from the remaining k equations to obtain:〈
T (f−1(O)),λ(U l)− λ(U0) 〉 + log Z(U0)
Z(U l) = 〈 T̃ (f̃−1(O)), λ̃(U l)− λ̃(U0) 〉 + log Z̃(U0)
Z̃(U l) , l = 1, · · · , k. (35)
Let L be defined as in assumption (iii) and L̃ defined similarly for λ̃. Note that L is invertible by assumption, but L̃ is not necessarily invertible. Letting b ∈ Rk in which bl = log Z̃(U 0)Z(U l)
Z̃(U l)Z(U0) , we
have LTT (f−1(O)) = L̃T T̃ (f̃−1(O)) + b. (36)
Left multiplying both sides of the above equation by L−T gives T (f−1(O)) = AT̃ (f̃−1(O)) + c, (37)
where A = L−T L̃ ∈ Rk×k and c = L−T b ∈ Rk.
Step III. To complete the proof, we need to show that A is invertible. Let Xl ∈ X , Ol = f(Xl), l = 0, · · · , k. We evaluate Eq. (37) at these k + 1 points to obtain k + 1 equations and subtract the first equation from the remaining k equations to obtain
[T (X1)− T (X0), · · · ,T (Xk)− T (X0)]︸ ︷︷ ︸ ,R∈Rk×k
= A [T̃ (f̃−1(O1))− T̃ (f̃−1(O0)), · · · , T̃ (f̃−1(Ok))− T̃ (f̃−1(O0))]︸ ︷︷ ︸ ,R̃∈Rk×k . (38)
We need to show that for a given X0 ∈ X , there exist k points X1, · · · ,Xk ∈ X such that the columns of R are linearly independent. Suppose, for contradiction, that the columns of R would never be linearly independent for any choice of X1, · · · ,Xk ∈ X . Then the function g(X) , T (X)− T (X0) would live in a k − 1 or lower dimensional subspace, and thus we could find a nonzero vector λ ∈ Rk orthogonal to that subspace. This would imply that 〈T (X)− T (X0),λ〉 = 0 and thus 〈T (X),λ〉 = 〈T (X0),λ〉 = const, ∀X ∈ X , which contradicts the assumption that the prior is strongly exponential (Definition 6). Therefore, we have shown that there exist k + 1 points X0,X1, · · · ,Xk ∈ X such that R is invertible. Since R = AR̃ and A is not a function of X , A must be invertible. This completes the proof.

heading: H.1.2 PART II
text: Theorem 5. Suppose that all assumptions in Theorem 4 hold. Let the sufficient statistics T (X) = [Tf (X) T,TNN (X) T]T given by the concatenation of a) the sufficient statistics Tf (X) = [T1(X1) T , · · · ,Tn(Xn)T ]T of a factorized exponential family, where all the Ti(Xi) have dimension larger or equal to 2, and b) the output TNN (X) of a neural network with ReLU activations. (note that a neural network with ReLU activation has universal approximation power and should be able to capture any dependencies of interest). Let k′ be the dimension of Tf and thus that k′ ≥ 2n. Assume that
(i) the sufficient statistics Tf have all second-order own derivatives;
(ii) the mixing function f has all second-order cross derivatives.

heading: Then the parameters {f ,T ,λ} are ∼P identifiable.
text: Proof. Let v = f̃−1 ◦ f : X → X . Since all assumptions in Theorem 4 hold, we have T (X) = AT̃ (v(X)) + c, (39)
where A ∈ Rk×k is invertible. We want to show that A is a block permutation matrix.
Step I. In this step, we show that v is a componentwise function. First we differentiate both sides of Eq. (39) with respect to Xs and Xt (s 6= t) to obtain
∂T (X)
∂Xs = A n∑ i=1 ∂T̃ (v(X)) ∂vi(X) · ∂vi(X) ∂Xs
(40)
∂2T (X) ∂Xs∂Xt = A n∑ i=1 n∑ j=1 ∂2T̃ (v(X)) ∂vi(X)∂vj(X) · ∂vj(X) ∂Xt · ∂vi(X) ∂Xs +A n∑ i=1 ∂T̃ (v(X)) ∂vi(X) · ∂ 2vi(X) ∂Xs∂Xt .
(41) By construction, the second-order cross derivatives of T and T̃ are all zero. Therefore, we have
0 = A n∑ i=1 ∂2T̃ (v(X)) ∂vi(X)2 · ∂vi(X) ∂Xt · ∂vi(X) ∂Xs +A n∑ i=1 ∂T̃ (v(X)) ∂vi(X) · ∂ 2vi(X) ∂Xs∂Xt . (42)
The above equation can be written in the matrix-vector form: 0 = AT̃ ′′(X)v′s,t(X) +AT̃ ′(X)v′′s,t(X), (43)
where
T̃ ′′(X) =
[ ∂2T̃ (v(X))
∂v1(X)2 , · · · , ∂
2T̃ (v(X))
∂vn(X)2
] ∈ Rk×n (44)
v′s,t(X) =
[ ∂v1(X)
∂Xt · ∂v1(X) ∂Xs , · · · , ∂vn(X) ∂Xt · ∂vn(X) ∂Xs
]T ∈ Rn, (45)
and
T̃ ′(X) =
[ ∂T̃ (v(X))
∂v1(X) , · · · , ∂T̃ (v(X)) ∂vn(X)
] ∈ Rk×n (46)
v′′s,t(X) =
[ ∂2v1(X)
∂Xs∂Xt , · · · , ∂
2vn(X)
∂Xs∂Xt
]T ∈ Rn. (47)
Now by concatenating T̃ ′′′(X) = [T̃ ′′(X), T̃ ′(X)] ∈ Rk×2n and v′′s,t(X) = [v′s,t(X)T ,v′′s,t(X)T ]T ∈ R2n, (48) we obtain 0 = AT̃ ′′′(X)v′′′s,t(X). (49) Finally, we take the rows of T̃ ′′′(X) that corresponds to the factorized strongly exponential family distribution part and denote them by T̃ ′′′f (X) ∈ Rk
′×2n. By Lemma 5 in the iVAE paper (Khemakhem et al., 2020a) and the assumption that k′ ≥ 2n, we have that the rank of T̃ ′′′f (X) is 2n. Since k ≥ k′ ≥ 2n, the rank of T̃ ′′′(X) is also 2n. Since the rank of A is k, the rank of AT̃ ′′′(X) ∈ Rk×2n is 2n. This implies that v′′′s,t(X) must be a zero vector. In particular, we have that v′s,t(X) = 0, ∀s 6= t. Therefore, we have shown that v is a componentwise function.
Step II. To complete the proof, we need to show that A is a block permutation matrix. Without loss of generality, we assume that the permutation in v is the identity. That is v(X) = [v1(X1), · · · , vn(Xn)]T for some nonlinear univariate scalar functions v1, · · · , vn. Since f and f̃ are bijective, we have that v is also bijective and v−1(X) = [v−11 (X1), · · · , v−1n (Xn)]T . We denote T̄ (v(X)) = T̃ (v(X)) + A−1c and plug it into Eq. (39) to obtain T (X) = AT̄ (v(X)). Applying v−1 to the variablesX at both sides gives T (v−1(X)) = AT̄ (X). (50) Let t be the index of an entry in the sufficient statistics T that corresponds to the the factorized strongly exponential family distribution part Tf . For all s 6= t, we have
0 = ∂T (v−1(X))t
∂Xs = k∑ j=1 atj ∂T̄ (X)j ∂Xs . (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 atj is zero for any j such that ∂T̄ (X)j∂Xs 6= 0. 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 Xt; and 2) the neural network part.
Therefore, when t is the index of an entry in the sufficient statistics T that corresponds to factor i in the factorized strongly exponential family distribution part Tf , the only non-zero atj are the ones that map between Ti(Xi) and T̄i(vi(Xi)), where Ti are the factors in Tf that only depend on Xi and T̄i is defined similarly. Therefore, we can construct an invertible submatrix A′i with all non-zero elements atj for all t that corresponds to factor i, such that
Ti(Xi) = A ′ iT̄i(vi(Xi)) = A ′ iT̃i(vi(Xi)) + ci, i = 1, · · · , n, (52)
where T̃i are the factors in T̃f that only depends on Xi, and ci are the corresponding elements of c. This means that the matrix A is a block permutation matrix. For each i = 1, · · · , n, the block A′i of A affinely transforms Ti(Xi) into T̃i(vi(Xi)). There is also an additional block A′ which affinely transforms TNN (X) into T̃NN (v(X)). This completes the proof.

heading: H.1.3 PART III
text: Now we combine Theorem 4 in Part I and Theorem 5 in Part II into one theorem, which completes the proof of Theorem 1.

heading: H.2 PROOF OF THEOREM 2
text: We recall that the loss function in Phase 1 is as follows: Lphase1(θ,φ) = LVAEphase1(f , T̂ , λ̂,φ)− LSMphase1(f̂ ,T ,λ, φ̂), (53) where LVAEphase1(θ,φ) :=EpD [ Eqφ(X|O,Y ,E) [log pf (O|X) + log pT ,λ(X|Y ,E)− log qφ(X|O,Y ,E)] ] , (54)
LSMphase1(T ,λ) :=EpD [ Eqφ(X|O,Y ,E) [ ||∇X log qφ(X|O,Y ,E)−∇X log pT ,λ(X|Y ,E)||2 ]] . (55)
Proof. If the family of qφ(X|O,Y ,E) is flexible enough to contain pθ(X|O,Y ,E), then by optimizing the loss over its parameter φ, we will minimize the score matching term LSMphase1 in Eq. (55), which will eventually reach zero. If we assume that the model is not degenerate and that qφ > 0 everywhere, then having that LSMphase1 = 0 implies that ∇X log qφ(X|O,Y ,E) and ∇X log pT ,λ(X|Y ,E) are equal. This implies that log qφ(X|O,Y ,E) = log pT ,λ(X|Y ,E) + c for some constant c. But c is necessarily 0 because both qφ(X|O,Y ,E) and pT ,λ(X|Y ,E) are pdf’s. Therefore, the VAE term LVAEphase1 in Eq. (54) will be equal to the log-likelihood, meaning that the loss Lphase1 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 a permutation and componentwise transformation, the consistency of MLE means that we converge to the true parameter θ∗ up to a 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 a permutation and componentwise transformation does not change the proof but only the conclusion.

heading: H.3 PROOF OF THEOREM 3
text: Proof. Theorem 1 and Theorem 2 guarantee that in the limit of infinite data, NF-iVAE can learn the true parameters θ∗ := (f∗,T ∗,λ∗) up to a permutation and componentwise transformation of the latent variables. Let (f̂ , T̂ , λ̂) be the parameters obtained by NF-iVAE. We, therefore, have (f̂ , T̂ , λ̂) ∼P (f∗,T ∗,λ∗), where ∼P denotes the equivalence up to a permutation and componentwise transformation. If there were no noise, this would mean that the learned f̂ transforms O into X̂ = f̂−1(O) that are equal toX∗ = (f∗)−1 (O) up to a permutation and componentwise transformation (Definition 9). If with noise, we obtain the posterior distribution of the latent variables up to an analogous indeterminacy.

heading: H.4 PROOF OF PROPOSITION 1
text: Proof. Theorem A.1 in (Arjovsky, 2021) has showed that i) any predictor w ◦ Φ with optimal OOD generalization uses only Pa(Y ) to compute Φ; ii) the classifier w in this optimal predictor can be estimated using data from any environment e for which the distribution of Pa(Y e) has full support; iii) the optimal predictor will be invariant across Eall. In iCaRL, the hypotheses of Theorems 1 and 2 and Assumption 1 guarantee that Pa(Y ) can be recovered by first identifying the latent variablesX fromO, Y andE 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 e) always has full support. Also, under Assumption 1 we have that p(Y |Pa(Y )) is invariant across Eall. Hence, the classifier w 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.

heading: H.5 PROOF OF PROPOSITION 2
text: 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 1 All the six structures in this group can be discovered only by performing conditional independence tests.
• Rule 1.1 If Xi ⊥ Y , Xi 6⊥ E, and E ⊥ Y , then Fig. 3d is discovered.
• Rule 1.2 If Xi 6⊥ Y , Xi ⊥ E, and E 6⊥ Y , then Fig. 3g is discovered.
• Rule 1.3 If Xi 6⊥ Y , Xi 6⊥ E, and E ⊥ Y , then Fig. 3f is discovered.
• Rule 1.4 If Xi 6⊥ Y , Xi 6⊥ E, E 6⊥ Y , and Xi ⊥ Y |E, then Fig. 3i is discovered.
• Rule 1.5 If Xi 6⊥ Y , Xi 6⊥ E, E 6⊥ Y , and Xi ⊥ E|Y , then Fig. 3h is discovered.
• Rule 1.6 If Xi 6⊥ Y , Xi 6⊥ E, E 6⊥ Y , and Y ⊥ E|Xi, then Fig. 3e is discovered.
Group 2 If Xi 6⊥ Y , Xi ⊥ E, and E ⊥ 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 Xi 6⊥ Y , Xi ⊥ E, and E ⊥ Y , and a chosen binary causal discovery algorithm prefers Xi → Y to Xi ← Y , then Fig. 3b is discovered.
• Rule 2.2 If Xi 6⊥ Y , Xi ⊥ E, and E ⊥ Y , and a chosen binary causal discovery algorithm prefers Xi ← Y to Xi → Y , then Fig. 3c is discovered.
Group 3 If Xi 6⊥ Y , Xi 6⊥ E, E 6⊥ Y , Xi 6⊥ Y |E, Xi 6⊥ E|Y , and Y 6⊥ E|Xi, then we can discover both Fig. 3j and Fig. 3k. These two structures cannot be further distinguished only by conditional independence tests, because they come from the same Markov equivalence class. They also cannot be distinguished by any binary causal discovery algorithm, since both Xi and Y are affected by E. Fortunately, Zhang et al. (2017) provided a heuristic solution to this case based on the invariance of causal mechanisms, i.e., P (cause) and P (effect|cause) change independently. The detailed description of their method is given in Section 4.2 of Zhang et al. (2017). For convenience, here we directly borrow their final result. Zhang et al. (2017) states that determining the causal direction between Xi and Y in Fig. 3j and Fig. 3k is finally reduced to calculating the following term:
∆Xi→Y =
〈 log
P̄ (Y |Xi) 〈P̂ (Y |Xi)〉
〉 , (56)
where 〈·〉 denotes the sample average, P̄ (Y |Xi) is the empirical estimate of P (Y |Xi) on all data points, and 〈P̂ (Y |Xi)〉 denotes the sample average of P̂ (Y |Xi), which is the estimate of P (Y |Xi) in each environment. We take the direction for which ∆ is smaller to be the causal direction.
• Rule 3.1 If Xi 6⊥ Y , Xi 6⊥ E, E 6⊥ Y , Xi 6⊥ Y |E, Xi 6⊥ E|Y , Y 6⊥ E|Xi, and ∆Xi→Y is smaller than ∆Y→Xi , then Fig. 3j is discovered.
• Rule 3.2 If Xi 6⊥ Y , Xi 6⊥ E, E 6⊥ Y , Xi 6⊥ Y |E, Xi 6⊥ E|Y , Y 6⊥ E|Xi, and ∆Y→Xi is smaller than ∆Xi→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 Xi∈Ip is either p(Xi|E) or p(Xi), depending on whether Xi is caused by E or not. By contrast, in the NF-iVAE model the prior is p(X|Y ,E). Is this going to affect the identifiability of the latent variables? Well, in practice not because the posterior distribution forX givenO, Y and E would be equivalent in both models (up to identifiability guarantees).

heading: J DATASETS
text: 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.

heading: J.1 SYNTHETIC DATA
text: For the nonlinear transformation, we use the MLP:
• Input layer: Input batch (batch size, input dimension) • Layer 1: Fully connected layer, output size = 6, activation = ReLU • Output layer: Fully connected layer, output size = 10

heading: J.2 COLORED MNIST DIGITS
text: We use the exact same environment as in Arjovsky et al. (2019). Arjovsky et al. (2019) propose to create an environment for training to classify digits in MNIST data9, 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 the preliminary label (ỹ = 0 if the digit is between 0-4 and ỹ = 1 if the digit is between 5-9) by flipping it with 25 percent probability to construct the final labels. We sample the color id z by flipping the final labels with probability pe, where pe is 0.2 in the first environment, 0.1 in the second environment, and 0.9 in the third environment. The third environment is the testing environment. We color the digit red if z = 1 or green if z = 0.

heading: J.3 COLORED FASHION MNIST
text: We modify the fashion MNIST dataset10 in a manner similar to the MNIST digits dataset. Fashion MNIST data has images from different categories: “t-shirt”, “trouser”, “pullover”, “dress”, “coat”, “sandal”, “shirt”, “sneaker”, “bag”, “ankle boots”. We add colors to the images in such a way that the colors correlate with the labels. The task is to classify whether the image is that of foot wear or a clothing item. There are three environments (two training, one test) We add noise to the preliminary label (ỹ = 0: “t-shirt”, “trouser”, “pullover”, “dress”, “coat”, “shirt” and ỹ = 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 is 0.2 in the first environment, 0.1 in the second environment, and 0.9 in the third environment, which is the test environment. We color the object red if z = 1 or green if z = 0.
K IN-DEPTH ANALYSIS ON SYNTHETIC DATA

heading: K.1 COMPARISONS WITH STATE-OF-THE-ART
text: We first conduct a series of experiments on synthetic data generated according to an extension of the SEM in Model 1. The extension is to map the variablesX := (X1, X2) into a 10 dimensional
9https://www.tensorflow.org/api_docs/python/tf/keras/datasets/mnist/ load_data
10https://www.tensorflow.org/api_docs/python/tf/keras/datasets/fashion_ mnist/load_data
observationO through a linear or nonlinear transformation. Our goal is to predict Y fromO, where O = g(X). We consider three transformations: (a) Identity: g(·) is the identity matrix I ∈ R2×2, i.e.,O = g(X) = X . (b) Linear: g(·) is a random matrix S ∈ R2×10, i.e.,O = g(X) = X ·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 a 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 X1, meaning that the conditional prior in Eq. (??) will not exhibit dependencies. Because of this, in this case we do not include a TNN (X) 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, we do include TNN (X) in the NF-iVAE prior.
We consider a simple scenario in which we fix σ1 = 1 and σ2 = 0 for all environments and only allow σ3 to vary across environments. In this case, σ3 controls the spurious correlations between X2 and Y . Each experiment draws 1000 samples from each of the three environments σ3 = {0.2, 2, 100}, where the first two are for training and the third for testing. We compare with the following baselines:11 ERM, and two variants of IRMG: F-IRM Game (with Φ fixed to the identity) and V-IRM Game (with variable Φ).
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 a simple nonlinear transformation renders the problem more difficult by converting the original identity or linear problem into a nonlinear problem. In the Nonlinear case, the gains of iCaRL are very clear.
We perform additional experiments to analyze the three phases of our approach, including a) an analysis of the importance of Assumptions 1&2 and of the necessity of NF-iVAE in Phase 1, b) how accurately the direct causes can be recovered in Phase 2, and c) how well the two optimization problems can be addressed in Phase 3, all of which can be found in Appendix K.
K.2 VERIFYING PHASE 1
Theorem 3 tells us that we can leverage NF-iVAE to learn the true latent variables up to a permutation and componentwise transformation. We empirically verify this point by comparing the raw dataX with the corresponding X̂ inferred through the learned inference model in NF-iVAE. Fig. 4 clearly shows that the inferred X̂ is equal toX up to a permutation and componentwise transformation.
11We also tried ICP, but ICP was unable to find any parent of Y even in the identity case.
To show the importance of NF-iVAE, we conduct an experiment in which we replace NF-iVAE with the original VAE in Phase 1. As shown in Table 4, the performance of iCaRL based on NF-iVAE significantly outperforms the one based on VAE (aka., iCaRL-VAE). It is worth noting that when VAE is instead used in Phase 1, it usually occurs in Phase 2 that either all the dimensions or no dimension of X̂ are identified as the parents of Y . This is because all components of X̂ are mixed together and will influence one another even when conditioning on Y and E.
K.3 VERIFYING PHASE 2
To show how well our method can identify the direct causes of Y in Phase 2, we compare the final performance when the identified direct cause (i.e., X1) and the identified non-cause (i.e., X2) are respectively used in Phase 3 to learn the predictor. More specifically, we respectively treat X1 and X2 as Pa(Y ) to solve Eq. (12) and Eq. (11) to learn the data representation Φ and the invariant classifier w. Note that, there might exist a permutation between the inferred {X̂1, X̂2} and the true {X1, X2}. Fig. 5a and Fig. 5b show the results on the regression task, from which we can obviously see that the predictor elicited by the identified cause has a much better generalization performance. The classification result (i.e., Y is binarized according to whether it is positive or negative) in Fig. 5c further demonstrates this point.
K.4 VISUALIZATION OF IDENTIFIABILITY OF NF-IVAE
To further verify identifiability of NF-iVAE, we conduct a series of experiments on synthetic data generated according to the data generating process (Fig. 6b) corresponding to the causal graph shown in Fig. 6a. 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., Xi 6⊥ Xj |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. We compare with the following baselines: VAE Kingma & Welling (2013) (without identifiability guarantees), iVAE Khemakhem et al. (2020a) (with a conditionally factorized prior for identifiability), and ICE-BeeM Khemakhem et al. (2020b). It is worth noting that in ICE-BeeM the primary assumption leading to identifiability is similar to that in iVAE, where the base measure Q(X) 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 5. Clearly, from the table we can see that our method has a more general assumption on the prior leading to identifiability. This is also demonstrated empirically in Fig. 7. Our method iCaRL can recover the original data X up to a permutation and a simple componentwise transformation, whereas all the other methods fail because they are unable to handle the case in which Xi 6⊥ Xj |Y ,E.
L IN-DEPTH ANALYSIS ON MORE REALISTIC DATA

heading: L.1 COLORED MNIST
text: 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 e (on each environment e), and ERM GRAYSCALE (on data with no spurious correlations), and 4) ROBUST MIN MAX (minimizing the maximum loss across the multiple environments). Table 1 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.

heading: L.2 PACS
text: 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 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-parameter tuning procedure: train domain validation. As shown in Table 7, iCaRL achieves state-of-the-art performance when compared to those most popular domain generalization alternatives.
M IMPLEMENTATION DETAILS

heading: M.1 JOINT TRAINING
text: As described in Section 4.1, we can jointly learn (θ,φ) by optimizing the following objective: Lphase1(θ,φ) =LVAEphase1(f , T̂ , λ̂,φ)− LSMphase1(f̂ ,T ,λ, φ̂) (62)
=EpD [ Eqφ(X|O,Y ,E) [ log pf (O|X) + log pT̂ ,λ̂(X|Y ,E)− log qφ(X|O,Y ,E) ]] − EpD [ Eq
φ̂ (X|O,Y ,E)
[ ||∇X log qφ̂(X|O,Y ,E)−∇X log pT ,λ(X|Y ,E)|| 2 ]] (63)
=EpD [ Eqφ(X|O,Y ,E) [ log pf (O|X) + log pT̂ ,λ̂(X|Y ,E)− log qφ(X|O,Y ,E) ]] − EpD [ Eq
φ̂ (X|O,Y ,E)
[ n∑
j=1
[ ∂2pT ,λ(X|Y ,E)
∂X2j +
1
2
( ∂pT ,λ(X|Y ,E)
∂Xj )2]]] + const. (64)
where the last equality is due to the equation in Appendix C.2, and f̂ , T̂ , λ̂, φ̂ are copies of f ,T ,λ,φ that are treated as 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 “stop_gradient” in TensorFlow Abadi et al. (2015).

heading: M.2 THE GENERAL NON-FACTORIZED PRIOR
text: In the experiments, the general non-factorized prior in Assumption 2 is implemented as follows: pT ,λ(X|Y ,E) = 〈 NN(X;param1)︸ ︷︷ ︸
TNN (X)
,NN(Y ,E;param2)︸ ︷︷ ︸ λNN (Y ,E)
〉 + 〈 concat(X,X2)︸ ︷︷ ︸
Tf (X)
,NN(Y ,E;param3)︸ ︷︷ ︸ λf (Y ,E)
〉 ,
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(X,X2) is a vector of the latent variables and their squared values, and NN(Y ,E;param3) is a deep neural network parameterized by param3 that computes a vector of natural parameters as a function of Y and E. Hence, the term〈 concat(X,X2),NN(Y ,E;param3) 〉 is equivalent to the factorized exponential family, which
also satisfies that each Ti(Xi) has dimension larger or equal to 2.
Secondly, NN(X;param1) is a neural network that receives as input a 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 a corresponding vector of natural parameters. Hence, the term 〈 NN(X;param1),NN(Y ,E;param2) 〉 will allow this prior to capture the dependencies between the latent variablesX .

heading: N HYPERPARAMETERS AND ARCHITECTURES
text: In this section, we describe the hyperparameters and architectures of different models used in different experiments. Unless stated otherwise, we have λ1 = 1 and λ2 = 1, both of which are selected on training/validation data.

heading: N.1 SYNTHETIC DATA
text: We used Adam optimizer for training with learning rate set to 1e-3 and batch size set to 128.
N.1.1 ERM

heading: Linear ERM
text: • Input layer: Input batch (batch size, input dimension) • Output layer: Fully connected layer, output size = 1

heading: Nonlinear ERM
text: • Input layer: Input batch (batch size, input dimension) • Layer 1: Fully connected layer, output size = 6, activation = ReLU • Output layer: Fully connected layer, output size = 1
N.1.2 IRM

heading: Linear Data Representation Φ
text: • Input layer: Input batch (batch size, input dimension) • Output layer: Fully connected layer, output size = 1

heading: Nonlinear Data Representation Φ
text: • Input layer: Input batch (batch size, input dimension) • Layer 1: Fully connected layer, output size = 6, activation = ReLU • Output layer: Fully connected layer, output size = 1

heading: N.1.3 F-IRM GAME

heading: Linear Classifier w
text: • Input layer: Input batch (batch size, input dimension) • Output layer: Fully connected layer, output size = 1

heading: Nonlinear Classifier w
text: • Input layer: Input batch (batch size, input dimension) • Layer 1: Fully connected layer, output size = 6, activation = ReLU • Output layer: Fully connected layer, output size = 1
N.1.4 V-IRM GAME

heading: Linear Data Representation Φ
text: • Input layer: Input batch (batch size, input dimension) • Output layer: Fully connected layer, output size = 2

heading: Nonlinear Data Representation Φ
text: • Input layer: Input batch (batch size, input dimension) • Layer 1: Fully connected layer, output size = 6, activation = ReLU • Output layer: Fully connected layer, output size = 2

heading: Linear Classifier w
text: • Input layer: Input batch (batch size, 2) • Output layer: Fully connected layer, output size = 1

heading: Nonlinear Classifier w
text: • Input layer: Input batch (batch size, 2) • Layer 1: Fully connected layer, output size = 6, activation = ReLU • Output layer: Fully connected layer, output size = 1
N.1.5 ICARL

heading: NF-iVAE λf -Linear Prior
text: • Input layer: Input batch (batch size, input dimension) • Output layer: Fully connected layer, output size = 2

heading: NF-iVAE λf -Nonlinear Prior
text: • Input layer: Input batch (batch size, input dimension) • Layer 1: Fully connected layer, output size = 6, activation = ReLU • Output layer: Fully connected layer, output size = 2

heading: NF-iVAE Linear Encoder
text: • Input layer: Input batch (batch size, input dimension) • Mean Output layer: Fully connected layer, output size = 2 • Log Variance Output layer: Fully connected layer, output size = 2

heading: NF-iVAE Nonlinear Encoder
text: • Input layer: Input batch (batch size, input dimension) • Layer 1: Fully connected layer, output size = 6, activation = ReLU • Mean Output layer: Fully connected layer, output size = 2 • Log Variance Output layer: Fully connected layer, output size = 2

heading: NF-iVAE Linear Decoder
text: • Input layer: Input batch (batch size, 2) • Mean Output layer: Fully connected layer, output size = output dimension • Variance Output layer: 0.01 × 1, where 1 is a vector full of 1 with the length of output
dimension

heading: NF-iVAE Nonlinear Decoder
text: • 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 × 1, where 1 is a vector full of 1 with the length of output
dimension

heading: Linear Classifier w
text: • Input layer: Input batch (batch size, 1)
• Output layer: Fully connected layer, output size = 1

heading: Nonlinear Classifier w
text: • Input layer: Input batch (batch size, 1)
• Layer 1: Fully connected layer, output size = 6, activation = ReLU
• Output layer: Fully connected layer, output size = 1

heading: N.2 CMNIST AND CFMNIST
text: Considering that the results of most baselines come from IRMG (Ahuja et al., 2020a), for a fair comparison, we follow the same setting of IRMG in terms of hyper-parameters and validation considerations. For example, the batch size is set to 256, and the learning rate is 10−4. We also did not use the test environment data for validation. Please find more details in Ahuja et al. (2020a).
NF-iVAE TNN -Prior
• Input layer: Input batch (batch size, input dimension)
• Layer 1: Fully connected layer, output size = 50, activation = ReLU
• Output layer: Fully connected layer, output size = 45

heading: NF-iVAE λNN -Prior
text: • Input layer: Input batch (batch size, input dimension)
• Layer 1: Fully connected layer, output size = 50, activation = ReLU
• Output layer: Fully connected layer, output size = 45

heading: NF-iVAE λf -Prior
text: • Input layer: Input batch (batch size, input dimension)
• Layer 1: Fully connected layer, output size = 50, activation = ReLU
• Output layer: Fully connected layer, output size = 20
NF-iVAEO-Encoder
• Input layer: Input batch (batch size, 2, 28, 28)
• Layer 1: Convolutional layer, output channels = 32, kernel size = 3, stride = 2, padding = 1, activation = ReLU
• Layer 2: Convolutional layer, output channels = 32, kernel size = 3, stride = 2, padding = 1, activation = ReLU
• Layer 3: Convolutional layer, output channels = 32, kernel size = 3, stride = 2, padding = 1, activation = ReLU
• Output layer: Flatten
NF-iVAE (Y ,E)-Encoder
• Input layer: Input batch (batch size, input dimension)
• Output layer: Fully connected layer, output size = 100, activation = ReLU
NF-iVAE (O,Y ,E)-Merger/Encoder
• Input layer: Input batch (batch size, input dimension) • Layer 1: 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

heading: NF-iVAE Decoder
text: • Input layer: Input batch (batch size, 10) • Layer 1: Fully connected layer, output size = 32× 4× 4, activation = ReLU • Layer 2: Reshape to (batch size, 32, 4, 4) • Layer 3: Deconvolutional layer, output channels = 32, kernel size = 3, stride = 2, padding =
1, outpadding = 0, activation = ReLU • Layer 4: Deconvolutional layer, output channels = 32, kernel size = 3, stride = 2, padding =
1, outpadding = 1, activation = ReLU • Layer 5: Deconvolutional layer, output channels = 2, kernel size = 3, stride = 2, padding = 1,
outpadding = 1 • Mean Output layer: activation = Sigmoid • Variance Output layer: 0.01× 1, where 1 is a matrix full of 1 with the size of 2× 28× 28.

heading: Classifier w
text: • Input layer: Input batch (batch size, 50) • Layer 1: Fully connected layer, output size = 100, activation = ReLU • Output layer: Fully connected layer, output size = 1, activation = Sigmoid
N.3 VLCS AND PACS
We used the exact experimental setting that is described in Gulrajani & Lopez-Paz (2020). Specifically, we trained our model over all possible train and test environment combination for one of the commonly used hyper-parameter tuning procedure: train domain validation. We use ResNet-50 as an encoder and reverse the architecture of ResNet-50 as a decoder. We set the number of the latent variables to n = 50. We do the hyperparameter search by exactly following the guides given in Gulrajani & Lopez-Paz (2020).

Review summaries:
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.

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.