Flow Matching with Uncertainty Quantification and Guidance
This addresses sample reliability issues in generative modeling for applications like image generation, but it is incremental as it extends existing flow matching methods.
The paper tackled the problem of inconsistent or degraded sample quality in flow matching generative models by proposing UA-Flow, which predicts velocity fields with heteroscedastic uncertainty to assess reliability and guide generation, resulting in uncertainty signals more correlated with sample fidelity and improved generation quality in image experiments.
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.