Flow Matching with Uncertainty Quantification and Guidance

The gist: Uncertainty-aware flow matching (UA-Flow) is a lightweight extension of flow matching that predicts velocity field uncertainty and uses it to steer generation via uncertainty-aware classifier guidance and classifier-free guidance, producing samples with higher fidelity

How it works

UA-Flow models heteroscedastic uncertainty in the velocity field by estimating both the mean, u¯θt(xt), and the diagonal variance, (σθt(xt))2, of the velocity field using a Gaussian negative log-likelihood (NLL) loss formulated as LUFM(θ). This loss is reformulated into a conditional uncertainty-aware flow matching loss, LCUFM(θ), to account for the conditional nature of the training data. The model estimates the variance element-wise, following common practice in heteroscedastic UQ for neural networks.

Uncertainty Propagation through Flow Dynamics

The paper propagates velocity uncertainty through the flow dynamics to estimate sample uncertainty, interpreting the resulting mean x¯1 and variance Var[x1] at the final time as the generated sample and its associated uncertainty. The mean dynamics are approximated by linearizing u¯θt(xt) around x¯t, yielding dx¯t/dt = ¯uθt(x¯t). Variance propagation uses Euler discretization, approximating the evolution of element-wise variance as Var[xt+∆t] ≈ Var[xt] + (σθt(x¯t)∆t)2 + 2∆txCov(xt, uθt(xt)). The covariance term Cov(xt, uθt(xt)) is approximated using Hutchinson’s diagonal estimator to avoid forming the Jacobian explicitly.

Uncertainty-Aware Guidance for Flow Matching

UA-Flow incorporates the predicted velocity uncertainty into guided sampling through two mechanisms. The first is an uncertainty-based pseudo-likelihood used as a classifier guidance term, leading to the U-CG mean velocity u¯θt,CG(x¯t) = ¯uθt(x¯t) + btw∇x¯t f((σθt(xbar t))2). The second is an adaptive choice of the classifier-free guidance (CFG) scale λ that reduces the predicted variance of the extrapolated velocity, leading to U-CFG. For U-CFG, the optimal scale λopt is found by minimizing the total predicted variance of the extrapolated velocity, and it is clamped by a maximum scale λmax to prevent divergence.

Experimental Validation and Results

Experiments validate UA-Flow's effectiveness in two ways: as a sample-level reliability signal and as a control signal for improved generation quality. Filtering out high-uncertainty samples yields better fidelity metrics than prior uncertainty-quantification baselines for sampling-based generative models. Furthermore, uncertainty reduction can be incorporated as guidance, where both U-CG and U-CFG lead to improved generation quality and. Qualitative results show that UA-Flow's uncertainty maps localize high-uncertainty regions, which correlates with sample fidelity, while BayesDiff produces noisy maps that fail to localize.

Generalization and Additional Findings

UA-Flow generalizes beyond image generation by showing meaningful reliability signals on non-image domains such as a 2D checkerboard density, where predicted uncertainty reliably identifies samples outside the target support. In time series forecasting, filtering by UA-Flow's uncertainty monotonically reduces prediction error, confirming its utility for sequential decision-making. Additionally, in robot policy environments like Push-T, chunk-level state uncertainty correlates negatively with episode success and can be used as an online abort signal to increase the retained success rate.

The work proposes that UA-Flow's uncertainty estimates are meaningfully calibrated and serve as a reliable measure of sample quality beyond their practical usefulness for filtering. The paper concludes that while the formulation relies on simplifying assumptions for tractability, it provides per-sample and spatially localized uncertainty estimates that can help practitioners assess the reliability of generated outputs.

--- Page 1 ---

Flow Matching with Uncertainty Quantification and Guidance

Juyeop Han, MIT, juyeop@mit.edu Lukas Lao Beyer, MIT, llb@mit.edu Sertac Karaman, MIT, sertac@mit.edu 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 generate higher-quality outputs, we propose uncertaintyaware flow matching (UA-Flow), a lightweight extension of flow matching that predicts the velocity field together with heteroscedastic uncertainty. UA-Flow estimates per-sample uncertainty by propagating velocity uncertainty through the flow dynamics. These uncertainty estimates act as a reliability signal for individual samples, and we further use them to steer generation via uncertainty-aware classifier guidance and classifier-free guidance. Experiments on image generation show that UA-Flow produces uncertainty signals more highly correlated with sample fidelity than baseline methods, and that uncertainty-guided sampling further improves generation quality.

1 Introduction In recent years, sampling-based generative models such as diffusion [Sohl-Dickstein et al., 2015, Ho et al., 2020, Song et al., 2021] and flow matching [Lipman et al., 2023, Liu et al., 2023, Albergo and Vanden-Eijnden, 2023] have achieved remarkable success across a wide range of domains, most particularly in image and video generation [Dhariwal and Nichol, 2021, Ho et al., 2022] as well as in sequential decision-making [Janner et al., 2022, Chi et al., 2025, Black et al., 2024]. Despite this progress, these models often produce samples of inconsistent quality. As a result, using reliably generated samples in downstream applications remains challenging. To address this issue, the uncertainty associated with each generated sample can be interpreted as a measure of the reliability of the generation process. Recently, several works have explored uncertainty estimation for diffusion-based generative models by adapting techniques from existing uncertainty quantification literature for neural networks [Lakshminarayanan et al., 2017, Kendall and Gal, 2017, Ritter et al., 2018]. In sampling-based generative modeling, uncertainty plays two central roles: (i) it provides a principled signal for assessing the quality or reliability of individual generated samples, and (ii) uncertainty can be leveraged during the generative process to actively improve sample quality through guided sampling. However, most prior works [Kou et al., 2023, Jazbec et al., 2025] primarily focus on the first role, using uncertainty in a post-hoc manner for sample filtering or selection. Moreover, existing approaches are often domain-specific [Sun et al., 2023, Franchi et al., 2025] and quantify uncertainty over conditional inputs [Berry et al., 2024]. We propose uncertainty-aware flow matching (UA-Flow), a lightweight extension of flow matching that models heteroscedastic uncertainty in the velocity field. By propagating this velocity uncertainty through the flow dynamics, UA-Flow provides principled sample uncertainty estimates with minimal additional overhead. Because uncertainty is modeled element-wise at the velocity level of general flow matching, our approach provides spatially localized uncertainty within each generated sample. Unlike methods that depend on domain-specific adaptations such as the CLIP encoder used by [Jazbec Preprint. arXiv:2602.10326v2 [cs.CV] 12 May 2026], this formulation is, in principle, not specialized to a particular data domain.

1. The system can produce samples of higher fidelity by employing uncertainty-aware classifier guidance (U-CG) which steers generation via uncertainty-aware classifier guidance and classifier-free guidance to steer the trajectory toward low-uncertainty regions.

2. The system can improve generation quality by utilizing uncertainty reduction as a control signal through uncertainty-aware classifier-free guidance (U-CFG), which uses an adaptive choice of the CFG scale that reduces the predicted variance of the extrapolated velocity.

3. The system can achieve better sample reliability by filtering generated images based on their per-sample uncertainty, as this approach yields lower FID and higher precision after filtering compared to baseline methods like BayesDiff and AU.

4. The system can be made more robust to overconfidence or excessively guided generation by using uncertainty estimates as a complementary reliability signal rather than a guarantee of correctness.

5. For sequential decision-making tasks, the system can be used for online abort signals by using chunk-level state uncertainty to terminate rollouts when it exceeds a threshold τ, which has been shown to raise the retained success rate from 66.5% to 83.3%.