Generative Flows for Potential Outcomes and Counterfactuals

Dongze Wu; David I. Inouye; Yao Xie

stat.MLMay 21, 2025

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

Abstract

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

Panel of Figure 1 — 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	$\mathbf{W_{1}}$	$\mathbf{W_{1,\mu}}$
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	$\mathbf{d_{Y}}$	Method	CF-RMSE	PEHE	PO-RMSE	$\mathbf{SW_{1}}$	$\mathbf{SW_{1,\mu}}$
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,\alpha)$	AUC	clipped	ESS/ $n$	max $w$	max SMD before $\to$ after
ACIC16	$(.99,.25)$	.6002	.073%	.9785	1.611	.2490 $\to$ .1903
ACIC18	$(.95,.75)$	.5575	.128%	.9796	1.260	.4517 $\to$ .3406
IHDP	$(.95,.10)$	.6659	3.585%	.9937	1.255	.4198 $\to$ .3980
IBM	$(.95,1.0)$	.9483	36.975%	.7828	2.650	1.2824 $\to$ 1.0483

Appendix

TABLE III: Overlap diagnostics for the final scalar IPW settings.

Data	$(q,\alpha)$	AUC	ESS/ $n$	max $w$	max SMD before $\to$ after
IHDP $d_{Y}=2$	$(.95,.50)$	.6991	.8728	2.220	.4384 $\to$ .3529
IHDP $d_{Y}=5$	$(.90,.75)$	.6991	.8541	2.006	.4384 $\to$ .3447
IHDP $d_{Y}=10$	$(.95,.50)$	.6991	.8728	2.220	.4384 $\to$ .3529
ACIC16	$(.95,.75)$	.7705	.7653	2.713	1.1332 $\to$ .7526
MorphoMNIST	$(.95,.50)$	.6663	.9811	1.332	.5055 $\to$ .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.

Counterfactual Prediction

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

See all 38 papers

Debiased Counterfactual Generation via Flow Matching from Observations
May 8, 2026Hugh Dance, Johnny Xi, Peter Orbanz +1Flow Matching Causal Counterfactual Generation
Gatsby Computational Neuroscience Unit University College London · Department of Statistics University of British Columbia
Counterfactual Generation via Flow Matching: Coupling-Sensitive End-to-End Rates
Oct 1, 2026Yunrui Guan, Krishnakumar Balasubramanian, Shiva Prasad KasiviswanathanFlow Matching Causal Counterfactual Generation
Department of Applied Mathematics and Statistics, Johns Hopkins University. · Amazon · Department of Statistics, University of California, Davis.
CAFE: Counterfactual Prediction via Fast Posterior Estimation
Sep 26, 2026Xinyan Han, Xiaoyu Lin, Hao Zou +3Counterfactual Prediction Structural Causal Models
Tsinghua University · Stable AI

May 8, 2026J/K move · Enter open · S save

Debiased Counterfactual Generation via Flow Matching from Observations

Hugh Dance, Johnny Xi, Peter Orbanz, Benjamin Bloem-Reddy

Gatsby Computational Neuroscience Unit University College London · Department of Statistics University of British Columbia

Flow Matching Causal Counterfactual Generation Causal Inference Counterfactual Learning

Open paper AskPDF arXiv

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.

Abstract

Counterfactual Prediction

Explore similar work

Debiased Counterfactual Generation via Flow Matching from Observations

Counterfactual Generation via Flow Matching: Coupling-Sensitive End-to-End Rates

CAFE: Counterfactual Prediction via Fast Posterior Estimation