Wasserstein Geometry: Diffusion vs Flow Matching, a 2026 paper

The paper "The Geometry Behind Diffusion and Flow Matching: Gradient Flows and Geodesics in Wasserstein Space", by Yian Yao and Weiwei Zhang, was submitted on 23 Jun 2026 as arXiv:2606.24157. It places diffusion models and optimal-transport Flow Matching on a single geometric stage, the Wasserstein manifold P2(R^d), and spells out how each family follows a different variational principle.

What does the paper claim?

The authors assert that diffusion models are gradient flows of a free energy on the Wasserstein manifold, while Flow Matching follows Wasserstein geodesics; the paper identifies precise correspondences between PDEs, variational schemes, and generative algorithms. Specifically, the free energy F(rho) = KL(rho || π) has a gradient flow that the paper says is "exactly the Fokker-Planck equation," and its implicit-Euler discretization is the JKO scheme, which the authors connect to denoising steps in diffusion models.

They also describe a second variational principle on the same manifold: the geodesics from the Benamou-Brenier formula are the minimum-action curves that Flow Matching learns. Fixing both endpoints and following that geodesic turns generation into a deterministic ordinary differential equation along a straight line, which the paper argues requires far fewer sampling steps than diffusion's path.

How do diffusion and Flow Matching map onto Wasserstein geometry?

Diffusion is cast as an initial-value gradient flow while Flow Matching is a boundary-value geodesic problem, and both evolve measures on P2(R^d) under the quadratic Wasserstein distance W_2. The paper lays out that diffusion's forward process "descends the free energy," with each denoising step realizing one JKO step; this recovery unifies DDPM, DDIM, NCSN/SMLD, and Energy Matching under one scheme. In contrast, Flow Matching follows optimal-transport paths: its geodesics are the Benamou-Brenier minimum-action curves, and following them between fixed endpoints yields a deterministic generation ODE.

The authors emphasize that these are not two disconnected theories but two behaviors on the same manifold: diffusion solves an initial-value problem, starting from noise and descending KL free energy; Flow Matching solves a boundary-value problem, directly interpolating between endpoints via optimal transport.

Why it matters

Putting both families on the same geometric manifold clarifies why diffusion methods typically need many denoising steps while Flow Matching can use a deterministic ODE along a geodesic to reach the same endpoint with fewer steps. That shifts comparisons between methods from heuristics about samplers to concrete geometric differences: one follows a free-energy descent, the other follows minimum-action transport paths. For researchers this reframing suggests new ways to translate tools across the two approaches, for instance adapting JKO-style discretizations or Benamou-Brenier perspectives to hybrid samplers.

What to watch

Look for implementations and empirical comparisons that test how closely practical Flow Matching ODEs follow Benamou-Brenier geodesics and how many sampling steps are needed in practice versus diffusion-based JKO discretizations. Also monitor follow-up work that uses the paper's P2(R^d) formulation to derive new discretizations or to bridge DDPM/DDIM and Flow Matching algorithms.

Details and provenance: the paper appears on arXiv as arXiv:2606.24157 and was submitted on 23 Jun 2026. The authors name the free energy explicitly as F(rho) = KL(rho || π), connect its gradient flow to the Fokker-Planck PDE, identify JKO as the implicit-Euler discretization, and link Flow Matching geodesics to the Benamou-Brenier formula. The generative methods mentioned by name are DDPM, DDIM, NCSN/SMLD, and Energy Matching.