PO-Flow: Flow-based Generative Models for Sampling Potential Outcomes and Counterfactuals
Dongze Wu, David I. Inouye, Yao Xie
TL;DR
PO-Flow formulates potential outcomes and counterfactual reasoning within a continuous normalizing flow, enabling joint PO prediction, CATE estimation, and conditional counterfactual generation directly from observational data. Trained with Flow Matching, it learns the full PO densities and provides likelihood-based evaluation, uncertainty quantification, and most likely PO selection via log-densities. The framework includes a counterfactual recovery guarantee under certain monotone SCM assumptions and does not rely on explicitly specifying a structural causal model. Empirically, PO-Flow achieves state-of-the-art performance across multiple benchmark datasets while offering efficient training, fast sampling, and integrated density-based analyses for robust clinical decision support.
Abstract
Predicting potential and counterfactual outcomes from observational data is central to clinical decision-making, where physicians must weigh treatments for an individual patient rather than relying solely on average effects at the population level. We propose PO-Flow, a continuous normalizing flow (CNF) framework for causal inference that jointly models potential outcomes and counterfactuals. Trained via flow matching, PO-Flow provides a unified approach to average treatment effect estimation, individualized potential outcome prediction, and counterfactual prediction. Besides, PO-Flow directly learns the densities of potential outcomes, enabling likelihood-based evaluation of predictions. Furthermore, PO-Flow explores counterfactual outcome generation conditioned on the observed factual in general observational datasets, with a supporting recovery result under certain assumptions. PO-Flow outperforms modern baselines across diverse datasets and causal tasks in the potential outcomes framework.