# Flow-based Generative Modeling of Potential Outcomes and Counterfactuals

arXiv: https://arxiv.org/abs/2505.16051 · HTML: https://mlsanity.com/paper/2505.16051/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 26 of its 45 references that are in the map.
- First posted on arXiv 21 May 2025.

## Builds on (most cited first, 26 of 26)

- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Building Normalizing Flows with Stochastic Interpolants](https://mlsanity.com/paper/2209.15571/) (2022)
- [Meta-learners for Estimating Heterogeneous Treatment Effects using Machine Learning](https://mlsanity.com/paper/1706.03461/) (2017)
- [Causal Effect Inference with Deep Latent-Variable Models](https://mlsanity.com/paper/1705.08821/) (2017)
- [Learning Representations for Counterfactual Inference](https://mlsanity.com/paper/1605.03661/) (2016)
- [Deep Structural Causal Models for Tractable Counterfactual Inference](https://mlsanity.com/paper/2006.06485/) (2020)
- [Bayesian Inference of Individualized Treatment Effects using Multi-task Gaussian Processes](https://mlsanity.com/paper/1704.02801/) (2017)
- [Conformal Inference of Counterfactuals and Individual Treatment Effects](https://mlsanity.com/paper/2006.06138/) (2020)
- [The probability flow ODE is provably fast](https://mlsanity.com/paper/2305.11798/) (2023)
- [Nonparametric Estimation of Heterogeneous Treatment Effects: From Theory to Learning Algorithms](https://mlsanity.com/paper/2101.10943/) (2021)
- [Foundations of Structural Causal Models with Cycles and Latent Variables](https://mlsanity.com/paper/1611.06221/) (2016)
- [Causal Transformer for Estimating Counterfactual Outcomes](https://mlsanity.com/paper/2204.07258/) (2022)
- [On Inductive Biases for Heterogeneous Treatment Effect Estimation](https://mlsanity.com/paper/2106.03765/) (2021)
- [Causal machine learning for predicting treatment outcomes](https://mlsanity.com/paper/2410.08770/) (2024)
- [Generalization Bounds and Representation Learning for Estimation of Potential Outcomes and Causal Effects](https://mlsanity.com/paper/2001.07426/) (2020)
- [Counterfactual Identifiability of Bijective Causal Models](https://mlsanity.com/paper/2302.02228/) (2023)
- [Normalizing flow neural networks by JKO scheme](https://mlsanity.com/paper/2212.14424/) (2022)
- [Identifying Patient-Specific Root Causes with the Heteroscedastic Noise Model](https://mlsanity.com/paper/2205.13085/) (2022)
- [Benchmarking Framework for Performance-Evaluation of Causal Inference Analysis](https://mlsanity.com/paper/1802.05046/) (2018)
- [Normalizing Flows for Interventional Density Estimation](https://mlsanity.com/paper/2209.06203/) (2022)
- [DiffPO: A causal diffusion model for learning distributions of potential outcomes](https://mlsanity.com/paper/2410.08924/) (2024)
- [Meta-Learners for Partially-Identified Treatment Effects Across Multiple Environments](https://mlsanity.com/paper/2406.02464/) (2024)
- [Annealing Flow Generative Models Towards Sampling High-Dimensional and Multi-Modal Distributions](https://mlsanity.com/paper/2409.20547/) (2024)
- [Comparison of meta-learners for estimating multi-valued treatment heterogeneous effects](https://mlsanity.com/paper/2205.14714/) (2022)

## Built on by (most cited first, 3 of 3)

- [Intervention-Based Time Series Causal Discovery via Simulator-Generated Interventional Distributions](https://mlsanity.com/paper/2605.09870/) (2026)
- [Extended Wasserstein-GAN Approach to Causal Distribution Learning: Density-Free Estimation and Minimax Optimality](https://mlsanity.com/paper/2605.10206/) (2026)
- [OrthoGen: A Generative Orthogonal Learner for Time-Varying Treatments](https://mlsanity.com/paper/2610.10210/) (2026)

## Abstract (arXiv)

Predicting potential and counterfactual outcomes from observational data is central to individualized decision-making, particularly in clinical settings where treatment choices must be tailored to each patient rather than guided solely by population averages. We propose PO-Flow, a continuous normalizing flow (CNF) framework for causal inference that jointly models potential outcome distributions and factual-conditioned counterfactual outcomes. Trained via flow matching, PO-Flow provides a unified approach to individualized potential outcome prediction, conditional average treatment effect estimation, and counterfactual prediction. By encoding an observed factual outcome and decoding under an alternative treatment, PO-Flow provides an encode-decode mechanism for factual-conditioned counterfactual prediction. In addition, PO-Flow supports likelihood-based evaluation of potential outcomes, enabling uncertainty-aware assessment of predictions. A supporting recovery guarantee is established under certain assumptions, and empirical results on benchmark datasets demonstrate strong performance across a range of causal inference tasks within the potential outcomes framework.