Flow Matching with Uncertainty Quantification and Guidance

TL;DR

This paper introduces UA-Flow, an extension of flow matching models that predicts uncertainty to improve sample reliability and quality in image generation.

Contribution

It proposes a lightweight method to estimate per-sample uncertainty in flow matching, guiding generation for higher fidelity outputs.

Findings

Uncertainty estimates correlate strongly with sample quality.

Uncertainty-guided sampling enhances image generation results.

UA-Flow outperforms baseline methods in reliability assessment.

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.