Causal Diffusion Models for Counterfactual Outcome Distributions in Longitudinal Data

Farbod Alinezhad; Jianfei Cao; Gary J. Young; Brady Post

stat.MLApr 14, 2026

Causal Diffusion Models for Counterfactual Outcome Distributions in Longitudinal Data

Authors: Farbod Alinezhad, Jianfei Cao, Gary J. Young, Brady Post

Organizations: Analysis Group · Department of Economics, Northeastern University · Center for Health Policy and Healthcare Research, Northeastern University · Department of Public Health and Health Sciences, Bouvé College of Health Sciences, Northeastern University

Abstract

Predicting counterfactual outcomes in longitudinal data, where sequential treatment decisions heavily depend on evolving patient states, is critical yet notoriously challenging due to complex time-dependent confounding and inadequate uncertainty quantification in existing methods. We introduce the Causal Diffusion Model (CDM), the first denoising diffusion probabilistic approach explicitly designed to generate full probabilistic distributions of counterfactual outcomes under sequential interventions. CDM employs a novel residual denoising architecture with relational self-attention, capturing intricate temporal dependencies and multimodal outcome trajectories without requiring explicit adjustments (e.g., inverse-probability weighting or adversarial balancing) for confounding. In rigorous evaluation on a pharmacokinetic-pharmacodynamic tumor-growth simulator widely adopted in prior work, CDM consistently outperforms state-of-the-art longitudinal causal inference methods, achieving a 15-30% relative improvement in distributional accuracy (1-Wasserstein distance) while maintaining competitive or superior point-estimate accuracy (RMSE) under high-confounding regimes. By unifying uncertainty quantification and robust counterfactual prediction in complex, sequentially confounded settings, without tailored deconfounding, CDM offers a flexible, high-impact tool for decision support in medicine, policy evaluation, and other longitudinal domains.

Figures & tables

Fig 1: Root Mean Squared Error (RMSE, %) and 1-Wasserstein distance (%) vs. confounding level $\gamma$ .

	Amount of Time-Dependent Confounding ( $\gamma$ )
	0	1	2	3	4	5	6	7	8	9	10	11	12	13	14	15
RMSE
CDM	0.88	0.89	0.95	1.61	1.50	1.78	2.14	2.41	2.69	2.95	3.62	3.43	3.45	4.09	3.50	3.73
CRN	0.78	0.70	0.86	0.93	1.37	1.74	2.07	2.64	3.01	3.33	3.89	4.27	4.44	4.42	4.55	4.47
CT	0.82	0.72	0.87	0.93	1.19	1.58	1.91	2.35	3.36	4.79	4.87	4.85	4.48	5.86	5.35	4.34
GNet	0.82	0.78	1.04	1.05	1.31	1.49	1.96	2.67	3.07	3.20	3.94	4.09	3.98	4.05	4.21	4.57

Table 1: Performance of CDM versus alternative methods (%)

Fig 2: RMSE (%) for the 2.5th and 97.5th percentiles of the predicted distribution. CDM (Causal Diffusion Model), CRN (Counterfactual Recurrent Network), CT (Causal Transformer), GNet, and RMSN (Recurrent Marginal Structural Network) are compared. CDM performs competitively with the baselines at low and moderate confounding levels, but outperforms most of the baselines at higher confounding levels.

	RMSE ${}_{\text{median}}$ (%)			1-Wasserstein (%)
Configuration	$\gamma{=}0$	$\gamma{=}5$	$\gamma{=}10$	$\gamma{=}0$	$\gamma{=}5$	$\gamma{=}10$
Full CDM (baseline)	0.94	1.64	4.29	0.08	0.26	1.04
Diffusion steps $=20$	1.12	3.37	3.24	0.09	0.68	0.83
Linear $\beta$ schedule	1.90	4.53	8.53	0.13	0.65	1.91
Inclusion of a residual layer	1.05	1.87	3.19	0.09	0.33	0.83
Embedding dimension $=8$	2.90	5.37	8.46	0.47	1.38	3.73

Table 2: Ablation study on CDM design choices at three confounding levels.

Hyperparameter	Value
Gradient Clipping	4
Epochs	25
Learning Rate (LR)	starting at $10^{-3}$ , decaying factor 0.9 on validation loss plateau
Minimum Learning Rate	$10^{-6}$

Table 3: Training Hyperparameters

Hyperparameter	Value
Embedding Dimension	32
Residual Layers	2
Diffusion Steps	5
Beta Schedule	cosine
Attention Heads	8
Kernel Size (for RSA convolution)	(3,7)

Table 4: Diffusion Imputation Model Hyperparameters

Counterfactual Prediction

Counterfactual prediction estimates outcomes under hypothetical interventions using observed data and assumptions about how systems generate outcomes.

See all 38 papers

Causal Longitudinal Prior-Fitted Networks for Counterfactual Outcome Prediction
Jun 4, 2026Amirhossein Zare, Amirhessam Zare, Herlock Rahimi +2Structural Causal Models Counterfactual Prediction
School of Medicine, Tehran University of Medical Sciences, Tehran, Iran. · Department of Electrical and Computer Engineering, Yale University, New Haven, CT, USA. · School of Medicine, Shiraz University of Medical Sciences, Shiraz, Iran. +2
Causal State-Space Model for Causal Inference: Estimating Longitudinal Individual Treatment Effects
Aug 8, 2026Abisoye Abidakun, Mingjun Zhong, Georgios LeontidisCausal Effect Estimation Selective SSMs
Department of Computing Science, University of Aberdeen, Aberdeen, UK · UiT The Arctic University of Norway
Robust Counterfactual Policy Optimisation via Nondeterministic Causal Models
Aug 3, 2026Jessica Lally, Milad Kazemi, Nicola Paoletti +2Markov Decision Processes Structural Causal Models
1King’s College London · 2Durham University · University College London

Jun 4, 2026J/K move · Enter open · S save

Causal Longitudinal Prior-Fitted Networks for Counterfactual Outcome Prediction

Amirhossein Zare, Amirhessam Zare, Herlock Rahimi, Reza Salarikia+1

School of Medicine, Tehran University of Medical Sciences, Tehran, Iran. · Department of Electrical and Computer Engineering, Yale University, New Haven, CT, USA. · School of Medicine, Shiraz University of Medical Sciences, Shiraz, Iran.+2

Structural Causal Models Counterfactual Prediction Prior-Data Fitted Networks Causal Inference Longitudinal EHR Modeling Counterfactual Learning

Open paper AskPDF arXiv

Longitudinal treatment decisions from multivariate time-series data require predicting potential outcomes under future treatment sequences in the presence of time-varying confounding, heterogeneous patient dynamics, and limited domain-specific data. Existing longitudinal causal estimators typically address this problem by training a new model for each cohort or simulator. We introduce Causal Longitudinal Prior-Fitted Networks (CausalLongPFN), a prior-fitted network for time-series causal inference in longitudinal treatment-response data and zero-shot in-context counterfactual outcome prediction. The model is pretrained entirely on synthetic episodes sampled from a broad prior over temporal structural causal models, exposing it to treatment-confounder feedback, latent heterogeneity, nonlinear state evolution, delayed effects, and cumulative treatment responses. At test time, CausalLongPFN remains frozen and is used zero-shot: it conditions on support trajectories, a query history, and a planned future treatment sequence, and returns a predictive distribution over future outcomes without gradient updates or propensity-model fitting. Multi-step predictions are obtained by recursively applying the one-step predictor under the specified treatment sequence. We evaluate the model on branchable cancer, HIV, and warfarin benchmarks with ground-truth counterfactual labels, and on factual-only rolling-origin prediction in MIMIC-III ICU trajectories. CausalLongPFN is competitive with domain-trained longitudinal baselines on counterfactual benchmarks and performs strongly on factual MIMIC-III prediction, suggesting that broad synthetic causal pretraining can provide a frozen, amortized alternative for zero-shot longitudinal treatment-response prediction when repeated domain-specific training is costly or impractical.

Abstract

Counterfactual Prediction

Explore similar work

Causal Longitudinal Prior-Fitted Networks for Counterfactual Outcome Prediction

Causal State-Space Model for Causal Inference: Estimating Longitudinal Individual Treatment Effects

Robust Counterfactual Policy Optimisation via Nondeterministic Causal Models