PO-Flow: CNF for Causal Inference

• PO-Flow is a continuous normalizing flow framework for causal inference that models potential outcomes and counterfactuals without making parametric assumptions.
• It employs a flow matching training paradigm with invertible ODEs to enable efficient nonparametric density estimation and uncertainty quantification.
• Empirical evaluations on benchmarks and visual datasets demonstrate state-of-the-art performance, with error reductions of 20-30% over traditional methods.

PO-Flow is a continuous normalizing flow (CNF) framework for causal inference that enables joint modeling of potential outcomes and counterfactuals in observational data. By leveraging flow matching to train invertible ODE-based generative models, PO-Flow provides nonparametric density estimation for individual-level potential outcomes, interventional predictions, and factual-conditioned counterfactual queries—all with explicit uncertainty quantification. Uniquely, PO-Flow imposes no parametric assumptions on the target distributions and scales to high-dimensional output spaces, including vision applications. This approach yields state-of-the-art empirical performance on standard causal inference benchmarks, consistently reducing error and better capturing the true distributional structure of counterfactuals (Wu et al., 21 May 2025).

1. Causal Inference and the Need for Flexible Modeling

The Rubin model of causal inference frames each observational unit with covariates $X$ as possessing a pair of potential outcomes, $Y^{(0)}$ and $Y^{(1)}$, corresponding to untreated and treated conditions, respectively. Only the factual outcome $Y = Y^{(A)}$ (where $A \in \{0,1\}$ indicates treatment assignment) is observed, while the counterfactual outcome $Y^{(1-A)}$ is inherently missing. Estimating both potential outcomes and their entire conditional distributions, particularly conditioned on observed factuals and covariates, underpins treatment effect estimation, individualized policy evaluation, and prediction in disciplines such as medicine, economics, and public policy.

Traditional methods (e.g., S-learners, T-learners) generally fit regression models to predict means, occasionally extending to simple parametric distributional forms (e.g., mixtures of Gaussians). These approaches are limited in their ability to capture complex and non-Gaussian uncertainty—especially in high-dimensional and multimodal generative settings.

2. Continuous Normalizing Flows (CNFs) as a Generative Backbone

CNFs define invertible, time-continuous mappings from a base density (often $q(z) = \mathcal N(0, I)$) to a target conditional distribution $p(Y \mid X, A)$ via learned neural ODEs: $\frac{dy(t)}{dt} = v_\theta(y(t), t)$ The evolution in log-density is governed by the Liouville equation: $$\frac{d}{dt}\log p(y(t)) = -\nabla_y\!\cdot v_\theta(y(t), t)$$ This formalism yields exact, tractable likelihoods, full conditional density estimation, and invertibility for encoding and decoding information—a crucial property for counterfactual modeling conditioned on observed outcomes. Unlike fixed-form parametric models, CNFs can fit arbitrarily complex output distributions, supporting rigorous uncertainty quantification and sampling.

3. Flow Matching for Efficient Training

CNFs are typically trained via maximum-likelihood estimation, necessitating costly ODE solves through the likelihood path. PO-Flow replaces this with a flow matching paradigm, where a known reference velocity field $u(y, t)$ transports samples between base and target distributions: $$\mathcal L(\theta) = \mathbb E_{t, y_0, y_1} \| v_\theta(\phi_t, t, x, a) - (y_1 - y_0) \|^2$$ with interpolation

$\phi_t = (1-t) y_0 + t y_1$

and target velocity $u(\phi_t, t) = y_1 - y_0$. This regression objective sidesteps repeated ODE solves during training, rendering the method more computationally efficient and scalable to both large datasets and high-dimensional outputs.

4. Joint Modeling of Potential Outcomes, Counterfactuals, and Uncertainty

PO-Flow parameterizes the CNF velocity as a function of outcome variables, time, covariates, and treatment: $\frac{dy(t)}{dt} = v_\theta(y(t), t, x, a)$ This enables various queries:

• Interventional sampling: To draw samples from $p(Y^{(a)} \mid X)$, run the flow backward from a base sample at $t=1$ to $t=0$ under treatment $a$.
• Counterfactual prediction: Encode an observed factual outcome $y^{(a)}$ forward to the latent space, inject Gaussian uncertainty ($\delta \sim \mathcal N(0, \Sigma(x))$), then decode backward under the alternate treatment, yielding

$$\hat{y}_{\text{cf}} = z'_y - \int_0^1 v_\theta(y(t), t, x, 1-a)\,dt$$

This mechanism induces a conditional predictive distribution for the counterfactual outcome given the factual and covariates.

As a CNF, PO-Flow provides per-sample log-probabilities along generative trajectories, supporting entropy-based confidence measures, credible interval estimation, and sample selection via highest-likelihood scoring. Divergence terms are efficiently estimated using Hutchinson's trace.

5. Empirical Performance on Causal Inference Benchmarks

PO-Flow has been comprehensively evaluated on the following datasets:

• ACIC 2016/2018 (semi-synthetic, high-dimensional tabular covariates)
• IHDP (747 units, 25 covariates)
• IBM Causal Inference Benchmark (1,000 units, 177 covariates)

Performance is measured by:

• Root mean squared error (RMSE) for mean potential outcome predictions
• Root PEHE (Precision in Estimation of Heterogeneous Effect) for conditional average treatment effect (CATE)
• RMSE for factual-conditioned counterfactuals
• KL divergence or Wasserstein-1 distance to true potential outcome distributions
• Absolute error in estimating the average treatment effect (ATE)

PO-Flow outperforms diffusion-based DiffPO, discrete-flow INFs, standard meta-learners (S-/T-learner), representation-based CFR, latent-variable CEVAE, and GANITE. It achieves ≈20–30% lower RMSE for POs, the lowest PEHE and ATE error, and significantly reduced divergence distances to ground-truth distributions across all tested scenarios.

6. Extension to High-Dimensional and Visual Counterfactuals

PO-Flow's scalability is demonstrated through high-dimensional counterfactual generation on the CelebA image dataset. Images ($178 \times 218$) are projected to 512-dimensional VAE latents, with binary semantic attributes ($X$) and categorical treatments ($A$). Training a lightweight MLP-based velocity field $v_\theta$, PO-Flow can synthesize realistic facial attribute transformations (e.g., male-to-female), with trajectory smoothness and diversity via latent space perturbations. This application illustrates the method's capacity to capture meaningful multimodal and nonparametric variations in complex real-world outputs.

7. Strengths, Limitations, and Future Directions

PO-Flow consolidates potential-outcome estimation, CATE computation, and factual-conditioned counterfactual prediction within a single principled generative framework. Its primary strengths include exact density learning without distributional assumptions, computationally efficient flow matching training, generalization to high-dimensional settings, and minimal hyperparameter sensitivity.

Limitations include the current reliance on Gaussian noise for uncertainty injection in counterfactual prediction; more expressive, possibly non-Gaussian perturbations may better capture correlations between potential outcomes. The framework does not recover explicit structural causal models (SCMs), and currently lacks formal support for time-varying treatments or multi-armed interventions. Prospective research avenues comprise the integration of SCM components, exploration of stochastic interpolants, alternative flow architectures, richer noise models, and multi-treatment extensions.

PO-Flow constitutes a methodological advancement by bringing the expressive power of continuous normalizing flows to the field of counterfactual inference, supporting accurate point estimation and principled, uncertainty-aware density prediction for a range of causal questions (Wu et al., 21 May 2025).