Uncertainty-Aware Distribution-to-Distribution Flow Matching for Scientific Imaging

Dongxia Wu; Yuhui Zhang; Serena Yeung-Levy; Emma Lundberg; Emily B. Fox

cs.LGMar 23, 2026

Uncertainty-Aware Distribution-to-Distribution Flow Matching for Scientific Imaging

Authors: Dongxia Wu, Yuhui Zhang, Serena Yeung-Levy, Emma Lundberg, Emily B. Fox

Organizations: Stanford University Stanford, CA

Abstract

Distribution-to-distribution generative models support scientific imaging tasks ranging from modeling cellular perturbation responses to translating medical images across conditions. Trustworthy generation requires reliability, or generalization across labs, devices, and experimental conditions, and accountability, or detecting out-of-distribution cases where predictions may be unreliable. We leverage Stochastic Flow Matching (SFM), a marginal-preserving stochastic extension of flow matching for improved generalization under distribution shift. SFM augments deterministic flows with a diffusion term together with a learned score-based drift correction, retaining the learned transport marginals while modeling conditional variability. Building on this SFM framework, we introduce Bayesian Stochastic Flow Matching (BSFM) as a companion uncertainty quantification mechanism and develop AVUQ (Antithetic Variance-reduction Uncertainty Quantification) to approximately estimate epistemic and aleatoric uncertainty via sample-efficient antithetic sampling with approximate posterior inference. We further use AVUQ to yield anomaly scores for unreliable generation detection. Experiments on cellular imaging (BBBC021, JUMP) and brain fMRI (Theory of Mind) across diverse unseen scenarios show that SFM improves generalization while AVUQ provides effective uncertainty-based anomaly scores under practical sampling budgets.

Figures & tables

Figure 1: Trustworthy distribution-to-distribution generative modeling. Left: In-domain samples are drawn from the source distribution $p_{0}$ , while unseen samples may arise from shifted source distributions $\{p_{0}^{\prime}\}$ . Right: Our approach centers on Stochastic Flow Matching (SFM) to improve model generalization, with Bayesian Stochastic Flow Matching (BSFM) and AVUQ providing approximate aleatoric and epistemic uncertainty estimates for unreliable generation detection.

Figure 2: Core components of the Bayesian Stochastic Flow Matching (BSFM) framework: (a) Stochastic Flow Matching (SFM) training process. (b) MC-Dropout as one approximate inference method that can be used by AVUQ. (c) AVUQ with antithetic sampling.

	BBBC021				JUMP				fMRI
Method	Unseen Pert.		Intensity Shift		Unseen Cell Lines		Unseen Plates		Low ToM
	FID $\downarrow$	KID $\downarrow$	FID $\downarrow$	KID $\downarrow$	FID $\downarrow$	KID $\downarrow$	FID $\downarrow$	KID $\downarrow$	FID $\downarrow$	KID $\downarrow$
BBDM	107.69	11.88	86.34	9.25	126.41	13.25	107.24	10.68	45.31	3.55
GOUB	51.89	4.71	36.94	3.46	62.47	6.62	57.66	6.14	33.49	2.63
UniDB	50.16	4.33	36.78	3.06	65.1	6.9	60.33	6.34	30.23	1.85
UNSB	88.31	6.27	56.95	4.7	57.24	4.19	45.5	4.04	75.54	8.89

Table 1: Generalization performance. We report FID and KID scores across unseen scenarios on BBBC021, JUMP, and fMRI datasets. Best results are in bold , second best are underlined .

	BBBC021				JUMP
Method	Unseen Pert.		Intensity Shift		Unseen Cell Lines		Intensity Shift
	AUROC $\uparrow$	AUPR $\uparrow$	AUROC $\uparrow$	AUPR $\uparrow$	AUROC $\uparrow$	AUPR $\uparrow$	AUROC $\uparrow$	AUPR $\uparrow$
non-nested
MAP Aleatoric	0.691 $\pm$ 0.033	0.482 $\pm$ 0.063	0.711 $\pm$ 0.018	0.637 $\pm$ 0.023	0.799 $\pm$ 0.021	0.746 $\pm$ 0.024	0.276 $\pm$ 0.07	0.325 $\pm$ 0.017
MCD-DFM Epistemic	0.621 $\pm$ 0.056	0.332 $\pm$ 0.047	0.654 $\pm$ 0.051	0.532 $\pm$ 0.05	0.645 $\pm$ 0.03	0.545 $\pm$ 0.02	0.425 $\pm$ 0.047	0.368 $\pm$ 0.017
nested-i.i.d.

Table 2: Unreliable generation detection performance. We report AUROC and AUPR scores as mean

\pm

standard error over 3 random seeds across 4 unreliable-generation scenarios on BBBC021 and JUMP datasets. We group methods into non-nested baselines, nested i.i.d. baselines, and nested-antithetic AVUQ.

Figure 3: Examples of generated images from different methods under various unseen scenarios, compared with source and ground-truth target images.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	BBBC021	BBBC021	JUMP	JUMP	JUMP	fMRI
OOD Case	Unseen Pert.	Intensity Shift	Unseen Cell Lines	Unseen Plates	Intensity Shift	Low ToM
# Channels	3	3	5	5	5	1
# Images	98K	98K	72K	72K	72K	12K
Dist Shift	Extreme	High	Medium	Low	High	-

Appendix

Table 3: Summary of datasets and OOD scenarios.

	BBBC021				JUMP
Method	Unseen Pert.		Intensity Shift		Unseen Cell Lines		Unseen Plates		Intensity Shift
	AUROC $\uparrow$	AUPR $\uparrow$	AUROC $\uparrow$	AUPR $\uparrow$	AUROC $\uparrow$	AUPR $\uparrow$	AUROC $\uparrow$	AUPR $\uparrow$	AUROC $\uparrow$	AUPR $\uparrow$
2x2 Aleatoric	0.673	0.404	0.679	0.572	0.838	0.792	0.735	0.662	0.303	0.331
2x2 Epistemic	0.752	0.551	0.785	0.715	0.743	0.678	0.628	0.595	0.826	0.710
2x4 Aleatoric	0.718	0.469	0.715	0.647	0.852	0.818	0.747	0.673	0.363	0.349
2x4 Epistemic	0.725	0.597	0.784	0.763	0.759	0.709	0.614	0.542	0.846	0.737

Appendix

Table 4: Sample-size sensitivity for AVUQ. Each row reports AUROC/AUPR for a posterior-sample

\times

SDE-sample budget.

	BBBC021				JUMP
Method	Unseen Pert.		Intensity Shift		Unseen Cell Lines		Unseen Plates		Intensity Shift
	AUROC $\uparrow$	AUPR $\uparrow$	AUROC $\uparrow$	AUPR $\uparrow$	AUROC $\uparrow$	AUPR $\uparrow$	AUROC $\uparrow$	AUPR $\uparrow$	AUROC $\uparrow$	AUPR $\uparrow$
MAP Aleatoric	0.616	0.443	0.636	0.626	0.670	0.627	0.517	0.519	0.227	0.351
SWAG Aleatoric	0.684	0.589	0.680	0.729	0.699	0.647	0.522	0.520	0.813	0.712
SWAG Epistemic	0.661	0.572	0.680	0.730	0.689	0.638	0.516	0.518	0.806	0.702
LA Aleatoric	0.662	0.574	0.671	0.727	0.684	0.646	0.527	0.528	0.804	0.707

Appendix

Table 5: Unreliable generation detection results without thresholding.

	BBBC021				JUMP
Method	Unseen Pert.		Intensity Shift		Unseen Cell Lines		Unseen Plates		Intensity Shift
	AUROC $\uparrow$	AUPR $\uparrow$	AUROC $\uparrow$	AUPR $\uparrow$	AUROC $\uparrow$	AUPR $\uparrow$	AUROC $\uparrow$	AUPR $\uparrow$	AUROC $\uparrow$	AUPR $\uparrow$
MAP Aleatoric	0.639	0.375	0.667	0.547	0.767	0.698	0.638	0.574	0.286	0.319
SWAG Aleatoric	0.701	0.550	0.670	0.683	0.794	0.728	0.615	0.562	0.813	0.712
SWAG Epistemic	0.676	0.523	0.678	0.684	0.782	0.714	0.627	0.574	0.879	0.797
LA Aleatoric	0.674	0.534	0.654	0.678	0.779	0.721	0.630	0.583	0.875	0.793

Appendix

Table 6: Unreliable generation detection results with mean threshold

\mu

Figure 4: Examples of generated images from different methods on BBBC021 under the Unseen Pert. OOD scenario, compared with source images and ground-truth target images.

Figure 5: Examples of generated images from different methods on BBBC021 under the Intensity Shift OOD scenario, compared with source images and ground-truth target images.

Figure 6: Examples of generated images from different methods on JUMP under the Unseen Cell Lines OOD scenario, compared with source images and ground-truth target images.

Figure 7: Examples of generated images from different methods on JUMP under the Unseen Plates OOD scenario, compared with source images and ground-truth target images.

Explore similar work

Flow Matching with Uncertainty Quantification and Guidance
Feb 10, 2026Juyeop Han, Lukas Lao Beyer, Sertac KaramanFlow Matching Image Generation
MIT
Divergence is Uncertainty: A Closed-Form Posterior Covariance for Flow Matching
May 1, 2026Jiarui Xing, Song Wang, Jian WangFlow Matching Uncertainty Quantification
Yale University · Harvard Medical School
Local Flow Matching Generative Models
Oct 3, 2024Chen Xu, Xiuyuan Cheng, Yao XieFlow Matching Generative Modeling
H. Milton Stewart School of Industrial and Systems Engineering, Georgia Institute of Technology. · Department of Mathematics, Duke University

Feb 10, 2026J/K move · Enter open · S save

Flow Matching with Uncertainty Quantification and Guidance

Juyeop Han, Lukas Lao Beyer, Sertac Karaman

MIT

Flow Matching Image Generation Uncertainty Quantification Classifier-Free Guidance

Open paper AskPDF arXiv

Despite the remarkable success of sampling-based generative models such as flow matching, they can still produce samples of inconsistent or degraded quality. To assess sample reliability and improve sample fidelity, we propose uncertainty-aware flow matching (UA-Flow), a lightweight extension of flow matching that predicts the velocity field together with heteroscedastic uncertainty. UA-Flow estimates spatially localized sample uncertainty by propagating velocity uncertainty through the flow dynamics. The resulting estimates serve as reliability signals that reflect data-manifold support for each generated sample or region. We further use them to steer generation via uncertainty-aware guidance (UG) and uncertainty-aware classifier-free guidance (UCFG). Experiments on ImageNet and CIFAR-10 show that UA-Flow's uncertainty tracks sample fidelity more consistently than element-wise baselines and that these signals also carry information about data-manifold support. On ImageNet, UG increases precision at the cost of recall and improves FID, while UCFG maintains lower FID than CFG at large guidance scales.

Abstract

Explore similar work

Flow Matching with Uncertainty Quantification and Guidance

Divergence is Uncertainty: A Closed-Form Posterior Covariance for Flow Matching

Local Flow Matching Generative Models