Learning Causally Disentangled Representations via the Principle of Independent Causal Mechanisms

1 Introduction

Disentangled representation learning aims to learn meaningful and compact representations that capture semantic aspects of data by structurally disentangling the factors of variation [Bengio et al., 2013]. Such representations have been shown to offer useful properties such as better interpretability, robustness to distribution shifts, efficient out-of-distribution sampling, and fairness [Locatello et al., 2019]. However, disentangled representation learning typically assumes that the underlying factors are independent, which is unrealistic in practice. The factors generating the data can contain correlations or even causal relationships that are disregarded when factors are assumed to be independent. Further, a generative model learning from an independent prior assumes that all combinations of the latent factors are equally likely to appear in the training data. Thus, disentangling the factors would yield a sub-optimal likelihood since the assumed support could be well outside the support of the training data.

Recently, there has been a growing interest in bridging causality [Pearl, 2009] and representation learning [Bengio et al., 2013]. The goal of causal representation learning (CRL) is to map unstructured low-level data to high-level abstract causal variables of interest [Scholkopf et al., 2021]. The key assumption is that high-dimensional observations are generated from a set of underlying low-dimensional causally related factors of variation. Causal representations have been shown to be useful for tasks involving reasoning and planning. Causal representations also adhere to the principle of independent causal mechanisms (ICM) [Parascandolo et al., 2018], which states that the mechanisms that generate each causal variable are independent such that a change in one mechanism does not affect another [Sch¨olkopf and von K¨ugelgen, 2022; Peters et al., 2017]. Learning a generative model that captures the causal structure among latent factors can be crucial for reasoning about the world under manipulation. For example, a pendulum, light source, and shadow, as seen in Figure 1, may be causally related but are separate entities in the world that can be independently manipulated. Particularly, manipulating the pendulum's angle will affect the shadow's position and length. These hypothetical scenarios could be counterfactually generated from a causal model.

Ensuring that causal representations are disentangled is useful for the causal controllable generation to generate counterfactual instances unseen during training. Such instances could be utilized as augmented data for robust learning in downstream tasks. The notion of disentanglement may be trivial when the factors are independent but becomes difficult to achieve when there are correlations or causal relationships among factors in the observed data. For highly correlated factors, it can be difficult to separate the factors of variation from their latent codes [Tr¨auble et al., 2021]. Recently, it was shown that it is impossible to learn a disentangled representation in an unsupervised manner without some form of inductive bias [Locatello et al., 2019]. Recent work proved that models with an independent prior are unidentifiable [Shen et al., 2022]. Further, most existing disentanglement methods fail to disentangle factors when correlations exist in the data [Tr¨auble et al., 2021]. However, results from large-scale empirical studies [Locatello et al., 2020] have indicated that supervision in the form of auxiliary labels or contrastive data can effectively disentangle correlated or causal factors.

Related Work. Our work builds upon the ideas presented in iVAE [Khemakhem et al., 2020] and causal variants [Yang et al., 2021; Komanduri et al., 2022] and extends them to consider a principled view of causal disentanglement in the label supervised setting. DIVA [Ilse et al., 2020] and CCVAE [Joy et al., 2021] are special case implementations of the iVAE framework. Yang et al. [2021] proposed CausalVAE, which uses a causal masking layer and is limited to linear SCMs. Komanduri et al. [2022] extended this to a nonlinear setting and proposed a causal prior. Both works proposed simplistic models to learn causal mechanisms under the strictly additive noise assumption and do not, from an empirical or theoretical perspective, focus on disentanglement. Shen et al. [2022] proposed learning causal representations supervised by a GAN loss. There has also been work focusing on learning causal representations from paired counterfactual data [Brehmer et al., 2022], temporal data [Lippe et al., 2023], in self-supervised learning [von K¨ugelgen et al., 2021], and when interventional data is available [Ahuja et al., 2023]. Unlike many previous works in label-supervised VAEbased CRL, we consider general nonlinear SCMs instead of restricting to additive noise models, propose a causal prior to causally factorize the latent space, and achieve disentanglement of causal mechanisms, as summarized in Table 1.

Contributions. (1)We propose ICM-VAE, a framework for causal representation learning under label supervision, where causal variables are derived from nonlinear flow-based diffeomorphic causal mechanisms. (2) Based on the ICM principle, we propose the notion of causal disentanglement for causal models from the perspective of mechanisms and design a causal disentanglement prior to causally factorize the learned distribution over causal variables. (3) Using the structure from our causal disentanglement prior, we theoretically show identifiability of the learned causal factors up to permutation and elementwise reparameterization. (4) We experimentally validate our method and show that our model can almost perfectly disentangle the causal factors, improve interventional robustness, and generate consistent counterfactuals.

2 Preliminaries

Let X ⊂ Rd denote the support of the observed data x generated from a set of causally related ground-truth factors. The observations are assumed to be explained by some latent causal factors of variation z with domain Z ⊂ Rn , where n ≪ d. z represents a set of causal factors while zi represents a single causal factor. We assume x can be decomposed as x = g(z) + ξ where ξ ∼ N (0, σ2 I) are mutually independent Gaussian noise terms for reconstruction. Let g : Z → X be the decoder (or mixing function). Each factor zi contains semantically meaningful information about the observation. In the traditional VAE [Kingma and Welling, 2013], we assume the observed data is generated by the latent generative model with the structure pθ(x, z) = pθ(x|z)pθ(z) where θ are the true but unknown parameters. We aim to learn the joint distribution p(x, z) to estimate the marginal density and a posterior p(z|x) to describe the underlying factors of variation given a prior p(z) over the latent variables.

2.1 Identifiability

The goal of learning a useful representation that recovers the true underlying data-generating factors is closely tied to the problem of blind source separation (BSS) and independent component analysis (ICA) [Hyvarinen et al., 2018; Comon, 1994; Hyv¨arinen and Pajunen, 1999]. Provably showing that a learning algorithm achieves this goal up to tolerable ambiguities under certain conditions is formalized as the identifiability of a model. In this section, we use the notion of ∼-equivalence [Khemakhem et al., 2020] to formulate the notion of identifiability.

Definition 1 (∼-identifiability). Let ∼ be an equivalence relation on θ. The generative model is ∼-identifiable if

If two different choices of model parameter θ and ˆ θ lead to the same marginal density pθ(x), then they must be equal and this implies that pθ(x, z) = pˆ θ (x, z), pθ(z) = pˆ θ (z), and pθ(z|x) = pˆ θ (z|x). However, recent work showed that it is impossible to achieve marginal density equivalence pθ(x) = pˆ θ (x) with an unconditional prior pθ(z) [Khemakhem et al., 2020]. Since the VAE is unidentifiable without some form of additional restriction on the function class of the mixing function or auxiliary information, the identifiable VAE (iVAE) was proposed to utilize auxiliary information in the form of a conditionally factorial prior for identifiability guarantees [Khemakhem et al., 2020]. In iVAE, each factor zi is assumed to have a univariate exponential family distribution (due to their universal approximation capabilities) given the conditioning variable u, where a function λ determines the natural parameters of the distribution. The general PDF of the conditional distribution is defined as follows:

where hi(zi) is the base measure, Ti : Z → Rk and Ti = (Ti,1, . . . , Ti,k) are the sufficient statistics, λi(u) = (λi,1(u), . . . , λi,k(u)) are the corresponding natural parameters, k is the dimension of each sufficient statistic, and the remaining term ψi(u) acts as a normalizing constant. A prior conditioned on auxiliary information u can guarantee that the joint densities pθ(x, z) = pˆ θ (x, z) are equivalent up to some equivalence class. The following two definitions describe the conditions necessary to achieve the identifiability of a learned model up to linear transformation and block permutation indeterminacies, respectively.

Definition 2 ([Khemakhem et al., 2020]). Let ∼ be an equivalence relation on θ, X = g(Z), and ˆ X = ˆ g(Z). We say that θ and ˆ θ are linearly-equivalent if and only if there exists an invertible matrix A ∈ Rnk×nk and vectors b, c ∈ Rnk such that ∀x ∈ X :

We denote this equivalence as θ ∼A ˆ θ.

Definition 3 ([Khemakhem et al., 2020]). We say θ and ˆ θ are permutation-equivalent, denoted θ ∼P ˆ θ, if and only if P is permutation matrix that has block-permutation structure respecting T. That is, there exist n invertible k × k matrices A1, . . . ,An and an n-permutation π such that for all z ∈ Rnk , P ˆ z = [zπ(1)AT1 , zπ(2)AT2 , . . . , zπ(n)ATn ]T .

Linear equivalence indicates the true representation is a linear transformation of the learned representation and only guarantees the learned representation captures the true representation. In general, linear-equivalent identifiability does not guarantee that the factors of variation are disentangled since the linear transformation can mix up the variables (i.e. one component of g-1 corresponds to multiple components of ˆ g-1 ). Permutation equivalence implies that the i-th factor zi of one representation corresponds to a unique factor in another representation, given the permutation π. To truly disentangle factors of variation, we must ensure that each coordinate of the learned representation is equal to the scaled and shifted coordinate of the ground truth up to some permutation. To this end, we define the notion of disentanglement as permutation equivalence [Lachapelle et al., 2022] as follows.

Definition 4 (Permutation Disentanglement). Given some ground-truth model, a learned model ˆ θ is said to be disentangled if θ and ˆ θ are permutation-equivalent.

2.2 Structural Causal Model

Henceforth, we assume z is described by a structural causal model (SCM) [Pearl, 2009], which is formally defined by a tuple M = ⟨Z, E, F⟩, where Z is the domain of the set of n endogenous causal variables z = {z1, . . . , zn}, E is the domain of the set of n exogenous noise variables ϵ = {ϵ1, . . . , ϵn}, which is learned as an intermediate latent variable, and F = {f1, . . . , fn} is a collection of n independent causal mechanisms of the form

where ∀i, fi : Ei × ∏j∈pai Zj → Zi are causal mechanisms that determine each causal variable as a function of the parents and noise, zpai are the parents of causal variable zi; and a probability measure pE(ϵ) = pE1 (ϵ1)pE2 (ϵ2) . . . pEn (ϵn), which admits a product distribution. For the purposes of this work, we assume a causally sufficient setting (no hidden confounding), no SCMmisspecification, and faithfulness is satisfied. We depict the causal structure of z by a directed acyclic graph (DAG) G z with adjacency matrix Gz ∈ {0, 1}n×n .

3 Causal MechanismEquivalence

Although the existing notions of disentanglement may be suitable for independent factors of variation [Khemakhem et al., 2020], they fail to capture important information in a causal model where the factors are causally related. As formulated in Def. 2 and Def. 3, linear equivalence or permutation equivalence cannot capture the causal mechanisms accurately or distinguish the mechanisms afflicted to factors. For a counterexample to the definitions, refer to the appendix in [Komanduri et al., 2024]. The framework of iVAE captures identifiability in the sense that the joint distributions of the latent variables of two different models are equivalent. However, for a causally factorized model, we have that pθ(z) = pˆ θ (z) does not imply pθ(zi|zpai ) = pˆ θ (zi|zpai ). That is, the ground-truth causal factors and the learned causal factors should entail the same causal conditional mechanisms, where the minimal conditioning set is the set of causal parents.

Based on the intuition that causal models are described by mechanisms, we define a new notion of disentanglement that takes into account conditional distributions of causal variables under the Markov factorization. The new causal conditional equivalence preserves information about the independent causal mechanisms (ICM), which is a unique formulation for a causal model and important for performing correct interventions. The following two definitions describe the conditions necessary to satisfy causal mechanism equivalence.

Definition 5 (Causal Mechanism Permutation Equivalence). Let ∼ be an equivalence relation between ˆ θ and θ, X = g(Z), and ˆ X = ˆ g(Z). If the factors z are causally related, we say that θ is causal mechanism permutation equivalent to ˆ θ iff:

1. There exists a permutation matrix P such that I = P · J where I and J are indices of z and ˆ z, respectively.
2. Given an equivalence pair (i, j), i.e., Pij ̸= 0, from this permutation matrix, one has Ti(zi|zpai ) = Dij ˆTj(zj|zpaj ), ∀zi ∈ Zi, ∀ˆ zj ∈ Zj, where Dij is a scaling coefficient.
3. For all i, j ∈ {1, . . . , n}, we have the mechanism equivalence λi(zpai , u) = Dij ˆλj (zpaj , u), where D is a diagonal scaling matrix.

Definition 6 (Causal Disentanglement). Given some groundtruth model θ, a learned model ˆ θ is said to be causally disentangled if θ and ˆ θ are causal mechanism permutationequivalent.

4 Proposed Framework

We design a framework to achieve causal disentanglement. We propose ICM-VAE, a VAE-based framework based on the independent causal mechanisms (ICM) principle that achieves disentanglement of causal mechanisms. Figure 2 shows the overall architecture of our proposed framework.

4.1 Structural Causal Flow

Rather than assuming the limiting linear causal graphical model (CGM), as done in CausalVAE [Yang et al., 2021], we consider causal mechanisms to be complex nonlinear functions. Diverging from the strictly additive noise model assumption, we propose to parameterize causal mechanisms with a more general diffeomorphic 1 function. Flow-based models [Papamakarios et al., 2021] are often quite expressive in low-dimensional settings, which makes them desirable for learning complex distributions due to efficient and exact evaluation of densities. We parameterize the causal mechanisms with a conditional flow, which we refer to as the latent structural causal flow (SCF), that learns to map the independent noise distribution to a distribution over causal variables. This module is inspired by the causal autoregressive flow [Khemakhem et al., 2021]. This type of model is more realistic and general to better capture the complex distribution over the latent causal variables compared to simple linear mappings and leads to counterfactual identifiability [Nasr-Esfahany et al., 2023]. The SCF, denoted as f RF , is the reduced form (RF) of a nonlinear SCM function that conceptually maps noise variables ϵ to causal variables z as follows

where f RF : E → Z is derived from the recursive substitution of causal mechanisms fi in topological order of the causal graph as follows

realized as a function of the noise term and parent variables. The noise encoding ϵi is exactly the SCM noise variable corresponding to the causal variable zi.

Similar to several prior works [Shen et al., 2022; Liang et al., 2023], we assume that the latent causal graph is known in the form of a binary adjacency matrix to focus on formulating the problem of causal disentanglement. To implement a diffeomorphic function f RF , we use flow-based models to parameterize the causal mechanisms. Specifically, this flow is implemented as an affine-form autoregressive flow, where we derive each causal variable one at a time in topological order such that each variable is dependent only on a subset of previously derived variables (i.e. parents). Thus, the change of variables can be computed quite easily for exact and efficient likelihood estimation. In general, one can parameterize causal mechanisms using any nonlinear diffeomorphic function, as long as it takes into account the topological ordering. Let's take the pendulum example in Figure 1 to illustrate. The causal structure is z1 → z3, z4 and z2 → z3, z4. Then, the SCF would be f RF : (ϵ1, ϵ2, ϵ3, ϵ4) ↦→ (z1 = f1(ϵ1), z2 = f2(ϵ2), z3 = f3(ϵ3, z1, z2), z4 = f4(ϵ4, z1, z2)), where zi = f RF i (ϵi; ϵpai ) = fi(ϵi; zpai ) and each fi is a diffeomorphic transformation of the form

1 A diffeomorphism is a differentiable bijection with a differentiable inverse.

where ai = r1(zpai ) and bi = r2(zpai ) are the slope and offset parameters of the affine transformation, respectively, learned via neural networks r1 and r2 that capture information about the causal parents. Since the Jacobian of the function will be triangular by construction and the slope parameter is learned for each variable, the slope is equivalent to the diagonal elements of the Jacobian matrix as follows

where ϵpai denotes the noise terms associated with the parents of causal variable zi. The structural causal flow can easily be generalized to multivariate scenarios by masking groups of latent codes corresponding to each causal variable.

4.2 Generative Model

To achieve an identifiable model, we leverage auxiliary information as a weak supervision signal [Khemakhem et al., 2020]. Let u ∈ Rn be the auxiliary observed labels corresponding to the causally related ground-truth factors with support U ⊂ Rn . We assume that the decoder g is diffeomorphic onto its image. Several prior works [Locatello et al., 2020; Khemakhem et al., 2020; Lachapelle et al., 2022] assume that the nonlinear mixing function mapping Z to X is a diffeomorphism. Consider the pendulum system from Figure 1 consisting of a light source, a pendulum, and a shadow. Given only the image, it is completely certain that we can identify where each object appears in the image. So, we find it reasonable to assume a diffeomorphic mixing function g for our exploration. Let θ = (g,T, λ,Gz ) be the parameters of the conditional generative model defined as follows

where

If we assume that the distribution over the noise ξ is Gaussian with infinitesimal variance, we can model non-noisy observations as a special case of Eq. (9). The prior distribution in the generative model is given by

where we choose p(ϵ) as a standard Gaussian base distribution and p(z|u) is assumed to be conditionally factorial. However, the conditional prior in Eq. (2) cannot properly capture causal mechanisms for causally related factors. We next define a causally factorized prior suitable to achieve causal disentanglement.

4.3 Causal Disentanglement Prior

We aim to use a structured prior and perform conditioning in the latent space, similar to previous work on nonlinear ICA [Khemakhem et al., 2020], to enforce z to be a disentangled causal representation. However, for a model incorporating causal structure, the form of the conditional prior in Eq. (2) needs to be modified and generalized to causally factorized distributions. To enforce the disentanglement of z, we parameterize the prior distribution to learn a mapping from u to z. That is, since the goal of causal disentanglement is to map each latent/mechanism to exactly one corresponding ground-truth factor/mechanism, we can explicitly incorporate this into the prior. Using u as our observational labels, we parameterize the factorized causal conditionals with a conditional flow between u and z to establish a bijective relationship. The goal is for the distribution over the causal variables to tend towards the learned prior. The prior over z is defined as follows

where λi(Gzi ⊙ z, ui) is the estimated parameter vector of the prior obtained via mechanism λi, Gzi is the ith column of the adjacency matrix of the causal graph of z, hi(z) is the base measure, and Ti(z) = (z, z2 ) is the sufficient statistic. The prior induces a causal factorization of z with causal conditionals pθ(zi|zpai , ui), where ui is introduced as a weak supervision signal for identifiability. Eq. (11) is reminiscent of temporal priors that define a distribution over a latent variable conditioned on the variable at a previous time step [Lippe et al., 2023]. In our case, we view the causal factors as derived autoregressively. With a slight abuse of notation, we define λ(z, u) to be the concatenation of all λi(Gzi ⊙ z, ui). The function λ(z, u) outputs the natural parameter vector for the causally factorized distribution. We further require λ : Z × U → Z to be a bijective map between u and learned representation z to encourage disentanglement of the causal mechanisms. In practice, we choose p(u) from a locationscale family such as Gaussian. The learned mechanism λi is defined as the following diffeomorphic map:

where ci = s1(zpai ) and di = s2(zpai ) are the slope and offset parameters of the flow, respectively, learned via neural networks. To obtain a causally factorized conditional prior over z, we map the base distribution p(u), which is known beforehand, to a distribution over z.

4.4 Learning Objective

Putting all the components together, ICM-VAE consists of a stochastic encoder qϕ(ϵ, z|x, u), a decoder pθ(x|ϵ, z), and diffeomorphic causal transformations fi(·; ϵ). All components are learnable and implemented as neural networks. Formally, we aim to optimize the following variational lower bound:

where β is the latent bottleneck parameter. We train the model by minimizing the negative of the ELBO loss and learn to map low-level pixel data to noise variables and map the noise variable distribution to a distribution over the causal variables. For a detailed derivation of the ELBO, refer to the appendix in [Komanduri et al., 2024]. Causal structure learning could be heuristically incorporated jointly with the learning objective by enforcing acyclicity and sparsity of the causal graph. However, most causal discovery methods, such as GraN-DAG [Lachapelle et al., 2020], require restricting parametric assumptions on the causal model (i.e., additive noise), to be practically applied. For an extended discussion on incorporating causal discovery and challenges, refer to the appendix in [Komanduri et al., 2024].

5 Identifiability Analysis

We design our framework to satisfy the conditions necessary to achieve causal mechanism equivalence and causally disentangle the factors of variation. The causally factorized prior in Eq. (11) induces disentanglement of causal mechanisms. Theorem 1 extends the identifiability theorem from iVAE [Khemakhem et al., 2020] to show causal mechanism equivalence identifiability when we have a causal model. We note that causal mechanism disentanglement implies the disentanglement of causal factors.

Theorem 1 (Identifiability of ICM-VAE). Suppose that we observe data sampled from a generative model defined according to (8)-(12) with two sets of model parameters θ = (g,T, λ,Gz ) and ˆ θ = (ˆ g, ˆT, ˆλ, ˆGz ). Suppose the following assumptions hold

1. The set {x ∈ X|ϕξ(x) = 0} has measure zero, where ϕξ is the characteristic function of the density pξ defined in Eq. (9).
2. The decoder g is diffeomorphic onto its image.
3. The sufficient statistics Ti are diffeomorphic.
4. [Sufficient Variability] There exists nk+1 distinct points u0, . . . , unk such that the matrix

of size nk × nk is invertible, the ground-truth function λ is affected sufficiently strongly by each individual label ui and previously derived variables zpai , and ∀i, λi(zpai , ui) ̸= 0.

Then θ and ˆ θ are causal mechanism permutation-equivalent, and the model ˆ θ is causally disentangled.

For the proof of Theorem 1, refer to the appendix in [Komanduri et al., 2024].

6 Experimental Evaluation

In this section, we empirically evaluate the effectiveness of ICM-VAE. We consider a setting where the causal variables can be multi-dimensional and fix the dimension of each causal variable, which has also been explored in previous work [Yang et al., 2021; Lippe et al., 2023]. This assumption enables the representation to learn more informative and specific latent codes to describe the factors of variation (i.e. x-position, y-position, etc.). We run experiments on datasets with continuous factors, but discrete flows [Tran et al., 2019] can be used for discrete factors. Our results suggest a component-wise correspondence between the learned and true causal factors.

Datasets. We show the performance of our framework on three datasets, each consisting of four real-valued causal variables. The Pendulum dataset [Yang et al., 2021] consists of causal variables with causal graph (pendulum angle → shadow length, shadow position) and (light position → shadow length, shadow position). The Flow dataset [Yang et al., 2021] consists of variables with causal graph (ball size → water height), (hole position → water flow), and (water height → water flow). We also show experiments on a more complex 3D dataset of a robot arm interacting with colored buttons called CausalCircuit [Brehmer et al., 2022], which consists of variables with causal graph (robot arm → blue light intensity, green light intensity, and red light intensity), (blue light intensity → red light intensity), and (green light intensity→ red light intensity).

Evaluation Metrics. The DCI metric [Eastwood and Williams, 2018] quantifies the degree to which ground-truth factors and learned latents are in one-to-one correspondence. We compute the DCI disentanglement (D) and completeness (C) scores, which are based on a feature importance matrix quantifying the degree to which each latent code is important for predicting each ground truth causal factor. The informativeness (I) score is the prediction error in the latent factors predicting the ground-truth generative factors and is constant (I = 0) throughout all datasets and models, so we omit it for brevity. We train models with 3 random seeds and select the median DCI score to report. We note that DCI is highly correlated with other disentanglement metrics, such as MCC, with strong connections to identifiability [Eastwood et al., 2023]. To evaluate how changes in the generative factors affect the latent factors, we compute the interventional robustness score (IRS) [Suter et al., 2019], which is similar to an R2 value.

Implementation. For the Pendulum (6K training and 1K testing) and Flow (6K training and 2K testing) datasets, we linearly increase the β parameter throughout training from 0 to 1. We train for 9 · 103 steps using a batch size of 64. We use a Gaussian encoder and decoder with mean and variance computed by fully connected neural networks. For the CausalCircuit dataset (35K training and 10K testing), we linearly increase β from 0 to 0.05. We train for 3.5 · 104 steps using a batch size of 100. We use a convolutional neural network architecture with 6 layers and ReLU activation followed by a fully connected layer to estimate the mean and variance. The noise level for the variance of the Gaussian distribution of the conditional prior is controlled by σ2 ∈ {0.01, 0.00001}. The structural causal flow and λ are implemented as affine form autoregressive flows with the slope and offset computed by fully connected 3-layer neural networks with 100 unit hidden layers and ReLU activation. We set the learning rate to 0.001 for all experiments. We set the dimension of each causal variable to 4 for all datasets. Our code is available at https://github.com/Akomand/ICM-VAE.

Baselines. We compare the performance of our approach, in terms of disentanglement and interventional robustness, with four baseline models: β-VAE [Higgins et al., 2017] (unsupervised and acausal), iVAE [Khemakhem et al., 2020] (acausal), CausalVAE [Yang et al., 2021] (causal), and SCMVAE [Komanduri et al., 2022] (causal). The causal baselines propose relatively simplistic models that do not necessarily guarantee the disentanglement of causal factors.

Causal Disentanglement. Our experiments show that learning diffeomorphic causal mechanisms, rather than linear SCM, and incorporating the causal structure to learn a bijective λ to estimate the parameters of the causally factorized distribution significantly improves the disentanglement and interventional robustness of learned causal factors compared with baselines, as shown in Table 2. Consistent with our intuition, iVAE fails to disentangle the causal factors. The results indicate that ICM-VAE disentangles the causal factors and mechanisms almost perfectly. A high DCI disentanglement score indicates a permutation matrix mapping the latent factors to ground-truth generative factors in an ideal one-toone mapping [Eastwood et al., 2023]. Further, our model improves the interventional robustness of the representation, where interventions on ground-truth factors map to interventions on the corresponding learned factors.

Counterfactual Generation. We show counterfactual generated results of intervening on learned latent factors. Figure 4 shows the CausalCircuit system and the result of intervening on the robot arm factor and propagating causal effects. We observe that the red light also turns on as the robot arm interacts with the blue or green lights. On the other hand, when the arm interacts with the red light, only the red light turns on and the other lights remain off. We observe a similar phenomenon in the Pendulum and Water Flow systems in Figure 3, which shows the result of intervening on causal factors and propagating effects. Intervening on the pendulum angle or light position has causal effects on the shadow. However, interventions on the shadow factors do not change the parent factors. For the counterfactual generation procedure, results from intervening on other causal factors, and iVAE latent traversals, refer to appendix in [Komanduri et al., 2024].

7 Conclusion

We propose ICM-VAE, a framework for causal representation learning under label supervision. We model causal mechanisms as flow-based transformations from noise to causal variables. We extend the idea of disentanglement to causal models and propose the notion of causal mechanism disentanglement. To this end, we design a causal disentanglement prior to causally factorize the distribution over causal variables. We theoretically show permutation-equivalent identifiability of the learned factors. Experimental results show that ICM-VAE almost perfectly disentangles the causal factors, improves interventional robustness, and generates consistent counterfactuals. Future work will incorporate causal discovery and disentanglement given partially observed labels.

Acknowledgements

This work is supported in part by National Science Foundation under awards 1910284, 1946391 and 2147375, the National Institute of General Medical Sciences of National Institutes of Health under award P20GM139768, and the Arkansas Integrative Metabolic Research Center at University of Arkansas.