The Geometry Behind Diffusion and Flow Matching: Gradient Flows and Geodesics in Wasserstein Space

A new paper reveals that diffusion models and flow matching models are unified under the geometry of Wasserstein space, where diffusion follows a gradient flow of free energy and flow matching follows a geodesic, explaining their different sampling behaviors.

arXiv:2606.24157v1 Announce Type: new Abstract: The space $\mathcal{P} 2 \mathbb{R}^d$ of probability measures with finite second moment carries a natural geometry: the quadratic Wasserstein distance W 2 makes it a complete metric space and, following Otto, a formal Riemannian manifold whose geodesics are the optimal-transport interpolations. On this manifold, the gradient flow of the free energy F rho = KL rho || \pi is exactly the Fokker-Planck equation, and its implicit-Euler discretization is the JKO scheme. This is the geometry underlying diffusion models: the forward process descends the free energy, and each denoising step realizes one JKO step, which recovers DDPM, DDIM, NCSN/SMLD, and Energy Matching; this is one scheme, not separate theories. The same manifold supports a second variational principle. Its geodesics - the minimum-action curves of the Benamou-Brenier formula - are precisely the optimal-transport paths that Flow Matching learns. Fixing both endpoints and following the geodesic, generation becomes a deterministic ODE along a straight line, hence far fewer sampling steps. Placing both families of models on one manifold makes their relationship exact: diffusion follows a free-energy gradient flow, an initial-value problem; optimal-transport Flow Matching follows a Wasserstein geodesic, a boundary-value problem. The two reach the same endpoints along different paths.