Uncertainty in a Single Pass: A Closed-Form Identity for One-Step Flow Matching

Abstract

Flow matching provides a highly effective framework for generative modeling, yet estimating the uncertainty of its generated samples remains a fundamental challenge. Existing methods rely on auxiliary variance heads, model ensembles, or iterative covariance propagation, which invariably introduce accumulated errors and structural interference along the noise-to-image trajectory. We resolve these limitations by adapting Tweedie's formula to the linear flow matching interpolant, yielding an exact, closed-form identity for the posterior covariance. Because this identity relies exclusively on the divergence of the learned velocity field, it can be evaluated post hoc on pretrained models without any architectural modifications. This analytical formulation proves exceptionally powerful for single-step generative frameworks. By computing the end-to-end posterior covariance in one forward pass, our approach completely bypasses the compounding errors inherent to sequential integration. This mechanism fundamentally enhances robustness, providing highly accurate and stable uncertainty estimates while preventing the numerical degradation typical of multi-step approximations. We validate our framework across benchmark datasets, including CIFAR-10 and real brain MRI scans. The resulting divergence-based uncertainty maps are highly interpretable and tightly correlated with empirical reconstruction errors. Crucially, in medical imaging, this precise localization of structural ambiguity provides essential decision support for high-stakes clinical tasks such as tumor boundary delineation and neurosurgical planning.