The Geometry Behind Diffusion and Flow Matching: Gradient Flows and Geodesics in Wasserstein Space
This work, authored by Yian Yao and Weiwei Zhang and published on arXiv, establishes a unified geometric framework for understanding and unifying two prominent classes of generative models: diffusion models and flow matching models. The core contribution is the demonstration that both families of models can be viewed as operating on the Wasserstein space $\mathcal{P}_2(\mathbb{R}^d$), a space of probability measures equipped with the quadratic Wasserstein distance. This geometric perspective reveals a deep, previously unarticulated relationship between them.
The problem this paper tackles is the perceived divergence in theoretical underpinnings and practical implementations of diffusion and flow matching. While both are powerful generative techniques, their underlying formalisms appear distinct. Diffusion models are often explained through stochastic differential equations and denoising processes, while flow matching focuses on deterministic optimal transport paths. This work bridges this gap by showing that both are manifestations of fundamental geometric principles on the Wasserstein manifold. Specifically, diffusion models are framed as gradient flows of a KL divergence free energy, mirroring solutions to the Fokker-Planck equation, with implicit-Euler discretizations corresponding to the JKO scheme. In contrast, flow matching is shown to learn geodesics – minimum-action curves – within the same Wasserstein space, governed by the Benamou-Brenier formula.
The most crucial technical insights revolve around this dual variational principle. First, the gradient flow interpretation of diffusion models explains why the forward process moves along a free energy landscape, akin to an initial-value problem. Second, the geodesic interpretation of flow matching highlights its nature as a boundary-value problem, where generation follows a deterministic ODE along a straight line in Wasserstein space, promising efficient sampling. Finally, the unified framework reveals that while diffusion and flow matching traverse the same endpoints in probability space, they do so via fundamentally different geometric paths: one descending an energy gradient, the other traversing a shortest path. This work is intended for researchers and engineers in machine learning, particularly those working on generative models, probabilistic modeling, and optimal transport.
This geometric unification has significant implications. It enables a more principled understanding of existing models and could inspire novel hybrid approaches. For instance, insights from gradient flow dynamics might inform the design of more efficient flow matching algorithms, or the deterministic sampling of flow matching could be leveraged to accelerate diffusion model inference. Going forward, this work suggests a research trajectory focused on exploring the richer geometric properties of Wasserstein space for generative modeling, potentially leading to more interpretable, controllable, and sample-efficient generative architectures. The abstract indicates this is a foundational theoretical contribution, serving as a conceptual roadmap for future model development.