Flow Matching with Uncertainty Quantification and Guidance

2026-02-10

AI Summary

This work addresses the issue of inconsistent or degraded sample quality in flow matching generative models and the lack of reliable metrics to assess generation fidelity. We propose Uncertainty-Aware Flow Matching (UA-Flow), which, for the first time, integrates heteroscedastic uncertainty modeling into the flow matching framework. UA-Flow jointly predicts the velocity field and estimates per-sample uncertainty, then propagates this uncertainty through the flow dynamics to derive a reliability measure. Leveraging this uncertainty estimate, we design both uncertainty-guided classifier-based and classifier-free sampling strategies. Experimental results demonstrate that the predicted uncertainty strongly correlates with sample fidelity, and that uncertainty-guided sampling significantly improves generation quality.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Uncertainty Representations
• Multiagent Systems: Multiagent Systems under Uncertainty

Application Category

• Economics, Online Markets and Human Computation: Trust and reliance of crowd workers and data experts on GenAI
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

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.

Problem

• flow matching
• uncertainty quantification
• sample quality
• generative models
• reliability assessment

Innovation

• flow matching
• uncertainty quantification
• heteroscedastic uncertainty
• classifier guidance
• generative modeling