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

Yunrui Guan; Krishnakumar Balasubramanian; Shiva Prasad Kasiviswanathan

cs.LGOct 1, 2026

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

Authors: Yunrui Guan, Krishnakumar Balasubramanian, Shiva Prasad Kasiviswanathan

Organizations: Department of Applied Mathematics and Statistics, Johns Hopkins University. · Amazon · Department of Statistics, University of California, Davis.

Abstract

Counterfactual generation seeks to sample outcomes under a hypothetical intervention or decision using observational data collected under the factual assignment mechanism. We develop a flow-matching approach that combines a sample-split, doubly robust training objective with a learned coupling between observed source outcomes and target outcomes drawn from a fitted conditional outcome model. To enable finite-step generation, we leverage a score-corrected stochastic sampler based on a Gaussian-smoothed interpolation. Our main theoretical contribution is a coupling-sensitive KL bound for constant-step Euler discretization: the error is controlled by moments of the source--target displacement under the chosen coupling, rather than by global uniform regularity of the velocity field, and has near-linear dependence on the ambient dimension. We also establish finite-sample non-parametric guarantees for the learned velocity and score fields when both the conditional outcome model and the source-target coupling are estimated from data. These bounds separate approximation, coupling-replacement, nuisance-estimation, generalization, and Monte Carlo errors and, combined with the sampler analysis, yield an end-to-end guarantee for counterfactual generation. Experiments on synthetic and semi-synthetic image benchmarks support the coupling-dependent theory and show that, at finite discretization budgets, the stochastic sampler can outperform the corresponding deterministic ODE sampler.

Figures & tables

Panel of Figure 1 — Figure 1: Angular-mode experiment. (a)–(b) Velocity and score errors across couplings ( $\sigma=0.12$ ). (c) Exact-field $8$ -step KL versus $D_{5}$ , using $2{,}000$ samples and $\tau=0.5\sqrt{\sigma D_{5}/d}$ . (d) Mode-mass allocation for $4$ and $64$ steps (EOT $\varepsilon_{\mathrm{OT}}=0.5$ , $\sigma=0.02$ ): $4$ -step ODE/SDE misallocate mass, while $64$ -step ODE/SDE match the target modes.

Panel of Figure 2 — Figure 2: ColorMNIST 1–6. (a) Clean-target SW2 for $\sigma\in\{0.01,0.015\}$ ( $5{,}000$ samples, five seeds; dashed: pure ODE). (b)–(d) Color evolution along the $\sigma=0.01$ SDE at $t=0,0.5,1$ ( $5{,}000$ samples; dashed: clean target).

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 4 — Figure 4: Fashion-MNIST shirts experiment. (a) Clean-target SW2 as a function of $\tau$ for the two smoothing levels; shaded regions show variation across five evaluation seeds, and the dashed line denotes the pure ODE baseline. (b)–(d) Color distributions at $t=0,0.5,1$ along the $\sigma=0.005$ SDE trajectory; dashed distributions denote the clean target.

Figure 5: Qualitative examples from the observed source distribution (top), samples generated by the $\sigma=0.005$ SDE (bottom). The generated samples preserve the boots shape and structure while showing clear, coherent color changes.

Figure 6: Qualitative examples from the observed source distribution (top), samples generated by the $\sigma=0.005$ SDE (bottom). The generated samples preserve the shirt shape and structure while showing clear, coherent color changes.

Work	Setting / Method	Metric	Dim.	Assumption	Coupling control
Stochastic interpolants and flow matching
This work Main theorem	constant- smoothing SI SDE, Euler	KL	$\widetilde{\mathcal{O}}(d)$	$L_{2}$ error	$D_{5}$
( Liu et al., 2025 ) Thm. 4.3; Sec. 5	SI SDE, Euler	KL	$\widetilde{\mathcal{O}}(d^{2})$	$L_{2}$ error	Partial
( Liu et al., 2026 ) Thm. 4.5	SI ODE, Euler	TV	$\widetilde{\mathcal{O}}(d^{2})$	drift/div. errors; unif. reg.	No
Diffusion-model baselines
( Benton et al., 2024 ) Cor. 1	diffusion exp. integrator	KL	$\widetilde{\mathcal{O}}(d)$	$L_{2}$ error	–

Appendix

Table 1: Comparison with related finite-step sampling results. Dimension dependence suppresses logarithmic factors and holds accuracy and the other problem parameters fixed; for our result this includes

\sigma,D_{5}

, with

\tau^{2}=\sigma D_{5}/d

and controlled field error.“Partial” means that source–target displacement enters the bound, but coupling-independent score terms remain.

Notation	Meaning	Notation	Meaning
$(C,A,Y)$	Observed covariate, treatment, and outcome.	$\mathcal{V}_{n},\mathcal{S}_{n}$	Clipped sparse-ReLU velocity class and score class induced by clipped sparse-ReLU denoisers.
$a$	Fixed target treatment value.	$R,\varepsilon$	Network input-clipping radius and approximation tolerance; field-specific choices appear in Proposition 1 .
$\mathbb{P}_{Y\mid C=c,A=a}$	Conditional outcome law at covariate value $c$ under treatment $a$ .	$\varepsilon_{\mathrm{vel}},\varepsilon_{\mathrm{score}}$	Integrated estimation errors of the learned velocity and score fields.
$\mu_{1}$	Target law $\int\mathbb{P}_{Y\mid C=c,A=a}\,d\mathbb{P}_{C}(c)$ .	$\tau$	Noise level of the score-corrected SDE sampler.
$\mathbb{P}_{Y\mid A=a}$	Observed outcome law among samples with treatment $a$ .	$N,h$	Number of sampling steps and Euler step size, with $h=1/N$ .
$\mathcal{D}_{\mathrm{tr}},\mathcal{D}_{\mathrm{aux}}$	Training and auxiliary splits, each containing $n$ observations.	$t_{k}$	Sampling grid point $t_{k}=kh$ .

Appendix

Table 2: Summary of notation used throughout the paper.

\displaystyle\widehat{L}_{\mathcal{B},\mathrm{DR}}(v)

\displaystyle=\frac{1}{|\mathcal{B}|}\sum_{i\in\mathcal{B}}\Bigg[\frac{\mathbf{1}\{A_{i}=a\}}{\widehat{\mathbb{P}}(A=a\mid C=C_{i})}\left\{\widehat{\ell}_{i}^{\mathrm{obs}}(v)-\frac{1}{M}\sum_{k=1}^{M}\widehat{\ell}_{i,k}^{\mathrm{plug}}(v)\right\}+\frac{1}{M}\sum_{k=1}^{M}\widehat{\ell}_{i,k}^{\mathrm{plug}}(v)\Bigg],

Appendix

Algorithm 2 Minibatch Doubly Robust Flow Training for a Fixed Treatment

a

Explore similar work

Debiased Counterfactual Generation via Flow Matching from Observations
May 8, 2026Hugh Dance, Johnny Xi, Peter Orbanz +1Counterfactual Generation Latent Confounders
Gatsby Computational Neuroscience Unit University College London · Department of Statistics University of British Columbia
Automatic, Debiased, and Invariant Counterfactual Generation under General Interventions
Jun 5, 2026Raphael C Kim, Jingsen Zhu, Ramin Zabih +1Counterfactual Generation Causal Inferences
Cornell Tech, Cornell University, New York, NY · Department of Biostatistics, Department of Population Health, New York University Grossman School of Medicine, New York, NY
Sharpen Your Flow: Sharpness-Aware Sampling for Flow Matching
May 12, 2026Aditi Gupta, Soon Hoe Lim, Annan Yu +1Conditional Flow Matching Accelerated Sampling
Lawrence Berkeley National Laboratory · International Computer Science Institute · 1Lawrence Berkeley National Laboratory, 2International Computer Science Institute +4

May 8, 2026, stat.MLJ/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

Counterfactual Generation Latent Confounders Influence Function

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

Explore similar work

Debiased Counterfactual Generation via Flow Matching from Observations

Automatic, Debiased, and Invariant Counterfactual Generation under General Interventions

Sharpen Your Flow: Sharpness-Aware Sampling for Flow Matching