Flow Matching with Uncertainty Quantification and Guidance
TL;DR
Uncertainty-aware flow matching (UA-Flow) is proposed, a lightweight extension of flow matching that predicts the velocity field together with heteroscedastic uncertainty that tracks sample fidelity more consistently than element-wise baselines and that these signals also carry information about data-manifold support.
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.
