PO-Flow: Flow-based Generative Models for Sampling Potential Outcomes and Counterfactuals

2025-05-21

AI Summary

This work addresses the challenge of jointly modeling the density of potential outcomes and counterfactuals without assuming explicit distributional forms. We propose a causal generative framework based on Continuous Normalizing Flows (CNFs), trained via flow matching and guided by causal graph–informed conditional density estimation. Our method enables end-to-end individualized prediction of potential outcomes, observation-conditional counterfactual inference, and uncertainty-aware density learning. Crucially, we extend CNFs to high-dimensional, image-level counterfactual generation—a first in the literature. Extensive evaluation on benchmark datasets (ACIC, IHDP, IBM) demonstrates consistent superiority over state-of-the-art methods in both accuracy and calibration. Furthermore, experiments on real-world image data confirm strong generalization capability and interpretability. The framework thus bridges rigorous density estimation with structural causal modeling, advancing both theoretical foundations and practical applicability of counterfactual reasoning in complex, high-dimensional domains.

Technology Category

• Machine Learning: Causal Learning
• Reasoning under Uncertainty: Causality
• Computer Vision: Generative Adversarial Networks (GANs) for Vision

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications

Abstract

We propose PO-Flow, a novel continuous normalizing flow (CNF) framework for causal inference that jointly models potential outcomes and counterfactuals. Trained via flow matching, PO-Flow provides a unified framework for individualized potential outcome prediction, counterfactual predictions, and uncertainty-aware density learning. Among generative models, it is the first to enable density learning of potential outcomes without requiring explicit distributional assumptions (e.g., Gaussian mixtures), while also supporting counterfactual prediction conditioned on factual outcomes in general observational datasets. On benchmarks such as ACIC, IHDP, and IBM, it consistently outperforms prior methods across a range of causal inference tasks. Beyond that, PO-Flow succeeds in high-dimensional settings, including counterfactual image generation, demonstrating its broad applicability.

Problem

• Jointly models potential outcomes and counterfactuals for causal inference
• Enables density learning without explicit distributional assumptions
• Supports high-dimensional settings like counterfactual image generation

Innovation

• Flow-based generative model for causal inference
• Unified framework for outcome and counterfactual prediction
• Density learning without explicit distributional assumptions