CDVAE: Causal Dynamic Variational Autoencoder

• CDVAE is a generative model for causal inference, combining deep latent variable techniques with dynamic treatment effect estimation.
• It fuses variational autoencoder architectures with causal adjustments to recover counterfactual outcomes and handle confounding in longitudinal data.
• The model employs weighted ELBO, IPM, and sparsity penalties, achieving state-of-the-art performance in both world modeling and causal representation recovery.

The Causal Dynamic Variational Autoencoder (CDVAE) is a family of generative models designed for causal inference and representation learning in high-dimensional, time-varying environments. CDVAE integrates deep latent variable modeling with causal adjustment, enabling robust estimation of individualized causal effects and facilitating mechanistic adaptation under interventions. This approach yields modular latent world models for complex dynamical systems and achieves state-of-the-art performance on counterfactual treatment effect estimation and causal representation recovery (Lei et al., 2022, Bouchattaoui et al., 2023, Bouchattaoui, 4 Dec 2025).

1. Problem Setting, Causal Assumptions, and Identifiability

CDVAE addresses estimation of Individual Treatment Effects (ITE) and Conditional Average Treatment Effects (CATE) in longitudinal panel data. For $N$ units observed over $T$ time-steps, with static covariates $V$, time-varying confounders $X_{i,t}$, treatment $W_{i,t}$, and outcome $Y_{i,t}$, the model posits potential outcome notation:

• $Y_{i,t}(w)$: response at time $t$ under intervention $W_{i,t}=w$
• ITE: $T_t(h)=E[Y_t(1) - Y_t(0)\mid H_t = h]$, $H_t$ being the observed history

Classical causal identification relies on:

• Consistency: $Y_t = Y_t(W_t)$
• Sequential ignorability: $Y_t(w)\perp W_t \mid H_t$
• Overlap: $p(W_t=w\mid H_t=h)>0$ for all $w,h$

CDVAE further augments with a static, unobserved adjustment variable $U$ that affects outcomes but not treatment assignment. This yields the "augmented CATE":

$T_t(h, u) = E[Y_t(1) - Y_t(0) | H_t = h, U = u]$

CDVAE infers a latent substitute $Z$ for $U$, such that $T_t(h, Z)$ remains identifiable, using a finite-order conditional Markov model (CMM($p$)) property:

$$p\bigl(Y_t\mid Y_{t-1:t-p}, X_t, W_t, Z\bigr) = p\bigl(Y_t\mid Y_{t-1:t-p}, X_t, W_t, U\bigr)$$

Underlying the architecture and theory are the results formalized in (Bouchattaoui, 4 Dec 2025), guaranteeing identifiability and uniqueness of $Z$ as a sufficient adjustment.

2. Model Architecture and Generative Process

CDVAE comprises two major architectural lines:

A. Variational Latent Dynamic Model for World Dynamics and Interventions (Lei et al., 2022)

• Observations $x^{0:T}$ (images, mixed state), actions $a^{0:T}$, latent dynamics $z^{0:T}$
• Generative model:

$$p_\theta(x^{0:T}, a^{0:T}) = \int\left[\prod_{t=0}^T p_\theta(x_t | z_t) p_\theta(z_t | z_{t-1}, a_{t-1})\right]dz^{0:T}$$

• Recognition model (encoder):

$$q_\phi(z_t | x_t) = \mathcal{N}(z_t; \mu_\phi(x_t), \operatorname{diag}\sigma^2_\phi(x_t))$$

• Structured transition model: Each latent dimension $z_{t,i}$ is treated as a causal variable in a causal DAG $G$, with factorized per-variable transitions:

$$p_\theta^{(k)}(z_t | z_{t-1}, a_{t-1}) = \prod_i \left[p_i^{(0)}(\cdot)\right]^{1 - R^I_{k,i}} [p_i^{(k)}(\cdot)]^{R^I_{k,i}}$$

where $R^I_{k,i}$ is an intervention mask.

B. Dynamic VAE with Propensity-Weighted Causal Adjustment (Bouchattaoui et al., 2023, Bouchattaoui, 4 Dec 2025)

• Encoder: RNN-based (GRU/LSTM) summarization of history to infer latent substitute $Z_i$, with recognition network $q_\phi(Z_i | D_{i,T})$
• Decoder: RNN plus MLP, using $Z_i$ and encoded history, to generate outcome sequences for both factual and counterfactual regimes
• Treatment assignment network, $e_\psi(r_{i,t})$, estimates $p(W_{it}=1|H_{it})$

3. Learning Objectives and Causal Regularization

CDVAE employs weighted variational inference and causal regularization to address selection bias and enforce latent validity:

Key elements:

• Weighted ELBO (W-ELBO):

$$\sum_{i=1}^n \sum_{t=1}^T E_{Z_i \sim q_\phi} [\alpha_{i,t} \log p_\theta(Y_{it} | \cdots)] - \beta\, \operatorname{KL}(q_\phi(Z_i|D_{i,T}) \Vert p(Z_i))$$

with overlap weights $\alpha_{i,t}$ derived from propensity scores

• Integral Probability Metric (IPM):

Enforces covariate balance in representation space across treated and control arms

• Posterior-consistency:

Penalty on Wasserstein distance between latent posteriors $q_\phi(Z|D_t)$ and $q_\phi(Z|D_{t-1})$ to ensure staticity of $Z$

• Sparsity Penalties:

Applied to learned graph and intervention masks to induce modularity in world dynamics (Lei et al., 2022)

• Moment-Matching Penalty:

$\sum_{t=2}^{T} \|g_{i,t} - g_{i,t-1}\|^2$ to further ensure $Z_i$ captures static heterogeneity

The overall loss combines the negative weighted-ELBO, IPM penalty, posterior-consistency regularizer, and cross-entropy for the propensity net.

4. Training Algorithms and Adaptation to Interventions

Model fitting follows alternating stochastic optimization via Adam/SGD:

• For world models (Lei et al., 2022):
  • Learn encoder/decoder/transition graph/interaction masks via reparameterized gradients and straight-through Gumbel-Softmax for discrete structures.
  • Adaptation to new environments by estimating shift masks $R'$ and training only changed mechanisms under the sparse-mechanism shift hypothesis.
• For treatment effect models (Bouchattaoui et al., 2023, Bouchattaoui, 4 Dec 2025):
  • Pretrain encoder on contrastive objectives (CPC, InfoMax)
  • Jointly optimize ELBO, regularization terms, and propensity net adversarially
  • Inference proceeds by encoding new histories, sampling $Z$, and forecasting outcomes under arbitrary treatments.

Pseudocode outlined in the cited works includes batch processing, overlap-weight sampling, intervention mask estimation, early stopping on factual validation losses, and estimation of Jacobian traces for scalable causal interpretability.

5. Theoretical Guarantees

CDVAE is supported by a suite of theoretical results, notably (Bouchattaoui, 4 Dec 2025):

• Identification of substitute $Z$: Theorems show that under CMM($p$), latent $Z$ suffices for valid adjustment as if true $U$ were observed.
• Minimality and uniqueness: If another variable $Z'$ satisfies CMM($p$), it must be a measurable function of $(Z, H_T)$
• Near-deterministic regime: As decoder variance $\sigma^2 \to 0$, posterior sampling collapses and any $Z$ sample yields the same causal estimate.
• Generalization bounds: Precision in Estimation of Heterogeneous Effects (PEHE) is bounded by empirical risk terms, IPM discrepancy, and sample complexity; uniform convergence achieves $O(1/\sqrt n \sqrt{\log d})$ rates.

This formal analysis provides guarantees for causal validity of estimated effects and adjustment.

6. Empirical Performance and Causal Representation Recovery

Empirical results span synthetic and real datasets:

• World modeling (modular dynamical systems) (Lei et al., 2022):
  • Accurately identifies axis-aligned ground-truth coordinates in latent space
  • Recovers sparse causal graphs and correct intervention patterns
  • Rapid, modular adaptation to environmental shifts (requiring few trajectories)
  • Outperforms RSSM, MultiRSSM on image/mixed-state settings
• Treatment effect estimation (Bouchattaoui et al., 2023, Bouchattaoui, 4 Dec 2025):
  • Demonstrated reduction in ITE error across synthetic autoregressive and tumor growth datasets
  • Ablations show the necessity of IPM and moment-matching for latent validity
  • Consistently outperforms Marginal Structural Recurrent Models, Counterfactual Recurrent Networks, Causal Forest DML, and Causal Transformer benchmarks
• Causal representation learning (Bouchattaoui, 4 Dec 2025):
  • Sparse self-expression of decoder Jacobian recovers known feature modularity
  • Overlapping groups identified even without anchor/single-parent assumptions; F1 and NMI metrics support clustering recovery

A plausible implication is that CDVAE provides robust, interpretable latent adjustment for ITE estimation and modular world-model adaptation.

7. Extensions and Generalizations

Emergent directions and enhancements include:

• Incorporation of causal-graph priors into latent dynamics, enabling SCM-to-GNN mapping for more expressive counterfactual modeling
• Invariance and equivariant decoders, supporting causal identifiability up to affine transformations
• Uncertainty quantification via conformal/sensitivity analysis for multi-horizon treatments
• Bayesian causal representation learning over latent graphs and decoder parameters
• Dynamic clustering of latent-to-observed mappings to allow adaptation for time-varying causal relations

Limitations noted include restriction to binary treatments, contemporaneous effects, and static risk factors; ongoing work focuses on continuous/multi-dose regimes and dynamic confounding.

Summary Table: Principal Elements of CDVAE Models

Element	World Model CDVAE (Lei et al., 2022)	Treatment Effect CDVAE (Bouchattaoui et al., 2023, Bouchattaoui, 4 Dec 2025)
Latent variable structure	Axis-aligned $z_t$ per timestep	Static $Z$ per subject/unit
Causal adjustment	DAG over latent dynamics + interventions	Static risk-factor substitute for unobserved confounders
Training objective	ELBO + sparsity via Gumbel-Softmax	Weighted ELBO + IPM + consistency + BCE
Adaptation mechanism	Sparse intervention mask re-learning	Latent static factor inference for new histories
Interpretation layer	Modular mechanisms in state-space	Causal representation, sparse Jacobian group recovery
Empirical benchmarks	RSSM, MultiRSSM	CRN, RMSM, CausalForestDML, Causal Transformer

In summary, Causal Dynamic Variational Autoencoders unify advances in latent world modeling, treatment effect adjustment, and interpretable causal representation learning, providing a flexible backbone for dynamic causal inference under complex confounding and environmental interventions (Lei et al., 2022, Bouchattaoui et al., 2023, Bouchattaoui, 4 Dec 2025).