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 uncertainty-aware 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.
| Subjects: | Computer Vision and Pattern Recognition (cs.CV); Machine Learning (cs.LG) |
| Cite as: | arXiv:2602.10326 [cs.CV] |
| (or arXiv:2602.10326v2 [cs.CV] for this version) | |
| https://doi.org/10.48550/arXiv.2602.10326 arXiv-issued DOI via DataCite |
Submission history
From: Juyeop Han [view email]
[v1]
Tue, 10 Feb 2026 22:03:13 UTC (36,577 KB)
[v2]
Tue, 12 May 2026 20:51:48 UTC (38,409 KB)