Diffusion Flow Matching: Dimension-Improved KL Bounds and Wasserstein Guarantees

Abstract: Diffusion Flow Matching (DFM) has recently emerged as a versatile framework for generative modeling, yet its theoretical convergence properties remain only partially understood. In this work, we provide refined and novel convergence guarantees for Brownian motion based DFMs, focusing on the discretization error. Our analysis is conducted under the Kullback-Leibler (KL) divergence and the 2-Wasserstein distance. Under finite-moment conditions and a mild score integrability assumption, we derive KL convergence bounds with improved dimensional dependence compared to prior work, achieving, up to our knowledge, state-of-the-art scaling under minimal conditions. We further extend the analysis to the 2-Wasserstein distance: under an additional first-order score integrability assumption and a weak log-concavity condition, we obtain convergence guarantees with dimensional dependence consistent with the KL case.

• The paper refines DFM analysis by establishing O(d^3) KL divergence bounds under finite moment conditions, enhancing previous O(d^4) results.
• It introduces both full and early-stopped DFM schemes that achieve sharp convergence guarantees with minimal integrability and regularity assumptions.
• The study derives Wasserstein-2 error bounds of O(√h) and O(√(d^3)), leveraging Brownian bridges and controlled drift approximations to ensure robust sample generation.

Dimension-Improved KL and Wasserstein Guarantees for Diffusion Flow Matching

Overview

This work provides a rigorous analysis of Diffusion Flow Matching (DFM) for generative modeling, with a focus on refining non-asymptotic convergence bounds under both Kullback-Leibler (KL) divergence and Wasserstein-2 ($W_2$) distance. Previous theoretical studies have been limited, often resulting in sub-optimal dimensional dependence and requiring restrictive assumptions. The authors address these gaps by deriving convergence bounds with improved dimension scaling, minimal integrability conditions, and guarantees for both full and early-stopped DFM procedures built on Brownian bridges.

Diffusion Flow Matching Framework

DFM generalizes flow-based generative modeling by learning stochastic transports from a simple base distribution to a complex data distribution using stochastic differential equations (SDEs). Brownian-bridge-based DFM explicitly constructs a stochastic interpolant between the source and target distributions, with marginal distributions matched via Markovian projections. The transport is parameterized through a drift estimated by neural network regression, and sample generation proceeds via Euler-Maruyama discretization.

A key technical aspect is deriving tractable approximations of the mimicking drift in the Markovian projection, which is inherently intractable due to its conditional expectation form. The learned drift surrogate is then used in time-discretized SDE simulation to produce approximate samples.

Main Theoretical Results

KL Divergence Bounds

Full DFM (No Early Stopping)

• Under standard moment conditions (finite 8th moment for marginals and coupling's score), and $L^2$ drift-approximation, the KL divergence between the DFM output and the target distribution enjoys an explicit upper bound with $O(d^3)$ dependence on dimension and $O(h)$ in the step size (Theorem 1).
• This refines the previous best result ([Gentiloni Silveri et al., 2024]) that scaled as $O(d^4)$, marking a strong improvement in the scaling regime while maintaining minimalistic assumptions.

Early-Stopped DFM

• By halting the simulation before terminal time ($t = 1 - \epsilon$), regularity demands on the coupling can be weakened (score integrability only required for a conditional coupling). The $O(d^3)$ scaling and $O(h)$ step-size dependence are preserved (Theorem 2).
• Under a novel non-uniform step-size schedule, accelerated rates for KL convergence are attainable without sacrificing dimension dependence, leading to overall improved complexity in the Fortet-Mourier metric (Theorem 3).

Wasserstein-2 Distance Bounds

• For the non-early-stopped regime, the authors derive $W_2$ bounds that depend as $O(\sqrt{h})$ in step-size and $L^2$0 in dimension, consistent with the corresponding KL scaling (Theorem 4).
• The analysis relies on weak log-concavity of the coupling and first-order score integrability, requiring only mild regularity on data and model.
• The derived rates match the KL setting; complexity is sharply characterized, and the requirements are shown to hold as well for independent couplings, provided the marginals satisfy analogous integrability and weak convexity conditions.

Methodology and Technical Advances

The technical core of the analysis involves:

• A refined handling of the mean acceleration (reciprocal characteristic) in the mimicking drift via integration-by-parts, allowing one fewer differentiation than previous approaches. This is pivotal in improving dimensional scaling from $L^2$1 to $L^2$2.
• For Wasserstein analysis, recursive error propagation is strictly controlled by bounding the Lipschitz constant of the Markovian drift through log-Sobolev and Poincaré inequalities, leveraging the structural properties of the interpolant and coupling.

In both metrics, error decompositions isolate the contributions of drift approximation and time discretization, making the results robust to practical estimation inexactness.

Comparison with Related Works

• The study advances KL convergence results of DFM both in the non-early-stopped ($L^2$3 vs $L^2$4) and early-stopped regimes, and substantially relaxes regularity assumptions for conditional scores. Complexity is established with tight dependence on both error tolerance and dimension.
• In Wasserstein-2, it improves upon bounds such as in [Xiangjun & Zhongjia, 2025], which require Gaussian priors and impose severe support assumptions on the target. Here, arbitrary priors are admissible and only weak convexity/integrability are required.
• The structural analysis does not suffer from the initialization bias issues found in certain two-sided early-stopping analyses ([Liu et al., 2025]), maintaining unbiasedness in the generative output.

Implications and Future Directions

The improved bounds provided in this work contribute nontrivial theoretical guarantees for DFM, supporting its deployment in high-dimensional generative tasks. By enabling precise characterization of both KL and Wasserstein convergence under minimal regularity, these results strengthen the theoretical foundations for Brownian-bridge-stochastic interpolant-based models.

Practically, sharper dimension dependence and relaxed regularity open the door to better scalability and robustness in generative modeling for large-scale or complex-structured data (e.g., high-res imagery, molecular graphs), suggesting enhanced reliability for practitioners.

Theoretically, the advances motivate further exploration in several directions:

• Further reductions in the dependence on dimension, possibly through leveraging symmetries or manifold structure in data
• Removing the remaining assumptions on score integrability, exploring models with heavy-tailed distributions or non-smooth densities
• Statistical analysis of DFM to quantify the propagation of finite-sample estimation errors in high-dimensional regimes or under out-of-distribution settings

Conclusion

The paper establishes dimension-improved, practically relevant convergence guarantees for DFM in both KL and Wasserstein metrics with minimalistic conditions. The analysis sharpens prior results, eliminates certain biases, and substantially advances the theoretical toolkit for DFM-based generative modeling. This work is foundational for future research on the statistical properties and scalable deployment of diffusion-driven generative models.

Reference:

"Diffusion Flow Matching: Dimension-Improved KL Bounds and Wasserstein Guarantees" (2606.16610)