CAFE: Counterfactual Prediction via Fast Posterior Estimation

Xinyan Han; Xiaoyu Lin; Hao Zou; Xingxuan Zhang; Bo Li; Peng Cui

cs.LGSep 26, 2026

CAFE: Counterfactual Prediction via Fast Posterior Estimation

Authors: Xinyan Han, Xiaoyu Lin, Hao Zou, Xingxuan Zhang, Bo Li, Peng Cui

Organizations: Tsinghua University · Stable AI

Abstract

Counterfactual prediction estimates an individual's outcome under an alternative intervention given their factual observations. Such outcomes are generally not identifiable from observational data without additional assumptions. Even within the class of fully observed additive noise models (ANMs), different causal graphs can generate the same observational distribution yet imply different individual counterfactual outcomes. Predictions based on a single estimated graph ignore this structural uncertainty. We therefore target a Bayesian counterfactual posterior predictive distribution that combines predictions from plausible SCMs. We introduce CAFE (CounterfActual Prediction via Fast Posterior Estimation), an amortized inference framework that directly approximates the Bayesian counterfactual posterior predictive distribution. We pretrain a transformer-based model on synthetic counterfactual tasks generated from a diverse prior over ANMs. Given an observational dataset, an individual's factual observations, and an intervention, CAFE approximates the corresponding posterior predictive distribution in a single forward pass. Experiments show that CAFE accurately predicts individual counterfactual outcomes in identifiable settings and approximates the posterior predictive distribution when structural uncertainty induced by observationally indistinguishable causal graphs exists. Strong performance in realistic manufacturing and viticulture settings further demonstrates its empirical robustness beyond the assumptions of the training prior.

Figures & tables

Figure 1: Overview of CAFE. Top: We sample SCMs with diverse DAG structures from the prior $\Pi$ and generate paired factual and counterfactual data with shared exogenous noise. Bottom: We train a Transformer on these tasks to approximate the counterfactual posterior predictive distribution from observational data $\mathcal{D}_{\mathrm{obs}}$ , an individual’s factual observations $\mathbf{v}$ , and an intervention $t$ .

Method	collider lingauss	fork lin nongauss	fork nonlin gauss	large backdoor	nonlin simpson	symprod simpson	weak arrows
y factual	$0.57$ $\pm 0.00$	$0.00$ $\pm 0.00$	$0.00$ $\pm 0.00$	$0.74$ $\pm 0.00$	$1.36$ $\pm 0.00$	$1.37$ $\pm 0.00$	$1.07$ $\pm 0.00$
CAREFL	$0.11$ $\pm 0.15$	$0.34$ $\pm 0.00$	$0.02$ $\pm 0.00$	$0.61$ $\pm 0.10$	$1.36$ $\pm 0.00$	$1.37$ $\pm 0.00$	$0.21$ $\pm 0.00$
CausalNF	$0.08$ $\pm 0.01$	$0.55$ $\pm 0.43$	$0.09$ $\pm 0.11$	$1.57$ $\pm 0.82$	$2.48$ $\pm 0.16$	$1.37$ $\pm 0.00$	$1.86$ $\pm 0.89$
VCI	$0.51$ $\pm 0.00$	$0.34$ $\pm 0.00$	$0.28$ $\pm 0.01$	$0.70$ $\pm 0.01$	$1.58$ $\pm 0.12$	$1.37$ $\pm 0.01$	$1.02$ $\pm 0.01$
SCIGAN	$0.81$ $\pm 0.01$	$0.62$ $\pm 0.04$	$0.47$ $\pm 0.00$	$1.20$ $\pm 0.06$	$1.07$ $\pm 0.09$	$1.51$ $\pm 0.05$	$0.98$ $\pm 0.03$
DECI	$0.39$ $\pm 0.30$	$0.60$ $\pm 0.46$	$0.03$ $\pm 0.07$	$0.42$ $\pm 0.23$	$3.04$ $\pm 1.07$	$1.45$ $\pm 0.10$	$0.42$ $\pm 0.17$

Table 1: Counterfactual prediction RMSE on 7 C-Suite datasets with additive noise (mean

\pm

std). Bold indicates the lowest mean RMSE in each column, excluding the factual-outcome baseline.

Panel of Figure 2 — Figure 2: Counterfactual uncertainty under non-identifiable causal structure. Top: six observationally indistinguishable SCMs. Bottom: predicted versus ground-truth counterfactual distributions for randomly selected individuals (left) and posterior variances across individuals (right) at $N=1000$ .

$N$	Mean RMSE	Variance RMSE	IQD
50	$0.464$ $\pm 0.120$	$0.558$ $\pm 0.132$	$0.161$ $\pm 0.050$
100	$0.342$ $\pm 0.191$	$0.565$ $\pm 0.267$	$0.104$ $\pm 0.092$
250	$0.153$ $\pm 0.048$	$0.085$ $\pm 0.017$	$0.031$ $\pm 0.017$
500	$0.110$ $\pm 0.020$	$0.073$ $\pm 0.015$	$0.020$ $\pm 0.008$
1000	$\mathbf{0.094}$ $\boldmath\mathbf{\pm 0.012}$	$\mathbf{0.055}$ $\boldmath\mathbf{\pm 0.010}$	$\mathbf{0.018}$ $\boldmath\mathbf{\pm 0.004}$

Table 2: Counterfactual posterior prediction under non-identifiable causal structure (mean

\pm

std).

Method	CausalMan Micro		CausalMan Small		Sangiovese
	Accuracy (%) $\uparrow$		Accuracy (%) $\uparrow$		MAPE (%) $\downarrow$
	$t=15000$	$t=30000$	$t=15000$	$t=30000$	Brix	pH
y factual	$92.85$ $\pm 0.00$	$12.15$ $\pm 0.00$	$94.05$ $\pm 0.00$	$6.70$ $\pm 0.00$	$6.83$ $\pm 0.00$	$2.07$ $\pm 0.00$
CAREFL	$92.85$ $\pm 0.00$	$12.15$ $\pm 0.00$	$94.05$ $\pm 0.00$	$6.70$ $\pm 0.00$	$2.52$ $\pm 0.14$	$1.88$ $\pm 0.23$
CausalNF	$92.98$ $\pm 0.26$	$20.86$ $\pm 17.42$	$94.40$ $\pm 0.54$	$24.30$ $\pm 21.61$	$4.20$ $\pm 0.25$	$1.93$ $\pm 0.17$
VCI	$92.79$ $\pm 0.06$	$12.09$ $\pm 0.06$	$94.04$ $\pm 0.02$	$6.69$ $\pm 0.02$	$5.70$ $\pm 0.25$	$1.83$ $\pm 0.07$

Table 3: Counterfactual prediction performance on CausalMan and Sangiovese (mean

\pm

std). We report accuracy (%, higher is better) on CausalMan under two press-fitting force interventions, and MAPE (%, lower is better) for Brix and pH on Sangiovese.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 3 — Figure 3: Counterfactual predictive distributions under identifiable causal structure. Predicted densities (blue) and ground-truth point masses (red) for four individuals from each of two C-Suite datasets: collider linear Gaussian (top) and fork nonlinear Gaussian (bottom). Panel titles report factual observations, and the corresponding causal graphs appear on the right. Each red marker represents unit probability mass at the individual’s true counterfactual outcome; its height is read on a separate scale from the predicted density.

Method	large backdoor binary treatment	mixed confounding	mixed simpson
y factual	$0.17$ $\pm 0.00$	$0.86$ $\pm 0.00$	$0.50$ $\pm 0.00$
CAREFL	$0.17$ $\pm 0.00$	$0.86$ $\pm 0.00$	$0.50$ $\pm 0.00$
CausalNF	$0.23$ $\pm 0.11$	$0.86$ $\pm 0.00$	$0.64$ $\pm 0.19$
VCI	$0.14$ $\pm 0.01$	$1.25$ $\pm 0.09$	$0.93$ $\pm 0.34$
SCIGAN	$1.16$ $\pm 0.04$	$0.71$ $\pm 0.07$	$0.43$ $\pm 0.06$
DECI	$0.14$ $\pm 0.03$	$0.73$ $\pm 0.23$	$1.13$ $\pm 0.32$

Appendix

Table 4: Counterfactual prediction RMSE on C-Suite datasets with non-additive noise (mean

\pm

standard deviation). Bold indicates the lowest mean RMSE in each column.

Method	Mean RMSE	Variance RMSE	IQD
DECI	$0.235$ $\pm 0.045$	$0.101$ $\pm 0.030$	$0.082$ $\pm 0.050$
DiBS	$0.142$ $\pm 0.050$	$0.175$ $\pm 0.125$	$0.034$ $\pm 0.024$
CAFE	$\mathbf{0.094}$ $\boldmath\mathbf{\pm 0.012}$	$\mathbf{0.055}$ $\boldmath\mathbf{\pm 0.010}$	$\mathbf{0.018}$ $\boldmath\mathbf{\pm 0.004}$

Appendix

Table 5: Counterfactual posterior prediction under non-identifiable causal structure (mean

\pm

standard deviation) at

N=1000

Panel of Figure 4 — Figure 4: Estimated and ground-truth counterfactual posteriors in the three-variable, six-DAG linear Gaussian system ( $N=1000$ ). Top: Counterfactual posterior estimates from DiBS. Bottom: Counterfactual posterior estimates from DECI.

Method	Sangiovese	CausalMan Micro	CausalMan Small
	$d=15$	$d=24$	$d=53$
CAREFL	$143.27$ $\pm 15.45$	$981.12$ $\pm 62.54$	$1064.83$ $\pm 526.84$
CausalNF	$21.13$ $\pm 0.42$	$1453.74$ $\pm 141.83$	$628.77$ $\pm 57.50$
VCI	$276.73$ $\pm 11.12$	$211.85$ $\pm 1.12$	$208.36$ $\pm 1.13$
SCIGAN	$233.44$ $\pm 9.16$	$117.91$ $\pm 2.24$	$122.23$ $\pm 1.94$
DECI	$78.33$ $\pm 1.04$	$1065.45$ $\pm 9.23$	$214.36$ $\pm 3.72$

Appendix

Table 6: Runtime comparison on Sangiovese and CausalMan (seconds, mean

\pm

standard deviation; lower is better). The number of observed variables is shown for each dataset.

Figure 5: Eight causal graphs used to evaluate CATE adjustment.

	Collider linear Gaussian		Fork linear non-Gaussian
Method	(1)	(2)	(1)	(2)
y factual	$0.81$ $\pm 0.00$	$0.00$ $\pm 0.00$	$0.00$ $\pm 0.00$	$0.00$ $\pm 0.00$
CausalForestDML	$0.16$ $\pm 0.01$	$0.72$ $\pm 0.02$	$0.51$ $\pm 0.01$	$0.30$ $\pm 0.03$
DRNet	$0.64$ $\pm 0.04$	$0.95$ $\pm 0.03$	$0.76$ $\pm 0.09$	$0.75$ $\pm 0.07$
DRNet cal	$0.13$ $\pm 0.01$	$0.56$ $\pm 0.09$	$0.65$ $\pm 0.10$	$0.46$ $\pm 0.13$
CCPFN	$0.58$ $\pm 0.00$	$0.98$ $\pm 0.00$	$0.56$ $\pm 0.00$	$0.57$ $\pm 0.00$

Appendix

Table 7: Counterfactual prediction RMSE on C-Suite datasets (mean

\pm

standard deviation). Bold indicates the lowest mean RMSE in each column, excluding the factual-outcome baseline. Numbers in parentheses denote outcome indices. The superscript “cal” denotes an adjustment variant. Light blue marks settings satisfying the two structural conditions in Proposition 4 .

	Collider linear Gaussian		Fork linear non-Gaussian
Method	(1)	(2)	(1)	(2)
y factual	$0.81$ $\pm 0.00$	$0.00$ $\pm 0.00$	$0.00$ $\pm 0.00$	$0.00$ $\pm 0.00$
Backdoor Prior	$\mathbf{0.02}$ $\boldmath\mathbf{\pm 0.00}$	$0.39$ $\pm 0.00$	$1.20$ $\pm 0.06$	$\mathbf{0.01}$ $\boldmath\mathbf{\pm 0.00}$
Full Prior	$0.03$ $\pm 0.00$	$\mathbf{0.06}$ $\boldmath\mathbf{\pm 0.01}$	$\mathbf{0.00}$ $\boldmath\mathbf{\pm 0.00}$	$\mathbf{0.01}$ $\boldmath\mathbf{\pm 0.00}$

Appendix

Table 8: Counterfactual prediction RMSE on C-Suite datasets (mean

\pm

standard deviation). Full Prior denotes CAFE trained with its original prior over diverse DAGs. Backdoor Prior denotes the ablation model trained only on back-door graphs. Bold indicates the lowest mean RMSE in each column, excluding the factual-outcome baseline. Numbers in parentheses denote outcome indices. Light blue marks settings satisfying the two structural conditions in Proposition 4 .

Explore similar work

Oct 9, 2026

Amortized Counterfactual Prediction with Graph Neural Networks

Counterfactual inference asks how an observed outcome might have differed under different circumstances, such as whether a patient would have recovered had they received a different treatment. Such questions arise whenever decisions depend on the outcomes of alternatives that were not observed. Answering counterfactual queries generally requires knowledge of the underlying causal structure, which is often unavailable in practice. A common approach is therefore to first estimate the causal graph from observational data and then perform counterfactual inference. However, errors in the estimated causal graph can propagate to downstream counterfactual predictions. To address this challenge, we introduce the Counterfactual Graph Neural Process (CFGNP), an end-to-end approach that answers counterfactual queries directly from observational data without first estimating a causal graph. CFGNP uses Graph Neural Networks (GNNs) to exploit the dual-graph structure by restricting cross-world message passing to corresponding factual--counterfactual variables. Across synthetic and image-based benchmarks, CFGNP achieves stronger predictive performance than existing approaches, including on the German Loan benchmark, where it reduces MSE from

4.51 \times 10^{-2}

0.19 \times 10^{-2}

. Furthermore, on synthetic causal graphs, a model trained only on 10-node graphs maintains strong predictive performance on graphs with up to 30 nodes, without retraining.

Hendrik Stumpf, Francesco Quinzan

Department of Computer Science, University of Oxford · Department of Engineering, University of Oxford

Counterfactual Prediction Causal Inference Graph Neural Networks

Ask PDF

Jun 25, 2026

A Causal Foundation Model for Structure and Outcome Prediction

We introduce TabPFN-CFM, a causal foundation model that can handle multiple causal problems. TabPFN-CFM predicts both causal structure and outcomes from observational data, supports queries on all three levels of Pearl's Causal Hierarchy and uses known graph structure when available to improve predictions. TabPFN-CFM is trained on synthetic datasets, and generalises to real datasets, demonstrating improved performance over both structural and outcome prediction baselines.

Max Zhu, Martino Mansoldo, Ching-Hao Wang +1

University of Cambridge, United Kingdom · GSK.ai.

Causal Foundation Models Causal Structure Learning Causal Inference

Ask PDF

Jun 5, 2026

Are Good Generators Good Decision-Makers? Policy Learning for General Interventions via Retargeted Counterfactual Generation

Generative models are increasingly used to support decision-making in complex systems, where interventions may be joint and high-dimensional, and outcomes are high-dimensional. However, using generators for these decision-making settings are challenged by three problems. First, they are often trained on noisy logs with limited intervention data. Second, generator learning can be noisy across different environments. Third, generators are not optimized for the decisions they support. We introduce policy learning via retargeted counterfactual generation, which trains a generator for the decisions it supports in three steps. We (1) learn a doubly-robust, invariant counterfactual generator for high-dimensional interventions and outcomes, (2) conduct policy-learning based on its rollouts, and (3) retarget the generator toward the learned policy's interventions and relearn the policy, so the generator is accurate where decisions are made. Theoretically, our generator's excess counterfactual risk has a doubly robust product remainder, and retargeting removes the worst-case density ratio between the logged and learned policies from the regret. We demonstrate the efficacy of our approach across synthetic data, Cell Painting images of SARS-CoV-2-infected cells, and physical video simulations.

Raphael C Kim, Jingsen Zhu, Ramin Zabih +1

Cornell Tech, Cornell University, New York, NY · Department of Biostatistics, Department of Population Health, New York University Grossman School of Medicine, New York, NY

Reinforcement Learning Causal Counterfactual Generation Policy Learning Conditional Generative Modeling Counterfactual Learning

Ask PDF