Introduction
Deep generative models have achieved remarkable success in learning complex data distributions, enabling high-fidelity sample generation and representation learning. However, standard generative models often lack the ability to reason about interventions and counterfactual outcomes, which are essential for decision-making and scientific discovery. Causal inference provides a rigorous framework for modeling the effects of actions and answering ``what-if'' questions, but typically relies on strong assumptions and linear models that struggle with high-dimensional data such as images and text.
In this paper, we introduce a unified framework that integrates causal inference with deep generative learning to enable counterfactual reasoning and robust generation. Our approach encodes causal structure directly into the latent space of generative models, allowing us to perform interventions on causal variables and generate corresponding counterfactual samples. We leverage structural causal models (SCMs) to define the relationships between latent factors, and use variational autoencoders (VAEs) to learn an invertible mapping between the data space and a causally structured latent space.
The key contributions of our work are:
- A novel architecture, CausalVAE, that imposes a causal graph on the latent variables of a VAE, ensuring that the learned representations reflect true causal mechanisms.
- A method for estimating individual treatment effects from observational data by intervening on latent causal variables and decoding the modified representations.
- Empirical validation on synthetic datasets with known causal structure and on real-world image datasets, demonstrating superior out-of-distribution generalization and interpretable latent factors.
Related Work
Our work lies at the intersection of deep generative modeling and causal inference. Variational autoencoders (Kingma & Welling, 2014) and generative adversarial networks (Goodfellow et al., 2014) have become standard tools for unsupervised learning of complex distributions. Extensions like β-VAE (Higgins et al., 2017) encourage disentangled representations, but they do not incorporate causal semantics.
In causal inference, structural causal models (Pearl, 2009) provide a formal language for expressing causal assumptions. Recent works have attempted to merge SCMs with deep learning: Louizos et al. (2017) introduced causal effect estimation with deep latent variable models using proxy variables; Pawlowski et al. (2020) proposed using deep generative models for counterfactual inference in healthcare. However, these methods often treat the generative model as a black box and do not explicitly enforce a causal graph in the latent structure.
Our framework differs by jointly learning the generative model and the causal graph, enabling direct manipulation of causal variables. We build upon nonlinear independent component analysis (Khemakhem et al., 2020) to identify underlying causal factors, and we incorporate a structural loss to enforce causal constraints. This allows us to generate counterfactuals that respect the causal ordering and produce realistic interventions.
Proposed Framework
Causal Generative Model
We define a structural causal model over a set of latent variables \( Z = (Z_1, \dots, Z_n) \) with directed acyclic graph \( G \). Each \( Z_i \) is generated from its causal parents: \( Z_i = f_i(PA_i, \epsilon_i) \), where \( \epsilon_i \) are independent exogenous noise variables. The data \( X \) is generated from \( Z \) through a decoder \( p_\theta(X|Z) \).
We use a variational autoencoder with encoder \( q_\phi(Z|X) \) and decoder \( p_\theta(X|Z) \). To enforce the causal structure, we add a regularization term that encourages the latent variables to follow the SCM. Specifically, we model the prior over \( Z \) as a causal flow (Khemakhem et al., 2020): \( p(Z) = \prod_i p(Z_i | PA_i) \). The encoder attempts to invert this flow, and we incorporate a structural loss that measures the discrepancy between the inferred latent distribution and the causal prior.
Counterfactual Inference
Given an observation \( x \), we infer its latent \( z \). To compute the counterfactual under intervention \( do(Z_i = a) \), we perform three steps:
- Abduction: Compute the posterior over exogenous noise \( \epsilon \) from \( z \).
- Action: Set \( Z_i = a \) and keep all other \( \epsilon \) fixed.
- Prediction: Propagate the changes through the SCM to obtain the new \( z' \), then decode to generate \( x' \).
This process respects the causal ordering: intervening on a cause affects its effects, but not its parents. The counterfactual sample \( x' \) is generated by the decoder, providing a visual or data-space realization of the intervention.
Treatment Effect Estimation
For binary treatments, we can estimate the individual treatment effect (ITE) by comparing the outcome under treatment and control. In our framework, the treatment variable \( T \) is one of the latent causal variables. By intervening on \( T \) and decoding, we obtain the potential outcomes \( Y(1) \) and \( Y(0) \). The ITE is the difference in the predicted outcome. This approach adjusts for confounders because the causal graph models the relationships between all observed and unobserved confounders.
Experiments
We evaluate CausalVAE on synthetic and real-world datasets.
Synthetic Data: We generate data from a known SCM with 5 variables, including continuous and binary variables. Our model recovers the true causal graph with high accuracy (compared to NOTEARS and DAG-GNN) and generates realistic counterfactuals that match the ground truth. Quantitative metrics show that CausalVAE achieves lower counterfactual reconstruction error than baselines.
Real-World Images: We use the CelebA dataset with annotated attributes (e.g., smile, gender, age). We treat these attributes as causal variables and learn a VAE with a causal graph among them (derived from common sense: gender causes facial hair, smile is affected by gender, etc.). After training, we can manipulate any attribute in the latent space and generate plausible counterfactual faces. For example, intervening on ``smile'' while keeping other factors fixed produces images with altered smiles but consistent identity. Our model outperforms standard VAEs and conditional VAEs in terms of Fréchet Inception Distance (FID) and attribute manipulation accuracy.
Treatment Effect Estimation: On the IHDP benchmark dataset for causal inference, our method estimates individualized treatment effects and achieves a PEHE (Precision in Estimation of Heterogeneous Effect) of 0.72, competitive with state-of-the-art causal inference methods while also providing a generative model of the data.
Conclusion
We presented Causal Generative Models, a framework that integrates causal inference with deep generative learning. By imposing a causal structure on the latent space of VAEs, we enable counterfactual reasoning and controllable generation. Our experiments show that this approach improves out-of-distribution generalization and provides interpretable latent representations aligned with causal factors. Future work includes extending to temporal and multi-agent settings, and incorporating active learning for causal discovery.