5.2.1 Causal Variational Autoencoder (CausalVAE)

Note: CausalVAE requires PyTorch. The CausalVAE estimator in pydeepcausalml.generative embeds a learnable DAG in the latent space for causal disentanglement.

This notebook provides a detailed guide to implementing and applying the Causal Variational Autoencoder (CausalVAE) in Python, following the architecture introduced in “CausalVAE: Disentangled Representation Learning via Neural Structural Causal Models” (Yang et al., CVPR 2021). We cover the theoretical motivation, a full torch-based implementation via the pydeepcausalml package, and an application to Average Treatment Effect (ATE) estimation on synthetic data with a known causal ground truth.

Overview

Standard VAEs — and even identifiable VAEs such as iVAEs — assume that the latent dimensions are statistically independent. This is a convenient mathematical assumption, but it conflicts with how the world actually generates data: causes produce effects, and latent factors are often related by causal mechanisms rather than being simply correlated or independent.

CausalVAE addresses this by embedding a Structural Causal Model (SCM) directly in the latent space. Where a standard VAE encodes \(x\) to a vector of independent noise variables \(\varepsilon\) and decodes back, CausalVAE adds a causal layer between the encoder and decoder that transforms the independent \(\varepsilon\) into causally structured latent variables \(z\) according to a learnable Directed Acyclic Graph (DAG). The result is a generative model whose latent space is not just disentangled — it is causally disentangled, meaning each latent dimension corresponds to a distinct node in a causal graph, and the learned edges describe how those nodes causally influence one another.

This structure enables operations that are impossible in standard VAEs:

Interventions: set \(\mathrm{do}(z_i = v)\) and propagate the effect through the graph to the other latents and to the reconstruction, without touching the unrelated dimensions.
Counterfactual reasoning: ask “what would \(z_2\) have been if \(z_1\) had taken a different value?”
Treatment effect estimation: because the causal graph is explicit, ATE can be estimated by intervening on the relevant latent node and reading off the downstream effect on the outcome.

Model Architecture

The Two-Stage Latent Space

CausalVAE separates the latent space into two layers with distinct roles:

Exogenous variables \(\varepsilon\): Independent noise sources with prior \(p(\varepsilon) = \mathcal{N}(0, I)\). These are what the encoder directly outputs — a representation that is, by construction, free of causal dependencies.
Endogenous variables \(z\): Causally generated from \(\varepsilon\) via the learned SCM. Each \(z_i\) is a function of its parents in the DAG plus its own exogenous noise \(\varepsilon_i\).

The forward pass therefore has an additional step compared to a standard VAE: after sampling \(\varepsilon\), the causal layer applies the structural equations to produce \(z\), and the decoder then reconstructs \(x\) from \(z\).

The Causal Layer (SCM)

The causal layer implements the structural equations of the SCM. In the linear case:

\[z = (I - A \odot M)^{-1} \varepsilon\]

where \(A\) is a matrix of learnable real-valued edge weights and \(M\) is a binary mask that determines which edges are structurally permitted. In practice, this inversion is replaced by an additive fixed-point iteration:

\[z = \varepsilon + (A \odot M)\, z\]

Because \(z\) must follow a DAG (no cycles), the matrix \(A \odot M\) must be nilpotent, which is enforced by a continuous differentiable DAG penalty on \(A\):

\[h(A) = \mathrm{tr}\!\left(\exp(A \circ A)\right) - d\]

where \(d\) is the number of latent dimensions. This penalty equals zero if and only if \(A \circ A\) has no positive-weight cycles, i.e., \(A\) encodes a DAG. It is added to the loss with a Lagrange multiplier \(\gamma\), giving the optimizer a continuous incentive to push the learned graph toward acyclicity without requiring discrete search.

Data Flow

The complete forward pass is:

\[x \;\xrightarrow{\text{encoder}}\; (\mu_\varepsilon, \log\sigma^2_\varepsilon) \;\xrightarrow{\text{reparam.}}\; \varepsilon \;\xrightarrow{\text{causal layer}}\; z \;\xrightarrow{\text{decoder}}\; (\mu_x, \log\sigma^2_x)\]

The encoder outputs parameters of a Gaussian over \(\varepsilon\) (not \(z\) directly). The decoder receives the causally structured \(z\) (not \(\varepsilon\)). This separation means the KL divergence is computed in the \(\varepsilon\) space against \(\mathcal{N}(0,I)\), keeping that term tractable, while the causal structure lives entirely in the transformation \(\varepsilon \to z\).

Loss Function

The training objective is a \(\beta\)-VAE ELBO augmented with causal regularizers:

\[\mathcal{L} = \underbrace{\mathbb{E}_{q(\varepsilon|x)}\!\left[\log p(x \mid z)\right]}_{\text{reconstruction}} - \underbrace{\beta\, D_{\mathrm{KL}}\!\left(q(\varepsilon \mid x) \,\|\, \mathcal{N}(0,I)\right)}_{\text{KL regularization}} - \underbrace{\gamma\!\left(h(A) + \lambda \|A\|_1\right)}_{\text{DAG + sparsity penalty}}\]

The three terms serve distinct purposes:

The reconstruction loss ensures the decoded output faithfully reproduces the input.
The KL term (weighted by \(\beta > 1\) for stronger disentanglement) regularizes the exogenous variables toward independence.
The DAG penalty \(h(A)\) enforces acyclicity of the learned causal graph.
The \(\ell_1\) sparsity term \(\lambda \|A\|_1\) encourages the learned graph to have few edges, promoting parsimonious causal explanations.

The causal layer deserves its own close-up showing how the DAG transforms ε into z.

Identifiability and Weak Supervision

CausalVAE, like iVAE, faces identifiability challenges. Without additional constraints, the latent space can be arbitrarily rotated. In practice, the paper addresses this through weak supervision: the training procedure uses partial knowledge of the causal graph (e.g., which edges are forbidden, or which interventional experiments have been performed) to anchor the latent space. In the implementation below, the synthetic data is generated from a known causal chain, giving us ground-truth latents against which to evaluate the learned representation.

Performing Interventions

Once trained, an intervention \(\mathrm{do}(z_i = v)\) is carried out by:

Encoding \(x\) to obtain \(\varepsilon\).
Setting the \(i\)-th component of \(z\) to \(v\), bypassing the structural equation for node \(i\).
Propagating \(v\) through the downstream edges in \(A\) to update the children of \(z_i\).
Decoding the modified \(z\) to produce the interventional reconstruction \(\hat{x}_{\mathrm{do}(z_i=v)}\).

This is the mechanism that enables ATE estimation: intervene on the treatment node, observe the change in the outcome node, and average over the data distribution.

ATE Estimation via Causal Interventions

Why CausalVAE Can Estimate ATE

Because the learned latent space has explicit causal structure, ATE estimation reduces to a simple intervention: set the treatment node \(T\) to 1 (treated) versus 0 (control), hold all exogenous noise \(\varepsilon\) fixed, propagate through the DAG, and read off the expected change in the outcome. This is the do-calculus \(\mathrm{do}(T = t)\) operationalized directly on the learned graph, without needing a separate propensity model or outcome regression.

In the synthetic setting below, treatment \(T\) acts on \(z_1\) and the outcome \(Y\) is generated as \(Y = z_3 + T \cdot 0.5 \cdot z_2 + \varepsilon_Y\), giving a true ATE of \(0.5 \cdot \mathbb{E}[z_2]\).

Generating Data with Treatment and Outcome

Training the Outcome-Augmented Model

CausalVAE_ATE extends the base model with a linear outcome head \(\mathbb{E}[Y \mid z, T]\). The outcome head is trained jointly with the VAE objective, so the latent representation learns to encode variation that is relevant for predicting \(Y\) — not just for reconstructing \(x\).

Implementation in Python

We fit CausalVAE with pydeepcausalml.generative.CausalVAE on synthetic data with a known causal chain.

Check and Install Required Python Packages

The following packages are used in this notebook. Install any missing packages with the cell below:

numpy, pandas, torch, matplotlib, seaborn, pydeepcausalml

Code

import importlib
import subprocess
import sys
from pathlib import Path

PACKAGES = [
    "numpy", "pandas", "scipy", "torch", "scikit-learn",
    "matplotlib", "seaborn",
]

for pkg in PACKAGES:
    mod = "sklearn" if pkg == "scikit-learn" else pkg
    try:
        importlib.import_module(mod)
    except ImportError:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", pkg])

try:
    import pydeepcausalml  # noqa: F401
except ImportError:
    repo_root = Path("..").resolve()
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "-e", str(repo_root)])

print("Packages ready.")

Verify imports

Code

import numpy as np
import pandas as pd
import torch
import matplotlib.pyplot as plt
import seaborn as sns
from pydeepcausalml import set_seed

print("torch:", torch.__version__)
print("cuda:", torch.cuda.is_available())

Code

set_seed(42)

Load data and prepare synthetic causal chain

Data Generation

We generate synthetic data from a known three-node causal chain \(z_1 \to z_2 \to z_3\) with nonlinear structural equations:

\[z_1 = \varepsilon_1, \qquad z_2 = z_1 + \varepsilon_2, \qquad z_3 = z_2^2 + \varepsilon_3\]

Observations \(x\) are a nonlinear mix of the causal latents, ensuring that the true causal structure cannot be read off from \(x\) directly and must be recovered from the learned representation. Having access to the ground-truth latents \(z\) allows us to directly evaluate whether the causal layer has learned the correct graph.

Data generation

Synthetic causal chain \(z_1 \to z_2 \to z_3\) with nonlinear observations.

Code

def generate_data(n_samples=5000, latent_dim=3, random_state=42):
    rng = np.random.default_rng(random_state)
    eps = rng.normal(size=(n_samples, latent_dim))
    z = np.zeros((n_samples, latent_dim))
    x = np.zeros((n_samples, latent_dim))
    z[:, 0] = eps[:, 0]
    z[:, 1] = z[:, 0] + eps[:, 1]
    z[:, 2] = z[:, 1] ** 2 + eps[:, 2]
    x[:, 0] = z[:, 0] * z[:, 2]
    x[:, 1] = np.sin(z[:, 1]) + z[:, 0]
    x[:, 2] = z[:, 2] ** 2 + rng.normal(scale=0.1, size=n_samples)
    return x, z, eps


x, true_z, true_eps = generate_data()
x_mean, x_std = x.mean(axis=0), x.std(axis=0) + 1e-8
x_norm = (x - x_mean) / x_std

n = len(x_norm)
train_n, val_n = int(0.8 * n), int(0.1 * n)
x_train = x_norm[:train_n]
x_val = x_norm[train_n : train_n + val_n]
x_test = x_norm[train_n + val_n :]
print(f"Train: {len(x_train)} | Val: {len(x_val)} | Test: {len(x_test)}")

Fit CausalVAE

Model Initialization

The CausalVAE() constructor (from pydeepcausalml) builds the encoder, causal layer, and decoder. The loss_function() computes the full augmented ELBO including the DAG penalty and sparsity regularizer.

We use a learning rate scheduler that halves the learning rate when validation loss plateaus, and a LR warmup phase in the first 10 epochs to prevent the DAG penalty from dominating before the reconstruction loss has stabilized.

Training

The training loop implements several practical stability measures:

LR warmup (first 10 epochs): the optimizer starts at a fraction of the target learning rate and ramps up linearly, preventing large early updates from disrupting the nascent causal structure.
Capped gamma schedule: the DAG penalty weight \(\gamma\) is introduced gradually over the first 50 epochs and capped at 0.35, preventing the acyclicity constraint from overwhelming reconstruction early in training.
Gradient clipping (max norm 0.5): prevents gradient explosions, which are especially common when the matrix exponential in \(h(A)\) produces large gradients.
NaN/Inf detection: batches with numerically unstable losses are skipped rather than allowed to corrupt the model weights.
Early stopping: training halts 10 epochs after the validation loss stops improving, and the best checkpoint is restored at the end.

Fit CausalVAE

Code

from pydeepcausalml.generative import CausalVAE

model = CausalVAE(
    latent_dim=3,
    hidden=64,
    beta_kl=1.0,
    lambda_dag=0.1,
    epochs=80,
    batch_size=128,
    lr=1e-3,
    random_state=42,
    verbose=True,
    log_every=10,
)
model.fit(x_train)
z_latent = model.transform(x_test)
print("Learned latent shape:", z_latent.shape)
adj = model.adjacency_matrix()
print("Adjacency (edge weights):\n", np.round(adj, 3))

Training and validation loss

Training and Validation Loss

The plot below shows the negative ELBO over training. A well-behaved run exhibits three phases: a rapid initial decline as the reconstruction loss improves; a mid-training inflection as the DAG penalty activates and temporarily increases the loss; and a final convergence as the graph structure stabilizes and the ELBO settles. Early stopping ensures that we restore the checkpoint from the first loss valley rather than continuing into a noisier regime.

Training loss

Code

hist = model.history_
if hist.get("loss"):
    plt.figure(figsize=(6, 4))
    plt.plot(range(1, len(hist["loss"]) + 1), hist["loss"], label="Train loss")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.title("CausalVAE: training loss")
    plt.legend()
    plt.tight_layout()
    plt.show()

ATE estimation via latent interventions

Training the Outcome-Augmented Model

ATE estimation via latent interventions

We extend the setting with treatment \(T\) and outcome \(Y\) on the synthetic SCM.

Code

def generate_data_with_treatment(n_samples=5000, latent_dim=3, random_state=42):
    rng = np.random.default_rng(random_state)
    eps = rng.normal(size=(n_samples, latent_dim))
    z = np.zeros((n_samples, latent_dim))
    z[:, 0] = eps[:, 0]
    z[:, 1] = z[:, 0] + eps[:, 1]
    z[:, 2] = z[:, 1] ** 2 + eps[:, 2]
    T = rng.binomial(1, 0.5, size=n_samples)
    Y = z[:, 2] + T * 0.5 * z[:, 1] + rng.normal(scale=0.1, size=n_samples)
    x = np.zeros((n_samples, latent_dim))
    x[:, 0] = z[:, 0] * z[:, 2]
    x[:, 1] = np.sin(z[:, 1]) + z[:, 0]
    x[:, 2] = z[:, 2] ** 2 + rng.normal(scale=0.1, size=n_samples)
    return x, z, T, Y


x_ate, z_ate, T_ate, Y_ate = generate_data_with_treatment()
x_ate = (x_ate - x_ate.mean(0)) / (x_ate.std(0) + 1e-8)
true_ate = 0.5 * z_ate[:, 1].mean()
print(f"True ATE (population): {true_ate:.4f}")

# Joint model: CausalVAE representation + linear outcome head on latents
from sklearn.linear_model import LinearRegression

cvae_ate = CausalVAE(latent_dim=3, hidden=64, epochs=60, batch_size=256, verbose=False, random_state=42)
cvae_ate.fit(x_ate)
z_hat = cvae_ate.transform(x_ate)
outcome_model = LinearRegression().fit(np.column_stack([z_hat, T_ate]), Y_ate)

rng = np.random.default_rng(0)
z_samples = rng.normal(size=(500, cvae_ate.latent_dim))
y1 = outcome_model.predict(np.column_stack([z_samples, np.ones(500)]))
y0 = outcome_model.predict(np.column_stack([z_samples, np.zeros(500)]))
pred_ate = y1.mean() - y0.mean()
print(f"Predicted ATE: {pred_ate:.4f}")
print(f"Absolute error: {abs(pred_ate - true_ate):.4f}")

Summary and Conclusions

CausalVAE combines the expressive power of variational autoencoders with the structural guarantees of causal graphical models. By embedding a learnable DAG directly in the latent space, it enables a form of representation learning that goes beyond statistical disentanglement to causal disentanglement: each latent dimension corresponds to a distinct node in a causal graph, and the learned edges describe how those nodes influence one another.

The practical consequences are significant. Interventions on the latent space are well-defined (via do-calculus on the learned DAG), counterfactuals can be computed by fixing \(\varepsilon\) and modifying \(z\), and ATE estimation reduces to a simple evaluation of the outcome head under two interventional distributions. These capabilities are unavailable in CEVAE, iVAE, or standard VAEs.

The main practical challenges are identifiability (which typically requires weak supervision from partial graph knowledge or interventional data), training stability (the DAG penalty requires careful scheduling to avoid disrupting reconstruction early in training), and scalability (the matrix exponential penalty is \(O(d^3)\) in the latent dimension \(d\), so very large latent spaces are expensive to regularize).

For production applications: use weak supervision whenever any partial graph knowledge is available; monitor the DAG penalty and reconstruction loss separately to diagnose convergence issues; and validate the learned graph against any available domain knowledge before trusting the inferred causal structure for downstream decisions.

References

Yang, M., Liu, F., Chen, Z., Shen, X., Hao, J., & Wang, J. (2021). CausalVAE: Disentangled Representation Learning via Neural Structural Causal Models. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR).
Zheng, X., Aragam, B., Ravikumar, P., & Xing, E. P. (2018). DAGs with NO TEARS: Continuous Optimization for Structure Learning. NeurIPS. (Source of the \(h(A)\) acyclicity penalty.)
Khemakhem, I., et al. (2020). Variational Autoencoders and Nonlinear ICA: A Unifying Framework. AISTATS. (iVAE, predecessor in the identifiable VAE family.)
Louizos, C., et al. (2017). Causal Effect Inference with Deep Latent Variable Models. NeurIPS. (CEVAE, complementary approach for proxy confounding.)
pydeepcausalml: R package for Machine Learning-based Causal Inference
Yang et al. Official CausalVAE implementation. Huawei Noah’s Ark Lab.