Flow Matching with Uncertainty Quantification and Guidance

Juyeop Han; Lukas Lao Beyer; Sertac Karaman

cs.CVFeb 10, 2026

Flow Matching with Uncertainty Quantification and Guidance

Authors: Juyeop Han, Lukas Lao Beyer, Sertac Karaman

Organizations: MIT

Abstract

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.

Figures & tables

Panel of Figure 1 — Figure 1: UG and UCFG sweep on ImageNet-256 (class guinea pig ). Left: generated samples; right: predicted latent pixel-wise uncertainty (brighter indicates higher uncertainty). Rows sweep the UG scale $w\in\{0,10,30,50\}$ , columns sweep the maximum UCFG scale $\lambda_{\max}\in\{0,1,2,5,10,20\}$ . Stronger guidance yields more class-consistent samples with lower predicted uncertainty.

$\bar{u}^{\theta}_{t}(\cdot)$ :	velocity mean	$\bar{\mathbf{x}}_{t}$ :	state mean
$\mathrm{Var}[\mathbf{x}_{t}]$ :	state variance	$S$ :	number of probes

Algorithm 1 Estimating

\operatorname{diag}(J^{\theta}_{t}(\bar{\mathbf{x}}_{t}))\odot\mathrm{Var}[\mathbf{x}_{t}]

Figure 2: Generative quality metrics vs. filtering ratio on ImageNet-256. At each filtering level, high-uncertainty samples are removed and 50k samples are drawn from the remainder for evaluation. Compared to the element-wise baselines AU ( De Vita and Belagiannis, 2025 ) and BayesDiff ( Kou et al., 2023 ) , UA-Flow achieves lower FID and higher precision after filtering. UA-Flow + CLIP further matches the CLIP-based scalar reference GenUnc ( Jazbec et al., 2025 ) .

Panel of Figure 3 — Figure 3: Uncertainty maps from UA-Flow, AU, and BayesDiff on generated samples. Brighter values indicate higher uncertainty; maps are normalized per image. UA-Flow localizes high-uncertainty regions, whereas AU often produces inverted patterns and BayesDiff produces noisy maps.

Method	CIFAR-10	ImageNet-128	ImageNet-256
Vanilla	4.044	1.112	4.601
AU	5.014	1.379	5.706
BayesDiff	9.201	2.479	10.26
GenUnc	24.32	7.655	31.39
UA-Flow	5.506	1.380	5.761
UA-Flow + CLIP	5.550	2.200	8.916

Table 1: Computational cost (TFLOPS) for sampling and uncertainty quantification per image. Vanilla reports sampling cost only.

		Sample-wise		Patch-wise
Dataset	Metric	GenUnc	UA-Flow + CLIP	AU	BayesDiff	UA-Flow
CIFAR-10	Dist. to manifold	$\mathbf{0.3462}$	$0.2421$	$-0.5654$	$0.5064$	$\mathbf{0.6114}$
CIFAR-10	Rarity	$\mathbf{0.3184}$	$0.2183$	$-0.5603$	$0.4803$	$\mathbf{0.6066}$
ImageNet-128	Dist. to manifold	$0.3579$	$\mathbf{0.3672}$	$-0.3187$	$0.1340$	$\mathbf{0.4609}$
ImageNet-128	Rarity	$0.3501$	$\mathbf{0.3901}$	$-0.2925$	$0.1078$	$\mathbf{0.4037}$
ImageNet-256	Dist. to manifold	$0.4190$	$\mathbf{0.4426}$	$-0.4719$	$0.1820$	$\mathbf{0.5404}$

Table 2: Spearman rank correlation

\rho

between predicted uncertainty and manifold support. Sample-wise methods produce a single CLIP-space scalar per image, whereas patch-wise methods produce element-wise uncertainty maps and are correlated at the patch level. Sample-wise, UA-Flow + CLIP is strongest on ImageNet-128 and ImageNet-256, while GenUnc leads on CIFAR-10. Patch-wise, UA-Flow attains the strongest correlation with both the distance to the data manifold (Dist. to manifold) and the rarity score on every dataset.

Figure 4: FID, precision, and recall as the UG scale

w

varies on ImageNet-256 under fixed CFG scales

\lambda\in\{0,0.25,0.5\}

. Increasing

w

generally improves precision while reducing recall. At nonzero CFG scales, FID reaches a minimum at an intermediate UG strength and then increases.

Panel of Figure 5 — Figure 5: ImageNet-256 samples under different UG scales $w$ at a fixed CFG scale $\lambda=0.5$ . Each panel uses a fixed class label and random seed; columns sweep $w\in\{0,10,30,50\}$ from left to right. UG trades diversity for fidelity as the guidance strength increases.

Table 3: FID, precision and recall under uncertainty-aware guidance (UG) across datasets. For each CFG scale

\lambda

, we report the best-performing UG setting (lowest FID). AU denotes uncertainty-aware guidance using aleatoric uncertainty ( De Vita and Belagiannis, 2025 ) and is included as a reference baseline. UG consistently improves precision compared to vanilla sampling under matched CFG settings, and it also yields better FID on ImageNet-128 and ImageNet-256.

Figure 6: FID, precision, and recall as a function of the fixed CFG scale

\lambda

or the maximum UCFG scale

\lambda_{\max}

on ImageNet-256. CFG degrades sharply at large

\lambda

, while UCFG remains more stable as

\lambda_{\max}

increases.

Table 4: FID, precision and recall under CFG variants on ImageNet. Best-performing scale (lowest FID) per method, averaged over 3 seeds.

Panel of Figure 7 — Figure 7: Distributions of the adaptive guidance scale $\lambda^{*}$ across sampling steps on ImageNet-256. Violin plots are computed from 1,000 samples for (a) UCFG and (b) APG+UCFG. In both cases, $\lambda^{*}$ tends to be smaller in early steps and larger in later steps.

Panel of Figure 8 — Figure 8: ImageNet-256 samples under increasing classifier-free guidance scales. Rows show CFG, APG, UCFG, and APG+UCFG from top to bottom; columns sweep the fixed scale $\lambda$ or adaptive-scale cap $\lambda_{\max}$ over $\{1,2,5,10,20\}$ from left to right. Compared to CFG and APG at large fixed scales $\lambda$ , UCFG and APG+UCFG preserves structure more robustly under large $\lambda_{\max}$ .

Batch size $B$	Average bias-correction $U_{t}$	Recovered
Reference	$-0.579$	—
2	$-0.100$	17.21%
16	$-0.273$	47.21%
64	$-0.349$	60.35%
256	$-0.407$	70.21%
1,024	$-0.454$	78.32%

Table 5: Batch-size sensitivity of the correction estimator on ImageNet-256 latents. Average bias-correction

U_{t}

measured over

t\in[0,0.1]

, where the correction is largest, together with the fraction of the reference value it recovers. The reference replaces the mini-batch estimate with a

65{,}536

-candidate estimate.

	Before		After
$t$	ECE $\downarrow$	Brier $\downarrow$	ECE $\downarrow$	Brier $\downarrow$
0.5	0.1328	0.4771	0.0066	0.0515
0.7	0.0992	0.3281	0.0038	0.0492
0.9	0.0277	0.1009	0.0008	0.0481

Table 6: Calibration of UA-Flow on ImageNet-256 validation images . Scores before/after post-hoc isotonic calibration at three noise levels

t

. Details and baselines in Section G.3 .

Threshold $\tau$	Excl. %	Success % $\uparrow$	Mean MaxR $\uparrow$
no removal	0.0%	62.7%	0.9514
0.0010	87.8%	91.0 %	0.9702
0.0025	62.9%	77.6%	0.9597
0.0050	29.2%	66.8%	0.9536
0.0075	14.6%	63.6%	0.9536
0.0100	7.9%	63.3%	0.9521

Table 7: Push-T early termination. Rollouts whose chunk uncertainty exceeds

\tau

are excluded. Excl. % is the excluded fraction of

1{,}000

rollouts, and Success % (maximum reward

>1.0

) and Mean MaxR are measured over those retained. Full table in Table 15 .

Appendix figures & tables25 assets

Supplementary material from the paper’s appendix.

Appendix

	CIFAR-10		ImageNet-128		ImageNet-256
	pretraining	finetuning	pretraining	finetuning	finetuning
Learning rate	$1e\text{-}4$	$(1e\text{-}5,1e\text{-}6)$	$2e\text{-}4$	$(2e\text{-}5,2e\text{-}5)$	$(1e\text{-}5,1e\text{-}5)$
AdamW $(\beta_{1},\beta_{2})$	$(0.9,0.95)$	$(0.9,0.999)$	$(0.9,0.999)$	$(0.9,0.999)$	$(0.9,0.999)$
Gradient warm-up step	–	10000	–	1000	1000
Gradient clipping	–	1.0	–	1.0	1.0
EMA decay rate	0.999	0.999	0.9999	0.9999	0.9999

Appendix

Table 8: Hyper-parameters. When two learning rates are reported as

(\eta_{\text{head}},\eta_{\text{backbone}})

, they denote the head and backbone learning rates, respectively.

	CIFAR-10
# of ResNet blocks per scale	4
Base channels	128
Channel multiplier per scale	(2, 2, 2)
Attention resolutions	2
Dropout	0.3

Appendix

Table 9: ADM configuration for CIFAR-10.

Panel of Figure 9 — Figure 9: Filtering high-uncertainty samples across datasets. Generative quality metrics as a function of the filtering ratio. On CIFAR-10, uncertainty-based filtering primarily induces a precision-recall trade-off that can negatively affect FID despite high overall sample quality. On ImageNet-128, the trends are consistent with ImageNet-256: UA-Flow achieves lower FID and higher precision after filtering compared to AU and BayesDiff; GenUnc is included as a reference baseline.

Panel of Figure 10 — Figure 10: Generation quality metrics as the UG scale $w$ varies on CIFAR-10 and ImageNet-128 under fixed classifier-free guidance (CFG) scales. Increasing $w$ induces a fidelity–diversity trade-off, improving precision while reducing recall. The FID response depends on the dataset and the fixed CFG scale.

Setting	FID $\downarrow$	Prec. $\uparrow$	Rec. $\uparrow$
CFG only ( $w\!=\!0$ )	$5.337_{\pm.042}$	$\mathbf{0.7132_{\pm.0019}}$	$\mathbf{0.5495_{\pm.0024}}$
CFG + UG ( $w\!=\!20$ )	$\mathbf{4.999_{\pm.047}}$	$0.7281_{\pm.0005}$	$0.5393_{\pm.0012}$

Appendix

Table 10: Multi-seed evaluation of UG on ImageNet-256 (

\lambda=0.5

). Mean

\pm

standard deviation over 3 random seeds. This corresponds to the

\lambda=0.5

rows of Table 3 (c); we report multi-seed statistics here because the FID gap at this operating point is small, and we wish to confirm that the improvement is statistically consistent.

Figure 11: FID, precision, and recall as a function of the fixed CFG scale

\lambda

or the maximum UCFG scale

\lambda_{\max}

on ImageNet-128. CFG degrades sharply at large

\lambda

, while UCFG remains more stable as

\lambda_{\max}

increases.

Panel of Figure 12 — Figure 12: Distributions of the adaptive guidance scale $\lambda^{*}$ across sampling steps on ImageNet-128. Violin plots are computed from 1,000 samples for (a) UCFG and (b) APG+UCFG. In both cases, $\lambda^{*}$ tends to be smaller in early steps and larger in later steps.

Panel of Figure 13 — Figure 13: Qualitative uncertainty evolution for BayesDiff. Images are sampled on ImageNet-256 without guidance (no classifier guidance or classifier-free guidance). Top : the generated sample. Middle : BayesDiff’s velocity uncertainty maps at intermediate sampling steps 12/24/36/48 (left to right). Bottom : state uncertainty maps obtained by propagating the velocity uncertainty through the sampling dynamics using the same variance propagation rule as UA-Flow. While the velocity uncertainty exhibits spatial correlation, the propagated state uncertainty becomes largely noise-like and fails to preserve coherent spatial structure.

Panel of Figure 14 — Figure 14: Qualitative uncertainty evolution for UA-Flow. Images are sampled on ImageNet-256 without guidance (no classifier guidance or classifier-free guidance). Top : the generated sample. Middle : UA-Flow’s velocity uncertainty maps at intermediate sampling steps 12/24/36/48 (left to right). Bottom : propagated state uncertainty maps computed from the velocity uncertainty via variance propagation. In contrast to BayesDiff, the state uncertainty remains spatially coherent and tracks structured regions throughout sampling.

Figure 15: Uncertainty scale over sampling steps for BayesDiff. We quantify the magnitude of uncertainty at each sampling step by summing the uncertainty map over all elements (spatial locations and channels) for each sample. The plot reports the mean

\pm

standard deviation of this scalar summary across the

5\times 5

sample grid shown in Figure 13 . The noisy state uncertainty maps in Figure 13 are consistent with the elevated uncertainty scale in the early sampling phase, suggesting that BayesDiff’s uncertainty is poorly calibrated across sampling time and can dominate variance propagation.

Panel of Figure 16 — Figure 16: Teddy bear (ImageNet-256, CFG $\lambda=0.5$ ): lowest vs. highest uncertainty samples. We generate 1,000 samples and select the 25 lowest-uncertainty images (top) and the 25 highest-uncertainty images (bottom) using the same uncertainty aggregation and ranking procedure as in the main paper. Low-uncertainty samples are visually clean and class-consistent, while high-uncertainty samples often contain clutter, distortions, or weak class identity.

Panel of Figure 17 — Figure 17: Hen (ImageNet-256, CFG $\lambda=0.5$ ): lowest vs. highest uncertainty samples. We generate 1,000 samples and select the 25 lowest-uncertainty images (top) and the 25 highest-uncertainty images (bottom). Low uncertainty corresponds to high-fidelity, easily recognizable hens (often in canonical views), whereas high uncertainty corresponds to less reliable generations with more frequent artifacts or reduced class salience, illustrating a qualitative fidelity–diversity trade-off.

Table 11: Calibration results on ImageNet-256 validation images. UA-Flow achieves lower ECE and Brier score than BayesDiff across all noise levels

t

, both before and after post-hoc calibration. AU is reported only at

t=0.9

Figure 18: Correlation between unconditional and conditional velocities. Left: Scatter plot of unconditional vs. conditional predicted mean velocities $(\bar{u}^{\theta}_{t,\varnothing},\bar{u}^{\theta}_{t,y})$ , using $10^{6}$ randomly sampled elements across all sampling steps and 1,000 generated samples without CFG ( $\lambda=0$ ). Right: Pearson correlation between $\bar{u}^{\theta}_{t,\varnothing}$ and $\bar{u}^{\theta}_{t,y}$ at each sampling step, aggregated over all elements and samples; the correlation exceeds $0.98$ at every step.

Panel of Figure 19 — Figure 19: Ablations on UA-Flow’s uncertainty estimation pipeline (ImageNet-256, UA-Flow, no guidance). Each panel sweeps a single design choice while keeping the rest of the filtering pipeline fixed, and reports FID, precision, and recall as a function of the filtering ratio. Panel (c) additionally shows a scatter plot of $\mathrm{Var}(\bar{u}^{\theta}_{t}(\mathbf{x}_{t}))$ (MC estimate with $K=10$ ) versus $(\sigma^{\theta}_{t}(\bar{\mathbf{x}}_{t}))^{2}$ across sampling steps and spatial locations. Continued on the next page.

Panel of Figure 20 — Figure 20: Ablations on UA-Flow’s uncertainty estimation pipeline (continued).

Panel of Figure 21 — Figure 21: Ablations on UG (continued; ImageNet-256, UA-Flow, $\lambda=0$ ).

Panel of Figure 22 — Figure 22: Velocity uncertainty magnitude ( $\lambda=0$ ) and adaptive CFG scale ( $\lambda_{\max}=20$ ) over sampling steps, with and without the bias-correction term $U_{t}$ (ImageNet-256, 1,000 images). Without correction, the variance head overestimates uncertainty at the very early steps, while the median of the adaptive CFG scale is lower during the intermediate phase where coarse image structure is formed. With correction, velocity uncertainty is lower at the early steps and the adaptive CFG scale is higher during the intermediate phase.

Task	FID $\downarrow$	Prec. $\uparrow$	Rec. $\uparrow$
—	22.87	0.5063	0.6683
Filtering (50%)	20.66	0.5158	0.6462
UG ( $\lambda{=}0$ , $w{=}50$ )	18.44	0.5296	0.6389
Best UCFG † ( $\lambda_{\max}{=}1.25$ )	$4.67_{\pm.05}$	$0.7578_{\pm.0040}$	$0.5063_{\pm.0029}$
Best fixed CFG † ( $\lambda{=}0.75$ )	$\mathbf{4.53_{\pm.04}}$	$\mathbf{0.7821_{\pm.0066}}$	$0.4920_{\pm.0037}$

Appendix

Table 12: Bias-correction ablation (ImageNet-256). Metrics for the UA-Flow model trained with

U_{t}=0

. † Results are averaged over 3 seeds.

Figure 23: Generated 2D checkerboard samples and uncertainty. Left : samples colored by log state variance. Middle : generated samples overlaid on the ground-truth checkerboard grid. Right : samples after removing the top 5% highest-uncertainty samples. High-uncertainty samples concentrate near or outside the grid boundaries, and removing them visibly improves coverage.

Filtering ratio	Coverage $\uparrow$
0% (no filtering)	0.9572
1%	0.9668
2.5%	0.9793
5%	0.9876
10%	0.9936

Appendix

Table 13: Coverage after uncertainty-based filtering on the 2D checkerboard. Coverage is the fraction of retained samples that fall inside the checkerboard grid. Filtering by UA-Flow’s uncertainty monotonically improves coverage.

Figure 24: Conditional time-series predictions and uncertainty-based filtering. Each panel shows generated predictions (colored curves) given an observed input (blue) and ground truth curve (orange dash). From left to right: no filtering, 10% filtering, and 25% filtering. Removing high-uncertainty predictions eliminates incorrect forecasts.

Filtering ratio	MSE ( $\times 10^{-4}$ ) $\downarrow$
0% (no filtering)	14.65
10%	3.87
25%	2.11
50%	1.77
90%	0.78

Appendix

Table 14: MSE after uncertainty-based filtering on the sine-wave time-series task. MSE is computed between each predicted trajectory and the ground-truth continuation, averaged over retained samples. Filtering by UA-Flow’s uncertainty monotonically reduces prediction error.

Threshold $\tau$	Included	Excluded	Excl. %	Success % $\uparrow$	Mean MaxR $\uparrow$	Std MaxR
no removal	1000	0	0.0%	62.7%	0.9514	0.1506
0.0010	122	878	87.8%	91.0 %	0.9702	0.1253
0.0025	371	629	62.9%	77.6%	0.9597	0.1424
0.0050	708	292	29.2%	66.8%	0.9536	0.1526
0.0075	854	146	14.6%	63.6%	0.9536	0.1504
0.0100	921	79	7.9%	63.3%	0.9521	0.1511

Appendix

Table 15: Uncertainty-triggered early termination on Push-T robot manipulation. During a rollout, if the state uncertainty of the generated action chunk exceeds a threshold

\tau

, the rollout is terminated and excluded from evaluation; otherwise the chunk is executed. Results are averaged over the rollouts that complete without early termination, out of

1{,}000

total. A rollout is successful when its maximum reward exceeds 1.0.

Panel of Figure 25 — Figure 25: Selected ImageNet-256 sample grids under uncertainty-aware guidance sweeps. Rows sweep the UG scale $w\in\{0,10,30,50\}$ and columns sweep the maximum UCFG scale $\lambda_{\max}\in\{0,1,2,5,10,20\}$ . Left: generated samples. Right: predicted latent pixel-wise uncertainty maps.

Explore similar work

May 1, 2026

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

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

10^4 \times

less total compute than ensembling or Monte Carlo dropout.

Jiarui Xing, Song Wang, Jian Wang

Yale University · Harvard Medical School

Flow Matching Uncertainty Quantification Bayesian Inference Flow-Based Generative Modeling Covariance Estimation Uncertainty Propagation

Ask PDF

Mar 23, 2026

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

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.

Dongxia Wu, Yuhui Zhang, Serena Yeung-Levy +2

Stanford University Stanford, CA

Generative Modeling Flow Matching Uncertainty Quantification OOD Generalization OOD Detection

Ask PDF

Feb 23, 2026

Momentum Guidance: Plug-and-Play Guidance for Flow Models

Flow-based generative methods offer a simple and effective framework for high-fidelity generation, yet pretrained flow models are rarely used in their vanilla conditional form: in image generation, samples without guidance often appear diffuse and lack fine-grained detail. Existing guidance techniques such as classifier-free guidance (CFG) improve fidelity but reduce sample diversity. We introduce Momentum Guidance (MG), a guidance method that improves sample quality by extrapolating the current velocity away from an exponential moving average of past velocities along the ODE trajectory, while preserving the standard one-evaluation-per-step cost. MG provides gains beyond CFG, improving the precision-recall Pareto frontier. Experiments demonstrate the effectiveness of MG across benchmarks. On ImageNet-256, MG improves FID by 36.54% without CFG and 25.42% with CFG on average across sampling settings, attaining an FID of 1.553 at 16 sampling steps. Evaluations on large flow-based models, including Stable Diffusion 3 and FLUX.1-dev, further confirm improvements across standard metrics.

Runlong Liao, Jian Yu, Baiyu Su +3

The University of Texas at Austin, Austin, TX 78712, USA

Momentum Methods Diffusion Guidance Image Generation Generative Modeling Flow-Based Generative Modeling Inference-Time Guidance

Ask PDF