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 p0 , while unseen samples may arise from shifted source distributions {p0′} . 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 ↓
KID ↓
FID ↓
KID ↓
FID ↓
KID ↓
FID ↓
KID ↓
FID ↓
KID ↓
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 ↑
AUPR ↑
AUROC ↑
AUPR ↑
AUROC ↑
AUPR ↑
AUROC ↑
AUPR ↑
non-nested
MAP Aleatoric
0.691 ± 0.033
0.482 ± 0.063
0.711 ± 0.018
0.637 ± 0.023
0.799 ± 0.021
0.746 ± 0.024
0.276 ± 0.07
0.325 ± 0.017
MCD-DFM Epistemic
0.621 ± 0.056
0.332 ± 0.047
0.654 ± 0.051
0.532 ± 0.05
0.645 ± 0.03
0.545 ± 0.02
0.425 ± 0.047
0.368 ± 0.017
nested-i.i.d.
Table 2: Unreliable generation detection performance. We report AUROC and AUPR scores as mean ± 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 ↑
AUPR ↑
AUROC ↑
AUPR ↑
AUROC ↑
AUPR ↑
AUROC ↑
AUPR ↑
AUROC ↑
AUPR ↑
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 × SDE-sample budget.
BBBC021
JUMP
Method
Unseen Pert.
Intensity Shift
Unseen Cell Lines
Unseen Plates
Intensity Shift
AUROC ↑
AUPR ↑
AUROC ↑
AUPR ↑
AUROC ↑
AUPR ↑
AUROC ↑
AUPR ↑
AUROC ↑
AUPR ↑
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 ↑
AUPR ↑
AUROC ↑
AUPR ↑
AUROC ↑
AUPR ↑
AUROC ↑
AUPR ↑
AUROC ↑
AUPR ↑
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 μ .
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.
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.
Flow matching has become a leading framework for generative modeling, but quantifying the uncertainty of its samples remains an open problem. Existing approaches retrain the model with auxiliary variance heads, maintain costly ensembles, or propagate approximate covariance through many integration steps, trading off training cost, inference cost, or accuracy. We show that none of these trade-offs is necessary. By extending Tweedie's formula from the denoising setting to the flow matching interpolant, we derive an exact, closed-form expression for the posterior covariance at every point along the generative trajectory. The result depends on a single quantity, namely the divergence of the learned velocity field, which can be computed post-hoc on any pre-trained flow matching model, requiring no retraining and no architectural modification. For one-step generators such as MeanFlow, the same formula yields the end-to-end generation uncertainty in a single forward pass, eliminating the multi-step variance propagation required by all prior methods. Experiments on MNIST confirm that the resulting per-pixel uncertainty maps are semantically meaningful, concentrating on digit boundaries where inter-sample variation is highest, and that the scalar uncertainty score tracks actual prediction error, all at roughly 104× less total compute than ensembling or Monte Carlo dropout.
Flow Matching (FM) is a simulation-free method for learning a continuous, invertible flow that interpolates between two distributions, and in particular generates data from noise. Inspired by the variational nature of the diffusion process as a gradient flow, we introduce a stepwise FM model, Local Flow Matching (LFM), which sequentially learns a sequence of FM submodels, each matching a diffusion process up to the time-step size in the data-to-noise direction. In each step, the two distributions to be interpolated by the sub-flow model are closer than those in the full-flow matching model, which interpolates data to noise distributions, enabling smaller models with more efficient training. This variational perspective also allows us to prove a theoretical generation guarantee for the proposed flow model in terms of the χ2-divergence between the generated and true data distributions, leveraging the contraction property of the diffusion process. In practice, the stepwise structure of LFM is naturally amenable to model distillation, and various distillation techniques can be applied to accelerate generation. We empirically demonstrate that LFM achieves competitive generative performance compared to FM on unconditional generation of tabular and image datasets, and on conditional generation of robotic manipulation policies.
Chen Xu, Xiuyuan Cheng, Yao Xie
H. Milton Stewart School of Industrial and Systems Engineering, Georgia Institute of Technology. · Department of Mathematics, Duke University