Divergence is Uncertainty: A Closed-Form Posterior Covariance for Flow Matching

We introduce a novel Wasserstein-1 ($W_1$) path-space divergence for stochastic and deterministic dynamics and establish a Wasserstein Uncertainty Propagation (WUP) theorem that bounds the $W_1$ distance between terminal distributions by the proposed divergence, equivalently characterized by a weighted $L^2$ discrepancy between the underlying drifts and the $W_1$ distance between their initial measures. A key ingredient is a probabilistic framework combining adjoint Feynman-Kac representations with synchronous coupling (and reflection coupling on bounded domains), yielding Wasserstein stability estimates beyond existing PDE- and Girsanov-based approaches. The framework accommodates time-varying and possibly degenerate diffusion coefficients, empirical and singular measures, and remains valid in the deterministic limit of flow matching. Unlike KL-based uncertainty quantification bounds, it does not require absolute continuity of path measures and therefore remains well-defined in singular settings. As consequences of the WUP theorem, we derive $W_1$ robustness and generalization bounds for score-based generative models and flow matching at both population and finite-sample levels. We further specialize the framework to group-symmetric targets, providing the first error analysis of equivariant flow-based models and the first quantitative comparison between data augmentation and equivariant inductive bias. Our analysis identifies a symmetry-aware Wasserstein path-space divergence that quantifies the model-form error induced by non-equivariant parametrizations. We prove that this error cannot be removed by additional data or training and vanishes only under equivariant architectures, establishing a precise theoretical advantage of equivariant inductive bias over data augmentation. Numerical experiments on group-symmetric Gaussian mixtures corroborate the theory.

Flow-based models learn a target distribution by modeling a marginal velocity field, defined as the average of sample-wise velocities connecting each sample from a simple prior to the target data. When sample-wise velocities conflict at the same intermediate state, however, this averaged velocity can misguide samples toward low-density regions, degrading generation quality. To address this issue, we propose the Flow Divergence Sampler (FDS), a training-free framework that refines intermediate states before each solver step. Our key finding reveals that the severity of this misguidance is quantified by the divergence of the marginal velocity field that is readily computable during inference with a well-optimized model. FDS exploits this signal to steer states toward less ambiguous regions. As a plug-and-play framework compatible with standard solvers and off-the-shelf flow backbones, FDS consistently improves fidelity across various generation tasks including text-to-image synthesis, and inverse problems.

One-step generative modeling has emerged as a leading approach for amortizing the inference cost of diffusion and flow-matching models. Among distillation-free methods, MeanFlow training is notoriously unstable, with non-decreasing loss and unbounded gradient variance. In this work, we establish a theory that attributes this pathology to a misuse of the conditional velocity field. We show that the conditional velocity plays two distinct statistical roles in the loss: both as an unbiased regression target and as a Monte Carlo control variate in a Jacobi-vector product, with the original MeanFlow loss assigning the wrong coefficient to the latter. We derive the optimal coefficient in closed form and show that a family of fixes in concurrent works corresponds to different practical realizations of the same optimum. A controlled sweep of this coefficient on two-dimensional benchmarks and on a latent Diffusion Transformer recovers the predicted bias-variance ordering. Our DiT experiment also reveals a quantitative FID-MSE landscape mismatch. Specifically, although the gradient-MSE is minimized at an interior coefficient value near $β\!=\!0.94$, the coefficient that minimizes FID prefers to use conditional velocity directly at the unbiased corner. Our analysis therefore explains why MeanFlow is unstable and unifies its concurrent remedies, and shows that the variance-optimal coefficient need not coincide with the quality-optimal one.