Recovering Latent Causal Factor for Generalization to Distributional Shifts

Abstract:
Distributional shifts between training and target domains may degrade the prediction accuracy of learned models, mainly because these models often learn features that possess only correlation rather than causal relation with the output. Such a correlation, which is known as “spurious correlation” statistically, is domaindependent hence may fail to generalize to unseen domains. To avoid such a spurious correlation, we propose Latent Causal Invariance Models (LaCIM) that specifies the underlying causal structure of the data and the source of distributional shifts, guiding us to pursue only causal factor for prediction. Specifically, the LaCIM introduces a pair of correlated latent factors: (a) causal factor and (b) others, while the extent of this correlation is governed by a domain variable that characterizes the distributional shifts. On the basis of this, we prove that the distribution of observed variables conditioning on latent variables is shift-invariant. Equipped with such an invariance, we prove that the causal factor can be recovered without mixing information from others, which induces the ground-truth predicting mechanism. We propose a Variational-Bayesian-based method to learn this invariance for prediction. The utility of our approach is verified by improved generalization to distributional shifts on various real-world data. Our code is freely available at https://github.com/wubotong/LaCIM.

heading: 1 Introduction
text: Current data-driven deep learning models, revolutionary in various tasks though, often exploit all types of correlations to fit data well. Among such correlations, there can be spurious ones corresponding to biases (e.g., confounding bias due to the presence of a third unseen factor) inherited from the data provided. Such data-dependent spurious correlations can erode the prediction power on unseen domains with distributional shifts, which can cause serious consequences especially in safety-critical tasks such as healthcare.
Recently, there is a Renaissance of causality in machine learning, expected to pursue causal relationships [59] to achieve stable generalization across domains. The so-called area of “causality” is pioneered by Structural Causal Models [51], as a mathematical formulation of this metaphysical concept grasped in the human mind. The incorporation of these human priors about cause and effect endows the model with the ability to identify the causal structure [51] which entails not only the data but also the underlying process of how they are generated. To achieve causal modeling, the old-school methods [52, 10] directly causally related the output label Y to a subset of covariates X , which is however not conceptually reasonable in applications with sensory-level data (e.g. model pixels as causal factors of the output does not make sense in image classification [11]). ∗Corresponding author †Work done during an internship at Microsoft Research Asia.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
For such applications, we rather adopt the manner of human visual perception [8, 9, 80] to causally relate the label Y to unobserved abstractions denoted by S, i.e., Y ← S. We further assume the existence of another non-causal latent factor (of Y ) denoted as Z, that together with S generate the input X: X ← (S,Z). Such an assumption is similarly adopted in the literature [25, 27, 35, 75, 71]. To model shifts across domains, we allow Z to be spuriously correlated with S (hence also the output), as marked by the bidirected arrow in Fig. 1 (a). Taking image classification as an example, the S and Z respectively refer to object-related abstractions (e.g., contour, texture) and contextual information (e.g., background, view). Due to this correlation, the model can learn contextual information into prediction, which may fail to generalize to the domain such that this correlation is broken.
We encapsulate above assumptions into the skeleton illustrated in Fig. 1 (a), in which the spurious correlation between S and Z varies across domains, as marked by the red bi-directed arrow in Fig. 1 (b). Taking a closer inspection, such a domain-dependent spurious correlation is governed by an auxiliary domain variable D in Fig. 1 (c), which causes the domain shifts. We call the set of causal models augmented with D as Latent Causal Invariance Models (LaCIM). Here, the “Causal Invariance” refers to P (Y |S), which together with P (X|S,Z), can be proved to be stable to the shifts across domains, under the assumptions embedded in the causal structure of LaCIM. Equipped with such an invariance, we prove that the S and the ground-truth predictor: P (Y |s?) for x generated from (s?, z?), are identifiable up to transformations that do not mix the non-causal information.
Under such an identifiability guarantee, we propose to learn the P (Y |S) and P (X|S,Z) by reformulating the Variational Auto-encoder (VAE) [37] to fit the joint distribution of the input and output variables from training domains. During the test stage, we first infer the value of S by optimizing the estimated P (X|S,Z) over latent space, followed by the learned P (Y |S) for prediction. We first use simulated data to verify the correctness of the identifiability claim. Then, to demonstrate the utility, we test our approach on real-world data, consistently achieving better generalization to the new distribution; besides, we find that our inferred causal factor can be concentrated in highly explainable semantic regions for the task of image classification.
We summarize our contribution as follows: Methodologically (in sec. 4.1), we propose LaCIM in which the causal assumptions of two latent factors and the distributional shifts are incorporated; Theoretically (in theorem 4.4), we prove the identifiability of the causal factor and the ground-truth predicting mechanism; Algorithmically (in sec. 4.3), guided by the identifiability, we reformulate Variational Bayesian method to learn P (X|S,Z), P (Y |S) for prediction; Experimentally (in sec. 5.2), our approach generalizes better to distributional shifts, compared with others.

heading: 2 Related Work
text: Causality for Domain Generalization. Due to its stable transferability, the concept of causality has been introduced in many recent works for domain generalization [39, 59, 52, 10, 40, 21, 68]. Most of these works learned the assumed (causal) invariance for generalizing to unseen domains. However, they suffer from either i) lacking explicit causal modeling; or ii) inappropriate causal relations made for the output. Specifically, for i), the [39, 59] are still data-driven methods to learn stable correlation (i.e., invariance) without incorporating causal assumptions [51] beyond data, which may impede its generalization to a broader set of domains; for ii), the [52, 10, 40, 21, 68] causally relate the output with covariates, which is inappropriate for sensory-level data.
Our Specification. We explicitly incorporate the causal assumptions. Specifically, we introduce i) latent factors and separate them into the causal and the non-causal factor; ii) the domain variable D, as a selecting mechanism to generate the varied S-Z correlation across domains. Such a causal modeling makes it possible to recover the causal factor S for generalization. In independent and concurrent works, [75] and [28] also explore latent variables in causal relation. As comparisons, [75] did not differentiate S from Z. The spurious correlation in [28] is limited in the correlation between domains and the output; while it is allowed in our setting to exist in a single domain, which is more aligned with real scenarios, e.g., the dog is more associated with grass than snow in a domain when most samples are collected in sunny morning.
Other Conceptually Related Works: i) transfer learning that leverages invariance in the context of domain adaptation [60, 81, 17] or domain generalization [43, 63]; (ii) causal inference [51, 53] which builds structural causal models and define intervention (a.k.a, “do-calculus”) for cause-effect reasoning and counterfactual learning; and (iii) latent generative model that assumes generation from latent space to observed data [37, 71] but aims at learning generator in the unsupervised scenario.

heading: 3 Preliminaries
text: Problem Setting. Let X,Y respectively denote the input and output variables. The training data {De}e∈Etrain are collected from multiple environments e ∈ Etrain, where each e is associated with a distribution Pe(X,Y ) over X × Y and De := {xei , yei }i∈[ne]
i.i.d∼ Pe with [k] := {1, ..., k} for any k ∈ Z+. Our goal is to learn a robust predictor f : X → Y that only exploit the causal factor for prediction and generalize well to all domains E ⊃ Etrain. We use respectively upper, lower case letter and Cursive letter to denote the random variable, the instance and the space, e.g., a is an instance in the spaceA of random variable A. ForA := f(X )∩B with B := Rp[i1]×Rp[i2]× ...×Rp[ik], the [f(x)]A denotes the f(x) restricted on dimensions of A, i.e., [f(x)]A := [fi1(x), ..., fik(x)]. The Sobolev space W k,p(A) contains all f such that ∫ A
∣∣∂Afα|A=a∣∣pdµ(a) <∞,∀α ≤ k. Structural Causal Model. The structural causal model (SCM) is defined as a triplet M := 〈G,F , P (ε)〉, in which i) the causal structure G := (V,E) (V,E respectively denote the node and edge set) is described by a directed acyclic graph (DAG); ii) the structural equations F := {fk}Vk∈V are autonomous, i.e., intervening on Vk does not affect others, based on which we can define the dooperator and calculate the causal effect; iii) the P (ε) are probability measure for exogenous variables {εk}k. By assuming independence among {εk}k, we obtain according to Causal Markov Condition that each P that is compatible with G has P({Vk = vk}Vk∈V ) = ΠkP(Vk = vk|Pa(k) = pa(k)). An acyclic directed mixed graph (ADMG) can further allow the existence of bidirectional arrows↔, meaning the spurious correlation between two variables connected.

heading: 4 Methodology
text: We first incorporate the causal assumptions into LaCIM in sec. 4.1. Under such assumptions, we identify the invariant distributions P (X|S,Z) and P (Y |S), which are repectively dubbed as generative invariance and causal invariance that are robust to domain shifts. Equipped with these invariances, we in sec. 4.2 show that the causal factor can be identified without mixing information from non-causal one during prediction. Finally, we introduce our learning method in sec. 4.3 to estimate the P (X|S,Z) and P (Y |S), which are respectively resorted in the inference and prediction that constitute a robust predictor during test stage.

heading: 4.1 Latent Causal Invariance Models
text: In this section, we introduce a set of structural causal models dubbed as Latent Causal Invariance Model (LaCIM), which incorporates the causal assumptions mentioned above and also the source of distributional shifts. The corresponding causal structure of LaCIM is illustrated in Fig. 1 (c), which we will introduce step-by-step from the skeleton in Fig. 1 (a).
Fig. 1 (a). Specifically, the ADMG in Fig. 1 (a) introduces latent factors V := {S,Z} to model the abstractions/concepts that generate the observed variables (X,Y ), as similarly assumed in
unsupervised latent generative models [37] for image generation. Further, we explicitly separate the V into S and Z, with only S causally related to the label Y . In image classification, such a causal factor refers to the (shape,contour) of the object need to be classified; while the image X is additionally affected by contextual factor such as light, view.
Fig. 1 (a)→ Fig. 1 (b). In addition, we assume that S is spuriously correlated with Z, as marked by the red “↔” in Fig. 1 (a). Such a spurious correlation corresponds to the bias inherited from data, e.g. the contextual information in image classification. Therefore, the magnitude of this correlation is distribution-dependent and thus can vary across domains. Statistically, the “spurious correlation" implicates the presence of a third unobserved (we use dot circle to represent unobserved variables) confounder, which is denoted as C in Fig. 1 (b). The unblocked path from Z to Y induced by C can lead to learning the non-causal factor during data-fitting, which can degrade the performance on unseen domains if the correlation between this non-causal factor and the output is broken.
Fig. 1 (b)→ Fig. 1 (c). Taking a further inspection in Fig. 1 (b), the varying degree of correlation can be either due to the distributional shift of S,Z|C or of the C itself across domains (we use red color to mean varied distributions). As both shifts are domain-dependent, we in Fig. 1 (c) ascribe them to a domain variable D, which causes the mutation of its children nodes’ distribution, i.e., S,Z and C. Such a domain variable has been similarly introduced in [69, 68] to generate mutable variables. In our scenario, we do not require D to be observed; rather, we only need the domain index d̃e (one-hot encoded vector with length m := |Etrain|). The set of SCMs augmented with D, with the SCM Markovian compatible to the DAG of C, S, Z,X, Y in Fig. 1 (c), is dubbed as Latent Causal Invariance Models (LaCIM) that is formally defined as follows: Definition 4.1 (LaCIM). The LaCIM denotes a set of SCMs augmented with the domain variable D, i.e., {〈Me, de〉}e∈E , in which de denotes the value of D and Me := 〈G,Fe, P (ε)〉 for e. The G denotes the DAG restricted on C, S, Z,X, Y . For each environment/domain e, the Fe := {fx, fy, fes , fez , fec } correspond to generating mechanism ofX,Y, S, Z,C, with fec (εc) := gc(εc, de), fes (c, εs) := gs(c, εs, d e) and fez (c, εz) := gz(c, εz, d e) from some gc, gs, gz .
Remark 1. Different from scenarios in which X generates [28] nor generated from Y [1], we consider the scenario when the X and Y are generated concurrently, which can widely exist but ignored in the literature. For example, the clinicians are recording the disease status while implementing the ultrasound test at the same time, during medical diagnosis.
As an illustration, we consider the following example, in which the distributional shifts caused by domain variable D can refer to sampling bias in data.
Example 4.1 (Sampling Bias). Consider the cat/dog classification, in which the animal in each image is either associated with the snow or grass. The D refers to the sampler, which generates the C that denotes time and weather to collect each sample. The S,Z respectively refer to the features of animals and context. Since each sampler may have a fixed sampling pattern (e.g. gets used to going out in the sunny morning (or in the snowy evening)), the data one collects may have sampling bias: dogs (cats) more associated with grass (snow) in the sunny morning (or snowy evening).
The Def. 4.1 specifies the generating mechanisms across environments and how they differ. Equipped with such a specification, we can identify the invariant mechanisms that are stable to domain shifts:
Proposition 4.2 (Causal Invariance & Generative Invariance). For LaCIM in Def. 4.1, the P (Y |S) and P (X|S,Z) are invariant to shifts across E , and are respectively denoted as Causal Invariance (CI) and Generative Invariance (GI). Remark 2. The generating process from latent variables to observed variables follows from physical law, e.g., the shape, contour, color, view, light should satisfy physical constraints to generate a reasonable image. Therefore, it is naturally hold that such generating processes are invariant.
The P (X|S,Z) and P (Y |S) can induce an invariant predicting mechanism. Specifically, for a new sample x← fx(s?, z?, εx), y ← fy(s?, εy), we can first infer the causal factor s? from pfx(x|s, z) by maximizing log-likelihood of pfx(x|s, z) over S ×Z and then feed the estimated s into pfy (y|s?) for prediction. To ensure the robustness of such a two-step invariant prediction, we need to answer two following identifiability questions:
1. Can the inferred causal factor S not mix the information of (disentangled from) others? 2. Can such an invariant predictor recover the ground-truth predictor P (Y |s?)?
We will answer these questions in the subsequent section, followed by our learning methods to identify the causal factor and the causal/generative invariance for prediction.

heading: 4.2 Identifiability Analysis
text: We present the identifiability results regarding (i) the disentanglement of inferred causal factor S from non-causal Z, and (ii) the induced true predicting mechanism P (Y |s?) for x← fx(s?, z?, εx), which respectively echo the two questions imposed in the last section.
Our main results are presented in theorem 4.4. To distinguish the causal factor S from others, our results require that the degree of diversity of S-Z correlation across environments is large enough, which has been similarly assumed in the literature of identifiability [52, 1]. Such a diversity condition implies the dramatical change of correlation between Z and Y , thus providing a clue to disentangle the S. Such a disentanglement analysis, is crucial to causal prediction but is ignored in existing literature about identifiability, such as those identifying the discrete latent confounders [32, 62], or those relying on Additive Noise Model (ANM) assumption [31], or linear Independent Component Analysis (ICA) [14, 35, 36, 75] (Please refer to supplement D.1 for more exhaustive reviews). More importantly, we will later in theorem 4.5 show the extension of above analysis from exponential family of P (S,Z|C) to Sobelev space; and from ANM for Y to categorical distribution for Y .
We assume the ANM for fx(s, z, εx)= f̂x(s, z) + εx (we replace f̂x with fx for simplicity), which has been widely adopted to identify the causal factor [30, 54, 35]. We assume the fx to be bijective and invertible (we will discuss it later). We first narrow our interest to a subset of LaCIM denoted as Pexp in which any model in Pexp satisfies that (i) the S,Z belong to the exponential family; and (ii) the Y is generated from the ANM:
Pexp = {
LaCIM with any m > 0| y = fy(s) + εy, pe(s, z|c) := Πt=s,zpTt,Γt c,de
(t|c),∀e } ,with
pTt,Γt c,de (t) = qt∏ i=1 exp ( kt∑ j=1 T ti,j(ti)Γ t c,de,i,j +Bi(ti)−Atc,de,i ) ,∀kt, qt (1)
for t = s, z and e ∈ E , with qt, kt respectively denoting the dimension of t = s, z and the number of natural parameters in each dimension. The {T ti,j(ti)}, {Γtc,de,i,j} denote the sufficient statistics and natural parameters, {Bi} and {Atc,de,i} denote the base measures and normalizing constants to ensure the integral of distribution equals to 1. Let Tt(t) := [Tt1(t1), ...,Ttqt(tqt)] ∈ Rkt×qt ( Tti(ti) := [T t i,1(ti), ..., T t i,kt(ti)], ∀i ∈ [qt] ) , Γtc,de := [ Γtc,de,1, ...,Γ t c,de,qt ] ∈ Rkt×qt ( Γtc,de,i := [Γtc,de,i,1, ...,Γ t c,de,i,kt ], ∀i ∈ [qt] ) . We further assume that the P e(C) serves to discrete distributions
on the set {c1, ..., cR}, with which the pe(s, z) := ∫ p(s|c)p(z|c)dP e(c) = ∑ r p
e(s, z|cr)pe(cr) can be regarded as the mixture of exponential family distributions. Rather than uniquely inference, we target on disentangling the S from Z and also recovering the ground-truth predictor, which is formally defined as ∼exp-identifiability as follows: Definition 4.3 (∼exp-identifiability). Suppose the X ⊇ fx(S × Z). We define a binary relation θ ∼exp θ̃ on the parameter space of X × Y: there exist two sets of permutation matrices and vectors, (Ms, as) and (Mz, az) for s and z respectively, such that for any (x, y) ∈ X ×Y , the following hold:
T̃s([f̃−1x ]S(x)) = MsT s([f−1x ]S(x)) + as, T̃ z([f̃−1x ]Z(x)) = MzT z([f−1x ]Z(x)) + az; (2)
pf̃y (y|[f̃ −1 x ]S(x)) = pfy (y|[f−1x ]S(x)). (3)
We then say that θ is∼exp-identifiable, if for any θ̃, peθ(x, y) = peθ̃(x, y) ∀e ∈ Etrain, implies θ ∼exp θ̃.
This definition is inspired by but beyond the scope of unsupervised scenario considered in nonlinear ICA [27, 35] in that, the former further disentangle S from Z (in Eq. (2)) and identify the true predicting mechanism (in Eq. (3)). To see disentanglement, note that for any clean (noise-free) sample x← fx(s?, z?), the Eq. (2) ensures that the inferred causal factor T̃s([f̃−1x ]S(x)) does not mix the information of others, unless the extreme case that there is a deterministic function between S and Z, in which it is impossible for S to be identified. With such an identification of s, the Eq. (3) further guarantees that the learned pf̃y (y|[f̃
−1]S(x)) can recover the ground-truth prediction probability density, i.e., pfy (y|[f−1x ]S(x)) = pfy (y|s?). With noise, the s? can be inferred with some indeterminacy. The formal result is presented in theorem 4.4.
Theorem 4.4 (∼exp-identifiability). For θ of Pexp in Def. 4.1 with m := |Etrain|, we have that the θ is ∼exp identifiable under following assumptions:
1. The characteristic functions of εx, εy are almost everywhere nonzero. 2. fx, f ′x, f ′′ x are continuous and fx, fy are bijective; 3. The {T ti,j}1≤j≤kt are linearly independent in S or Z for each i ∈ [qt] for any t = s, z; and T ti,j are twice differentiable for any t = s, z, i ∈ [qt], j ∈ [kt];
4. The { ( Ts([f−1]S(x)),T z([f−1]Z(x)) )
;B(x) > 0} contains a non-empty open set in Rqs×ks+qz×kz , with B(x) := ∏ is∈[qs]Bis([f −1]is(x)) ∏ iz∈[qz ]Biz ([f
−1]iz (x)). 5. The L := [P e1(C)T, ..., P em(C)T]T ∈ Rm×R and [ [Γt=s,zc2,de1 − Γ t=s,z c1,de1
]T, ..., [Γt=s,zcR,dem − Γt=s,zc1,de1 ] T ]T ∈ R(R×m)×(qt×kt) have full column rank.
The assumptions 1-3 are trivial and easy to satisfy. The characteristics functions of εx, εy can be almost everywhere non-zero for most continuous variables, such as Gaussian, exponential, beta, gamma distribution. This assumption can ensure the identifiability of p(f−1(x), as will be shown in the appendix. The bijectivity of fx and fy have been widely assumed in [30, 54, 53, 35, 75] as a basic condition for identifiability. It naturally holds for fx to be bijective since it has been empirically proven in auto-encoder [38] that the low-dimension embeddings (i.e., qs + qz < qx) can recover the original input well and also that the variational auto-encoder can extract meaningful representations from x. For the θ with categorical Y such that p(y = k|s) = [fy]k(s)/ ( ∑ k[fy]k(s)), the fy may not satisfy the bijectivity condition. We will shown identifiability for such a categorical case later in theorem 4.5. The assumption 3 can be uniformly satisfied for all distributions in the strongly exponential family. The containment of an open set in assumption (4) for { ( Ts([f−1]S(x)),T z([f−1]Z(x)) )
;B(x) > 0} implies that space expanded by sufficient statistics are dense in some open set, as a sufficient condition for the mixture distribution P e(C) and also P e(X,Y |c) to be identified. The diversity assumption (5) implies that i) m ≥ R and m ∗ R ≥ max(kz ∗ qz, ks ∗ qs) + 1; and that ii) different environments are diverse enough in terms of S-Z correlation, as an almost a necessary for the invariant one to be identified (a different version is assumed in [1]). In supplement B.2, we will show that the ii) can hold unless the space of Γ belong to a zero-(Lebesgue) measure set. As indicated by the formulation, a larger m would be easier to satisfy the condition, which agrees with the intuition that more environments can provide more complementary information. Besides, our result can be extended to non-independent case among {s1, ..., sqs} (or {z1, ..., zqz}), i.e., pTt,Γt
c,de (t) = exp(〈Tt(t),Γtc,de〉+B(t)−Atc,de) (t = s, z), which will shown in supplement B.2.
Extension to the general forms of LaCIM. We extend to general forms of LaCIM in theorem 4.5 as long as its P(S,Z|C = c) ∈W r,2(S × Z) (for some r ≥ 2) and categorical Y , in the following theorem. This is accomplished by proving that any model in LaCIM can be approximated by a sequence of distributions with parameterization in Pexp, motivated by [3] that the exponential family is dense in the set of distributions with bounded support, and in [44] that the continuous variable with multinomial logit model can be approximated by a series of distributions with i.i.d Gumbel noise as the temperature converges to infinity. The proof is left in the supplement.
Theorem 4.5 (Asymptotic∼exp-identifiability). Suppose the LaCIM satisfy that p(x|s, z) and p(y|s) are smooth w.r.t s, z and s respectively. For each e and c ∈ C, suppose Pe(S,Z|c) ∈W r,2(S×Z) for some r ≥ 2, we have that the LaCIM is asymptotically∼exp-identifiable: ∀ > 0, ∃ ∼exp-identifiable P̃θ ∈ Pexp, s.t. dPok(pe(X,Y ), p̃eθ(X,Y )) < ,∀e ∈ Etrain 3.
Our proof is built on [3] that any probability in Sobolev space can be approximated by a sequence of distribution with the number of natural paramters going to infinity, i.e., kt →∞.

heading: 4.3 Learning and Inference
text: Guided by the identifiability result, we propose to learn P (X|S,Z) and P (Y |S) via generative modeling following from Fig. 1 (c). Then to predict the label for a new sample x generated from (s?, z?), we first leverage the learned p(x|s, z) to infer s? that is ensured to be able to not mix the non-causal information, followed by learned P (y|s̃?) for prediction.
3The dPok(µ1, µ2) denotes the Pokorov distance between µ1 and µ2, with limn→∞ dPok(µn, µ) → 0 ⇐⇒ µn d→ µ.

heading: 4.3.1 Learning Method
text: To learn the P (X|S,Z), P (Y |S) for invariant prediction, we reformulate the objective function of Variational Auto-Encoder (VAE) in the supervised scenario, in order to fit {pe(x, y)}e∈Etrain . As a latent generative model, the VAE was originally proposed for unsupervised generation from latent variables V to high-dimensional input variable X . To make such a generation tractable, the VAE introduced a variational distribution qψ parameterized by ψ to approximate the intractable posterior by
maximizing the following Evidence Lower Bound (ELBO):−Lθ,ψ = Ep(x) [ Eqψ(v|x) log pθ(x,v) qψ(v|x) ] ≤ Ep(x)[log pθ(x)], where the equality is achieved when qψ(v|x) = pθ(v|x). Therefore, maximizing the ELBO over pθ and qψ will drive (i) qψ(v|x) to approximate pθ(v|x); (ii) pθ to estimate the ground-truth model p.
To adapt the above surrogate loss to our DAG in Fig. 1 (c), we introduce the variational distribution qeψ(s, z|x, y) for each environment e. The corresponding ELBO for e is
−Leθ,ψ ∆ =Epe(x,y) [ Eqeψ(s,z|x,y) log peθ(x, y, s, z) qeψ(s, z|x, y) ] , where peθ(x, y, s, z) = pθ(x|s, z)pθ(y|s)pe(s, z).
Similarly, minimizing Leθ,ψ can drive pθ(x|s, z), pθ(y|s) to approximate the p(x|s, z), p(y|s), and also qeψ(s, z|x, y) to estimate peθ(s, z|x, y). Therefore, the qψ can inherit the properties of pθ. As peθ(s, z|x, y)= peθ(s,z|x)pθ(y|s) peθ(y|x) for our DAG in Fig. 1 (c), we can similarly reparameterize qeψ(s, z|x, y) as qeψ(s,z|x)pθ(y|s)
qeψ(y|x) with qψ(y|s) replaced by pθ(y|s) (since the goal of qψ is to mimic the behavior of
pθ). Then, the Leθ,ψ can be rewritten as:
Leθ,ψ = Epe(x,y) [ − log qeψ(y|x)− Eqeψ(s,z|x)
pθ(y|s) qeψ(y|x) log pθ(x|s, z)peθ(s, z) qeψ(s, z|x)
] , (4)
where qeψ(y|x) = ∫ S q e ψ(s|x)pθ(y|s)ds. We correspondingly parameterize the prior model peθ(s, z) and inference model qeψ(s, z|x) as pθ(s, z|d̃e) and qψ(s, z|x, d̃e), in which d̃e (of environment e) denotes the domain index that can be represented by the one-hot encoded vector with length m := |Etrain|. The overall loss function is:
Lθ,ψ ∆ = ∑ e∈Etrain Leθ,ψ. (5)
The training datasets {De}e∈Etrain are applied to optimize the prior models {p(s, z|d̃e)}e, inference models {qψ(s, z|x, d̃e)}e, generative model pθ(x|s, z) and predictive model pθ(y|s). Particularly, the parameters of pθ(x|s, z) and pθ(y|s) are shared among all environments, motivated by the the invariance property of P (X|S,Z) and P (Y |S) across all domains.

heading: 4.3.2 Inference & Prediction.
text: We leverage the learned P (X|S,Z), P (Y |S) for prediction. According to Prop. 4.2 and Eq. (3) in theorem 4.4, the induced predictor via P (X|S,Z), P (Y |S) can recover the true predicting mechanism for any distributional shifts from E . Specifically, for any x generated by (s?, z?), we first optimize the following log-likelihood of pθ(x|s, z) over S × Z to infer s?, z?,
max s,z
log pθ(x|s, z) + λs‖s‖22 + λz‖z‖22, (6)
with hyperparameters λs > 0 and λz > 0 in order to control the learned s, z in a reasonable scale. Note that Eq. Eq. (6) is different from the maximum a posterior estimation since the posterior qeψ(s, z|x) is parameterized differently for different e while the pθ(x|s, z) is invariantly parameterized for E (this is because p(x|s, z) is invariant). For optimization, we adopt the strategy in [61] that first sample some candidate points from N (0, I) and select the optimal one in terms of Eq. (6) as initial point; then use Adam to optimize for another T iterations. The implementation details and optimization effect are shown in supplement E.2. Finally, with estimated s̃?, z̃?, we implement the learned pθ(y|s̃?) for prediction: ỹ := arg maxy pθ(y|s̃?).

heading: 5 Experiments
text: We first verify the identifiability claims of theorem 4.4 in sec. 5.1. Then we evaluate LaCIM on real-world data in sec. 5.2: Non-I.I.D. Image dataset with Contexts (NICO); Colored MNIST (CMNIST); Alzheimer’s Disease Neuroimaging Initiative (ADNI www.loni.ucla.edu/ADNI for early prediction of Alzheimer’s Disease), to verify the generalization ability of our method on the target domain with distributional shifts.

heading: 5.1 Simulation
text: To verify the identifiability claims, we implement LaCIM on synthetic data. We generate C, S, Z,X, Y following Fig. 1 (with details left in supplementary). We choose m = 3, 5 with the same total number of samples. To verify the advantage of learning on multiple diverse domains (m > 1), we compare with pool-LaCIM: minimizing the loss Eq. (4) on the pooled data from all m domains. We compute the mean correlation coefficient (MCC) adopted in [35], which measures the goodness of identifiability under permutation by introducing cost optimization to assign each learned component to the source component. We run all methods for 100 times, with the average recorded in Fig. 2a. The superiority of LaCIM over pool-LaCIM, together with the fact that LaICM with m = 5 performs better than m = 3, verify the benefit of more domains to satisfy the diversity condition. To illustrate the learning effect, we visualize the learned Z (with S left in supplement E.1) in Fig. 2b.

heading: 5.2 Real-world Data
text: We verify the generalization ability of LaCIM on three data: NICO, CMNIST and ADNI.
Dataset. We describe the datasets as follows (X,Y denotes the input and output; D is unobserved):
• NICO. We consider the cat/dog classification in “Animal” dataset in NICO, a benchmark for non-i.i.d problem in [20]. Each animal is associated with “grass”,“snow” contexts. The D denotes the attributes of the sampler. The C denotes the time and weather of sampling, which generates the S,Z that respectively denote the semantic and contextual features. We split the dataset into m training domains and the test domain, in which each domain has different proportions of contexts associated with each animal, i.e., (%cat in grass, %cat in snow, %dog in grass, %dog in snow), due to different sampling strategies determined by D. The proportion vectors of all domains are introduced in Tab. 3. The distributional shift refers to the spurious correlation between the context and the label.
• CMNIST: We relabel the digits 0-4 and 5-9 as y = 0 and y = 1, based on MNIST. Then we color pe (1 − pe) of images with y = 0 (y = 1) as green and others as red. We set m = 2 with pe1 = 0.95, pe2 = 0.99; while the petest for the test domain is set to 0.1. The D denotes the attributes of the painter. The Z, S respectively represent the features related to the color and the digit. Their confounder C denotes the time and weather for which the painter D draws the number and color, e.g., the painter tends to draw red 0 more often than green 1 in the sunny morning. In this regard, the distributional shift refers to the spurious correlation between the color and the label.
• ADNI. The Y := {0, 1, 2}, with 0,1,2 respectively denoting Normal Control, Mild Cognitive Impairment and AD. The X is structural Magnetic resonance imaging. We split the data into m = 2 training domains and the test domain, with different values of D that denotes Age, TAU (a biomarker [24]). The C, S (Z) respectively denote the hormone level that
affects the brain structure development and the disease-related (-unrelated) brain regions. The distributional shifts among all domains are due to different values of D.
Compared Baselines & Implementation Details. We compare with (i) Empirical Risk Mnimization from X → Y (ERM), (ii) domain-adversarial neural network (DANN) [15], (iii) Maximum Mean Discrepancy with Adversarial Auto-Encoder (MMD-AAE) [43], (iv) Domain Invariant Variational Autoencoders (DIVA) [29], (v) Invariant Risk Mnimization (IRM) [1], (vi) Supervised VAE (sVAE): our LaCIM implemented by VAE without disentangling S,Z. For all methods, the network structures of qeψ(s, z|x), pθ(x|s, z) and pθ(y|s) for CMNIST, NICO and ADNI are shared (details introduced in supplement E.4, E.5, E.6, Tab. 7, 8). We implement SGD as optimizer, with learning rate (lr) 0.5 and weight decay (wd) 1e-5 for CMNIST; lr 0.01 with decaying 0.2× every 60 epochs, wd 5e-5 for NICO and ADNI (wd is 2e-4). The batch-size are set to 256, 30 and 4 for CMNIST, NICO, ADNI.
Main Results & Discussions. We report accuracy over 10 runs for each method. As shown in Tab. 1, our LaCIM consistently outperforms others on all data. Specifically, the advantage over IRM and ERM may due to the incorporation of causal assumptions embedded in Fig. 1 (c). Further, the improvement over sVAE is benefited from the separation of S from others to avoid spurious correlation. Besides, a larger m (with the total sample size fixed) can bring further benefit on NICO, which may due to the easier satisfaction of the diversity condition in theorem 4.4.
Interpretability. We visualize the learned S andZ on CMNIST and NICO. Specifically, for CMNIST, we visualize the generated image (with only digit “0” among all classes that belong to Y = 0 and digit “7” among all classes that belong to Y = 1) by interpolating S (and Z) with fixed Z (and S); for NICO, we adopt the gradient method [67], which visualizes the derivatives of the S? (i.e., dimension of S that has the highest correlation with Y ) with respect to each image.
As shown in Fig. 3a, the generated sequential images in the 1st and 2nd row look more like “7” from “0” as s increases; while the sequential images in the 2nd-row change from red to green as z increases. Besides, different dimensions of S can learn different differentiating semantic information. For example, the first dimension can learn to add the dash in the hand-writing "7"; while the second dimension can learn to remove the left part of "0" to "7" as interpolated. For the dimension of Z, it learned other non-differentiating factors such as width, color. This result reflects that the learned S and Z correspond to the digit (causal factor of Y ) and color-related features. For NICO, Fig. 3b shows the ability of identifying more explainable semantic features of LaCIM than ERM in which the learned features can mix the background information. Supplement E.5 provides more results.

heading: 6 Conclusions & Discussions
text: We propose recovering latent causal factor that is robust to distributional shifts caused by a domain variable. We introduce the causal and non-causal latent factors that are spuriously correlated with each other, and generate the input and the output via invariant mechanisms. Under this invariance, the causal factor is guaranteed to be disentangled from the non-causal one, which induces the groundtruth predictor that holds on all domains. A reformulated generative model is proposed for inferring the causal factor and prediction. A possible drawback of our model lies in our requirement of the number of environments for identifiability, the relaxation of which is left in future work.
Broader Impact
We claim that this work does not present any foreseeable negative social impact.

Review summaries:
This paper introduces the "Latent Causal Invariance Model" (LaCIM), a structural variational autoencoder-like approach for learning to isolate causal factors from "spurious" ones in prediction to improve generalization to new environments. The authors assume a latent generative structure for the data, with separate latent constructs for causal and contextual factors. Under strong parametric assumptions and requirements on the number of source environments the authors prove identifiability of the causal factors. Experiments on imaging datasets (2 benchmarks and a small medical dataset) show LaCIM outperforms some existing domain adaptation baselines as well as Invariant Risk Minimization (IRM). Further, qualitative assessment of images generated by ablating the causal and contextual factors show disentanglement of these factors.

This work proposes an identifiable causal model with latent variables to model spurious correlations between labels and irrelevant context (such as background or light setting) in order to construct a classifier which is robust to certain distribution shifts. The authors show under which conditions the irrelevant context and important features can be disentangled from only input observations and labels. They propose a VAE-like approach to learn the model which they use to come up with an invariant classifier. The identifiability of the model is assessed on synthetic data and the robustness to dataset shift is evaluated on multiple real world data sets, drawing a favorable picture of their contribution. 

To deal with the issue of distributional shifts between training and testing domains, the authors propose Latent Causal Invariance Models (LaCIM) consisting of both causal latent factors and non-causal latent factors and the extent of the correlations between them is governed by a domain variable. They theoretically show the identifiability of the causal latent factors and the ground-truth predicting mechanism in the proposed LaCIM. Based on the identifiability, they learn the model by reformulating VAE and then verify it on various real-world data. 

In this paper, the authors propose a method to estimate the invariant causal mechanism. By modeling the process as the causal graph in Fig.1 (c), the authors first provide the condition when disentangling causal factors is feasible. Based on the theoretical result, they present the method to infer causal factor and non-causal factor with X based on VAE.