Causal Diffusion Autoencoders: Toward Counterfactual Generation via Diffusion Probabilistic Models

A structural causal model (SCM) is formally defined by a tuple $\mathcal{M}=\langle\mathcal{Z},\mathcal{U},F\rangle$, where $\mathcal{Z}$ is the domain of the set of $n$ endogenous causal variables $\mathbf{z}=\{z_{1},\dots,z_{n}\}$, $\mathcal{U}$ is the domain of the set of $n$ exogenous noise variables $\mathbf{u}=\{u_{1},\dots,u_{n}\}$, which is learned as an intermediate latent variable, and $F=\{f_{1},\dots,f_{n}\}$ is a collection of $n$ independent causal mechanisms of the form

where $\forall i$, $f_{i}:\mathcal{U}_{i}\times\prod_{j\in\textbf{pa}_{i}}\mathcal{Z}_{j}\to\mathcal{Z}_{i}$ are causal mechanisms that determine each causal variable as a function of the parents and noise, $z_{\textbf{pa}_{i}}$ are the parents of causal variable $z_{i}$; and a probability measure $p_{\mathcal{U}}(\mathbf{u})=p_{\mathcal{U}_{1}}(u_{1})p_{\mathcal{U}_{2}}(u_{2})\dots p_{\mathcal{U}_{n}}(u_{n})$, which admits a product distribution. For the purposes of this work, we assume a causally sufficient setting (no hidden confounding), no SCM misspecification, and faithfulness is satisfied.

Diffusion Probabilistic Models (DPMs) [11, 20] have shown impressive results in image generation tasks, even beating out GANs in many cases [5]. The idea of the denoising diffusion probabilistic model (DDPM) [11] is to define a Markov chain of diffusion steps to slowly destroy the structure in a data distribution through a forward diffusion process by adding noise [11] and learn a reverse diffusion process that restores the structure of the data. Some proposed methods, such as denoising diffusion implicit model (DDIM) [32], break the Markov assumption to speed up the sampling in the diffusion process by carrying out a deterministic encoding of the noise.

Forward Diffusion. Given some input data sampled from a distribution $\mathbf{x}_{0}\sim q(\mathbf{x})$, the forward diffusion process is defined by adding small amounts of Gaussian noise to the sample in $T$ steps thereby producing noisy samples $\mathbf{x}_{1},\dots,\mathbf{x}_{T}$. The distribution of the noisy sample at time step $t$ is defined as a conditional distribution as follows:

where $\beta_{t}\in(0,1)$ is a variance parameter that controls the step size of noise. As $t\to\infty$, the input sample $\mathbf{x}_{0}$ loses its distinguishable features. In the end, when $t=T$, $\mathbf{x}_{T}$ follows an isotropic Gaussian. From Eq (2), we can then define a closed-form tractable posterior over all time steps factorized as follows:

Now, $\mathbf{x}_{t}$ can be sampled at any arbitrary time step $t$ using the reparameterization trick. Let $\alpha_{t}=\prod_{i=1}^{t}1-\beta_{i}$:

Reverse Diffusion. In the reverse process, to sample from $q(\mathbf{x}_{t-1}|\mathbf{x}_{t})$, the goal is to recreate the true sample $\mathbf{x}_{0}$ from a Gaussian noise input $\mathbf{x}_{T}\sim\mathcal{N}(\mathbf{0},\mathbf{I})$. Unlike the forward diffusion, $q(\mathbf{x}_{t-1}|\mathbf{x}_{t})$ is not analytically tractable and thus requires learning a model $p_{\theta}$ to approximate the conditional distributions as follows:

where $\mu_{\theta}$ and $\Sigma_{\theta}$ are learned via neural networks. It turns out that conditioning on the input $\mathbf{x}_{0}$ yields a tractable reverse conditional probability

where $\tilde{\mu}$ and $\tilde{\beta}_{t}$ are the true mean and variance. The learning objective is then formulated as a simplified objective of the ELBO via reparameterization to minimize the following mean squared error loss

where $\epsilon_{t}\sim\mathcal{N}(\mathbf{0},\mathbf{I})$ is the noise that takes an analytical form via a reparameterization from $\mathbf{x}_{0}$, as shown in [11].

DPMs produce latent variables $\mathbf{x}_{1:T}$ through the forward process. However, these variables are stochastic [24]. Song et al. proposed a DPM called Denoising Diffusion Implicit model (DDIM), which enables a deterministic process as follows:

with the following deterministic decoding process

which keeps the DDPM marginal distribution $q(\mathbf{x}_{t}|\mathbf{x}_{0})=\mathcal{N}(\sqrt{\alpha_{t-1}}\mathbf{x}_{0},(1-\alpha_{t})\mathbf{I})$. It turns out that this formulation shares the same objective and solution of DDPM and only differs in the sampling procedure. Thus, we can deterministically obtain the noise map $\mathbf{x}_{T}$ corresponding to a given image $\mathbf{x}_{0}$.

Let $\mathbf{x}_{0}\in\mathbb{R}^{d}$ be the observed input image. We carry out the forward diffusion process until we have a set of $T$ perturbed samples $\{\mathbf{x}_{1},\mathbf{x}_{2},\dots,\mathbf{x}_{T}\}$, each at a different noise scale. Suppose there are $n$ abstract causal variables that describe the high-level semantics of the observed image. To learn a meaningful representation, we propose to encode the input image $\mathbf{x}_{0}$ to a low-dimensional noise encoding $\mathbf{u}\in\mathbb{R}^{n}$. We then map the noise encoding to latent causal factors $\mathbf{z}_{\text{causal}}\in\mathbb{R}^{n}$ corresponding to the abstract causal variables. In this formulation, each noise term $u_{i}$ is the exogenous noise term for causal variable $z_{i}$ in the SCM. Let $\mathbf{A}$ be the adjacency matrix encoding the causal graph among the underlying factors where $A_{ji}$ implies $z_{j}$ is a cause of $z_{i}$. Then, we parameterize the mechanisms among causal variables as follows

where $f_{i}$ is the causal mechanism generating causal variable $z_{i}$ as a function of its parents and exogenous noise term and $z_{\textbf{pa}_{i}}$ denotes the causal parents of factor $z_{i}$. In practice, we can implement $f_{i}$ as a post-nonlinear additive noise model such that

where $\nu_{i}$ are the parameters of the neural network parameterizing each mechanism, $\odot$ is the elementwise product, and $\mathbf{z}_{\text{causal}}=\{z_{1},\dots,z_{n}\}$. This module captures the causal relations between latent variables using neural structural causal models. For the purposes of this work, we assume that the causal graph is known since we focus on counterfactual generation. However, a more end-to-end framework may include a causal discovery component. See Appendix C for a more detailed discussion.

Let $\mathbf{x}_{0}$ denote the high-dimensional input image and $\mathbf{y}\in\mathbb{R}^{n}$ denote an auxiliary weak supervision signal. Then, the CausalDiffAE generative process can be factorized as follows:

where $\theta$ are the parameters of the reverse process of the causal diffusion decoder (will discuss in Section 4.3), $p(\mathbf{u},\mathbf{z}_{\text{causal}}|\mathbf{y})=p(\mathbf{u})p(\mathbf{z}_{\text{causal}}|\mathbf{y})$, $p(\mathbf{u})=\mathcal{N}(\mathbf{0},\mathbf{I})$, and $p(\mathbf{z}_{\text{causal}}|\mathbf{y})$ is the alignment prior defined in Eq. (19). The joint posterior distribution $p(\mathbf{x}_{1:T},\mathbf{u},\mathbf{z}_{\text{causal}}|\mathbf{x}_{0},\mathbf{y})$ is intractable, so we approximate it using a variational distribution $q(\mathbf{x}_{1:T},\mathbf{u},\mathbf{z}_{\text{causal}}|\mathbf{x}_{0},\mathbf{y})$ which can be factorized into the following conditional distributions

where $\phi$ are the parameters of the variational encoder network parameterizing the joint distribution over the noise $\mathbf{u}$ and causal factors $\mathbf{z}_{\text{causal}}$. We can remove the dependence on $\mathbf{y}$ for the second conditional term in the decomposition of Eq. (13) since $\mathbf{x}_{1:T}$ is independent of the auxiliary label $\mathbf{y}$. We note that $q_{\phi}(\mathbf{z}_{\text{causal}},\mathbf{u}|\mathbf{x}_{0},\mathbf{y})$ can be factorized as $q_{\phi}(\mathbf{z}_{\text{causal}}|\mathbf{x}_{0},\mathbf{y})q_{\phi}(\mathbf{u}|\mathbf{x}_{0})$ since $\mathbf{u}$ and $\mathbf{z}_{\text{causal}}$ have a one-to-one correspondence.

We use a conditional DDIM decoder that takes as input the pair of latent variables $(\mathbf{z}_{\text{causal}},\mathbf{x}_{T})$ to generate the output image. We approximate the inference distribution $q(\mathbf{x}_{t-1}|\mathbf{x}_{t},\mathbf{x}_{0})$ by parameterizing the probabilistic decoder via a conditional DDIM $p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_{t},\mathbf{z}_{\text{causal}})$. With DDIM, the forward process becomes completely deterministic except for $t=1$. Similar to [24], we define the joint distribution of the reverse generative process as follows:

where $\mathbf{f}_{\theta}$ is parameterized by a noise prediction network $\epsilon_{\theta}$ (i.e., UNet [5]) as follows:

Note that in Eq. (14), $\mathbf{u}$ is omitted since $\mathbf{z}_{\text{causal}}$ already captures all the information about the noise. By leveraging the reparameterization trick, we can optimize the following mean squared error between noise terms

where $\epsilon_{t}\sim\mathcal{N}(\mathbf{0},\mathbf{I})$ and $\mathbf{x}_{t}=\sqrt{\alpha_{t}}\mathbf{x}_{0}+\sqrt{1-\alpha_{t}}\epsilon_{t}$.

To ensure the causal representation is disentangled, we incorporate label information $\mathbf{y}\in\mathbb{R}^{n}$ as a prior in the variational objective to aid in learning semantic factors and for identifiability guarantees [13]. We define the following joint loss objective:

where $\gamma$ is a regularization hyperparameter similar to the bottleneck parameter in $\beta$-VAEs [9], and the alignment prior over latent variables is defined as the following exponential family distribution

where $\mu_{\nu}$ and $\sigma^{2}_{\nu}$ are functions that estimate the mean and variance of the Gaussian, respectively. Intuitively, this prior ensures that the learned factors are one-to-one mapped to an indicator of the underlying ground truth factors. DiffAE requires training a latent DDIM in the latent space of the pre-trained autoencoder to enable sampling of latent semantic representation. However, CausalDiffAE is formulated as a variational objective with a stochastic encoder. Thus, we can sample the representation from the defined prior directly without having to train a separate diffusion model in the latent space. The training procedure for CausalDiffAE is outlined in Algorithm 1. See Appendix A for a derivation of the ELBO. For a detailed discussion on the connection of our diffusion objective to score-based generative models [33], see Appendix B.

A fundamental property of causal models is the ability to perform interventions and observe changes to a system. In generative models, this enables the sampling of counterfactual data. Given a pre-trained CausalDiffAE, we can controllably manipulate any factor of variation, propagate the causal effects to descendants, and perform reverse diffusion to sample from the counterfactual distribution. Algorithm 2 shows the process of generating counterfactuals from a trained CausalDiffAE, where $\mathbf{x}_{0}$ refers to the factual observation and $\mathbf{x}_{0}^{CF}$ refers to the generated counterfactual sample. To generate counterfactual instances, we first encode the high dimensional observation $\mathbf{x}_{0}$ to a noise encoding $\mathbf{u}$ (abduction) and transform it to causal latent variables $\mathbf{z}_{\text{causal}}$. Then, we intervene on a desired variable and propagate the causal effects via neural mechanisms to yield the intervened representation $\bar{\mathbf{z}}_{\text{causal}}$. We utilize the DDIM sampling algorithm to ensure the stochastic noise $\mathbf{x}_{T}$ is a deterministic encoding to enable semantic manipulations. Finally, we decode using DDIM conditioned on $(\bar{\mathbf{z}}_{\text{causal}},\mathbf{x}_{T})$ to obtain a counterfactual $\mathbf{x}_{0}^{CF}$. In lines 12-13, we use the DDIM non-Markovian deterministic generative process to generate counterfactual instances as follows:

Conditioning vs. Intervening. When we study causal generative models, we utilize the intervention operation, which is a fundamentally different operation than conditioning. When we condition, we narrow our scope to a specific subgroup of the data based on the conditioning variable. Interventions are population-level operations that fix a variable’s value (rendering it independent of its parents) to determine causal effects downstream. We emphasize that, under this intervention operation, causal models are robust to distribution shifts and can generate data outside the support of the training distribution.

To reduce the reliance on labeled data, inspired by classifier-free [10] guidance, we train a CausalDiffAE with a weak supervision guidance paradigm on the representation level [5].

Training. In the limited labeled-data regime, we train two models: an unconditional denoising diffusion model $p_{\theta}(\mathbf{x})$ parameterized by the score estimator $\epsilon_{\theta}(\mathbf{x}_{t},t)$ and a representation-conditioned model $p_{\theta}(\mathbf{x}|\mathbf{z}_{\text{causal}})$ parameterized through $\epsilon_{\theta}(\mathbf{x}_{t},t,\mathbf{z}_{\text{causal}})$. We use a single neural network to parameterize both models, where for the unconditional model we use only the unlabeled data for predicting the score (i.e., $\epsilon_{\theta}(\mathbf{x}_{t},t)$).

Generation. The counterfactual generation procedure in lines 12-13 of Algorithm 2 can be modified to generate counterfactuals with a guidance strength $\omega$, which can be interpreted as controlling the strength of the intervention on the causal variable to generate the counterfactual in our case. The overall modified score estimation during generation can be performed using the following linear combination of conditional and unconditional score estimates

where $\bar{\mathbf{z}}_{\text{causal}}$ is the set of latent causal factors after an intervention. The original utility of the classifier-free paradigm was to decrease the generation of diverse data in favor of higher-quality image samples without needing classifier gradients. So, $\omega$ controls the trade-off between higher quality and diverse samples. In our case, we care about generating high-quality counterfactual data. Intuitively, a higher $\omega$ implies a stronger effect of the intervention on the generated counterfactual since the conditional model $\epsilon_{\theta}(\mathbf{x}_{t},t,\mathbf{z}_{\text{causal}})$ is sensitive to interventions. So, as $\omega$ decreases, the unconditional model dilutes the effect of the intervention-sensitive model. In this sense, the sampling mechanism can be used to evaluate the causal strength of interventions. We find that the weak supervision paradigm enables (1) more efficient training with a weaker supervision signal, and (2) fine-grained control over generated counterfactuals.

Datasets. We experiment on three datasets. We use the MorphoMNIST dataset [4] created by imposing a 2-variable SCM to generate morphological transformations on the original MNIST dataset, where thickness is the cause of the intensity of the digit [22], as shown in Figure 2(a). The Pendulum dataset [35] consists of images of a causal system consisting of a light source, pendulum, and shadow. The light source and the pendulum angle determine the length and position of the shadow, as shown in the causal graph in Figure 2(b). We also use CausalCircuit [3], a complex 3D robotics dataset where a robot arm moves around to turn on red, green, or blue lights. The causal graph of this system is shown in Figure 3.

Baselines. CausalVAE [35] is a VAE-based causal representation learning framework that models causal variables using a linear SCM and enables counterfactual generation during inference time through interventions on causal variables. Class-conditional diffusion model (CCDM) [20] is a conditional diffusion model that utilizes class labels as the conditioning signal in reverse diffusion. Thus, this model is capable of generating new samples determined by a discrete or continuous set of labels $\mathbf{y}$. DiffAE [24] is a diffusion model that aims to learn manipulable and semantically meaningful latent codes. However, this approach learns an arbitrary representation in an unsupervised fashion and does not disentangle the latent space. Manipulations are performed using a post-hoc classifier for linear interpolation. Thus, the learned representation would not be ideal to perform causal interventions. For a fair comparison in counterfactual generation, we modify the objective to disentangle the latent space by incorporating label information in a prior to regularize the posterior. We call this extension DisDiffAE. We use DisDiffAE as a baseline to evaluate counterfactual generation and the DiffAE to evaluate disentanglement.

Metrics. We primarily use two quantitative metrics to evaluate the performance of our approach. To evaluate the disentanglement of the learned representations, we use the DCI disentanglement metric [7]. A high DCI score also suggests the effectiveness of controllable generation. In the context of a causal representation, this means that we can intervene on latent codes in an isolated fashion without any entanglements (i.e., two factors are encoded in the same latent code). To quantitatively evaluate generated counterfactuals, we adopt the Effectiveness metric from Melistas et al [18], which evaluates how successful the performed intervention was at generating the counterfactual. We train anti-causal predictors via convolutional regressors on the training dataset for each continuous causal variable $z_{i}$. Then, we report the mean absolute error (MAE) loss between the predicted values from the generated counterfactual and the true values of the counterfactual. This metric captures how controllable the factors are and the accuracy of the generated counterfactuals.

For details about the datasets, implementation, metrics, and computational requirements, see Appendix D. Our code is available at https://github.com/Akomand/CausalDiffAE.

We compare the disentanglement of CausalDiffAE with other baseline models, as shown in Table 1. We observe that diffusion-based representation learning objectives coupled with a suitable prior can better disentangle latent variables compared to VAE-based models. We do not include CCDM as a baseline here since it does not produce a representation to be evaluated. Compared to CausalVAE, the diffusion-based decoder in CausalDiffAE disentangles the semantic factors of variation to a much greater degree. Thus, we can perform interventions on causal variables in isolation and observe their downstream effects. We also note that DiffAE [24] does not learn a disentangled latent space since the semantic representation learned is arbitrary. To perform controllable manipulations with DiffAE, a post-hoc classifier must be trained to guide the sampling process. CausalDiffAE offers more precise control over learned factors through the disentanglement objective without the need to train additional classifiers.

Qualitative Evaluation. We show that CausalDiffAE produces much more realistic counterfactual samples compared to other acausal baselines and its VAE counterpart, CausalVAE. We attribute this to the diffusion process, which is better capable of capturing causally relevant information along with low-level stochastic variation.

Figure 2(a) shows the counterfactual generation results for the MorphoMNIST dataset. CausalVAE can generate counterfactual images after intervening on either thickness or intensity, but the accuracy and quality of the generated counterfactuals is far lower than CausalDiffAE. For instance, lower thickness does not lower the intensity and lower intensity intervention seems to change the thickness of the digit. CCDM fails to produce samples consistent with the underlying causal model. For example, intervening on the intensity produces a sample that increases in thickness. From a conditioning perspective, high-intensity digits tend to be thicker in the training distribution. For DisDiffAE, increasing the thickness does not influence the intensity.

Figure 2(b) shows counterfactual generation results for the Pendulum causal system. Upon interventions, images generated by CCDM are not consistent with the causal model. For example, intervening on the light position changes both the light position and the pendulum angle. DisDiffAE produces images where we can control one factor at a time, but does not reflect causal effects. For example, changing the angle of the pendulum does not accurately change the shadow length and position. CausalDiffAE generates higher-quality counterfactuals that are consistent with the causal model. Specifically, intervention on the pendulum angle or light position changes the shadow length and position accurately. On the other hand, interventions on children variables leave the parents unchanged.

Figure 3 shows counterfactuals from the CausalCircuit dataset. CausalVAE generates inaccurate counterfactuals for many scenarios (e.g., intervention on the blue light intensity changes the intensity of all other lights). For CCDM, moving the robot arm over the green light fails to turn the light on. Furthermore, manipulating the light intensity of other lights affects the position of the robot arm. DisDiffAE enables control over the generative factors, but does not consider causal effects. For example, moving the robot arm over the green light does not turn it on. Counterfactuals generated from CausalDiffAE are consistent with the causal system. For example, moving the robot arm over the green button turns the light on and as a result also turns on the red light, which is a downstream child variable. Intervening on the blue or green light slightly increases the intensity of the red light. Intervening on the red light leaves all parent variables unchanged. For additional counterfactual generation results, see Appendix D.5.

Quantitative Evaluation. We quantitatively show using the effectiveness metric that CausalDiffAE generates counterfactuals that are both accurate and realistic. We perform random interventions from a uniform distribution over the test dataset for each causal variable. We find that CausalDiffAE almost always outperforms other baselines in the effectiveness metric, as shown in Tables 2 and 3, for all causal factors. Specifically, for the MorphoMNIST dataset, we observe that interventions on thickness produce counterfactuals that accurately reflect both the thickness and intensity values. In the scenario where we intervene on thickness, the intensity MAE is lower for CausalDiffAE than other baselines, which indicates that the generated counterfactual has an accurate intensity value consistent with the causal effect of thickness on intensity. When we intervene on intensity, the thickness MAE is lower for CausalDiffAE than baselines, which suggests that the generated counterfactual retains its original thickness value upon intervention on intensity. For the Pendulum dataset, we see a similar phenomenon, where interventions on causal factors along with their downstream effects are accurately captured in the generated counterfactuals. We do not evaluate effectiveness for the CausalCircuit dataset since we do not have access to the generative process used to obtain the factors. We compute the average effectiveness value over $5$ runs with different random seeds. Our results strongly imply that the generated counterfactuals closely match the true counterfactuals.

Unlike VAE-based approaches, the weak supervision paradigm of diffusion models reduces the full-label supervision. We study the weak supervision scenario with the MorphoMNIST dataset. We jointly train a representation-conditioned and unconditional model, where the conditioned split is far less than the unconditioned split. We have two main motivations for doing this: (1) it greatly reduces the need for fully labeled datasets, and (2) it enables granular control over generated counterfactuals. We denote the proportion of unlabeled data by $p_{\text{unlabeled}}$. The CausalDiffAE model trained with $p_{\text{unlabeled}}=0.8$ on the MorphoMNIST dataset yields a DCI score of $0.9964$, which suggests that even under strictly limited label supervision, CausalDiffAE learns disentangled representations. We also empirically show that changing the $\omega$ parameter controls the strength of the intervention on the generated counterfactual. Figure 4 shows MNIST digits generated using the joint estimated score from the reduced supervision version of CausalDiffAE. We observe that interventions have virtually no effect when sampling using the joint score with $\omega=0.2$. For $\omega=0.5$, we see a stronger effect of the intervention on the thickness and intensity of the digit. Finally, for the fully-supervised score $\omega=1.0$, the intervention acts the strongest. Thus, varying $\omega$ in the range $(0,1)$ can be interpreted as generating a range of different counterfactuals.