Generative Flows for Potential Outcomes and Counterfactuals
Authors: Dongze Wu, David I. Inouye, Yao Xie
Organizations: School of Industrial and Systems Engineering Georgia Institute of Technology Atlanta, GA · Elmore Family School of Electrical and Computer Engineering Purdue University West Lafayette, IN
Average treatment effects summarize how interventions change outcomes, but individualized decisions require predictions and outcome distributions of what would have happened to the same individual under an alternative treatment. We develop PO-Flow, a continuous normalizing flow framework that represents treatment-specific outcome distributions through invertible transport maps. It supports potential-outcome generation, conditional average treatment-effect estimation, factual-conditioned counterfactual prediction, and density evaluation. To address treatment imbalance in observational data, we incorporate cross-fitted inverse-probability weighting into the flow-matching loss. We establish sufficient conditions under which encoding a factual outcome and decoding under an alternative treatment recovers the individual counterfactual. Under shared exogenous noise and a common encoded reference distribution, scalar recovery follows from increasing noise-to-latent maps. For multivariate outcomes, we assume triangularity or a convex-gradient structure. These conditions distinguish fitting outcome distributions from recovering their causal correspondence. Experiments on four scalar benchmarks and seven multivariate or image settings, with dimensions up to 784, compare against ten scalar and five multivariate baselines. Controlled ablations show IPW improves all reported scalar counterfactual RMSE, PEHE, and potential-outcome RMSE. PO-Flow achieves the lowest PEHE and potential-outcome RMSE on all four scalar benchmarks and attains or ties the lowest mean in 26 of 35 multivariate dataset--metric comparisons.
Figures & tables
Fig. 1: From average effects to individual counterfactuals. (a) ATE and CATE summarize treatment effects on average, overall and for people with the same covariates. (b) Given a person’s observed outcome under one treatment, we ask what would have happened to that same person under the other. The dashed person marks the unobserved outcome; the arrow shows a pairing. Outcome distributions describe variation and tail risk, but identifying this pairing requires additional structural assumptions.
Data
Method
CF-RMSE
PEHE
PO-RMSE
W1
W1,μ
ACIC16
PO-Flow
1.52 (.01)
.41 (.02)
1.25 (.04)
.17 (.01)
.17 (.02)
PO-Flow (no IPW)
1.53 (.17)
.82 (.09)
1.37 (.09)
.42 (.12)
.45 (.10)
TARNet
1.53 (.10)
.58 (.18)
1.26 (.08)
–
.27 (.10)
DragonNet
1.34 (.12)
.60 (.14)
1.27 (.08)
–
.22 (.08)
DiffPO
1.78 (.20)
1.12 (.26)
1.42 (.08)
.51 (.18)
–
INFs
–
1.42 (.30)
1.36 (.25)
.67 (.14)
–
TABLE I: Scalar-outcome comparison over five seeds (mean with standard deviation); lower is better.
Data
dY
Method
CF-RMSE
PEHE
PO-RMSE
SW1
SW1,μ
IHDP
2
PO-Flow
.69 (.00)
.20 (.01)
.48 (.00)
.21 (.00)
.05 (.00)
TARNet
.64 (.02)
.55 (.03)
.56 (.01)
–
.16 (.01)
DragonNet
.64 (.01)
.59 (.03)
.57 (.01)
–
.19 (.01)
CFR
.63 (.02)
.55 (.03)
.56 (.01)
–
.16 (.01)
GANITE
1.15 (.37)
1.03 (.47)
.76 (.11)
–
.44 (.21)
CEVAE
1.03 (.16)
.39 (.02)
.54 (.01)
.46 (.01)
.11 (.02)
TABLE II: Vector- and image-outcome comparison over five seeds (mean with standard deviation); lower is better.
Appendix figures & tables3 assets
Supplementary material from the paper’s appendix.
Appendix
Data
(q,α)
AUC
clipped
ESS/ n
max w
max SMD before → after
ACIC16
(.99,.25)
.6002
.073%
.9785
1.611
.2490 → .1903
ACIC18
(.95,.75)
.5575
.128%
.9796
1.260
.4517 → .3406
IHDP
(.95,.10)
.6659
3.585%
.9937
1.255
.4198 → .3980
IBM
(.95,1.0)
.9483
36.975%
.7828
2.650
1.2824 → 1.0483
Appendix
TABLE III: Overlap diagnostics for the final scalar IPW settings.
Data
(q,α)
AUC
ESS/ n
max w
max SMD before → after
IHDP dY=2
(.95,.50)
.6991
.8728
2.220
.4384 → .3529
IHDP dY=5
(.90,.75)
.6991
.8541
2.006
.4384 → .3447
IHDP dY=10
(.95,.50)
.6991
.8728
2.220
.4384 → .3529
ACIC16
(.95,.75)
.7705
.7653
2.713
1.1332 → .7526
MorphoMNIST
(.95,.50)
.6663
.9811
1.332
.5055 → .2875
Appendix
TABLE IV: IPW overlap diagnostics for the vector families. IHDP and ACIC16 share X,A across output dimensions, but IHDP uses dimension-specific weight strength.
Method
ACIC16
ACIC18
IHDP
IBM
PO-Flow
0.14 (0.04)
0.08 (0.02)
0.09 (0.03)
0.08 (0.03)
INF
0.22 (0.09)
0.15 (0.06)
0.14 (0.07)
0.10 (0.04)
GP-ITE
1.73 (0.34)
0.73 (0.19)
1.42 (0.30)
0.68 (0.17)
Appendix
TABLE V: KL divergence of learned potential outcome densities (10-fold cross-validation). Parentheses indicate standard deviation.
Estimating counterfactual distributions under interventions is central to treatment risk assessment and counterfactual generation tasks. Existing approaches model the counterfactual distribution as a standalone generative target, without exploiting its relationship to the observational data. In this work, we show that under standard assumptions, observational and counterfactual outcome distributions are tightly linked: they have identical support and tail behavior, remain statistically close under weak confounding, and share any features of high-dimensional outcomes which are invariant to confounders. These properties motivate learning counterfactual distributions not from scratch, but via a deconfounding flow from the observational distribution. We formulate this problem via flow-matching and derive a semiparametrically efficient estimator based on a novel efficient influence function correction. We subsequently extend our estimator to target minimal-energy flows in high-dimensions, which we show can be especially simple targets between observational and counterfactual distributions. In experiments, deconfounding flows outperform existing debiased counterfactual distribution estimators, while also mitigating known failure modes of flow-based methods.
Hugh Dance, Johnny Xi, Peter Orbanz +1
Gatsby Computational Neuroscience Unit University College London · Department of Statistics University of British Columbia
Counterfactual generation seeks to sample outcomes under a hypothetical intervention or decision using observational data collected under the factual assignment mechanism. We develop a flow-matching approach that combines a sample-split, doubly robust training objective with a learned coupling between observed source outcomes and target outcomes drawn from a fitted conditional outcome model. To enable finite-step generation, we leverage a score-corrected stochastic sampler based on a Gaussian-smoothed interpolation. Our main theoretical contribution is a coupling-sensitive KL bound for constant-step Euler discretization: the error is controlled by moments of the source--target displacement under the chosen coupling, rather than by global uniform regularity of the velocity field, and has near-linear dependence on the ambient dimension. We also establish finite-sample non-parametric guarantees for the learned velocity and score fields when both the conditional outcome model and the source-target coupling are estimated from data. These bounds separate approximation, coupling-replacement, nuisance-estimation, generalization, and Monte Carlo errors and, combined with the sampler analysis, yield an end-to-end guarantee for counterfactual generation. Experiments on synthetic and semi-synthetic image benchmarks support the coupling-dependent theory and show that, at finite discretization budgets, the stochastic sampler can outperform the corresponding deterministic ODE sampler.
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.