🤖 AI Summary
This work unifies the theoretical foundations of diffusion models and flow matching within the Riemannian geometric framework of Wasserstein space. It interprets diffusion processes as gradient flows of a free energy functional—formulated as an initial value problem—whereas flow matching corresponds to geodesics with fixed endpoints, cast as a boundary value problem. By leveraging Otto calculus, the Fokker–Planck equation, the Jordan–Kinderlehrer–Otto (JKO) scheme, and the Benamou–Brenier principle, the study rigorously characterizes both the equivalence and distinctions between these two paradigms. This geometric perspective places DDPM, DDIM, NCSN/SMLD, and Energy Matching within a common variational framework and elucidates how flow matching achieves efficient deterministic generation through Wasserstein geodesics.
📝 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 || π) 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.