LATENT CAUSAL INVARIANT MODEL
Year: 2020

Abstract:
Current supervised learning can learn spurious correlation during the data-fitting process, imposing issues regarding interpretability, out-of-distribution (OOD) generalization, and robustness. To avoid spurious correlation, we propose a Latent Causal Invariance Model (LaCIM) which pursues causal prediction. Specifically, we introduce latent variables that are separated into (a) output-causative factors and (b) others that are spuriously correlated to the output via confounders, to model the underlying causal factors. We further assume the generating mechanisms from latent space to observed data to be causally invariant. We give the identifiable claim of such invariance, particularly the disentanglement of output-causative factors from others, as a theoretical guarantee for precise inference and avoiding spurious correlation. We propose a Variational-Bayesian-based method for estimation and to optimize over the latent space for prediction. The utility of our approach is verified by improved interpretability, prediction power on various OOD scenarios (including healthcare) and robustness on security.

heading: 1 INTRODUCTION
text: Current data-driven deep learning models, revolutionary in various tasks though, heavily rely on i.i.d data to exploit all types of correlations to fit data well. Among such correlations, there can be spurious ones corresponding to biases (e.g., selection or confounding bias due to coincidence of the presence of the third factor) inherited from the data provided. Such data-dependent spurious correlations can erode the (i) interpretability of decision-making, (ii) ability of out-of-distribution (OOD) generalization, i.e., extrapolation from observed to new environments, which is crucial especially in safety-critical tasks such as healthcare, and (iii) robustness to small perturbation (Goodfellow et al., 2014).
Recently, there is a Renaissance of causality in machine learning, expected to pursue causal prediction (Schölkopf, 2019). The so-called “causality” is pioneered by Judea Pearl (Pearl, 2009), as a mathematical formulation of this metaphysical concept grasped in the human mind. The incorporation of a priori about cause and effect endows the model with the ability to identify the causal structure which entails not only the data but also the underlying process of how they are generated. For causal prediction, the old-school methods (Peters et al., 2016; Bühlmann, 2018) causally related the output label Y to the observed input X , which however is NOT conceptually reasonable in scenarios with sensory-level observed data (e.g. modeling pixels as causal factors of Y does not make much sense).
For such applications, we rather adopt the manner in Bengio et al. (2013); Biederman (1987) to relate the causal factors of Y to unobserved abstractions denoted by S, i.e., Y ← fy(S, εy) via mechanism fy. We further assume existence of additional latent components denoted as Z, that together with S generates the input X via mechanism fx as X ← fx(S,Z, εx). Taking image classification as an example, the S and Z respectively refer to object-related abstractions (e.g., contour, texture, color) and contextual information (e.g., light, view). Such an assumption is similarly adopted in the literature of nonlinear Independent Components Analysis (ICA) (Hyvarinen and Morioka, 2016; Hyvärinen et al., 2019; Khemakhem, Kingma and Hyvärinen, 2020; Teshima et al., 2020) and latent generative models (Suter et al., 2019), which are however without separation of output (y)-causative factors (a.k.a, S) and other correlating factors (a.k.a, Z) that can both be learned in data-fitting process.
We encapsulate these assumptions into a novel causal model, namely Latent Causal Invariance Model (LaCIM) as illustrated in Fig. 1, in which we assume the structural equations fx (associated with S,Z → X), fy (associated with S → Y ) to be the Causal Invariant Mechanisms (CIMe) that hold under any circumstances with P(S,Z) allowed to be varied across domains. The incorporation of these
priories can explain the spurious correlation embedded in the back-door path from Z to Y (contextual information to the class label in image classification). To avoid learning spurious correlations, our goal is to identify the intrinsic CIMe fx, fy. Specifically, we first prove the identifiability (i.e., the possibility to be precisely inferred up to an equivalence relation) of the CIMe. Notably, far beyond the scope in existing literature (Khemakhem, Kingma and Hyvärinen, 2020), our results can implicitly, and are the first to disentangle the output-causative factors (a.k.a, S) from others (a.k.a, Z) for prediction, to ensure the isolation of undesired spurious correlation. Guaranteed by such, we propose to estimate the CIMe by extending the Variational Auto-encoder (VAE) (Kingma and Welling, 2014) to the supervised scenario. For OOD prediction, we propose to optimize over latent space under the identified CIMe. To verify the correctness of our identifiability claim, we conduct a simulation experiment. We further demonstrate the utility of our LaCIM via high explainable learned semantic features, improved prediction power on various OOD scenarios (including tasks with confounding and selection bias, healthcare), and robustness on security.
We summarize our contribution as follows: (i) Methodologically, we propose in section 4.1 a latent causal model in which only a subset of latent components are causally related to the output, to avoid spurious correlation and benefit OOD generalization; (ii) Theoretically, we prove the identifiability (in theorem 4.3) of CIMe fx, fy from latent variables to observed data, which disentangles outputcausative factors from others; (iii) Algorithmically, guided by the identifiability, we in section 4.3 reformulate Variational Bayesian method to estimate CIMe during training and optimize over latent space during the test; (iv) Experimentally, LaCIM outperforms others in terms of prediction power on OOD tasks and interpretability in section 5.2, and robustness to tiny perturbation in section 5.3.

heading: 2 RELATED WORK
text: The invariance/causal learning proposes to learn the assumed invariance for transferring. For the invariance learning methods in Krueger et al. (2020) and Schölkopf (2019), the “invariance” can refer to stable correlation rather than causation, which lacks the interpretability and impedes its generalization to a broader set of domains. For causal learning, Peters et al. (2016); Bühlmann (2018); Kuang et al. (2018); Heinze-Deml and Meinshausen (2017) assume causal factors as observed input, which is inappropriate for sensory-level observational data. In contrast, our LaCIM introduces latent components as causal factors of the input; more importantly, we explicitly separate them into the output-causative features and others, to avoid spurious correlation. Further, we provide the identifiability claim of causal invariant mechanisms. In independent and concurrent works, Teshima et al. (2020) and Ilse et al. (2020) also explore latent variables in causal relation. As comparisons, Teshima et al. (2020) did not differentiate S from Z; and Ilse et al. (2020) proposed to augment intervened data, which can be intractable in real cases.
Other works which are conceptually related to us, as a non-exhaustive review, include (i) transfer learning which also leverages invariance in the context of domain adaptation (Schölkopf et al., 2011; Zhang et al., 2013; Gong et al., 2016) or domain generalization (Li et al., 2018; Shankar et al., 2018); and (ii) causal inference (Pearl, 2009; Peters et al., 2017) which proposes a structural causal model to incorporate intervention via “do-calculus” for cause-effect reasoning and counterfactual learning; (iii) latent generative model which also assumes generation from latent space to observed data (Kingma and Welling, 2014; Suter et al., 2019) but aims at learning generator in the unsupervised scenario.

heading: 3 PRELIMINARIES
text: Problem Setup & Notation LetX,Y respectively denote the input and output variables. The training data {De}e∈Etrain are collected from the set of multiple environments Etrain, where each domain e is associated with a distribution Pe(X,Y ) over X × Y and De = {xei , yei , de}i∈[ne]
i.i.d∼ Pe with [k] := {1, ..., k} for any k ∈ Z+. The de∈{0, 1}m denotes the one-hot encoded domain index for e, where 1 ≤ m := ∣∣Etrain∣∣ ≤ n := ∑e∈Etrain ne. Our goal is to learn a model f : X 7→ Y that learns output (y)-causative factors for prediction and performs well on the set of all environments E ⊃ Etrain, which is aligned with existing OOD generalization works (Arjovsky et al., 2019; Krueger et al., 2020). 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 space A of random variable A. The [f ]A denotes the f restricted on dimensions of A. The Sobolev space W k,p(A) contains all f such that ∫ A ∣∣∂Afα∣∣A=a∣∣pda <∞,∀α ≤ k.
Structural Causal Model. The structural causal model (SCM) is defined as the causal graph assigned with structural equations. The causal graph encodes the assumptions in missing arrows in a directed acylic graph (DAG): G = (V,E) with V,E respectively denoting the node set and the edge set. The Pa(k) denotes the set of parent nodes of Vk for each Vk ∈ V and the X → Y ∈ E indicates the causal effect of X on Y . The structural equations {Vk ← fk(Pa(k), εk)}Vk∈V , quantify the causal effects shown in the causal graph G. By assuming independence among exogenous variables {εk}k, the Causal Markov Condition states that P({Vk = vk}Vk∈V ) = ΠkP(Vk = vk|Pa(k) = pa(k)). A back-door path from Va to Vb is defined as a path that ends with an arrow pointing to Va (Pearl, 2009).

heading: 4 METHODOLOGY
text: We build our causal model associated with Causal Invariant Mechanism (CIMe, i.e., fx, fy) and a priori about the generating process in section 4.1, followed by our identifiability result for CIMe in section 4.2. Finially, we introduce our learning method to estimate CIMe in section 4.3.
4.1 LATENT CAUSAL INVARIANCE MODEL
We introduce latent variables to model the abstractions/concepts that play as causal factors that generate the observed variables (X,Y ), which is more reasonable than assuming the X as the direct cause of Y in scenarios with sensory-level data. We explicitly separate the latent variables into two parts: the S and Z that respectively denote the y (output)-causative and y-non-causative factors, as shown by the arrow S → Y in Fig. 1. Besides, the X and Y are respectively generated by S,Z and S, via structural equations (with noise) fx, fy, which are denoted as Causal Invariant Mechanisms (CIMe) that hold across all domains. The output Y denotes the label generated by human knowledge, e.g., the semantic shape, the contour to discern the object, etc. Hence, we assume the Y as the outcome/effect of these high-level abstractions (Biederman, 1987) rather than the cause (detailed comparison with Y → S is left in supplementary 7.7.1). We call the model associated with the causal
graph in Fig. 1 as Latent Causal Invariance Model (LaCIM), with formal definition given in Def. 4.1.
As an illustration, we consider the image classification in which X,Y denote the image and the class label. Instead of X , i.e., the pixels, it is more reasonable to assume the causal factors (of X,Y ) as latent concepts (S,Z) that can denote light, angle, the shape of the object to generate X following the physical mechanisms. Among these concepts, only the ones that are causally related to the object, i.e., S (e.g., shape) are causal factors of the object label, i.e., Y . Following the physical or natural law, the mechanisms S,Z → X,S → Y invariantly hold across domains. The S := Rqs ,Z := Rqz denote the space of S,Z, with Pe(S,Z) (that characterizes the correlation between S and Z) varying across E (e.g., the object is more associated with a specific scene than others). We assume that the y-non-causative factor (i.e., Z) is associated with (but not causally related to) S,Y through the confounder C, which is allowed to take a specific value for each sample unit. Therefore, the back-door path Z←C→S→Y induces the correlation between Z and Y in each single domain. Rather than invariant causation, this correlation is data-dependent and can vary across domains, which is known as “spurious correlation”. In real applications, this spurious correlation corresponds to the bias inherited from data, e.g. the contextual information in object classification. This domain-specific S-Z correlation, can be explained by the source variable D, which takes a specific and fixed value for each domain and functions the prior of distribution of the confounder C, as illustrated in Fig. 1. This source variable D can refer to attributes/parameters that characterize the distribution of S,Z in each domain. When such attributes are unobserved, we use the domain index as a substitute. Consider the cat/dog classification task as an illustration, the animal in each image is either associated with the snow or grass. The S,Z respectively denote the concepts of animals and scenes. The D denotes the sampler, which can be described by the proportions of scenes associated with the cat and those associated with the dog. The D generates the C that denotes the (time, weather) to go outside and
collect samples. Since each sampler may have a fixed pattern (e.g. gets used to going out in the sunny morning (or in the snowy evening)), the data he/she collects, may have sample selection bias (e.g. with dogs (cats) more associated with grass (snow) in the sunny morning (or snowy evening) ). In this regard, the scene concepts Z can be correlated with the animal concepts S, and also the label Y .
Definition 4.1 (LaCIM). The Latent Causal Invariance Model (LaCIM) for e ∈ E is defined as a SCM characterized by (i) the causal graph, i.e., the G = (V,E) with V = {C, S, Z,X, Y } and E = {C → S,C → Z,Z → X,S → X,S → Y }; and (ii) structural equations with causal mechanisms {fc, fz, fs, fx, fy} embodying the quantitative causal information: c ← fc(de, εc), z ← fz(c, εz), s← fs(c, εs);x← fx(s, z, εx); y ← fy(s, εy), in which {εc, εz, εs, εx, εy} are independent exogenous variables that induce pfc(c|de), pfz (z|c), pfs(s|c), pfx(x|s, z), pfy (y|s). The CIMe fx, fy are assumed to be invariant across E . We call the environment-dependent parts: Pe(S,Z) and Pe(S,Z|X) as S,Z-prior and S,Z-inference in the following. Remark 1. We denote LaCIM-ds and LaCIM-d as two versions of LaCIM, with the source variable ds with practical meaning (e.g. attributes or parameters of P(S,Z)) observed or not. The observation of ds can be possible in some applications (e.g., age, gender that characterizes population in medical diagnosis). As for the LaCIM-d with ds unobserved, we use domain index D as a substitute.
Denote C as the space of C. We assume that the C is finite union of disjoint sets {Cr}Rr=1, i.e. C := ∪Rr=1Cr, such that for any cr,i 6= cr,j ∈ Cr, it holds that p(s, z|cr,i) = p(s, z|cr,j) for any (s, z). Returning to the cat/dog classification example, the C denotes the range of time to collect samples, i.e., 00 : 00-24 : 00. The C can be divided into several time periods C1, ..., CR, such that the proportion of concepts of (animal,scene) given any c in the same period is unchanged, e.g., the dog often comes up on the grass in the morning. Further, since p(x, y|s, z) = p(x|s, z)p(y|s) is invariant, we have for each Cr that p(x, y|cr,i) = ∫ p(x, y|s, z)p(s, z|cr,i)dsdz = ∫ p(x, y|s, z)p(s, z|cr,j)dsdz = p(x, y|cr,j) for any (x, y). That is, the {p(x, y|cr}cr∈Cr for each (x, y) collapse to a single point, namely p(x, y|cr). In this regard, we have pe(x, y) := p(x, y|de) = ∑R r=1 p(x, y|cr)p(cr|de). Besides, we assume the Additive Noise Model (ANM) for X , i.e., fx(s, z, εx) = f̂x(s, z) + εx (we replace f̂x with fx without loss of generality), which has been widely adopted to identify the causal factors (Janzing et al., 2009; Peters et al., 2014; Khemakhem, Kingma and Hyvärinen, 2020). We need to identify the CIMe (i.e., fx, fy), guaranteed by the identifiability that ensures the learning method to distinguish S from Z to avoid spurious correlation, as presented in section 4.2. Traditionally speaking, the identifiability means the parameter giving rise to the observational distribution pθ?(x, y|de) can be uniquely determined, i.e., pθ(x, y|de) = pθ̃(x, y|de) =⇒ θ = θ̃. Instead of strict uniqueness, we rather identify an equivalent class of θ? (in Def. 4.2) that suffices to disentangle the y-causative features S from Z to avoid learning spurious correlation. To achieve this goal, we first narrow our interest in case when p(s, z|c) is exponential family in Eq. (1), in which we can respectively identify the S,Z up to linear and point-wise transformations given by theorem 4.3; then we generalize to any p(s, z|c) as long as it belongs to Sobolev space, as explained in theorem 4.4. A reformulated VAE is proposed to learn the CIMe practically. For generalization, note that the gap between two environments in terms of prediction given x, i.e.,
∣∣Epe2 [Y |X = x] − Epe1 [Y |X = x] ∣∣ = ∫S ∣∣pe2(s|x) − pe1(s|x)∣∣pfy (y|s)ds, is mainly due to the inconsistency of S,Z-inference, i.e., pe(s, z|x) 6= pe′(s, z|x) for e′ 6= e (for details please refer to theorem 7.1 in supplement 7.1). Therefore, one cannot directly apply the trained {pe(s, z|x), pe(y|x)}e∈Etrain to the inference model of new environment, i.e. pe ′ (s, z|x), pe′(y|x) for e′ /∈ Etrain. To solve this problem and generalize to new environment, we note that since pfx(x|s, z) and pfy (y|s) are shared among all environments, we propose to inference s, z that give rise to the test sample x via maximizing the identified pfx(x|s, z), as a pseudo-likelihood of x given (s, z), rather than using S,Z-inference model which is inconsistent among environments. Then, we feed estimated s into invariant predictor pfy (y|s) for prediction.

heading: 4.2 IDENTIFIABILITY OF CAUSAL INVARIANT MECHANISMS
text: We present the identifiability claim about the CIMe fx, fy, which implicitly distinguishes the ycausative factors (a.k.a, S) from others (a.k.a, Z) for prediction, to provide a theoretical guarantee for avoiding spurious correlations. Notably, the S and Z play “asymmetric roles” in terms of generating process, as reflected in additional generating flow from S to Y . This “information intersection” property of S, i.e., f−1y (ȳ) = [f −1 x ]S(x̄) for any (x̄, ȳ) ∈ fx(S,Z) × fy(S) if y = fy(s) + εy, is exploited to disentangle S from Z. Such a disentanglement analysis, is crucial to causal prediction
but lacked in existing literature about identifiability, such as those identifying the discrete latent confounders (Janzing, Sgouritsa, Stegle, Peters and Schölkopf, 2012; Sgouritsa et al., 2013); or those relying on ANM assumption (Janzing, Peters, Mooij and Schölkopf, 2012); linear ICA (Eriksson and Koivunen, 2003); (Khemakhem, Kingma and Hyvärinen, 2020; Khemakhem, Monti, Kingma and Hyvärinen, 2020; Teshima et al., 2020) (Please refer to supplement 7.6 for more broad reviews). Besides, our analysis extends the scope of Khemakhem, Kingma and Hyvärinen (2020) to categorical Y and general forms of P(S,Z|C = c) that belongs to Sobolev space, in theorem 4.4. Note that our analysis does NOT require observing the original source variable ds.
We first narrow our interest to a family class of LaCIM denoted as Pexp in which any p ∈ Pexp satisfies that (i) the S,Z belong to the exponential family; and that (ii) the Y is generated from the ANM. We will show later that Pexp can approximate any P(S,Z|c) ∈W r,2(S × Z) for some r ≥ 2:
Pexp = { LaCIM with y = fy(s) + εy, p(s, z|c) := pTz,Γzc (z|c)pTs,Γsc(s|c) }
with pTt,Γtc(t) := qt∏ i=1 exp ( kt∑ j=1 T ti,j(ti)Γ t c,i,j +Bi(ti)−Atc,i ) for t = s, z, and e ∈ E ,
(1)
where {T ti,j(ti)}, {Γtc,i,j} denote the sufficient statistics and natural parameters, {Bi}, {Atc,i} denote the base measures and normalizing constants to ensure the integral of distribution equals to 1. Let Tt(t) := [Tt1(t1), ...,Ttqt(tqt)] ∈ R kt×qt ( Tti(ti) := [T t i,1(ti), ..., T t i,kt(ti)], ∀i ∈ [qt] ) ,
Γtc := [ Γtc,1, ...,Γ t c,qt ] ∈Rkt×qt ( Γtc,i := [Γ t c,i,1, ...,Γ t c,i,kt ],∀i ∈ [qt] ) . We define the ∼p-identifiability for θ := {fx, fy,Ts,Tz} as: Definition 4.2 (∼p-identifiability). We define a binary relation on the parameter space of X × Y: θ ∼p θ̃ if 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 ,
Ts([f−1x ]S(x)) = MsT̃ s([f̃−1x ]S(x)) + as, T z([f−1x ]Z(x)) = MzT̃ z([f̃−1x ]Z(x)) + az, pfy (y|[f−1x ]S(x)) = pf̃y (y|[f̃ −1 x ]S(x)),
We say that θ is ∼p-identifiable, if for any θ̃, peθ(x, y) = peθ̃(x, y) ∀e ∈ Etrain, implies θ ∼p θ̃.
It can be shown that∼p satisfies the reflective property (θ∼p θ), the symmetric property (if θ∼p θ̃ then θ̃ ∼p θ), and the transitive property (if θ1 ∼p θ2 and θ2 ∼p θ3, then θ1 ∼p θ3), and hence is an equivalence relation (details in supplement 7.2). This definition states that the S,Z can be identified up to permutation and point-wise transformation, which is sufficient for disentanglement of S and identifying the predicting mechanism pfy (y|[f−1x ]S(x)). Specifically, the definition regarding fx implies the separation of S and Z unless the extreme case when S can be represented by Z, i.e., there exists a function h : S → Z such that [f−1x ]S(x) = h([f−1x ]Z(x)). This definition is inspired by but beyond the scope of unsupervised scenario considered in nonlinear ICA (Hyvärinen et al., 2019; Khemakhem, Kingma and Hyvärinen, 2020) to further distinguish of S from Z. Besides, the pfy (y|[f−1x ]S(x)) = pf̃y (y|[f̃ −1 x ]S(x)) further guarantees the identifiability of prediction: predict using fy(s) with s obtained from fx. The following theorem presents the ∼p-identifiability for Pexp: Theorem 4.3 (∼p-identifiability). For θ in the LaCIM peθ(x, y) ∈ Pexp for any e ∈ Etrain, we assume that i) CIMe satisfies that fx, f ′x and f ′′x are continuous and that fx, fy are bijective; ii) the T ti,j are twice differentiable for any t = s, z, i ∈ [qt], j ∈ [kt]; iii) the exogenous variables satisfy that the characteristic functions of εx, εy are almost everywhere nonzero. Under the diversity condition on A := [P>de1 , ..., P > dem ]
> ∈ Rm×R with Pde := [p(c1|de), ..., p(cR|de)] that the A and[ [Γt=s,zc2 − Γ t=s,z c1 ] T, ..., [Γt=s,zcR − Γ t=s,z c1 ] T ]T
have full column rank for both t = s and t = z, we have that the θ := {fx, fy,Ts,Tz} are ∼p identifiable.
The bijectivity of fx and fy have been widely assumed in Janzing et al. (2009); Peters et al. (2014; 2017); Khemakhem, Kingma and Hyvärinen (2020); Teshima et al. (2020) as a basic condition for identifiability. It naturally holds for fx to be bijective since the latent components S,Z, as high-level abstractions which can be viewed as embeddings in auto-encoder (Kramer, 1991), lies in lower-dimensional space compared with input X which is supposed to have more variations, i.e., (qs + qz < qx). For categorical Y , the fy which generates the classification result, i.e., p(y = k|s) = [fy]k(s)/ ( ∑ k[fy]k(s)), will be shown later to be identifiable.
The diversity condition implies that i) m ≥ R ≥ max(kz ∗ qz, ks ∗ qs) + 1; and that ii) different environments are variant enough in terms of S-Z correlation (which is also assumed in Arjovsky et al. (2019)), as a necessary for the invariant one to be identified. As noted in 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 for the identification of the invariant mechanisms. Remark 2. The dimensions of the ground-truth S,Z are unknown, making the check about whether m is large enough impossible. Besides, in some real applications, the training environments are passively observed and may not satisfy the condition. However, we empirically find the improvement of LaCIM in terms of both OOD prediction and interpretability, if the multiple environments provided are diverse enough. Besides, a training environment can be the mixture of many sub-environments, which motivates to splitting the data according to their source ID or clustering results (Teney et al., 2020) to obtain more environments, making the condition easier to satisfy.
Extension to the general form of LaCIM. We generalize the identifiable result in theorem 4.3 to any LaCIM 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 showing that any such LaCIM can be approximated by a sequence of distributions in Pexp, motivated by the facts in Barron and Sheu (1991) that the exponential family is dense in the set of distributions with bounded support, and in Maddison et al. (2016) 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. Theorem 4.4 (Asymptotic ∼p-identifiability). Consider a LaCIM satisfying that pfx(x|s, z) and pfy (y|s) are smooth w.r.t s, z and s respectively. For each e and c ∈ C, suppose Pe(S,Z|C = c) ∈ W r,2(S × Z) for some r ≥ 2, we have that P is asymptotically ∼p-identifiable defined as: ∀ > 0, ∃ ∼p-identifiable P̃θ ∈ Pexp, s.t. dPok(pe(x, y), p̃eθ(x, y)) < ,∀e ∈ Etrain, (x, y) ∈ X × Y 1.

heading: 4.3 CAUSAL SUPERVISED VARIATIONAL AUTO-ENCODER
text: Guided by identifiability, we first provide the training method to learn fx, fy by reformulating VAE in a supervised scenario, followed by optimization over latent space for inference and test.
Training. To learn the CIMe and pfx(x|s, z), pfy (y|s) for invariant prediction, we implement the generative model to fit {pe(x, y)}e∈Etrain , which has been guaranteed by theorem 4.3, 4.4 to be able to identify the ground-truth predicting mechanism. Specifically, we reformulate the objective of VAE, as a generative model proposed in (Kingma and Welling, 2014), in supervised scenario. For unsupervised learning, the VAE introduces the variational distribution qψ parameterized by ψ to approximate the intractable posterior by maximizing the following Evidence Lower Bound (ELBO): −Lφ,ψ = Ep(x) [ Eqψ(z|x) log pφ(x,z)
qψ(z|x)
] , as a tractable surrogate of maximum likelihood Ep(x) log pφ(x).
Specifically, the ELBO is less than and equal to Ep(x) [ log pφ(x) ]
and the equality can only be achieved when qψ(z|x)=pφ(z|x). Therefore, maximizing the ELBO over pφ and qψ will drive (i) qψ(z|x) to learn pφ(z|x); (ii) pφ to learn the ground-truth model p (including pφ(x|z) to learn p(x|z)). In our supervised scenario, we introduce the variational distribution qeψ(s, z|x, y) and the corresponding ELBO for any e is −Leφ,ψ=Epe(x,y) [ Eqe
ψ (s,z|x,y) log
peφ(x,y,s,z) qe ψ (s,z|x,y) ] . Similarly, minimizing Leφ,ψ can
drive pφ(x|s, z), pφ(y|s) to learn the CIMe (i.e. pfx(x|s, z), pfy (y|s)), and also qeψ(s, z|x, y) to learn peφ(s, z|x, y). In other words, 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, we can similarly reparameterize qeψ(s, z|x, y) as qeψ(s,z|x)qψ(y|s)
qeψ(y|x) . According
to Causal Markov Condition, we have that peφ(x, y, s, z) = pφ(x|s, z)peφ(s, z)pφ(y|s). Substituting the above reparameterizations into the ELBO with qψ(y|s) replaced by pφ(y|s), 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)
] , (2)
where qeψ(y|x) = ∫ S q e ψ(s|x)pφ(y|s)ds. The overall loss function is: Lφ,ψ = ∑ e∈Etrain L e φ,ψ. The training datasets {De}e∈Etrain are applied to optimize prior model peφ(s, z), inference
1The dPok denotes the Pokorov distance and limn→∞ dPok(µn, µ)→ 0⇐⇒ µn d→ µ.
model qeψ(s, z|x) and generative models pφ(x|s, z), pφ(y|s) in Eq. (2). The generative models pφ(x|s, z), pφ(y|s) are shared among all environments, while the peφ(s, z), qeψ(s, z|x) are respectively pφ(s, z|des), qψ(s, z|x, des) and pφ(s, z|de), qψ(s, z|x, de) for LaCIM-ds and LaCIM-d.
Inference & Test. When de ′ s can be acquired during test for e ′ ∈ Etest, we can predict y as arg maxy pφ(y|x, de ′ s ) = ∫ qψ(s|x, de ′
s )pφ(y|s)ds. Otherwise, for LaCIM-d with ds unobserved, we first optimize s, z via (s?, z?) := arg maxs,z log pφ(x|s, z) and predict y as arg maxy qψ(y|s?). Specifically, we adopt the strategy for optimization in Schott et al. (2018) that we first sample initial points and select the one with the maximum log pφ(x|s, z), then we optimize for 50 iterations using Adam. The implementation details and optimization effect are shown in supplement 7.9.

heading: 5 EXPERIMENTS
text: We evaluate LaCIM on (I) synthetic data to verify the identifiability in theorem 4.3; (II) OOD challenges: object classification with sample selection bias (Non-I.I.D. Image dataset with Contexts (NICO)); Hand-Writing Recognition with confounding bias (Colored MNIST (CMNIST)); prediction of Alzheimer’s Disease (Alzheimer’s Disease Neuroimaging Initiative (ADNI www.loni.ucla. edu/ADNI); (III) Robustness on detecting images with small perturbation (FaceForensics++).

heading: 5.1 SIMULATION
text: To verify the identifiability claim and effectiveness of our learning method, we implement LaCIM on synthetic data. The data generating process is provided in Supplement 7.8. The domain index D ∈ Rm is denoted as a one-hot encoded vector with m = 5. To verify the utility of training on multiple domains (m > 1), we also conduct LaCIM by pooling data from all m domains together, namely pool-LaCIM for comparison. We randomly generate m = 5 datasets and run 20 times for each. We compute the metric mean correlation coefficient (MCC) adopted in Khemakhem, Kingma and Hyvärinen (2020) to measure the goodness of identifiability under permutation by introducing cost optimization to assign each learned component to the source component. This measurement is aligned with the goal of ∼p-identifiability, which allows us to distinguish S from Z. Table 5.1 shows the superiority of our LaCIM-d, LaCIM-ds over pool-LaCIM in terms of the CIMe relating to S,Z under permutation, by means of multiple diverse experiments. Besides, we consider LaCIM-d on m = 3, 5, 7 with the same total number of samples. It yields that more environments can perform better; and that even m = 3 still performs much better than pool-LaCIM. To illustrate the learning effect, we visualize the learned Z in Fig. 7.8, with S left in supplement 7.8 due to space limit.

heading: 5.2 REAL-WORLD OOD CHALLENGE
text: We present our LaCIM’s results on three OOD tasks, with different environments associated with different values of ds. We implement both versions of LaCIM, i.e., LaCIM-ds and LaCIM-d, with task-dependent definition of ds. In CMNIST, the ds (digit color) is a fully observed confounder, and LaCIM-ds in this case is the ceiling of LaCIM-d under the same implementation. In NICO and ADNI, the LaCIM-d even outperform LaCIM-ds, when the source variables are only partially observed.
Dataset. We describe the datasets as follows (the X denote image; the Y denote label):
NICO: we evaluate the cat/dog classification in “Animal” dataset in NICO, a benchmark for non-i.i.d problem in He et al. (2019). Each animal is associated with “grass”,“snow” contexts with different proportions, denoted as ds ∈ R4 (cat,dog in grass,snow). We set m = 8 and m = 14. The C,Z, S respectively denote the (time,whether) of sampling, the context and semantic shape of cat/dog.
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 color others as red. We set m = 2 with pe1 = 0.9, pe2 = 0.8. The des is p
e to describe the intensity of spursiou correlation caused by color. We do not flip y with 25% like Arjovsky et al. (2019) 2, since doing so will cause the digit correlated rather than causally related to the label, which is beyond our scope. The Z, S respectively represent the color and number. The C can also denote (time,whether) for which the painter draws the number and color, e.g., the painter tends to draw red 0 more often than green 1 in the sunny morning.
ADNI. The data are obtained from the ADNI databaset, the Y := {0, 1, 2} with 0,1,2 respectively denoting AD, Mild Cognitive Impairment (MCI) and Normal Control (NC). The X is Magnetic resonance imaging (sMRI). We set m = 2. We consider two types of ds: Age and TAU (a biomarker Humpel and Hochstrasser (2011)). The S (Z) denote the disease-related (-unrelated) brain regions. The C denotes the hormone level that can affect the brain structure development.
Compared Baselines. We compare with (i) Cross-Entropy (CE) from X → Y (CE X → Y ), (ii) domain-adversarial neural network (DANN) for domain adaptation Ganin et al. (2016), (iii) Maximum Mean Discrepancy with Adversarial Auto-Encoder (MMD-AAE) for domain generalization Li et al. (2018), (iv) Domain Invariant Variational Autoencoders (DIVA) Ilse et al. (2019) (v) Selecting Data Augmentation (SDA) Ilse et al. (2020), (vi) Invariant Risk Mnimization (IRM) Arjovsky et al. (2019), (vii) CE (X, ds) → Y , (viii) VAE with causal graph C → V → {X,Y } with V mixing S,Z and we call it sVAE for simplicity. We only implement SDA on CMNIST, since the intervened-data generation of SDA requires explicitly extracting the S,Z, which is intractable in ADNI and NICO. For fair comparison, we keep the model capacity (numer of parameters) in the same level.
Implementation Details. For each domain e, we implement the reparameterization with ρes, ρez: s′, z′ = ρes(s), ρ e z(z), to transform the p
e(s, z) into isotropic Gaussian; then the generative models are correspondingly modified as {pφ(x|(ρes)−1(s), (ρez)−1(z)), pφ(y|(ρes)−1(s))} according to rule of change of variables. The optimized parameters are {{qeψ(s, z|x)}e, pφ(x|s, z), pφ(y|s), {ρet=s,z}e}, with the encoder qeψ(s, z|x) being sequentially composed of: i) the sequential of Conv-BN-ReLUMaxPool blocks that shared among Etrain, followed by ii) the sequential of ReLU-FC for the mean and log-variance of S,Z that are specific to e. The structure of ρet=s,z is FC-ReLU-FC. The decoder pφ(x|s, z) is the sequential of upsampling, several TConv-BN-ReLU blocks and Sigmoid. The predictor pφ(y|s) is sequential of FC→BN→ReLU blocks, followed by Softmax (or Sigmoid) for classification. The network structure and the output channel size for CMNIST, NICO and ADNI are introduced in supplement 7.11, 7.12, 7.13, Tab. 13, 14. 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. The “FC”, ”BN” stand for Fully-Connected, Batch-Normalization.
Results. We report accuracy over three runs for each method. As shown in Tab. 2 3 our LaCIM-d performs comparable and better than others on all applications, except the 99.3 achieved by SDA on CMNIST, which is comparable to the result on the original MNIST. This is because during training,
2We also conduct this experiment with flipping y in supplementary 7.11. 3On NICO, we implement ConvNet with Batch Balancing as a specifically benchmark in He et al. (2019).
The results are 60± 1 on m = 8 and 62.33± 3.06 on m = 14.
the SDA implemented data augmentation with random colors, which decorrelate the color-label. When S cannot be explicitly extracted in general case, the SDA is not tractable.
Method Dataset NICO CMNIST ADNI (m = 2) m = 8 m = 14 m = 2 C : Age C : TAU CE X → Y 60.67± 2.52 59.00± 1.73 97.87± 0.19 63.06± 2.26 64.58± 0.90 DANN 59.33± 4.93 60.00± 2.65 97.42± 0.13 60.84± 1.83 64.58± 0.90 MMD-AAE 61.33± 2.89 66.33± 3.21 81.23± 7.80 62.43± 2.42 65.62± 0.00 DIVA 60.67± 2.08 58.67± 1.53 97.97± 0.19 61.37± 3.30 65.10± 0.90 SDA - - 99.37 ± 0.03 - - IRM 61.67± 4.16 65.00± 3.00 98.18± 0.22 63.49± 1.59 65.10± 0.90 CE X, ds → Y 57.33± 6.03 64.00± 1.00 98.03± 0.27 62.43± 0.92 65.62± 0.00 sVAE 59.67± 3.79 64.33± 0.58 97.89± 0.61 63.67± 1.87 66.67± 0.91 LaCIM-ds (Ours) 62.00± 1.73 68.00± 2.64 98.81± 0.14 65.08± 1.59 66.14± 0.91 LaCIM-d (Ours) 62.67 ± 0.58 68.67± 2.64 98.78± 0.20 64.44± 0.96 68.23± 0.90
NICO, a larger m (with the total number of samples n fixed) can bring further benefit, which may due to the easier satisfaction of the diversity condition in theorem 4.3. One thing worth particular mention is that on NICO and ADNI (when ds denotes TAU), our LaCIM-d performs comparable and even better than LaCIM-ds, due to the existence of unobserved partial variables. For example, each ds only contains one attribute each time in ADNI. For completeness, we conduct experiments with fully observed confounders in supplement 7.13. Besides, we apply our method on intervened data, the result of which can validate more robustness of LaCIM, as shown in supplement 7.12.
Interpretability. We visualize learned S as side proof of interpretability. Specifically, we select the s∗ that has the highest correlation with y among all dimension of S, and visualize the derivatives of s∗ with respect to the image. For CE x → y and CE (x, ds) → y, we visualize the derivatives of predicted class scores with respect to the image. As shown in Fig. 5.2, LaCIM (the 4th column) can identify more explainable semantic features, which verifies the identifiability and effectiveness of the learning method. Supplement 7.12 provides more results.
5.3 ROBUSTNESS ON SECURITY
We consider the DeepFake-related security problem, which targets on detecting small perturbed fake images that can spread fake news. The Rossler et al. (2019) provides FaceForensics++ dataset from 1000 Youtube videos for training and 1,000 benchmark images from other sources (OOD) for testing. We split the train data into m = 2 environments accord-
ing to video ID. The considerable result in Tab. 5.3 verifies potential value on security.

heading: 6 CONCLUSIONS & DISCUSSIONS
text: We incorporate the causal structure as prior knowledge in proposed LaCIM, by introducing: (i) latent variables and explicitly separate them into y-causative factors (a.k.a, S) and others (a.k.a, Z) which are spuriously correlated with the output; (ii) the source variable ds that explains the distributional inconsistency among domains. When the environments are diverse and much enough, we can successfully identify the causal invariant mechanisms, and also y-causative factors for prediction without a mix of others. Our LaCIM shows potential value regarding robustness to OOD tasks with confounding bias, selection bias and others such as healthcare and security. A possible drawback of our model lies in our requirement of the number of environments (which may be not satisfied in some scenarios) for identifiability, and the relaxation of which is left in the future work.

heading: 7 SUPPLEMENTARY MATERIALS

heading: 7.1 O.O.D GENERALIZATION ERROR BOUND
text: Denote Ep[y|x] := ∫ Y yp(y|x)dy for any x, y ∈ X×Y . We have Epe [y|s] = ∫ Y yp(y|s)dy according to that p(y|s) is invariant across E , we can omit pe in Epe [y|s] and denote g(S) := E[Y |S]. Then, the OOD bound
∣∣Epe1 (y|x)− Epe2 (y|x)∣∣, ∀(x, y) is bounded as follows: Theorem 7.1 (OOD genearlization error). Consider two LaCIM Pe1 and Pe2 , suppose that their densities , i.e., pe1(s|x) and pe2(s|x) are absolutely continuous having support (−∞,∞). For any (x, y) ∈ X × Y , assume that
• g(S) is a Lipschitz-continuous function; • πx(s) := p e2 (s|x) pe1 (s|x) is differentiable and Epe1 [ πx(S) ∣∣g(S)− µ1∣∣] < ∞ with µ1 := Epe1 [g(S)|X = x] = ∫ S g(s)p e1(s|x)ds;
then we have ∣∣Epe1 (y|x)− Epe2 (y|x)∣∣ ≤ ‖g′‖∞‖π′x‖∞Varpe1 (S|X = x).
When e1 ∈ Etrain and e2 ∈ Etest, the theorem 7.1 describes the error during generalization on e2 for the strategy that trained on e1. The bound is mainly affected by: (i) the Lipschitz constant of g, i.e., ‖g‖∞; (ii) ‖π′x‖∞ which measures the difference between pe1(s, z) and pe2(s, z); and (iii) the Varpe1 (S|x) that measures the intensity of x→ (s, z). These terms can be roughly categorized into two classes: (i),(iii) which are related to the property of CIMe and gave few space for improvement; and the (ii) that describes the distributional change between two environments. Specifically for the first class, the (i) measures the smoothness of E(y|s) with respect to s. The smaller value of ‖g′‖∞ implies that the flatter regions give rise to the same prediction result, hence easier transfer from e1 to e2 and vice versa. For the term (iii), consider the deterministic setting that εx = 0 (leads to Varpe1 (S|x) = 0), then s can be determined from x for generalization if the f is bijective function. The term (ii) measures the distributional change between posterior distributions pe1(s|x) and pe2(s|x), which contributes to the difference during prediction:
∣∣Epe1 (y|x) − Epe2 (y|x)∣∣ = ∫S(pe1(s|x) − pe1(s|x))pfy (y|s)ds. Such a change is due to the inconsistency between priors pe1(s, z) and pe2(s, z), which is caused by different value of the confounder ds.
Proof. In the following, we will derive the upper bound∣∣Epe1 [Y |X=x]− Epe2 [Y |X=x] ∣∣ ≤ ‖g′‖∞‖π′x‖∞Varpe1 (S|X = x) , where πx(s) =: pe2 (s|x) pe1 (s|x) and g(s) is assumed to be Lipschitz-continuous.
To begin with, note that E[Y |X] = E[E(Y |X,S)|X] = E[g(S)|X] = ∫ g(s)p(s|x)ds.
Let p1(s|x) = pe1(s|x), p2(s|x) = pe2(s|x). For ease of notations, we use P1 and P2 denote the distributions with densities p1(s|x) and p2(s|x) and suppose S1 ∼ P1 and S2 ∼ P2, where x is omitted as the following analysis is conditional on a fixed X=x.
Then we may rewrite the difference of conditional expectations as
Epe2 [Y |X = x]− Epe1 [Y |X = x] = E(g(S2))− E(g(S1)), where E[g(Sj))] = ∫ g(s)pj(s|x)ds denotes the expectation over Pj .
Let µ1 := Epe1 [g(S)|X = x] = E[g(S1)] = ∫ g(s)p1(s|x)ds. Then
Epe2 [Y |X = x]− Epe1 [Y |X = x] = E(g(S2))− E(g(S1)) = E [g(S2)− µ1] .
Further, we have the following transformation E [g(S2)− µ1] = ∫ (g(s)− µ1)πx(s)p1(s|x)ds = E [(g(S1)− µ1)πx(S1)] . (3)
In the following, we will use the results of the Stein kernel function. Please refer to Definition 7.2 for a general definition. Particularly, for the distribution P1 ∼ p1(s|x), the Stein kernel τ1(s) is
τ1(s) = 1
p1(s|x) ∫ s −∞ (E(S1)− t)p1(t|x)dt, (4)
where E(S1) = ∫ s · p1(s|x)ds. Further, we define (τ1 ◦ g)(s) as
(τ1 ◦ g)(s)= 1
p1(s|x) ∫ s −∞ (E(g(S1))− g(t))p1(t|x)dt= 1 p1(s|x) ∫ s −∞ (µ1 − g(t))p1(t|x)dt. (5)
Under the second condition listed in Theorem 7.1, we may apply the result of Lemma 7.3. Specifically, by the equation (8), we have
E [(g(S1)− µ1)πx(S1)] = E [(τ1 ◦ g)(S1)π′x(S1)] .
Then under the first condition in Theorem 7.1, we can obtain the following inequality by Lemma 7.4, E [(τ1 ◦ g)(S1)π′x(S1)]= E [(
(τ1 ◦ g) τ1 π′xτ1
) (S1) ] ≤ E [∣∣∣ (τ1 ◦ g) τ1 (S1) ∣∣∣ · ∣∣∣π′xτ1(S1)∣∣∣]
≤ ‖g′‖∞E [| (π′xτ1) (S1)|] ≤ ‖g′‖∞‖π′x‖∞E [|τ1(S1)|] . (6)
In the following, we show that the Stein kernel is non-negative, which enables E [|τ1(S1)|] = E [τ1(S1)]. According to the definition, τ1(s) = 1p1(s|x) ∫ s −∞(E(S1)− t)p1(t|x)dt, where E(S1)=∫∞
−∞ t · p1(t|x)dt. Let F1(s) = ∫ s −∞ p1(t|x)dt be the distribution function for P1. Note that∫ s
−∞ E(S1)p1(t|x)dt = F1(s)E(S1) = F1(s) E(S1),∫ s −∞ tp1(t|x)dt = F1(s) ∫ s −∞ t p1(t|x) F1(s) dt = F1(s) E(S1|S1 ≤ s) ≤ F1(s) E(S1),
The last inequality is based on E(S1|S1 ≤ s)− E(S1) ≤ 0 that can be proved as the following∫ s −∞ t p1(t|x) F1(s) dt− ∫ ∞ −∞ tp1(t|x)dt = ∫ s −∞ t ( 1 F1(s) − 1 ) p1(t|x)dt− ∫ ∞ s tp1(t|x)dt
≤ s ∫ s −∞ ( 1 F1(s) − 1 ) p1(t|x)dt− s ∫ ∞ s p1(t|x) = 0.
Therefore, τ1(s) ≥ 0 and hence E [|τ1(S1)|] = E [τ1(S1)] in (6). Besides, by equation (9), the special case of Lemma 7.3, we have
E [τ1(S1)] = Var(S1) = Varpe1 (S|X = x).
To sum up,
E [(τ1 ◦ g)(S1)π′x(S1)] ≤ ‖g′‖∞‖πx‖∞E [τ1(S1)] = ‖g′‖∞‖π′x‖∞Varpe1 (S|X = x).
Definition 7.2 (the Stein Kernel τP of distribution P ). Suppose X∼P with density p. The Stein kernel of P is the function x 7→ τP (x) defined by
τP (x) = 1
p(x) ∫ x −∞ (E(X)− y)p(y)dy, (7)
where Id is the identity function for Id(x) = x. More generally, for a function h satisfying E[|h(X)|] <∞, define (τP ◦ h)(x) as
(τP ◦ h)(x) = 1
p(x) ∫ x −∞ (E(h(X))− h(y))p(y)dy.
Lemma 7.3. For a differentiable function ϕ such that E[|(τP ◦ h)(x)ϕ′(X)|] <∞, we have E [(τP ◦ h)(x)ϕ′(X)] = E[(h(X)− E(h(X))ϕ(X)]. (8)
Proof. Let µh =: E(h(X)). As E(h(X)− µh) = 0,
(τP ◦ h)(x) = 1
p(x) ∫ x −∞ (µh − h(y))p(y)dy = −1 p(x) ∫ ∞ x (µh − h(y))p(y)dy.
Then E [(τP ◦ h)(x)ϕ′(X)]= ∫ 0 −∞ (τP ◦ h)(x)ϕ′(x)p(x)dx+ ∫ ∞ 0 (τP ◦ h)(x)ϕ′(x)p(x)dx
= ∫ 0 −∞ ∫ x −∞ (µh − h(y))p(y)ϕ′(x)dydx− ∫ ∞ 0 ∫ ∞ x (µh − h(y))p(y)ϕ′(x)dydx
= ∫ 0 −∞ ∫ 0 y (µh − h(y))p(y)ϕ′(x)dxdy − ∫ ∞ 0 ∫ y 0 (µh − h(y))p(y)ϕ′(x)dxdy
= ∫ 0 −∞ ∫ y 0 (h(y)− µh)p(y)ϕ′(x)dxdy + ∫ ∞ 0 ∫ y 0 (h(y)− µh)p(y)ϕ′(x)dxdy
= ∫ ∞ −∞ (h(y)− µh)p(y) (∫ y 0 ϕ′(x)dx ) dy= ∫ ∞ −∞ (h(y)− µh)p(y)(ϕ(y)− ϕ(0))dy
= ∫ ∞ −∞ (h(y)− µh)p(y)(ϕ(y))dy=E[(h(X)− E(h(X))ϕ(X)]
Particularly, taking h(X) = X and ϕ(X) = X − E(X), we immediately have E(τP (X)) = Var(X) (9)
Lemma 7.4. Assume that E(|X|) < ∞ and the density p is locally absolutely continuous on (−∞,∞) and h is a Lipschitz continuous function. Then we have |fh| ≤ ‖h′‖∞ for
fh(x) = (τP ◦ h)(x) τP (x) = ∫ x −∞(E(h(X))− h(y))p(y)dy∫ x −∞(E(X)− y)p(y)dy .
Proof. This is a special case of Corollary 3.15 in Döbler et al. (2015), taking the constant c = 1.

heading: 7.2 PROOF OF THE EQUIVALENCE OF DEFINITION 4.2
text: Proposition 7.5. The binary relation ∼p defined in Def. 4.2 is an equivalence relation.
Proof. The equivalence relation should satisfy three properties as follows:
• Reflexive property: The θ ∼p θ with Mz , Ms being identity matrix and as, az being 0.
• Symmtric property: If θ ∼p θ̃, then there exists block permutation matrices Mz and Ms such that
Ts([fx] −1 S (x)) = MsT̃ s([f̃x] −1 S (x)) + as, T z([fx] −1 Z (x)) = MzT̃ z([f̃x] −1 Z (x)) + az, pfy (y|[fx]−1S (x)) = pf̃y (y|[f̃x] −1 S (x)).
The we have M−1s and M −1 z are also block permutation matrices and such that:
T̃s([f̃x] −1 S (x)) = M −1 s T s([fx] −1 S (x)) + (−as), T̃ s([f̃x] −1 Z (x)) = M −1 z T s([fx] −1 Z (x)) + (−az), pf̃y (y|[f̃x] −1 S (x)) = pfy (y|[fx] −1 S (x)).
Therefore, we have θ̃ ∼p θ.
• Transitive property: if θ1 ∼p θ2 and θ2 ∼p θ3 with θi := {f ix, f iy,Ts,1,Tz,1,Γs,i,Γz,i}, then we have
Ts,1((f1x,s) −1(x)) = M1sT s,2((f2x,s) −1(x)) + a1s, Tz,1((f1x,z) −1(x)) = M1zT z,2((f2x,z) −1(x)) + a2z, Ts,2((f2x,s) −1(x)) = M2sT s,3((f3x,s) −1(x)) + a2s, Tz,2((f2x,z) −1(x)) = M2zT z,3((f3z ) −1(x)) + a3x,z
for block permutation matrices M1s ,M 1 z ,M 2 s ,M 2 z and vectors a 1 s, a 2 s, a 1 z, a 2 z . Then we have
Ts,1((f1x,s) −1(x)) = M2sM 1 sT s,3((f3x,s) −1(x)) + (M2s a 1 s) + a 2 s, Tz,1((f1x,z) −1(x)) = M2zM 1 zT z,3((f3x,z) −1(x)) + (M2z a 1 z) + a 2 z.
Besides, it is apparent that
pf1y (y|(f 1 x) −1 s (x)) = pf2y (y|(f 2 x) −1 s (x)) = pf3y (y|(f 3 x) −1 s (x)). (10)
Therefore, we have θ1 ∼p θ3 since M2sM1s and M2zM1z are also permutation matrices.
With above three properties satisfied, we have that ∼p is a equivalence relation.

heading: 7.3 PROOF OF THEOREM 4.3
text: In the following, we write pe(x, y) as p(x, y|de) and also Γt=s,zc := Γt=s,z(de), Sc,i = Si(d e), Zc,i = Zi(d e). To prove the theorem 4.3, we first prove the theorem 7.6 for the simplest case when c|de = de, then we generalize to the case when C := ∪rCr. The overall roadmap is as follows: we first prove the ∼A-identifiability in theorem 7.9, and the combination of which with lemma 7.12, 7.11 give theorem 7.6 in the simplest case when c|de = de. Then we generalize the case considered in theorem 7.6 to the more general case when C := ∪rCr. Theorem 7.6 (∼p-identifiability). For θ in the LaCIM peθ(x, y) ∈ Pexp for any e ∈ Etrain, we assume that (1) the CIMe satisfies that fx, f ′x and f ′′x are continuous and that fx, fy are bijective; (2) that the T ti,j are twice differentiable for any t = s, z, i ∈ [qt], j ∈ [kt]; (3) the exogenous variables satisfy that the characteristic functions of εx, εy are almost everywhere nonzero; (4) the number of environments, i.e., m ≥ max(qs ∗ ks, qz ∗ kz) + 1 and [ Γt=s,zde2 − Γ t=s,z de1 , ...,Γ t=s,z dem − Γ t=s,z de1 ] have full column rank for both t = s and t = z, we have that the parameters θ := {fx, fy,Ts,Tz} are ∼p identifiable.
To prove theorem 7.6, We first prove the ∼A-identifiability that is defined as follows: Definition 7.7 (∼A-identifiability). The definition is the same with the one defined in 4.2, with Ms,Mz being invertible matrices which are not necessarily to be the permutation matrices in Def. 4.2. Proposition 7.8. The binary relation ∼A defined in Def. 7.7 is an equivalence relation.
Proof. The proof is similar to that of proposition 7.5.
The following theorem states that any LaCIM that belongs to Pexp is ∼A-identifiable. Theorem 7.9 (∼A-identifiability). For θ in the LaCIM peθ(x, y) ∈ Pexp for any e ∈ Etrain, we assume (1) the CIMe satisfies that fx, fy are bijective; (2) the T ti,j are twice differentiable for any t = s, z, i ∈ [qt], j ∈ [kt]; (3) the exogenous variables satisfy that the characteristic functions of εx, εy are almost everywhere nonzero; (4) the number of environments, i.e., m ≥ max(qs ∗ ks, qz ∗ kz) + 1 and [ [Γtde2 − Γtde1 ]T, ..., [Γtdem − Γtde1 ]T ]T have full column rank for t = s, z, we have that the parameters {fx, fy,Ts,Tz} are ∼p identifiable.
Proof. Suppose that θ = {fx, fy,Ts,Tz} and θ̃ = {f̃x, g̃y, T̃s, T̃z} share the same observational distribution for each environment e ∈ Etrain, i.e.,
pfx,fy,Ts,Γs,Tz,Γz (x, y|de) = pf̃x,f̃y,T̃s,Γ̃s,T̃z,Γ̃z (x, y|d e). (11)
Then we have
pfx,fy,Ts,Γs,Tz,Γz (x|de) = pf̃x,f̃y,T̃s,Γ̃s,T̃z,Γ̃z (x|d e) (12) =⇒ ∫ S×Z pfx(x|s, z)pTs,Γs,Tz,Γz (s, z|de)dsdz = ∫ S×Z pf̃x(x|s, z)pT̃s,Γ̃s,T̃z,Γ̃z (s, z|d e)dsdz
(13) =⇒ ∫ X pεx(x− x̄)pTs,Γs,Tz,Γz (f−1x (x̄)|de)volJf−1x (x̄)dx̄ (14)
= ∫ X pεx(x− x̄)pT̃s,Γ̃s,T̃z,Γ̃z (f̃ −1 x (x̄)|de)volJf̃−1x (x̄)dx̄ (15)
=⇒ ∫ X p̃Ts,Γs,Tz,Γz,fx(x̄|de)pεx(x− x̄)dx̄ = ∫ X p̃T̃s,Γ̃s,T̃z,Γ̃z,f̃x(x̄|d e)pεx(x− x̄)dx̄ (16) =⇒ (p̃Ts,Γs,Tz,Γz,fx ∗ pεx)(x|de) = (p̃T̃s,Γ̃s,T̃z,Γ̃z,f̃x) ∗ pεx(x|d e) (17) =⇒ F [p̃Ts,Γs,Tz,Γz,fx ](ω)ϕεx(ω) = F [p̃T̃s,Γ̃s,T̃z,Γ̃z,f̃x ](ω)ϕεx(ω) (18) =⇒ F [p̃Ts,Γs,Tz,Γz,fx ](ω) = F [p̃T̃s,Γ̃s,T̃z,Γ̃z,f̃x ](ω) (19) =⇒ p̃Ts,Γs,Tz,Γz,fx(x|de) = p̃T̃s,Γ̃s,T̃z,Γ̃z,f̃x(x|d e) (20)
where volJf (X) := det(Jf (X)) for any square matrix X and function f with “J” standing for the Jacobian. The p̃Ts,Γs,Tz,Γz,fx(x) in Eq. (16) is denoted as pTs,Γs,Tz,Γz (f −1 x (x|de)volJf−1(x). The ’*’ in Eq. (17) denotes the convolution operator. The F [·] in Eq. (18) denotes the Fourier transform, where φεx(ω) = F [pεx ](ω). Since we assume that the ϕεx(ω) is non-zero almost everywhere, we can drop it to get Eq. (20). Similarly, we have that:
pfy,Ts,Γs(y|de) = pf̃y,T̃s,Γ̃s(y|d e) (21) =⇒ ∫ S pfy (y|s)pTs,Γs(s|de)ds = ∫ S pf̃y (y|s)pT̃s,Γ̃s(s|d e)ds (22)
=⇒ ∫ Y pεy (y − ȳ)pTs,Γs(f−1y (ȳ)|de)volJf−1y (ȳ)dȳ (23)
= ∫ Y pεy (y − ȳ)pT̃s,Γ̃s(f̃ −1 y (ȳ)|de)volJg̃−1(ȳ)dȳ (24)
=⇒ ∫ S p̃Ts,Γs,fy (ȳ|de)pεy (y − ȳ)dȳ = ∫ S p̃T̃s,Γ̃s,f̃y (ȳ|d e)pεy (y − ȳ)dȳ (25) =⇒ (p̃Ts,Γs,fy ∗ pεy )(y|de) = (p̃T̃s,Γ̃s,f̃y ∗ pεy )(y|d e) (26)
=⇒ F [p̃Ts,Γs,fy ](ω)ϕεy (ω) = F [p̃T̃s,Γ̃s,f̃y ](ω)ϕεy (ω) (27) =⇒ F [p̃Ts,Γs,fy ](ω) = F [p̃T̃s,Γ̃s,f̃y ](ω) (28) =⇒ p̃Ts,Γs,fy (y) = p̃Ts,Γs,f̃y (y), (29)
and that
pfx,fyTs,Γs,Tz,Γz (x, y|de) = pf̃x,f̃y,T̃s,Γ̃s,T̃z,Γ̃z (x, y|d e) (30) =⇒ ∫ S×Z pfx(x|s, z)pfy (y|s)pTs,Γs,Tz,Γz (s, z|de)dsdz
= ∫ S×Z pf̃ (x|s, z)pf̃y (y|s)pT̃s,Γ̃s,T̃z,Γ̃z (s, z|d e)dsdz (31)
=⇒ ∫ V pε(v − v̄)pTs,Γs,Tz,Γz (h−1(v̄)|de)volJh−1(v̄)dv̄ (32)
= ∫ V pε(v − v̄)pT̃s,Γ̃s,T̃z,Γ̃z (h̃ −1(v̄)|de)volJh̃−1(v̄)dv̄ (33)
=⇒ ∫ S×Z p̃Ts,Γs,Tz,Γz,h,c(v̄|d)pε(v − v̄)dv̄ = ∫ S×Z p̃T̃s,Γ̃s,T̃z,Γ̃z,h̃,de(v̄|d e)pε(v − v̄)dv̄ (34)
=⇒ (p̃Ts,Γs,Tz,Γz,h ∗ pε)(v) = (p̃T̃s,Γ̃s,T̃z,Γ̃z,h̃ ∗ pε)(v) (35)
=⇒ F [p̃Ts,Γs,Tz,Γz,h](ω)ϕε(ω) = F [p̃T̃s,Γ̃s,T̃z,Γ̃z,h̃](ω)ϕε(ω) (36) =⇒ F [p̃Ts,Γs,Tz,Γz,h](ω) = F [p̃T̃s,Γ̃s,T̃z,Γ̃z,h̃](ω) (37) =⇒ p̃Ts,Γs,Tz,Γz,h(v) = p̃Ts,Γs,Tz,Γz,h(v), (38)
where v := [x>, y>]>, ε := [ε>x , ε > y ] >, h(v) = [[fx]−1Z (x) >, f−1y (y) >]>. According to Eq. (29), we have
log volJfy (y) + qs∑ i=1 logBi(f−1y,i (y))− logAi(de) + ks∑ j=1 T si,j(f −1 y,i (y))Γ s i,j(d e)  = log volJf̃y (y) + qs∑ i=1 log B̃i(f̃−1y,i (y))− log Ãi(de) + ks∑ j=1 T̃ si,j(f̃ −1 y,i (y))Γ̃ s i,j(d e)  (39)
Suppose that the assumption (4) holds, then we have
〈Ts(f−1y (y)),Γ s (dek)〉+ ∑ i log Ai(d e1) Ai(dek) = 〈T̃s(f̃−1y (y)), Γ̃ s (dek)〉+ ∑ i log Ãi(d e1) Ãi(dek) (40)
for all k ∈ [m], where Γ̄(d) = Γ(d)−Γ(de1). Denote b̃s(k) = ∑ i Ãi(d e1 )Ai(d ek )
Ãi(d ek )Ai(de1 ) for k ∈ [m], then we have
Γ s,> Ts(f−1y (y)) = Γ̃ s,> T̃s(f̃−1y (y)) + b̃s, (41)
Similarly, from Eq. (20) and Eq. (38), there exists b̃z, b̃s such that
Γ s,> Ts([fx] −1 S (x)) + Γ z,> Tz([fx] −1 Z (x)) = Γ̃ s,> T̃s([f̃x] −1 S (x)) + Γ̃ z,> T̃z([f̃x] −1 Z (x)) + b̃z + b̃s,
(42)
where b̃z(k) = ∑ i Z̃i(d e1 )Zi(d ek )
Z̃i(d ek )Zi(de1 )
for k ∈ [m]; and that,
Γ s,> Ts(f−1y (y)) + Γ z,> Tz([f−1x ]Z(x)) = Γ̃ s,> T̃s(f̃−1y (y)) + Γ̃ z,> T̃z([f̃−1x ]Z(x)) + b̃z + b̃s.
(43)
Substituting Eq. (41) to Eq. (42) and Eq. (43), we have that
Γ z,> Tz([f−1x ]Z(y)) = Γ̃ z,> T̃z([f̃−1x ]Z(y)) + b̃z, Γ s,> Ts([f−1x ]S(y)) = Γ̃ s,>
T̃s([f̃−1x ]S(y)) + b̃s. (44)
According to assumption (4), the Γ s,> and Γ z,> have full column rank. Therefore, we have that
Tz([f−1x ]Z(x)) = ( Γ z Γ z,>)−1 Γ̃ z,> T̃z([f̃−1x ]Z(x)) + ( Γ z Γ z,>)−1 b̃z (45)
Ts([f−1x ]S(x)) = ( Γ s Γ s,>)−1 Γ̃ s,> T̃s([f̃−1x ]S(x)) + ( Γ s Γ s,>)−1 b̃s. (46)
Ts(f−1y (y)) = ( Γ s Γ s,>)−1 Γ̃ s,> T̃s(f̃−1y (y)) + ( Γ s Γ s,>)−1 b̃s. (47)
Denote Mz := ( Γ z Γ z,>)−1 Γ̃ z,> , Ms := ( Γ s Γ s,>)−1 Γ̃ s,> and as = ( Γ s Γ s,>)−1
b̃s, az =( Γ z Γ z,>)−1
b̃z . The left is to prove that Mz and Ms are invertible matrices. Denote x̄ = f−1(x). Applying the (Khemakhem, Kingma and Hyvärinen, 2020, Lemma 3) we have that there exists ks points x̄1, ..., x̄ks , ˜̄x1, ..., ˜̄xkz such that ( (Ts)′i([f −1 x ]Si(x 1 i )), ..., (T s)′i([f −1 x ]Si(x ks i )) ) for each i ∈
[qs] and ( (Tz)′i([f −1 x ]Zi(x̃ 1 i )), ..., (T z)′i([f −1 x ]Si(x̃ kz i )) ) for each i ∈ [qt] are linearly independent.
By differentiating Eq. (45) and Eq. (46) for each x̄i with i ∈ [qs] and ˜̄xi with i ∈ [qz] respectively, we have that (
JTs(x̄ 1), ..., JTs(x̄ ks) ) = Ms ( JTs◦f̃−1x ◦fx(x̄ 1), ..., JTs◦f̃−1x ◦f (x̄ ks) )
(48)( JTz (˜̄x 1), ..., JTz (˜̄x kz ) ) = Mz ( JTz◦f̃−1x ◦fx( ˜̄x1), ..., JTz◦f̃−1x ◦fx( ˜̄xkz ) ) . (49)
The linearly independence of (
(Ts)′i([f −1 x ]Si(x 1 i )), ..., (T s)′i([f −1 x ]Si(x ks i )) ) and(
(Tz)′i([f −1 x ]Zi(x̃ 1 i )), ..., (T z)′i([f −1 x ]Si(x̃ kz i ))
) imply that the ( JTs(x̄ 1), ..., JTs(x̄ ks) )
and( JTz (˜̄x 1), ..., JTz (˜̄x kz ) )
are invertible, which implies the invertibility of matrix Ms and Mz . The rest is to prove pfy (y|[fx]−1S (x)) = pf̃y (y|[f̃x] −1 S (x)). This can be shown by applying Eq. (31) again. Specifically, according to Eq. (31), we have that∫ X pεx(x− x̄)p(y|[fx]−1S (x̄))pTs,Γs,Tz,Γz (f −1(x̄)|de)volJf−1(x̄)dx̄
= ∫ X pεx(x− x̄)p(y|[f̃x]−1S (x̄))pTs,Γs,Tz,Γz (f̃ −1(x̄)|de)volJf̃−1(x̄)dx̄. (50)
Denote lTs,Γs,Tz,Γz,fy,fx,y(x) := pfy (y|[fx]−1S (x̄))pTs,Γs,Tz,Γz (f−1(x̄)|de)volJf−1x (x̄), we have∫ X pεx(x− x̄)lTs,Γs,Tz,Γz,fy,fx,y(x̄)dx̄ = ∫ X pεx(x− x̄)lT̃s,Γ̃s,T̃z,Γ̃z,f̃y,f̃x,y(x̄)dx̄ (51)
=⇒(lTs,Γs,Tz,Γz,fy,fx,y ∗ pεx)(x|de) = (lT̃s,Γ̃s,T̃z,Γ̃z,f̃y,f̃x,y ∗ pεx)(x|d e) (52)
=⇒F [lT̃s,Γ̃s,T̃z,Γ̃z,f̃y,f̃x,y](ω)ϕεx(ω) = F [lTs,Γs,Tz,Γz,fy,fx,y](ω)ϕεx(ω) (53) =⇒F [lTs,Γs,Tz,Γz,fy,fx,y](ω) = F [lT̃s,Γ̃s,T̃z,Γ̃z,f̃y,f̃x,y](ω) (54) =⇒lTs,Γs,Tz,Γz,fy,fx,y(x) = lT̃s,Γ̃s,T̃z,Γ̃z,f̃y,f̃x,y(x) (55) =⇒pfy (y|[fx]−1S (x))pTs,Γs,Tz,Γz (f −1(x)|de)volJf−1x (x)
= pf̃y (y|[f̃x] −1 S (x))pT̃s,Γ̃s,T̃z,Γ̃z (f̃ −1(x)|de)volJf̃−1x (x). (56)
Taking the log transformation on both sides of Eq. (56), we have that
log pfy (y|[fx]−1S (x)) + log pTs,Γs,Tz,Γz (f −1(x)|de) + log volJf−1x (x)
= log pf̃y (y|[f̃x] −1 S (x)) + log pT̃s,Γ̃s,T̃z,Γ̃z (f̃ −1(x)|de) + log volJf̃−1x (x). (57)
Subtracting Eq. (57) with y2 from Eq. (57) with y1, we have
pfy (y2|[fx]−1S (x)) pfy (y1|[fx]−1S (x)) = pf̃y (y2|[f̃x] −1 S (x)) pf̃y (y1|[f̃x] −1 S (x))
(58)
=⇒ ∫ Y pfy (y2|[fx]−1S (x)) pfy (y1|[fx]−1S (x)) dy2 = ∫ Y pf̃y (y2|[f̃x] −1 S (x)) pf̃y (y1|[f̃x] −1 S (x)) dy2 (59) =⇒pfy (y1|[fx]−1S (x)) = pf̃y (y1|[f̃x] −1 S (x)), (60)
for any y1 ∈ Y . This completes the proof.
Understanding the assumption (4) in Theorem 7.9 and 7.6. Recall that we assume the confounder ds in LaCIM is the source variable for generating data in corresponding domain. Here we also use the C to denote the space of ds (since ds := c), then we have the following theoretical conclusion that the as long as the image set of C is not included in any sets with Lebesgue measure 0, the assumption (4) holds. This conclusion means that the assumption (4) holds generically.
Theorem 7.10. Denote ht=s,z(d) := (
Γt1,1(d)− Γt1,1(de1), ...,Γtqt,kt(d)− Γ t 1,1(d
e1) )>
, h(C) := hs(S) ⊕ hz(Z) ⊂ Rqz∗kz ⊕ Rqs∗ks , then assumption (4) holds if h(C) is not included in any zero-measure set of Rqz∗kz ⊕ Rqs∗ks . Denote rs := qs ∗ ks and rz := qz ∗ kz .
Proof. With loss of generality, we assume that rs ≤ rz . Denote Q as the set of integers q such that there exists de2 , ..., dq+1 that the rank([hz(de2), ..., hz(deq+1)]) = min(q, rz) and rank([hs(de2), ..., hs(deq+1)]) = min(q, rs). Denote u := max(Q). We discuss two possible cases for u, respectively:
• Case 1. u < rs ≤ rz . Then there exists de2 , ..., deu+1 s.t. hz(de2), ..., hz(deu+1) and hs(de2), ..., hs(deu+1) are linearly independent. Then ∀c, we have hz(d) ∈ L(hz(de2), ..., hz(deu+1)) or hs(d) ∈ L(hs(de2), ..., hs(deu+1)). Therefore, so we have hz(d) ⊕ hs(d) ∈ [L(hz(de2), ..., hz(deu+1))⊕ Rrs ] ∪ [Rrz ⊕ L(hs(de2), ..., hs(deu+1))], which has measure 0 in Rrz ⊕ Rrs .
• Case 2. rs ≤ u < rz . Then there exists de2 , ..., deu+1 s.t. hz(de2), ..., hz(deu+1) are linearly independent and rank([hs(de1), ..., hs(deu)]) = rs. Then ∀c, we have hz(d) ∈ L(hz(de1), ..., hz(deu+1)), which means that hz(d)⊕hs(d) ∈ L(hz(de1), ..., hz(deu+1))⊕ Rrs , which has measure 0 in Rrz ⊕ Rrs .
The above two cases are contradict to the assumption that h(C) is not included in any zero-measure set of Rrz ⊕ Rrs .
Lemma 7.11. Consider the cases when ks ≥ 2. Then suppose the assumptions in theorem 7.9 are satisfied. Further assumed that
• The sufficient statistics Tsi,j are twice differentiable for each i ∈ [qs] and j ∈ [ks].
• fy is twice differentiable.
Then we have Ms in theorem 7.9 is block permutation matrix.
Proof. Directly applying (Khemakhem, Kingma and Hyvärinen, 2020, Theorem 2) with fx, A, b,T, x replaced by fy,Ms, as,Ts, y.
Lemma 7.12. Consider the cases when ks = 1. Then suppose the assumptions in theorem 7.9 are satisfied. Further assumed that
• The sufficient statistics Tsi are not monotonic for i ∈ [qs].
• g is smooth.
Then we have Ms in theorem 7.9 is block permutation matrix.
Proof. Directly applying (Khemakhem, Kingma and Hyvärinen, 2020, Theorem 3) with fx, A, b,T, x replaced by fy,Ms, as,Ts, y.
Proof of Theorem 7.6. According to theorem 7.9, there exist invertible matrices Ms and Mz such that
T(f−1x (x)) = AT̃(f̃ −1 x (x)) + b Ts([f−1x ]S(x)) = MsT̃ s([f̃−1x ]S(x)) + as. Ts(f−1y (y)) = MsT̃ s(f̃−1y (y)) + as,
where T = [Ts,>,Tz,>]>, and
A = ( Ms 0 0 Mz ) . (61)
By further assuming that the sufficient statistics Tsi,j are twice differentiable for each i ∈ [qs] and j ∈ [ks] for ks ≥ 2 and not monotonic for ks = 1. Then we have that Ms is block permutation matrix. By further assuming that Tzi,j are twice differentiable for each i ∈ [nz] and j ∈ [kz] for kz ≥ 2 and not monotonic for kz = 1 and applying the lemma 7.11 and 7.12 respectively, we have that A is block permutation matrix. Therefore, Mz is also a block permutation matrix.
Proof of Theorem 4.3. We consider the general case when C := ∪Rr=1Cr, in which each Cr can be simplified as a representative point cr. For environment de, let Pde = [P(C = c1|de), · · · ,P(C = cR|de)] be the vector of probability mass of C in the environment de. And Etrain hasm environments with indexes de1 , · · · , dem . The latent factors (S,Z) belongs to the exponential family distribution p(s, z|c) = pTz,Γz(d)(z)pTs,Γs(d)(s). Suppose that θ = {fx, fy,Ts,Tz} and θ̃ = {f̃x, g̃y, T̃s, T̃z} share the same observational distribution for each environment, i.e., pθ(x, y|de) = pθ̃(x, y|de), then we have that
R∑ r=1 pθ(x, y|cr) P(C=cr|de) = R∑ r=1 pθ̃(x, y|cr) P(C=cR|d e). (62)
Let ∆x,y = [pθ(x, y|c1)−pθ̃(x, y|c1), · · · , pθ(x, y|cm)−pθ̃(x, y|cm)]T, then Eq. (62) can be written as A∆x,y = 0. Denote A := P>de1 ∈ Rm×R. According the diversity condition, we have that A and the [[Γt(c2)−Γt(c1)]T, ..., [Γt(cm)−Γt(c1)]T]T have full column rank, therefore we have that ∆x,y = 0, i.e. pθ(x, y|cr) = pθ̃(x, y|cr) for each r ∈ [R]. The left proof is the same with the one in theorem 7.6.

heading: 7.4 PROOF OF THEOREM 4.4
text: Proof of Theorem 4.4. Due to Eq. (62), it is suffices to prove the conclusion for every cr ∈ {cr}r∈[R]. Motivated by Barron and Sheu (1991, Theorem 2) that the distribution pe(s, z) defined on bounded set can be approximated by a sequence of exponential family with sufficient statistics denoted as polynomial terms, therefore the Tt=s,z are twice differentiable hence satisfies the assumption (2) in theorem 4.3 and assumption (1) in lemma 7.11. Besides, the lemma 4 in Barron and Sheu (1991) informs us that the KL divergence between pθ0(s, z|cr) (θ0 := (fx, fy,T z,T s,Γz0,Γs0) and pθ1(s, z|cr) (θ1 := (fx, fy,T z,T s,Γz1,Γs1) (the pθ0(s, z|cr), pθ1(s, z|cr) belong to exponential family with polynomial sufficient statistics terms) can be bounded by the `2 norm of [(Γs(cr) − Γs1(cr))
>, (Γz0(cr) − Γz1(cr))>]>. Therefore, ∀ > 0, there exists a open set of Γ(cr) such that the DKL(p(s, z|cr), pθ(s, z|cr)) < . Such an open set is with non-zero Lebesgue measurement therefore can satisfy the assumption (4) in theorem 4.3, according to result in theorem 7.10. The left is to prove that for any p defined by a LaCIM following Def. 4.1, there is a sequence of {pm}n ∈ Pexp such that the dPok(p, pn)→ 0 that is equivalent to pn d→ p. For any A,B, we consider to prove that
In ∆ = ∣∣∣∣p(x ∈ A, y ∈ B|cr)− pn(x ∈ A, yn ∈ B|cr)∣∣∣∣→ 0, (63) where pn(x ∈ A, yn ∈ B|cr) = ∫ S ∫ Z p(x ∈ A|s, z)p(yn ∈ B|s)pn(s, z|cr)dsdz with
yn(i) = exp((fy,i(s) + εy,i)/Tn)∑ i exp((fy,i(s) + εy,i)/Tn) , i = 1, ..., k, (64)
for y ∈ Rk denoting the k-dimensional one-hot vector for categorical variable and εy,1,...,k are Gumbel i.i.d. According to (Maddison et al., 2016, Proposition 1) that the yn(i) d→ y(i) with
p(y(i) = 1) = exp(fy,i(s))∑ i exp((fy,i(s)) , as Tn → 0. (65)
As long as fy is smooth, we have that the p(yn|s) is continuous. We have that In = ∣∣∣p(x ∈ A, y ∈ B|cr)− ∫
S×Z p(x ∈ A|s, z)p(yn ∈ B|s)pn(s, z|cr)dsdz ∣∣∣ ≤ ∣∣∣p(x ∈ A, y ∈ B|cr)− p(x ∈ A, yn ∈ B|cr)∣∣∣
+ ∣∣∣p(x ∈ A, yn ∈ B|cr)− ∫
S×Z p(x ∈ A|s, z)p(yn ∈ B|s)pn(s, z|cr)dsdz ∣∣∣ = ∣∣∣ ∫ S×Z p(x ∈ A|s, z) (p(y ∈ B|s)− p(yn ∈ B|s)) p(s, z|cr)dsdz ∣∣∣
+ ∣∣∣ ∫ S×Z p(x ∈ A|s, z)p(yn ∈ B|s) (p(s, z|cr)− pn(s, z|cr)) ∣∣∣
≤ ∣∣∣ ∫ Ms×Mz p(x ∈ A|s, z) (p(y ∈ B|s)− p(yn ∈ B|s)) p(s, z|cr)dsdz ∣∣∣︸ ︷︷ ︸
In,1 + ∣∣∣ ∫
(Ms×Mz)cr p(x ∈ A|s, z) (p(y ∈ B|s)− p(yn ∈ B|s)) p(s, z|cr)dsdz ∣∣∣︸ ︷︷ ︸ In,2
+ ∣∣∣ ∫ Ms×Mz p(x ∈ A|s, z)p(yn ∈ B|s) (p(s, z|cr)− pn(s, z|cr)) ∣∣∣︸ ︷︷ ︸
In,3 + ∣∣∣ ∫
(Ms×Mz)cr p(x ∈ A|s, z)p(yn ∈ B|s) (p(s, z|cr)− pn(s, z|cr)) ∣∣∣︸ ︷︷ ︸ In,4 . (66)
For In,1, if y is itself additive model with y = fy(s) + εy, then we just set yn d = y, then we have that In,1 = 0. Therefore, we only consider the case when y denotes the categorical variable with softmax distribution, i.e., Eq. (65). ∀cr ∈ C := {c1, ..., cR} and ∀ > 0, there exists M crs and M crz such that p(s, z ∈ M crs ×M crz |cr) ≤ ; Denote Ms ∆ = ∪mk=1M crs and Mz ∆ = ∪mk=1M crz , we have that p(s, z ∈Ms ×Mz|c) ≤ 2 for all cr ∈ C. Since ∀s1 ∈Ms, ∃Ns1 such that ∀n ≥ Ns1 , we have that
∣∣∣p(y ∈ B|s1) − p(y ∈ B|s1)| ≤ from that yn d→ y. Besides, there exists open set Os1 such that ∀s ∈ Os1 and∣∣∣p(y ∈ B|s1)− p(y ∈ B|s1)| ≤ , ∣∣∣p(yn ∈ B|s1)− p(yn ∈ B|s1)| ≤ . Again, according to Heine–Borel theorem, there exists finite s, namely s1, ..., sl such that Ms ⊂ ∪li=1O(si). Then there exists N
∆ = max{Ns1 , ..., Nsl} such that ∀n ≥ N , we have that∣∣p(y ∈ B|s)− p(yn ∈ B|s)∣∣ ≤ 3 , ∀s ∈Ms. (67)
Therefore, In,1 ≤ ∫ Ms×Mz 3 p(x ∈ A|s, z)p(s, z|c)dsdz ≤ 3 . Hence, In,1 → 0 as n → ∞.
Besides, we have that In,2 ≤ ∫ Ms×Mz 2 p(s, z|cr)dsdz ≤ 2 . Therefore, we have that ∣∣ ∫ S×Z p(x ∈
A|s, z) (p(y ∈ B|s)− p(yn ∈ B|s)) p(s, z|cr)dsdz ∣∣→ 0 as n→∞. For In,3, we have that
In,3 = ∣∣∣∣ ∫ Ms×Mz p(x ∈ A|s, z)p(yn ∈ B|s)1(s, z ∈Ms ×Mz) (p(s, z|cr)− pn(s, z|cr)) dsdz ∣∣∣∣
≤ ∣∣∣∣ ∫ Ms×Mz p(x ∈ A|s, z)p(yn ∈ B|s)p(s, z|cr) (
1 p(s, z ∈Ms ×Mz|cr) − 1 ) dsdz ∣∣∣∣︸ ︷︷ ︸ In,3,1
+ ∣∣∣∣ ∫ Ms×Mz p(x ∈ A|s, z)p(yn ∈ B|s)p(s, z|cr) (
1 p(s, z ∈Ms ×Mz|cr) − 1 ) dsdz ∣∣∣∣︸ ︷︷ ︸ In,3,2 .
(68)
The In,3,1 ≤ 1− . Denote p̃(s, z|cr) := p(s,z|cr)1(s,z∈Ms×Mz)
p(s,z∈Ms×Mz|cr) , according to (Barron and Sheu, 1991, Theorem 2), there exists a sequence of pn(s, z|c) defined on a compact support Ms ×Mz such that ∀cr ∈ C, we have that
pn(s, z|cr) d→ p(s, z|cr).
Applying again the Heine–Borel theorem, we have that ∀ , ∃N such that ∀n ≥ N , we have∣∣∣p̃(s, z|cr)− pn(s, z|cr)∣∣∣ ≤ , (69)
which implies that In,3,2 → 0 as n→∞ combining with the fact that p(x, y|s, z) is continuous with respect to s, z. For In,4, we have that
In,4 = ∣∣∣∣ ∫ Ms×Mz p(x ∈ A|s, z)p(yn ∈ B|s)p(s, z|cr) ∣∣∣∣ ≤ ∣∣∣∣ ∫ Ms×Mz p(s, z|cr) ∣∣∣∣ ≤ , (70) where the first equality is from that the pn(s, z|cr) is defined on Ms ×Mz . Then we have that∣∣∣∣ ∫
S×Z p(x ∈ A|s, z)p(yn ∈ B|s) (p(s, z|cr)− pn(s, z|c)) ∣∣∣∣→ 0, as n→∞. (71) The proof is completed.
7.5 REPARAMETERIZATION FOR LACIM-d
We provide an alternative training method to avoid parameterization of prior p(s, z|de) to increase the diversity of generative models in different environments. Specifically, motivated by Hyvärinen and Pajunen (1999) that any distribution can be transformed to isotropic Gaussian with the density denoted by pGau, we have that for any e ∈ Etrain, we have
pe(x, y) = ∫ S×Z pfx(x|s, z)pfy (y|s)p(s, z|de)dsdz
= ∫ S×Z p(x|(ρes)−1(s′), (ρez)−1(z′))p(y|ρs(s′))pGau(s′, z′)ds′dz′,
with s′, z′ := ρes(s), ρ e z(z) ∼ N (0, I). We can then rewrite ELBO for LaCIM-d for environment e as:
Leφ,ψ,ρe = Epe(x,y) [ − log qeψ(y|x) ] + Epe(x,y) [ −Eqeψ(s,z|x)
qψ(y|(ρes)−1(s)) qeψ(y|x) log pφ((ρ e s) −1(s), (ρez) −1(z))pGau(s, z) qeψ(s, z|x)
] .
(72)

heading: 7.6 IDENTIFIABILITY
text: Earlier works that identify the latent confounders rely on strong assumptions regarding the causal structure, such as the linear model from latent to observed variable or ICA in which the latent component are independent Silva et al. (2006), or noise-free model Shimizu et al. (2009); Davies (2004). The Hoyer et al. (2008); Janzing, Peters, Mooij and Schölkopf (2012) extend to the additive noise model (ANM) and other causal discovery assumptions. Although the Lee et al. (2019) relaxed the constraints put on the causal structure, it required the latent noise is with small strength, which does not match with many realistic scenarios, such as the structural MRI of Alzheimer’s Disease considered in our experiment. The works which also based on the independent component analysis (ICA), i.e., the latent variables are (conditionally) independent, include Davies (2004); Eriksson and Koivunen (2003); recently, a series of works extend the above results to deep nonlinear ICA (Hyvarinen and Morioka, 2016; Hyvärinen et al., 2019; Khemakhem, Kingma and Hyvärinen, 2020; Khemakhem, Monti, Kingma and Hyvärinen, 2020; Teshima et al., 2020). However, these works require that the value of confounder of these latent variables is fixed, which cannot explain the spurious correlation in a single dataset. In contrast, our result can incorporate these scenarios by assuming that each sample has a specific value of the confounder. Other works assume discrete distribution for latent variables, such as Janzing, Sgouritsa, Stegle, Peters and Schölkopf (2012); Kocaoglu et al. (2018); Sgouritsa et al. (2013). However, in the literature, no existing works can disentangle the prediction-causative features from others, in the scenario of avoiding spurious correlation in order for OOD generalization.

heading: 7.7 COMPARISON WITH EXISTING WORKS

heading: 7.7.1 Y → S OR S → Y ?
text: Many existing works Rojas-Carulla et al. (2018); Khemakhem, Monti, Kingma and Hyvärinen (2020); Ilse et al. (2020; 2019) assumed Y → S(X) as the causal direction. Such an difference from ours
can mainly be contributed to the generating process of Y . Different understanding leads to different causal graph. The example of digital hand-writing in Peters et al. (2017) provides a good explanation. Consider the case that the writer is provided with a label first (such as ”2”) before writing the digit (denoted as X), then it should be Y → X . Consider another case, when the writing is based on the incentive (denoted as S) of which digit to write, then the writer record the label Y and the digit X concurrently, in which case it should be X ← S → Y . For Y → S, the Y is thought to be the source variable that generates the latent components and is observed before X . In contrast, we define Y as ground-truth labels given by humans. Taking image classification as an example, it is the human that give the classification of all things such as animals. In this case, it can be assumed that the label given by humans are ground-truth labels. This assumption can be based by the work Biederman (1987) in the field of psychology that humans can factorize the image X by many components due to the powerful perception learning ability of human beings. These components which denoted as S, can be accurately detected by humans, therefore we can approximately assume that it is the S generating the label Y . Consider the task of early prediction in Alzheimer’s Disease, the disease label is given based on the pathological analysis and observed after the MRI X . Such a labelling outcome can be regarded as the ground-truth which itself is defined by medical science. The corresponding pathology features, as the evidences for labelling, can also thought as the generators of X . In these cases, it is more appropriate to assume the Y as the outcome than the cause. For example, the Peters et al. (2016); Kuang et al. (2018) assumed XS → Y . As an adaptation to sensory-level data such as image, we assume S → Y with S are latent variables to model high-level explanatory factors, which coincides with existing literature Teshima et al. (2020). Another difference lies in the definition of Y . The Invariant Risk Minimization (we will give a detailed comparison later) Arjovsky et al. (2019) assumes that X → S̃ → Y by defining the Y as the label with noise. The S̃ denoted as the extracted hidden components by observer.

heading: 7.7.2 COMPARISONS WITH DATA AUGMENTATION & ARCHITECTURE DESIGN
text: The goal of data augmentation Shorten and Khoshgoftaar (2019) is increase the variety of the data distribution, such as geometrical transformation Kang et al. (2017); Taylor and Nitschke (2017), flipping, style transfer Gatys et al. (2015), adversarial robustness Madry et al. (2017). On the other way round, an alternative kind of approaches is to integrate into the model corresponding modules that improve the robustness to some types of variations, such as Worrall et al. (2017); Marcos et al. (2016).
However, these techniques can only make effect because they are included in the training data for neural network to memorize Zhang et al. (2016); besides, the improvement is only limited to some specific types of variation considered. As analyzed in Xie et al. (2020); Krueger et al. (2020), the data augmentation trained with empirical risk minimization or robust optimization Ben-Tal et al. (2009) such as adversarial training Madry et al. (2017); Sagawa et al. (2019) can only achieve robustness on interpolation (convex hull) rather than extrapolation of training environments.

heading: 7.7.3 COMPARISONS WITH EXISTING WORKS IN DOMAIN ADAPTATION
text: Apparently, the main difference lies in the problem setting that (i) the domain adaptation (DA) can access the input data of the target domain while ours cannot; and (ii) our methods need multiple training data while the DA only needs one source domain. For methodology, our LaCIM shares insights but different with DA. Specifically, both methods assume some types of invariance that relates the training domains to the target domain. For DA, one stream is to assume the same conditional distribution shared between the source and the target domain, such as covariate shift Huang et al. (2007); Ben-David et al. (2007); Johansson et al. (2019); Sugiyama et al. (2008) in which P (Y |X) are assumed to be the same across domains, concept shift Zhang et al. (2013) in which the P (X|Y ) is assumed to be invariant. Such an invariance is related to representation, such as Φ(X) in Zhao et al. (2019) and P (Y |Φ(X)) in Pan et al. (2010); Ganin et al. (2016); Magliacane et al. (2018). However, these assumptions are only distribution-level rather than the underlying causation which takes the data-generating process into account. Taking the image classification again as an example, our method first propose a causal graph in which the latent factors are introduced as the explanatory/causal factors of the observed variables. These are supported by the framework of generative model Khemakhem, Kingma and Hyvärinen (2020); Khemakhem, Monti, Kingma and Hyvärinen (2020); Kingma and Welling (2014); Suter et al. (2019) which has natural connection with the causal
graph Schölkopf (2019) that the edge in the causal graph reflects both the causal effect and also the generating process. Until now, perhaps the most similar work to us are Romeijn and Williamson (2018) and Teshima et al. (2020) which also need multiple training domains and get access to a few samples in the target domain. Both work assumes the similar causal graph with us but unlike our LaCIM, they do not separate the latent factors which can not explain the spurious correlation learned by supervised learning Ilse et al. (2020). Besides, the multiple training datasets in Romeijn and Williamson (2018) refer to intervened data which may hard to obtain in some applications. We have verified in our experiments that explicitly disentangle the latent variables into two parts can result in better OOD prediction power than mixing them together.

heading: 7.7.4 COMPARISONS WITH DOMAIN GENERALIZATION
text: For domain generalization (DG), similar to the invariance assumption in DA, a series of work proposed to align the representation Φ(X) that assumed to be invariant across domains Li et al. (2017; 2018); Muandet et al. (2013). As discussed above, these methods lack the deep delving of the underlying causal structure and precludes the variations of unseen domains.
Recently, a series of works leverage causal invariance to enable OOD generalization on unseen domains, such as Ilse et al. (2019) which learns the representation that is domain-invariant. Notably, the Invariant Causal Prediction Peters et al. (2016) formulates the assumption in the definition of Structural Causal Model and assumes that Y = XSβ?S + εY where εY satisfies Gaussian distribution and S denotes the subset of covariates ofX . The Rojas-Carulla et al. (2018); Bühlmann (2018) relaxes such an assumption by assuming the invariance of fy and noise distribution εy in Y ← fy(XS , εy) which induces P (Y |XS). The similar assumption is also adopted in Kuang et al. (2018). However, these works causally related the output to the observed input, which may not hold in many real applications in which the observed data is sensory-level, such as audio waves and pixels. It has been discussed in Bengio et al. (2013); Bengio (2017) that the causal factors should be high-level abstractions/concepts. The Heinze-Deml and Meinshausen (2017) considers the style transfer setting in which each image is linear combination of shape-related variable and contextual-related variable, which respectively correspond to S and Z in our LaCIM in which the nonlinear mechanism (rather than linear combination in Heinze-Deml and Meinshausen (2017)) is allowed. Besides, during testing, our method can generalize to the OOD sample with intervention such as adversarial noise and contextual intervention.
Recently, the most notable work is Invariant Risk Minimization Arjovsky et al. (2019), which will be discussed in detail in the subsequent section.

heading: 7.7.5 COMPARISONS WITH INVARIANT RISK MINIMIZATION ARJOVSKY ET AL. (2019) AND REFERENCES THERE IN
text: The Invariant Risk Minimization (IRM) Arjovsky et al. (2019) assumes the existence of invariant representation Φ(X) that induces the optimal classifier for all domains, i.e., the E[Y |Pa(Y )] is domain-independent in the formulation of SCM. Similar to our LaCIM, the Pa(Y ) can refer to latent variables. Besides, to identify the invariance and the optimal classifier, the training environments also need to be diverse enough. As aforementioned, this assumption is almost necessary to differentiate the invariance mechanism from the variant ones. To learn such an invariance, a regularization function is proposed.
The difference of our LaCIM with IRM lies in two aspects: the direction of causal relation and the methodology. For the direction, as aforementioned in section 7.7.1, the IRM assumes X → S rather than the S,Z → X in our LaCIM. This is because the IRM defines Y as label with noise while ours definie the Y as the ground-truth label hence should be generated by the ground-truth hidden components that generating S. Such an inconsistency can be reflected by experiment regarding to the CMNIST in which the number is the causal factors of the label Y , rather than only invariant correlation. Besides, in terms of methodology, the theoretical claim of IRM only holds in linear case; in contrast, the CIMe fx, fy are allowed to be nonlinear.
Some other works share the similar spirit with or based on IRM. The Risk-Extrapolation (REx) Krueger et al. (2020) proposed to enforce the similar behavior of m classifiers with variance of which proposed as the regularization function. The work in Xie et al. (2020) proposed a Quasidistribution framework that can incorporate empirical risk minimization, robust optimization and
REx. It can be concluded that the robust optimization only generalizes the convex hull of training environments (defined as interpolation) and the REx can generalize extrapolated combinations of training environments. This work lacks model of underlying causal structure, although it performs similarly to IRM experimentally. Besides, the Teney et al. (2020) proposed to unpool the training data into several domains with different environment and leverages Arjovsky et al. (2019) to learn invariant information for classifier. Recently, the Bellot and van der Schaar (2020) also assumes the invariance to be generating mechanisms and can generalize the capability of IRM when unobserved confounder exist. However, this work also lacks the analysis of identifiability result.
We finish this section with the following summary of methods in section 7.7.4 and the IRM, in terms of causal factor, invariance type, direction of causal relation, theoretical judgement and the ability to generalize to intervened data.

heading: 7.8 IMPLEMENTATION DETAILS AND MORE RESULTS FOR SIMULATION
text: Data Generation We set m = 5, ne = 1000 for each e. The generating process of ds ∈ Rqds , Z ∈ Rqz , S ∈ Rqs , X ∈ Rqx and Y ∈ Rqy is introduced in the supplement 7.8. We set qds = qs = qz = qy = 2 and qx = 4. For each environment e ∈ [m] with m = 5, we generate 1000 samples De = {xi, yi} i.i.d∼ ∫ pfx(x|s, z)pfy (y|s)pe(s, z|des)dsdz. The des = ( N (0, Iqds×qds ) + 5 ∗ e ) ∗ 2;
the s, z|des ∼ N ( µφ?s,z (s, z|d e s), σ 2 φ?s,z (s, z|des) ) with µφ?s,z = A µ s,z ∗ des and log σφ?s,z = A σ s,z ∗ des
(Aµs,z , A σ s,z are random matrices); the x|s, z ∼ N ( µφ?x(x|s, z), σ 2 φ?x (x|s, z) ) with µφ?s,z = h(A µ,3 x ∗ h(Aµ,2x ∗ h(Aµ,2x ∗ [s>, z>]>]))) and log σφ?s,z = h(A σ,3 x ∗ h(Aσ,2x ∗ h(Aσ,2x ∗ [s>, z>]>]))) (h is LeakyReLU activation function with slope = 0.5 and Aµ,i=1,2,3x ,A σ,i=1,2,3 x are random matrices); the y|s is similarly to x|s, z with Aµ,i=1,2,3x ,Aσ,i=1,2,3x respectively replaced by Aµ,i=1,2,3y ,Aσ,i=1,2,3y .
Implementation Details We parameterize pθ(s, z|d), qφ(s, z|x, y, d), pθ(x|s, z) and pθ(y|s) as 3- layer MLP with the LeakyReLU activation function. The Adam with learning rate 5 × 10−4 is implemented for optimization. We set the batch size as 512 and run for 2,000 iterations in each trial.
Visualization. As shown from the visualization of S is shown in Fig. 7.8, our LaCIM can identify the causal factor S.
The setting when C can take a value in a sample-level. We consider the generation process of De
as De = {xi, yi} i.i.d∼ ∫ pfx(x|s, z)pfy (y|s)p(s, z|c)p(c|de)dsdzdc, with qc := 2. The generation is the same except that the after obtaining ds, we additionally generate c with c := N (ds, I). The results are summarized in Tab. 7.8.
7.9 IMPLEMENTATION DETAILS FOR OPTIMIZATION OVER S,Z
Recall that we first optimize s∗, z∗ according to
s∗, z∗ = arg max s,z log pφ(x|s, z).
We first sample some initial points from each posterior distribution qeψ(s|x) and then optimize for 50 iterations. We using Adam as optimizer, with learning rate as 0.002 and weight decay 0.0002. The
Fig. 7.9 shows the optimization effect of one run in CMNIST. As shown, the test accuracy keeps growing as iterates. For time saving, we chose to optimize for 50 iterations.
Figure 5: The optimization effect in CMNIST, starting from the point with initial sampling from inference model q of each branch. As shown, the test accuracy increases as iterates.

heading: 7.10 IMPLEMENTATIONS FOR BASELINE
text: For the CE X → Y and the CE X, ds → Y , they both composed of two parts: (i) feature extractor, followed by (ii) classifier. The network structure of the feature extractor for CE X → Y is the same with that of our encoder; while the extracted features for CE X, d→ Y is the concatenation of the features encoded from X → S,Z via the network with the same network structure of our encoder; and the network with the same structure of our prior network for LaCIM-d. The network structures of the classifier for both methods are the same to that of our pφ(y|s). The IRM and SDA adopt the same structure as CE X → Y . DANN adopt the same structure of CE X → Y and a additional domain classifier which is the same as that of pφ(y|s). sVAE adopt the same structure as LaCIM-ds with the exception that the pφ(y|s) is replaced by pφ(y|z, s). MMD-AAE adopt the same structure of encoder, decoder and classifier as LaCIM-d and a additional 2-layer MLP with channel 256-256-dimz is used to extract latent z. The detailed number of parameters and channel size on each dataset for each method are summarized in Tab. 13, 14.

heading: 7.11 SUPPLEMENTARY FOR COLORED MNIST
text: Implementation details The network structure for inference model is composed of two parts, with the first part shared among all environments and multiple branches corresponding to each environment
for the second part. The network structure of the first-part encoder is composed of four blocks, each block is the sequential of Convolutional Layer (Conv), Batch Normalization (BN), ReLU and maxpooling with stride 2. The output number of feature map is accordingly 32, 64, 128, 256. The second part network structure that output the mean and log-variance of S,Z is Conv-bn-ReLU(256) → Adaptive (1)→ FC(256, 256)→ ReLU→ FC(256, qt=s,z) with FC stands for fully-connected layer. The structure of ρt=s,z in Eq. (72) is FC(qt, 256)→ ReLU→ FC(256, qt). The network structure for generative model pφ(x|s, z) is the sequential of three modules: (i) Upsampling with stride 2; (ii) four blocks of Transpose-Convolution (TConv), BN and ReLU with respective output dimension being 128, 64, 32, 16; (iii) Conv-BN-ReLU-Sigmoid with number of channels in the output as 3, followed by cropping step in order to make the image with the same size as input dimension, i.e., 3× 28× 28. The network structure for generative model pφ(y|s) is commposed of FC (512)→ BN→ ReLU→ FC (256)→ BN→ ReLU→ FC (|Y|). The qt=s,z is set to 32. We implement SGD as optimizer with learning rate 0.5, weight decay 1e− 5 and we set batch size as 256. The total training epoch is 80. We first explain why we do not flip y with 25% in the manuscript, and then provide further exploration of our method for the setting with flipping y.
Invariant Causation v.s. Invariant Correlation by Flipping y in Arjovsky et al. (2019) The y is further flipped with 25% to obtain the final label in IRM setting and this step is omitted in ours. The difference lies in the definition of invariance. Our LaCIM defines invariance as the causal relation between S and the label Y , while the one in IRM can be correlation. As illustrated in Handwritting Sample Form in Fig. 7.11 in Grother (1995), the generting direction should be Y → X . If we denote the variable by flipping Y as Ỹ (a.k.a, the final label in IRM), then the causal graph should be X ← Y → Ỹ . In this case, the Ỹ is correlated rather than causally related to the digit X . For our LaCIM, we define the label as interpretable human label (which can approximate to y for any image x) and represented by Y in our experiments. The reason why we do not define the Y as ground-truth label is that (i) the prediction is only based on the extracted components of image which may be determined not only by the ground-truth label; (ii) the learning of ground-truth is interpretable that relevant to human. For example, if a writer is provided with digit “2” but he wrote it mistakenly as “4”, then it is more interpretable that we can predict the digit as “4” rather than “2”. For the digit with ambiguous label from the perspective of image, even if we predict it mistakenly, it is also interpretable in terms of prediction given the information of only digit. Returning back to the IRM setting, the label is flipping without reference to the semantic shape of digit. Therefore, the flipping may happen to noiseless digits rather than noisy and unsure ones, making the shape of number less semantically related to the label.
Experiment with IRM setting We further conduct the experiment on IRM setting, with the final label y defined by flipping original label with 25%, and further color pe proportions of digits with corresponding color-label mapping. If we assume the original ground-truth label to be the effect of the digit number of S, then the anti-causal relation with Z and Y can make the identifiability of S difficult in this flipping scenario. Note that the causal effect between S and Y is invariant across domains, therefore we adopt to regularize the branch of inferring S to be shared among inference models for multiple environments. Besides, we regularize the causal effect between S and Z to be shared among different environments via pairwise regularization. The combined loss is formulated as:
L̃ψ,φ = Lψ,φ + Γ
2m2 m∑ i=1 m∑ j=1 ‖E(x,y)∼pei (x,y)[y|x]− E(x,y)∼pej (x,y)[y|x]‖22,
with qeψ(s, z|x) in Eq. (72) factorized as qψez (z)qψs(s) and ρs shared among m environments. The appended loss is coincide with recent study Risk-Extropolation (REx) in Krueger et al. (2020), with the difference of separating y-causative factors S from others. We name such a training method as LaCIM-REx. For implementation details, in addition to shared encoder regarding S, we set learning rate as 0.1, weight decay as 0.0002, batch size as 256. we have that p(y|x) = ∫ S qψs(s|x)pφ(y|ρs(s)) for any x. We consider two settings: setting#1 with m2 and pe1 = 0.9, pe2 = 0.8; and setting#2 with m = 4 with pe1 = 0.9, pe2 = 0.8, pe3 = 0.7, pe4 = 0.6. We only report the number of IRM since the cross entropy performs poorly in both settings. As shown, our model performs comparably with LaCIM-ds and better than IRM Arjovsky et al. (2019) due to separation of S znd Z.

heading: 7.12 SUPPLEMENTARY FOR NICO
text: Implementation Details Due to size difference among images, we resize each image into 256×256. The network structure of pθ(z, s|ds), qφ(z, s|x, ds), pθ(x|z, s), pθ(y|s) for cat/dog classification is the same with the one implemented in early prediction of Alzheimer’s Disease with exception of 3D convolution/Deconvolution replaced by 2D ones. For each model, we train for 200 epochs using sgd, with learning rate (lr) set to 0.01, and after every 60 epochs the learning rate is multiplied by lr decay parameter that is set to 0.2. The weight decay coefficients parameter is set to 5× 10−4. The
batch size is set to 30. The training environments which is characterized by c can be referenced in
Table 7.12. For visualization, we implemented the gradient-based method Simonyan et al. (2013) to visualize the neuron (in fully connected layer for both CE x→ y and CE (x, ds)→ y; in s layer for LaCIM-ds) that is most correlated to label y.
The ds form environments We summarize the ds ofm = 8 andm = 14 environments in Table 7.12. As shown, the value of ds in the test domain is the extrapolation of the training environments, i.e., the dtests is not included in the convex hull of {dei}14i=1. More Visualization Results Fig. 7 shows more visualization results.
Results on Intervened Data. We test our model and the baseline on intervened data, in which each image is generated by intervention on Z, i.e., taking a specific value of Z. This intervention breaks the correlation between S and Z, thus the distribution of which can be regarded as a specific type of OOD. Specifically, we replace the scene of an image with the scene from the another image, as shown in Fig. 8. We generate 120 images, including 30 images of types: cat on grass, dog on grass, cat on snow, and dog on grass. We evaluate LaCIM-d, CE X → Y , IRM, DANN, NCBB, MMD-AAE, and DIVA methods on this intervened dataset. As shown in Tab 9, our LaCIM-d can performs the best among all methods, which validate the robustness of our LaCIM.

heading: 7.13 DISEASE PREDICTION OF ALZHEIMER’S DISEASE
text: Dataset Description. The dataset contains in total 317 samples with 48 AD, 75 NC, and 194 MCI.
Denotation of Attributes ds. The C ∈ R9 includes personal attributes (e.g., age Guerreiro and Bras (2015), gender Vina and Lloret (2010) and education years Mortimer (1997) that play as potential
risks of AD), gene (ε4 allele), and biomarkers (e.g., changes of CSF, TAU, PTAU, amyloidβ , cortical amyloid deposition (AV45) Humpel and Hochstrasser (2011)).
Implementation Details For LaCIM-ds, we parameterize inference model qψ(s, z|x, ds), pφ(s, z|ds), pφ(x|z, s) and pφ(y|s) and S,Z ∈ R64. For qψ(s, z|x, ds), we concatenate outputs of feature extractors of X and ds: the feature extractor for x is composed of four Convolution-Batch
Normalization-ReLU (CBNR) blocks and four Convolution-Batch Normalization-ReLU-MaxPooling (CBNR-MP) blocks with structure 64 BNR→ 128 CBNR-MP→ 128 CBNR→ 256 CBNR-MP → 256 CBNR→ 512 CBNR-MP→ 512 CBNR→ 1024 CBNR-MP; the feature extractor of c is composed of three Fully Connection-Batch Normalization-ReLU (FC-BNR) blocks with structure 128→ 256→ 512. The concatenated features are further transformed by four 64 FC-BNR to generate µs,z(x, ds) and log σs,z(x, ds). For the prior model pθ(s, z|ds), it shares the same structure without feature extractor of x. For pφ(x|s, z), the network is composed of three DeConvolution-Batch Normalization-ReLU (DCBNR) blocks and three Convolution-Batch Normalization-ReLU (CBNR) blocks, followed by a convolutional layer, with structure 256 DCBNR→ 256 CBNR→ 128 DCBNR → 128 CBNR→ 64 DCBNR→ 64 CBNR→ 48 Conv. For pφ(y|s), the network is composed of 256 FC-BNR→ 512 FC-BNR→ 3 FC-BNR. For prior model pφ(s, z|ds)N (µs,z(ds),diag(σ2s,z(ds))) the µs,z(x, ds) and log σs,z(x, ds) are parameterized by Multi Perceptron Neural Network (MLP). The decoders pφ(x|s, z) are pφ(y|s) parameterized by Deconvolutional neural network. For all methods, we train for 200 epochs using SGD with weight decay 2× 10−4 and learning rate 0.01 and is multiplied by 0.2 after every 60 epochs. The batch size is set to 4. For each variable in biomarker vector C ∈ R9, each person may have multiple records, and we take its median as representative to avoid extreme values due to device abnormality.
As for LaCIM-d, we adopt the same decoder pφ(x|z, s) and classifier pφ(y|s). For qψ(s, z|x, d), we adopt the same network for the shared part; for the part specific to each domain, µs,z(x, d) and log σs,z(x, d) are generated by the sub-network which is composed of 1024 FC-BNR→ 1024 FC-BNR→ qz,s FC-BNR. The z, s can be reparameterized by µs,z(x, d) and log σs,z(x, d) are fed into a sub-network which is composed of qz,s FC-BNR→ 1024 FC-BNR→ qz,s FC-BNR to get rid of the constraint of Gaussian distribution. Then the reconstructed images and predicted label are computed by pφ(x|z, s) and pφ(y|s) which have the same network structure of LaCIM-C with the z, s.
The ds variable in training and test. The selected attributes include Education Years, Age, Gender (0 denotes male and 1 denotes female), AV45, amyloidβ and TAU. We split the data into m = 2 training environments and test according to different value of ds. The Tab. 7.13 describes the data distribution in terms of number of samples, the value of ds (Age and TAU).

heading: 7.13.1 EXPERIMENTS WITH COMPLETE OBSERVABLE SOURCE VARIABLE
text: In image-based diagnosis, the personal attributes, genes and biomarkers are often available. Therefore, we consider the setting when ds can be fully observed. In this case, the value of ds is person-byperson. Therefore, the number of environments m is equal to the number of samples. In this case, the dataset turns to {xi, yi, dis}ni=1. The expected risk turns to:
Lψ,φ = Ep(x,y|ds) [ − log qψ(y|x, ds)− Eqψ(s,z|x,ds) [ pφ(y|s)
qψ(y|x, ds) log pφ(x|s, z)pφ(s, z|ds) qψ(s, z|x, ds)
]] .
(73)
And the corresponding empirical risk is:
L̃ψ,φ = 1
n
[ − log qψ(yi|xi, ds,i)− Eqψ(s,z|xi,ds,i) [ qψ(yi|s)
qψ(yi|xi, ds,i) log pφ(xi|s, z)pφ(s, z|ds,i) qψ(s, z|xi, ds,i)
]] .
(74)
The ds here is re-defined as the 9-dimensional vector that includes all attributes, genes and biomarkers mentioned above. We re-split the data into 80% train and 20% test, according to different average value of specific variable in the whole vector ds.
The ds variable in training and test We implemented OOD tasks in which the value of ds is different between training and test. Specifically, we repeatedly split the dataset into the training and the test according to a selected attribute in ds for three times. The average value of these attributes in train and test are recorded in Table 7.13.1.
Experimental Results We conduct OOD experiments with source variables Age, Gender, amyloidβ and TAU different between training data and the test. The results are shown in Table 12.

heading: 7.14 SUPPLEMENTARY FOR DEEPFAKE
text: Implementation Details. We implement data augmentations, specifically images with 30 angle rotation, with flipping horizontally with 50% probability. We additionally apply random compressing techniques, such as JpegCompression. For inference model, we adopt Efficient-B5 Tan and Le (2019), with the detailed network structure as: FC(2048, 2048)→ BN→ ReLU→ FC(2048, 2048)→ BN → ReLU→ FC(2048, qt=s,z). The structure of reparameterization, i.e., ρt=s,z is FC(qt=s,z , 2048) → BN→ ReLU→ FC(2048, 2048)→ BN→ ReLU→ FC(2048, qt=s,z). The network structure for generative model, i.e., pψ(x|s, z) is TConv-BN-ReLU(qt=s,z , 256)→ TConv-BN-ReLU(256, 128)→ TConv-BN-ReLU(128, 64)→ TConv-BN-ReLU(64, 32)→ TConv-BN-ReLU(32, 32)→ TConv-BN-ReLU(32, 16) → TConv-BN-ReLU(16, 16) → Conv-BN-ReLU(16, 3) → Sigmoid,
followed by cropping the image to the same size 3×224×224. We set qt=s,z as 1024. We implement SGD as optimizer, with learning rate 0.02, weight decay 0.00005, and run for 9 epochs.
Ta bl
e 13
:G en
er al
fr am
ew or
k ta
bl e
fo ro
ur m
et ho
d an
d ba
se lin
es on
D a ta ∈ {C
M N
IS T ,N
IC O ,A
D N
I, D
ee pF
ak e}
D at
as et
.W e
de no
te th
e di
m en
si on
of z
or z s as di m z ,z s .W e lis tt he ou tp ut di m en si on (e .g .t he ch an ne ln um be r) of ea ch m od ul e, if it is di ff er en tf ro m th e on e in Ta b. 14 .
D at
as et
M et
ho d
C E X → Y
C E X , d → Y
M M
D -A
A E
D A
N N
D IV
A L
aC IM
-d s
L aC
IM -d
D at
a: C
M N
IS T
E nc
D a t a
x FC (2 56
,d im z ) D ec -C E D a t a y
E nc
D a t a x ;E
nc D
a t a
d FC (5 12 ,d im z ) D ec -C E D a t a y
E nc
D a t a
x FC -B N -R eL
U (2
56 ,2
56 )
FC (2
56 ,2
56 )→
z
D ec
D a t a y ;D
ec D
a t a
x
E nc
D a t a
x D A N N -C L SD a t a
y ;D
A N
N -C
L SD
a t a
y
p D
a t a θ (x |z d , z x , z y
)
p D
a t a
θ d
(z d |d
)
p D
a t a
θ y
(z y |y
)
q D
a t a
φ d
(z d |x
)
q D
a t a
φ x
(z x |x
)
q D
a t a
φ y
(z y |x
)
E nc
D a t a x ;E
nc D
a t a
d FC -B N -R eL U (5
12 ,2
56 )
FC (2
56 ,d
im z s ) D ec D a t a y ;D ec D
a t a
x
pr io
r: E
nc D
a t a
d
E nc
D a t a
x E nc D
a t a z ,s × m
Φ D
a t a z ,s × m
D ec
D a t a y ;D
ec D
a t a
x
# of
Pa ra
m s
1. 12
M 1.
16 M
1. 23
M 1.
1M 1.
69 M
1. 09
M 0.
92 M
hy pe
rPa
ra m
s lr
:0 .1
w d:
0. 00
00 5
lr :0 .2 w d: 0.
00 05
lr :0
.0 1
w d:
0. 00
01 lr
:0 .1
w d:
0. 00
02 lr
:0 .0 01 w d: 0. 00 00
1 lr
:0 .1
w d:
0. 00
01 lr
:0 .0 1 w d: 0.
00 02
D at
a: N
IC O
E nc
D a t a
x FC (1 02
4, di
m z ) D ec -C E D a t a y
E nc
D a t a x ;E
nc D
a t a
d FC (5 12 ,d im z ) D ec -C E D a t a y
E nc
D a t a
x FC -B N -R eL
U (1
02 4,
10 24 ) FC (1 02 4, 10 24 )→ z D ec D a t a y ;D ec D a t a x
E nc
D a t a
x D A N N -C L SD a t a
y ;D
A N
N -C
L SD
a t a
y
p D
a t a θ (x |z d , z x , z y
)
p D
a t a
θ d
(z d |d
)
p D
a t a
θ y
(z y |y
)
q D
a t a
φ d
(z d |x
)
q D
a t a
φ x
(z x |x
)
q D
a t a
φ y
(z y |x
)
E nc
D a t a x ;E
nc D
a t a
d FC (1 53 6, di
m z s ) D ec D a t a y ;D ec D a t a x pr io r: E nc D a t a d
E nc
D a t a
x E nc D
a t a z ,s × m
Φ D
a t a z ,s × m
D ec
D a t a y ;D
ec D
a t a
x
# of
Pa ra
m s
(m =
8 )
18 .0
8M 19
.0 1M
19 .7
0M 19
.1 3M
14 .8
6M 16
.3 1M
18 .2 5M # of Pa ra m s (m = 1 4 ) 18 .0 8M 19 .0 1M 19 .7 0M 26 .4 9M 14 .8 7M 18 .0 8M 19 .7 0M hy pe rPa ra m s lr :0 .0 1 w d: 0. 00 02 lr :0 .0 1 w d: 0. 00 02 lr :0 .2 w d: 0. 00 01 lr :0 .0 5 w d: 0. 00 05 lr :0 .0 01 w d: 0. 00 01 lr :0 .0 1 w d: 0. 00 05 lr :0
.0 1
w d:
0. 00
01
D at
a: A
D N
I E
nc D
a t a
x FC (1 02
4, di
m z ) D ec -C E D a t a y
E nc
D a t a x ;E
nc D
a t a
d FC (1 53 6, di
m z ) D ec -C E D a t a y
E nc
D a t a
x FC -B N -R eL
U (1
02 4,
10 24 ) FC (1 02 4, 10 24 )→ z D ec D a t a y ;D ec D a t a x
E nc
D a t a
x D A N N -C L SD a t a
y ;D
A N
N -C
L SD
a t a
y
p D
a t a θ (x |z d , z x , z y
)
p D
a t a
θ d
(z d |d
)
p D
a t a
θ y
(z y |y
)
q D
a t a
φ d
(z d |x
)
q D
a t a
φ x
(z x |x
)
q D
a t a
φ y
(z y |x
)
E nc
D a t a x ;E
nc D
a t a
d FC (1 53 6, di
m z s ) D ec D a t a y ;D ec D a t a x pr io r: E nc D a t a d
E nc
D a t a
x E nc D
a t a z ,s × m
Φ D
a t a z ,s × m
D ec
D a t a y ;D
ec D
a t a
x
# of
Pa ra
m s
28 .2
7M 28
.2 7M
36 .6
8M 30
.2 1M
33 .2
2M 33
.0 7M
37 .7
8M
hy pe
rPa
ra m
s lr
:0 .0 1 w d: 0.
00 02
lr :0
.0 1
w d:
0. 00
02 lr
:0 .0 05 w d: 0.
00 02
lr :0
.0 1
w d:
0. 00
02 lr
:0 .0 05 w d: 0.
00 01
lr :0
.0 05
w d:
0. 00
02 lr
:0 .0 1 w d: 0.
00 02

Review summaries:
This paper proposes a VAE based model for learning latent causal factors given data from multiple domains. Similar to [Kingma and Hyv¨arinen, 2020], it utilizes additional labels as supervision signals and learns the model using a Bayesian optimization approach given a fixed hypothetical causal structure. The identifiability is obtained by assuming the casual mechanism to be domain invariant, which is partially supported by some empirical experiments.