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

# The Geometry Behind Diffusion and Flow Matching Gradient Flows and Geodesics in Wasserstein Space
Source: [https://arxiv.org/html/2606.24157](https://arxiv.org/html/2606.24157)
###### Abstract

The space𝒫2​\(ℝd\)\\mathcal\{P\}\_\{2\}\(\\mathbb\{R\}^\{d\}\)of probability measures with finite second moment carries a natural geometry: the quadratic Wasserstein distanceW2W\_\{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ℱ​\(ρ\)=KL​\(ρ∥π\)\\mathcal\{F\}\(\\rho\)=\\mathrm\{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—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\.

Wasserstein distance→endows geometry\\xrightarrow\{\\text\{endows geometry\}\}probability space becomes a Riemannian manifold,carryingtwo complementary variational principles:\(A\) Diffusion— free\-energy*gradient flow*\(*initial\-value*problem: givenρk\\rho\_\{k\}, descendℱ\\mathcal\{F\}while staying close inW2W\_\{2\}\)free energy→\\toFokker\-Planck→time discretization\\xrightarrow\{\\text\{time discretization\}\}JKO scheme\(B\) Flow Matching— Benamou–Brenier*minimum action*\(*boundary\-value*problem: given bothρ0\\rho\_\{0\}andρ1\\rho\_\{1\}, minimize kinetic energy\)minimum action→\\toWasserstein geodesics→\\toOT transport pathsEssential difference:*same endpointρ1=p*data*\\rho\_\{1\}=p\_\{\\text\{data\}\}, different path\.*

Outline

- •Prologuestarts from the continuity and Fokker–Planck equations as they appear in diffusion models and Flow Matching, and isolates the connections those treatments leave implicit\.
- •Section 1introduces the Wasserstein distance—the transport cost between two distributions\.
- •Section 2derives the continuity equation—the local form of mass conservation\.
- •Section 3establishes the Riemannian structure: the Benamou–Brenier formula identifiesW2W\_\{2\}as a geodesic distance, and its geodesics as the optimal\-transport paths learned by Flow Matching\.
- •Section 4derives the Fokker–Planck equation from a stochastic differential equation and recasts it as a continuity equation\.
- •Section 5proves the central identity: the Fokker–Planck equation is the Wasserstein gradient flow of the free energy\.
- •Section 6discretizes the flow into the JKO scheme, proves convergence, and recovers DDPM, DDIM, NCSN/SMLD, Flow Matching, and Energy Matching from a single variational template\.

Prerequisites\.We assume multivariable calculus, linear algebra, and elementary probability; measure theory, differential geometry, and convex analysis are developed in self\-contained appendices\.

###### Contents

1. [Notation and Conventions](https://arxiv.org/html/2606.24157#Sx1)
2. [Prologue: Things You May Already Know](https://arxiv.org/html/2606.24157#Sx2)1. [Scenario 1: How you encounter Fokker\-Planck in Diffusion Models](https://arxiv.org/html/2606.24157#Sx2.SSx1) 2. [Scenario 2: How you encounter the continuity equation in Flow Matching](https://arxiv.org/html/2606.24157#Sx2.SSx2) 3. [Scenario 3: Probability Flow ODE—where the two equations meet](https://arxiv.org/html/2606.24157#Sx2.SSx3) 4. [How deep do these connections run?](https://arxiv.org/html/2606.24157#Sx2.SSx4) 5. [This article’s mission](https://arxiv.org/html/2606.24157#Sx2.SSx5) 6. [Related perspectives](https://arxiv.org/html/2606.24157#Sx2.SSx6)
3. [1Wasserstein Distance: The Cost of Moving Dirt](https://arxiv.org/html/2606.24157#S1)1. [1\.1Why do we need a new distance?](https://arxiv.org/html/2606.24157#S1.SS1) 2. [1\.2The Kantorovich formulation](https://arxiv.org/html/2606.24157#S1.SS2) 3. [1\.3The Monge formulation and Brenier’s theorem](https://arxiv.org/html/2606.24157#S1.SS3) 4. [1\.4A concrete example: Gaussians](https://arxiv.org/html/2606.24157#S1.SS4) 5. [1\.5Key properties](https://arxiv.org/html/2606.24157#S1.SS5)
4. [2The Continuity Equation: Conservation of Mass](https://arxiv.org/html/2606.24157#S2)1. [2\.1Physical picture](https://arxiv.org/html/2606.24157#S2.SS1) 2. [2\.2Derivation from first principles](https://arxiv.org/html/2606.24157#S2.SS2) 3. [2\.3Weak formulation](https://arxiv.org/html/2606.24157#S2.SS3) 4. [2\.4Connection to particle ODEs](https://arxiv.org/html/2606.24157#S2.SS4)
5. [3The Riemannian Structure of Wasserstein Space](https://arxiv.org/html/2606.24157#S3)1. [3\.1The Benamou–Brenier formula](https://arxiv.org/html/2606.24157#S3.SS1) 2. [3\.2Why continuity equation \+ minimum kinetic energy gives the OT path](https://arxiv.org/html/2606.24157#S3.SS2) 3. [3\.3Otto’s Riemannian interpretation](https://arxiv.org/html/2606.24157#S3.SS3) 4. [3\.4Geodesics = displacement interpolation](https://arxiv.org/html/2606.24157#S3.SS4)
6. [4The Fokker–Planck Equation](https://arxiv.org/html/2606.24157#S4)1. [4\.1Overview](https://arxiv.org/html/2606.24157#S4.SS1) 2. [4\.2Motivation: the physical setup](https://arxiv.org/html/2606.24157#S4.SS2) 3. [4\.3Prerequisite: Itô’s formula](https://arxiv.org/html/2606.24157#S4.SS3) 4. [4\.4Derivation of Fokker–Planck: complete details](https://arxiv.org/html/2606.24157#S4.SS4) 5. [4\.5Rewriting as a continuity equation](https://arxiv.org/html/2606.24157#S4.SS5) 6. [4\.6Stationary distribution](https://arxiv.org/html/2606.24157#S4.SS6) 7. [4\.7Connection to modern generative models](https://arxiv.org/html/2606.24157#S4.SS7) 8. [4\.8Conceptual clarification: what needs an energy, and what does not?](https://arxiv.org/html/2606.24157#S4.SS8) 9. [4\.9Where is the modeling choice in diffusion?](https://arxiv.org/html/2606.24157#S4.SS9)
7. [5Free Energy and the Wasserstein Gradient](https://arxiv.org/html/2606.24157#S5)1. [5\.1The free energy functional](https://arxiv.org/html/2606.24157#S5.SS1) 2. [5\.2First variation \(functional derivative\)](https://arxiv.org/html/2606.24157#S5.SS2) 3. [5\.3Computing the Wasserstein gradient](https://arxiv.org/html/2606.24157#S5.SS3) 4. [5\.4The main result: Fokker–Planck as a Wasserstein gradient flow](https://arxiv.org/html/2606.24157#S5.SS4) 5. [5\.5End\-of\-section comparison: Euclidean vs\. Wasserstein gradient\-flow logic](https://arxiv.org/html/2606.24157#S5.SS5)
8. [6The JKO Scheme](https://arxiv.org/html/2606.24157#S6)1. [6\.1Implicit Euler method: from gradient flow to optimization problem](https://arxiv.org/html/2606.24157#S6.SS1) 2. [6\.2JKO scheme: Implicit Euler in Wasserstein space](https://arxiv.org/html/2606.24157#S6.SS2) 3. [6\.3Deriving JKO from the gradient flow](https://arxiv.org/html/2606.24157#S6.SS3) 4. [6\.4First\-order optimality conditions of JKO](https://arxiv.org/html/2606.24157#S6.SS4) 5. [6\.5Rigorous convergence theorem](https://arxiv.org/html/2606.24157#S6.SS5) 6. [6\.6Structural advantages of JKO](https://arxiv.org/html/2606.24157#S6.SS6) 7. [6\.7Example: Heat equation](https://arxiv.org/html/2606.24157#S6.SS7) 8. [6\.8Application: Energy Matching \(Balcerak et al\., 2025\)](https://arxiv.org/html/2606.24157#S6.SS8) 9. [6\.9Unified derivation from JKO to mainstream generative algorithms](https://arxiv.org/html/2606.24157#S6.SS9)
9. [7The Big Picture](https://arxiv.org/html/2606.24157#S7)
10. [AMeasures and Couplings](https://arxiv.org/html/2606.24157#A1)1. [A\.1What is a measure?](https://arxiv.org/html/2606.24157#A1.SS1) 2. [A\.2Absolutely continuous: what does it really mean?](https://arxiv.org/html/2606.24157#A1.SS2) 3. [A\.3Integration against a measure](https://arxiv.org/html/2606.24157#A1.SS3) 4. [A\.4Pushforward measure](https://arxiv.org/html/2606.24157#A1.SS4) 5. [A\.5Coupling: the key concept for optimal transport](https://arxiv.org/html/2606.24157#A1.SS5)
11. [BDifferential Geometry Prerequisites](https://arxiv.org/html/2606.24157#A2)
12. [CAnalysis Tools: Integration by Parts and du Bois\-Reymond Lemma](https://arxiv.org/html/2606.24157#A3)1. [C\.1Integration by Parts](https://arxiv.org/html/2606.24157#A3.SS1)1. [C\.1\.1One\-dimensional case](https://arxiv.org/html/2606.24157#A3.SS1.SSS1) 2. [C\.1\.2Higher\-dimensional case: divergence theorem version](https://arxiv.org/html/2606.24157#A3.SS1.SSS2) 3. [C\.1\.3Transferring the Laplacian \(two integrations by parts\)](https://arxiv.org/html/2606.24157#A3.SS1.SSS3) 4. [C\.1\.4Integration by parts in time](https://arxiv.org/html/2606.24157#A3.SS1.SSS4) 5. [C\.1\.5Formula summary table](https://arxiv.org/html/2606.24157#A3.SS1.SSS5) 2. [C\.2Du Bois\-Reymond Lemma \(Fundamental Lemma of the Calculus of Variations\)](https://arxiv.org/html/2606.24157#A3.SS2)1. [C\.2\.1Intuitive understanding](https://arxiv.org/html/2606.24157#A3.SS2.SSS1) 2. [C\.2\.2Proof sketch](https://arxiv.org/html/2606.24157#A3.SS2.SSS2) 3. [C\.2\.3Constrained version \(used in the main text\)](https://arxiv.org/html/2606.24157#A3.SS2.SSS3) 4. [C\.2\.4Where it is used in the main text](https://arxiv.org/html/2606.24157#A3.SS2.SSS4)
13. [DGibbs–Boltzmann Distribution](https://arxiv.org/html/2606.24157#A4)1. [D\.1Core intuition](https://arxiv.org/html/2606.24157#A4.SS1) 2. [D\.2The role of temperatureTT—the single control knob](https://arxiv.org/html/2606.24157#A4.SS2) 3. [D\.3Applications across disciplines](https://arxiv.org/html/2606.24157#A4.SS3) 4. [D\.4Core connection to the main text](https://arxiv.org/html/2606.24157#A4.SS4) 5. [D\.5Why the exponential form?](https://arxiv.org/html/2606.24157#A4.SS5)
14. [EFisher Information](https://arxiv.org/html/2606.24157#A5)1. [E\.1Fisher Information in Statistics \(Parametric Version\)](https://arxiv.org/html/2606.24157#A5.SS1) 2. [E\.2Fisher Information in Information Theory/PDE \(Distribution Version\)](https://arxiv.org/html/2606.24157#A5.SS2) 3. [E\.3Unification of the two versions](https://arxiv.org/html/2606.24157#A5.SS3) 4. [E\.4Relative Fisher Information](https://arxiv.org/html/2606.24157#A5.SS4) 5. [E\.5Three identities](https://arxiv.org/html/2606.24157#A5.SS5) 6. [E\.6Information inequalities](https://arxiv.org/html/2606.24157#A5.SS6) 7. [E\.7Example](https://arxiv.org/html/2606.24157#A5.SS7)
15. [FGenerative Models: Score, Velocity, and Diffusion](https://arxiv.org/html/2606.24157#A6)1. [F\.1The Complete Relationship among Score, Velocity, and Drift](https://arxiv.org/html/2606.24157#A6.SS1) 2. [F\.2Geometric meaning of the score: the Wasserstein gradient of entropy](https://arxiv.org/html/2606.24157#A6.SS2) 3. [F\.3A unified framework for mainstream models](https://arxiv.org/html/2606.24157#A6.SS3)1. [F\.3\.1DDPM \(Denoising Diffusion Probabilistic Models\)](https://arxiv.org/html/2606.24157#A6.SS3.SSS1) 2. [F\.3\.2NCSN / SMLD](https://arxiv.org/html/2606.24157#A6.SS3.SSS2) 3. [F\.3\.3VE\-SDE \(Variance Exploding\)](https://arxiv.org/html/2606.24157#A6.SS3.SSS3) 4. [F\.3\.4VP\-SDE \(Variance Preserving\)](https://arxiv.org/html/2606.24157#A6.SS3.SSS4) 5. [F\.3\.5Flow Matching](https://arxiv.org/html/2606.24157#A6.SS3.SSS5) 4. [F\.4Unified comparison table](https://arxiv.org/html/2606.24157#A6.SS4) 5. [F\.5A unified understanding from the Fokker\-Planck/JKO perspective](https://arxiv.org/html/2606.24157#A6.SS5)
16. [GConvex Analysis Quick Review](https://arxiv.org/html/2606.24157#A7)1. [G\.1Convex functions](https://arxiv.org/html/2606.24157#A7.SS1) 2. [G\.2The Hessian matrix](https://arxiv.org/html/2606.24157#A7.SS2) 3. [G\.3Positive definiteness and positive semidefiniteness](https://arxiv.org/html/2606.24157#A7.SS3) 4. [G\.4Strong convexity andλ\\lambda\-convexity](https://arxiv.org/html/2606.24157#A7.SS4)
17. [HSupplementary Concepts](https://arxiv.org/html/2606.24157#A8)1. [H\.1Langevin dynamics](https://arxiv.org/html/2606.24157#A8.SS1) 2. [H\.2Kolmogorov forward equation](https://arxiv.org/html/2606.24157#A8.SS2) 3. [H\.3Ornstein\-Uhlenbeck \(OU\) process](https://arxiv.org/html/2606.24157#A8.SS3) 4. [H\.4Otto Calculus](https://arxiv.org/html/2606.24157#A8.SS4) 5. [H\.5Geodesics](https://arxiv.org/html/2606.24157#A8.SS5) 6. [H\.6Riesz representation theorem](https://arxiv.org/html/2606.24157#A8.SS6) 7. [H\.7Lyapunov functions](https://arxiv.org/html/2606.24157#A8.SS7) 8. [H\.8KKT conditions](https://arxiv.org/html/2606.24157#A8.SS8)

## Notation and Conventions

The following notation is used frequently throughout this article and is collected here for easy reference\.

### Sets and Spaces

- •ℝd\\mathbb\{R\}^\{d\}:dd\-dimensional Euclidean space\.
- •Borel sets: Starting from all open sets inℝd\\mathbb\{R\}^\{d\}and closing under countable unions, intersections, and complements, one obtains theBorelσ\\sigma\-algebra, denotedℬ​\(ℝd\)\\mathcal\{B\}\(\\mathbb\{R\}^\{d\}\)\. Its elements are called Borel sets\. Intuitively, all “reasonable” geometric sets \(open, closed, countable unions/intersections\) are Borel sets; virtually every set you encounter in practice is a Borel set\.
- •𝒫2​\(ℝd\)\\mathcal\{P\}\_\{2\}\(\\mathbb\{R\}^\{d\}\): the set of all probability measures onℝd\\mathbb\{R\}^\{d\}with finite second moment, i\.e\., probability measuresμ\\musatisfying∫ℝd\|x\|2​𝑑μ​\(x\)<∞\\int\_\{\\mathbb\{R\}^\{d\}\}\|x\|^\{2\}\\,d\\mu\(x\)<\\infty\.
- •L2​\(ℝd\)L^\{2\}\(\\mathbb\{R\}^\{d\}\)\(orL2​\(ρ\)L^\{2\}\(\\rho\)\):space of square\-integrable functions\.L2​\(ℝd\)L^\{2\}\(\\mathbb\{R\}^\{d\}\)consists of all functionsffwith∫ℝd\|f​\(x\)\|2​𝑑x<∞\\int\_\{\\mathbb\{R\}^\{d\}\}\|f\(x\)\|^\{2\}\\,dx<\\infty\.L2​\(ρ\)L^\{2\}\(\\rho\)is the weighted version with weightρ\\rho:∫\|f\|2​ρ​𝑑x<∞\\int\|f\|^\{2\}\\rho\\,dx<\\infty\. The inner product is⟨f,g⟩L2​\(ρ\)=∫f​\(x\)​g​\(x\)​ρ​\(x\)​𝑑x\\langle f,g\\rangle\_\{L^\{2\}\(\\rho\)\}=\\int f\(x\)g\(x\)\\rho\(x\)\\,dx\. This is aHilbert space\(a complete inner product space\)\.
- •Cc∞​\(ℝd\)C\_\{c\}^\{\\infty\}\(\\mathbb\{R\}^\{d\}\):space of compactly supported smooth functions\. It contains all infinitely differentiable functions that vanish outside some bounded region\. Here: - –Smooth\(C∞C^\{\\infty\}\) = derivatives of all orders exist and are continuous - –Compactly supported\(ccfor compact support\) = there exists a bounded closed setKKsuch thatf​\(x\)=0f\(x\)=0for allx∉Kx\\notin K\. Intuitively, the function has “finite range of activity” and is zero far away\.

### Operators and Symbols

- •infA\\inf A\(infimum\): the greatest lower bound of a setAA\. For example,inf\{1/n:n∈ℕ\}=0\\inf\\\{1/n:n\\in\\mathbb\{N\}\\\}=0\. Similar tomin\\min, but allows the minimum not to be attained \(limit case\)\.
- •supA\\sup A\(supremum\): the least upper bound of a setAA, a generalization ofmax\\max\.
- •arg⁡minx⁡f​\(x\)\\arg\\min\_\{x\}f\(x\): the value ofxxat whichf​\(x\)f\(x\)attains its minimum \(not the minimum value itself, but the “location” where the minimum is achieved\)\.
- •∇f\\nabla f: gradient \(the vector of first partial derivatives offfwith respect to spatial variables\)\.
- •∇⋅v\\nabla\\cdot v: divergence \(the sum of partial derivatives of the components of a vector fieldvv\)\.
- •Δ​f=∇⋅\(∇f\)\\Delta f=\\nabla\\cdot\(\\nabla f\): Laplacian\.

### Key Terminology

- •a\.e\.\(almost everywhere\): a property “holds a\.e\.” means it may fail at some points, but the set of failure points has measure zero \(“occupies zero volume”\)\. For example, functionsffandggare equal a\.e\. under Lebesgue measure means\{x:f​\(x\)≠g​\(x\)\}\\\{x:f\(x\)\\neq g\(x\)\\\}has zero volume\.
- •Test function: a function inCc∞​\(ℝd\)C\_\{c\}^\{\\infty\}\(\\mathbb\{R\}^\{d\}\)\. The name comes from the fact that in the weak formulation of PDEs, we use them to “probe” whether the equation holds—multiply both sides of the PDE by a test function and integrate, converting the PDE into an integral identity that must hold for all test functions\. This avoids requiring smoothness of the solution itself\.

Note:For a detailed review of foundational concepts such as measures, absolute continuity, pushforward, and couplings, see Appendix[A](https://arxiv.org/html/2606.24157#A1)\. For basic concepts in convex analysis, see Appendix[G](https://arxiv.org/html/2606.24157#A7)\.

## Prologue: Things You May Already Know

For anyone who has worked on generative modeling, the continuity equation and the Fokker–Planck equation are familiar objects—one works with them constantly, often without naming them as such\.

### Scenario 1: How you encounter Fokker\-Planck in Diffusion Models

Open Song et al\. \(2021\)“Score\-Based Generative Modeling through SDEs”, and you will see the following setup\.

#### Discrete version: DDPM and SMLD

Beforethe unified framework of Song et al\. \(2021\), DDPM and SMLD each defined forward processes using discrete Markov chains:

DDPM forward process\(Ho et al\., 2020\): given noise scheduleβ1,…,βN\\beta\_\{1\},\\ldots,\\beta\_\{N\}:

𝐱i=1−βi​𝐱i−1\+βi​𝐳i−1,𝐳i−1∼𝒩​\(0,I\),i=1,…,N\\mathbf\{x\}\_\{i\}=\\sqrt\{1\-\\beta\_\{i\}\}\\,\\mathbf\{x\}\_\{i\-1\}\+\\sqrt\{\\beta\_\{i\}\}\\,\\mathbf\{z\}\_\{i\-1\},\\quad\\mathbf\{z\}\_\{i\-1\}\\sim\\mathcal\{N\}\(0,I\),\\quad i=1,\\ldots,NEach step: shrink the previous𝐱i−1\\mathbf\{x\}\_\{i\-1\}slightly \(multiply by1−βi<1\\sqrt\{1\-\\beta\_\{i\}\}<1\), then add a bit of noise \(βi​𝐳\\sqrt\{\\beta\_\{i\}\}\\,\\mathbf\{z\}\)\. AfterNNsteps,𝐱N≈𝒩​\(0,I\)\\mathbf\{x\}\_\{N\}\\approx\\mathcal\{N\}\(0,I\)\. Closed\-form expression:

𝐱i=α¯i​𝐱0\+1−α¯i​ϵ,ϵ∼𝒩​\(0,I\),α¯i:=∏j=1i\(1−βj\)\\mathbf\{x\}\_\{i\}=\\sqrt\{\\bar\{\\alpha\}\_\{i\}\}\\,\\mathbf\{x\}\_\{0\}\+\\sqrt\{1\-\\bar\{\\alpha\}\_\{i\}\}\\,\\boldsymbol\{\\epsilon\},\\quad\\boldsymbol\{\\epsilon\}\\sim\\mathcal\{N\}\(0,I\),\\quad\\bar\{\\alpha\}\_\{i\}:=\\prod\_\{j=1\}^\{i\}\(1\-\\beta\_\{j\}\)
SMLD forward process\(Song & Ermon, 2019\): given noise scalesσ1<σ2<⋯<σN\\sigma\_\{1\}<\\sigma\_\{2\}<\\cdots<\\sigma\_\{N\}:

𝐱i=𝐱i−1\+σi2−σi−12​𝐳i−1,i=1,…,N\\mathbf\{x\}\_\{i\}=\\mathbf\{x\}\_\{i\-1\}\+\\sqrt\{\\sigma\_\{i\}^\{2\}\-\\sigma\_\{i\-1\}^\{2\}\}\\,\\mathbf\{z\}\_\{i\-1\},\\quad i=1,\\ldots,NEach step only adds noise without rescaling\. The variance increases progressively toσN2\\sigma\_\{N\}^\{2\}\.

#### Continuous version: whenN→∞N\\to\\infty

The key insight of Song et al\. \(2021\): as the number of stepsN→∞N\\to\\infty, the discrete chain becomes a continuous SDE\.

DDPM→\\toVP\-SDE: letβi=β​\(i/N\)/N\\beta\_\{i\}=\\beta\(i/N\)/N, takeN→∞N\\to\\infty, the discrete update

𝐱i=1−βi​𝐱i−1\+βi​𝐳i−1≈𝐱i−1−βi2​𝐱i−1\+βi​𝐳i−1\\mathbf\{x\}\_\{i\}=\\sqrt\{1\-\\beta\_\{i\}\}\\,\\mathbf\{x\}\_\{i\-1\}\+\\sqrt\{\\beta\_\{i\}\}\\,\\mathbf\{z\}\_\{i\-1\}\\approx\\mathbf\{x\}\_\{i\-1\}\-\\frac\{\\beta\_\{i\}\}\{2\}\\mathbf\{x\}\_\{i\-1\}\+\\sqrt\{\\beta\_\{i\}\}\\,\\mathbf\{z\}\_\{i\-1\}becomes an SDE \(d​t=1/Ndt=1/N,𝐳​d​t=d​𝐰\\mathbf\{z\}\\sqrt\{dt\}=d\\mathbf\{w\}\):

d​𝐱=−β​\(t\)2​𝐱​d​t\+β​\(t\)​d​𝐰\(VP\-SDE\)d\\mathbf\{x\}=\-\\frac\{\\beta\(t\)\}\{2\}\\mathbf\{x\}\\,dt\+\\sqrt\{\\beta\(t\)\}\\,d\\mathbf\{w\}\\qquad\\text\{\(VP\-SDE\)\}\(1\)This is an OU process—the spring force−β2​𝐱\-\\frac\{\\beta\}\{2\}\\mathbf\{x\}pulls particles back to the origin, while noiseβ​d​𝐰\\sqrt\{\\beta\}\\,d\\mathbf\{w\}disperses them\.

SMLD→\\toVE\-SDE: similarly,𝐱i=𝐱i−1\+σi2−σi−12​𝐳i−1\\mathbf\{x\}\_\{i\}=\\mathbf\{x\}\_\{i\-1\}\+\\sqrt\{\\sigma\_\{i\}^\{2\}\-\\sigma\_\{i\-1\}^\{2\}\}\\,\\mathbf\{z\}\_\{i\-1\}in the continuous limit becomes:

d​𝐱=d​\[σ2​\(t\)\]d​t​d​𝐰\(VE\-SDE\)d\\mathbf\{x\}=\\sqrt\{\\frac\{d\[\\sigma^\{2\}\(t\)\]\}\{dt\}\}\\,d\\mathbf\{w\}\\qquad\\text\{\(VE\-SDE\)\}\(2\)Pure diffusion, no restoring force\. The variance grows monotonically \(“explodes”\), ultimately yielding a Gaussian with very large variance\.

Unified notation:d​𝐱=𝐟​\(𝐱,t\)​d​t\+g​\(t\)​d​𝐰d\\mathbf\{x\}=\\mathbf\{f\}\(\\mathbf\{x\},t\)\\,dt\+g\(t\)\\,d\\mathbf\{w\}, where

#### From SDE to the Fokker\-Planck equation

The probability densityρt​\(x\)\\rho\_\{t\}\(x\)induced by these SDEs satisfies theFokker\-Planck equation\(derivation in Section 4\):

∂tρt=−∇⋅\(𝐟​ρt\)\+g22​Δ​ρt\\partial\_\{t\}\\rho\_\{t\}=\-\\nabla\\cdot\(\\mathbf\{f\}\\,\\rho\_\{t\}\)\+\\frac\{g^\{2\}\}\{2\}\\Delta\\rho\_\{t\}\(3\)Substituting𝐟=−β2​𝐱\\mathbf\{f\}=\-\\frac\{\\beta\}\{2\}\\mathbf\{x\},g2=βg^\{2\}=\\betafor VP\-SDE:

∂tρt=β2​∇⋅\(𝐱​ρt\)\+β2​Δ​ρt\\partial\_\{t\}\\rho\_\{t\}=\\frac\{\\beta\}\{2\}\\nabla\\cdot\(\\mathbf\{x\}\\,\\rho\_\{t\}\)\+\\frac\{\\beta\}\{2\}\\Delta\\rho\_\{t\}

#### How Fokker\-Planck is used in the paper: deriving the Reverse\-time SDE

In the Diffusion Model paper, thesole useof Fokker\-Planck is to derive thereverse\-time SDE\. Here is the complete derivation\.

Goal: the forward processd​𝐱=𝐟​d​t\+g​d​𝐰d\\mathbf\{x\}=\\mathbf\{f\}\\,dt\+g\\,d\\mathbf\{w\}runs fromt=0t=0\(data\) tot=Tt=T\(noise\)\. We want to find an SDE running fromTTto0whose marginal distributions are exactlyρt\\rho\_\{t\}in reversed time\.

Supposethe reverse SDE \(lett¯=T−t\\bar\{t\}=T\-tbe the reverse time\) takes the form:

d​𝐱=𝐟¯​d​t¯\+g​d​𝐰¯d\\mathbf\{x\}=\\bar\{\\mathbf\{f\}\}\\,d\\bar\{t\}\+g\\,d\\bar\{\\mathbf\{w\}\}Its marginal distributionρ¯t¯:=ρT−t¯\\bar\{\\rho\}\_\{\\bar\{t\}\}:=\\rho\_\{T\-\\bar\{t\}\}should satisfy thereverse Fokker\-Planck:

∂t¯ρ¯t¯=−∇⋅\(𝐟¯​ρ¯t¯\)\+g22​Δ​ρ¯t¯\\partial\_\{\\bar\{t\}\}\\bar\{\\rho\}\_\{\\bar\{t\}\}=\-\\nabla\\cdot\(\\bar\{\\mathbf\{f\}\}\\,\\bar\{\\rho\}\_\{\\bar\{t\}\}\)\+\\frac\{g^\{2\}\}\{2\}\\Delta\\bar\{\\rho\}\_\{\\bar\{t\}\}\(4\)
Key step: fromρ¯t¯=ρT−t¯\\bar\{\\rho\}\_\{\\bar\{t\}\}=\\rho\_\{T\-\\bar\{t\}\}, we get∂t¯ρ¯t¯=−∂tρt\\partial\_\{\\bar\{t\}\}\\bar\{\\rho\}\_\{\\bar\{t\}\}=\-\\partial\_\{t\}\\rho\_\{t\}\.

Substituting the forward Fokker\-Planck \([3](https://arxiv.org/html/2606.24157#Sx2.E3)\):

−∂tρt=−\[−∇⋅\(𝐟​ρt\)\+g22​Δ​ρt\]=∇⋅\(𝐟​ρt\)−g22​Δ​ρt\-\\partial\_\{t\}\\rho\_\{t\}=\-\\left\[\-\\nabla\\cdot\(\\mathbf\{f\}\\rho\_\{t\}\)\+\\tfrac\{g^\{2\}\}\{2\}\\Delta\\rho\_\{t\}\\right\]=\\nabla\\cdot\(\\mathbf\{f\}\\rho\_\{t\}\)\-\\tfrac\{g^\{2\}\}\{2\}\\Delta\\rho\_\{t\}
Setting this equal to the right\-hand side of \([4](https://arxiv.org/html/2606.24157#Sx2.E4)\) \(both expressed in terms ofρt\\rho\_\{t\}\):

∇⋅\(𝐟​ρt\)−g22​Δ​ρt=−∇⋅\(𝐟¯​ρt\)\+g22​Δ​ρt\\nabla\\cdot\(\\mathbf\{f\}\\rho\_\{t\}\)\-\\tfrac\{g^\{2\}\}\{2\}\\Delta\\rho\_\{t\}=\-\\nabla\\cdot\(\\bar\{\\mathbf\{f\}\}\\,\\rho\_\{t\}\)\+\\tfrac\{g^\{2\}\}\{2\}\\Delta\\rho\_\{t\}
Rearranging:

∇⋅\(𝐟​ρt\)=−∇⋅\(𝐟¯​ρt\)\+g2​Δ​ρt\\nabla\\cdot\(\\mathbf\{f\}\\rho\_\{t\}\)=\-\\nabla\\cdot\(\\bar\{\\mathbf\{f\}\}\\,\\rho\_\{t\}\)\+g^\{2\}\\Delta\\rho\_\{t\}
Using the identityΔ​ρ=∇⋅\(∇ρ\)=∇⋅\(ρ​∇log⁡ρ\)\\Delta\\rho=\\nabla\\cdot\(\\nabla\\rho\)=\\nabla\\cdot\(\\rho\\,\\nabla\\log\\rho\):

∇⋅\(𝐟​ρt\)=∇⋅\[−𝐟¯​ρt\+g2​ρt​∇log⁡ρt\]\\nabla\\cdot\(\\mathbf\{f\}\\rho\_\{t\}\)=\\nabla\\cdot\\\!\\left\[\-\\bar\{\\mathbf\{f\}\}\\,\\rho\_\{t\}\+g^\{2\}\\rho\_\{t\}\\,\\nabla\\log\\rho\_\{t\}\\right\]
Comparing inside the∇⋅\\nabla\\cdot, solving for𝐟¯\\bar\{\\mathbf\{f\}\}:

𝐟¯=−𝐟\+g2​∇log⁡ρt\\boxed\{\\bar\{\\mathbf\{f\}\}=\-\\mathbf\{f\}\+g^\{2\}\\,\\nabla\\log\\rho\_\{t\}\}
Therefore, the reverse SDE \(returning to the original time parametertt, running backward fromTTto0\) is:

d​𝐱=\[𝐟​\(𝐱,t\)−g​\(t\)2​∇log⁡ρt​\(𝐱\)\]​d​t\+g​\(t\)​d​𝐰¯d\\mathbf\{x\}=\\left\[\\mathbf\{f\}\(\\mathbf\{x\},t\)\-g\(t\)^\{2\}\\,\\nabla\\log\\rho\_\{t\}\(\\mathbf\{x\}\)\\right\]dt\+g\(t\)\\,d\\bar\{\\mathbf\{w\}\}\(5\)
This is the Anderson \(1982\) theorem\. Note that∇log⁡ρt\\nabla\\log\\rho\_\{t\}—the score—is the only unknown quantity that needs to be learned\.

#### Discrete counterpart: DDPM denoising steps

ApplyEuler\-Maruyama discretizationto the reverse SDE \([5](https://arxiv.org/html/2606.24157#Sx2.E5)\)\. Recall: for a general SDEd​X=a​\(X,t\)​d​t\+b​\(t\)​d​WdX=a\(X,t\)\\,dt\+b\(t\)\\,dW, the Euler\-Maruyama method simply “replacesd​tdtwithΔ​t\\Delta tandd​WdWwithΔ​t​𝐳\\sqrt\{\\Delta t\}\\,\\mathbf\{z\}” \(𝐳∼𝒩​\(0,I\)\\mathbf\{z\}\\sim\\mathcal\{N\}\(0,I\)\):

Xn\+1=Xn\+a​\(Xn,tn\)​Δ​t\+b​\(tn\)​Δ​t​𝐳nX\_\{n\+1\}=X\_\{n\}\+a\(X\_\{n\},t\_\{n\}\)\\,\\Delta t\+b\(t\_\{n\}\)\\,\\sqrt\{\\Delta t\}\\,\\mathbf\{z\}\_\{n\}This is the simplest numerical scheme for SDEs, analogous to the Euler method for ODEs, but with the additional random termΔ​t​𝐳\\sqrt\{\\Delta t\}\\,\\mathbf\{z\}\.

For the reverse SDE \(a=𝐟−g2​∇log⁡ρta=\\mathbf\{f\}\-g^\{2\}\\nabla\\log\\rho\_\{t\},b=gb=g\), takingΔ​t=1/N\\Delta t=1/N, we obtain the DDPM denoising formula:

𝐱i−1=11−βi​\(𝐱i\+βi​sθ​\(𝐱i,i\)\)\+βi​𝐳,𝐳∼𝒩​\(0,I\)\\mathbf\{x\}\_\{i\-1\}=\\frac\{1\}\{\\sqrt\{1\-\\beta\_\{i\}\}\}\\left\(\\mathbf\{x\}\_\{i\}\+\\beta\_\{i\}\\,s\_\{\\theta\}\(\\mathbf\{x\}\_\{i\},i\)\\right\)\+\\sqrt\{\\beta\_\{i\}\}\\,\\mathbf\{z\},\\quad\\mathbf\{z\}\\sim\\mathcal\{N\}\(0,I\)wheresθ≈∇log⁡ρis\_\{\\theta\}\\approx\\nabla\\log\\rho\_\{i\}is the score learned by the neural network\. Training objective \(denoising score matching; Hyvärinen, 2005; Vincent, 2011\):

ℒ=𝔼t,𝐱0,ϵ​\[‖sθ​\(α¯t​𝐱0\+1−α¯t​ϵ,t\)−\(−ϵ1−α¯t\)‖2\]\\mathcal\{L\}=\\mathbb\{E\}\_\{t,\\mathbf\{x\}\_\{0\},\\boldsymbol\{\\epsilon\}\}\\left\[\\\|s\_\{\\theta\}\(\\sqrt\{\\bar\{\\alpha\}\_\{t\}\}\\,\\mathbf\{x\}\_\{0\}\+\\sqrt\{1\-\\bar\{\\alpha\}\_\{t\}\}\\,\\boldsymbol\{\\epsilon\},\\;t\)\-\\left\(\-\\frac\{\\boldsymbol\{\\epsilon\}\}\{\\sqrt\{1\-\\bar\{\\alpha\}\_\{t\}\}\}\\right\)\\\|^\{2\}\\right\]Note∇𝐱log⁡pt\|0​\(𝐱\|𝐱0\)=−𝐱−α¯t​𝐱01−α¯t=−ϵ1−α¯t\\nabla\_\{\\mathbf\{x\}\}\\log p\_\{t\|0\}\(\\mathbf\{x\}\|\\mathbf\{x\}\_\{0\}\)=\-\\frac\{\\mathbf\{x\}\-\\sqrt\{\\bar\{\\alpha\}\_\{t\}\}\\mathbf\{x\}\_\{0\}\}\{1\-\\bar\{\\alpha\}\_\{t\}\}=\-\\frac\{\\boldsymbol\{\\epsilon\}\}\{\\sqrt\{1\-\\bar\{\\alpha\}\_\{t\}\}\}, so score matching is essentially training the network to predict the direction of the noiseϵ\\boldsymbol\{\\epsilon\}\.

#### Connection to this article

In the diffusion paper, the derivation chain is:

SDE→Fokker\-Planckdensity evolution​ρt→time reversalreverse SDE with score→learn scoregeneration\\text\{SDE\}\\xrightarrow\{\\text\{Fokker\-Planck\}\}\\text\{density evolution\}\\rho\_\{t\}\\xrightarrow\{\\text\{time reversal\}\}\\text\{reverse SDE with score\}\\xrightarrow\{\\text\{learn score\}\}\\text\{generation\}Here Fokker–Planck appears only as an intermediate step in deriving the score from the SDE, and is set aside once the derivation is complete\.

This article takes the equation more seriously\. Rewriting Fokker\-Planck in the form of a continuity equation∂tρ=−∇⋅\(ρ​v\)\\partial\_\{t\}\\rho=\-\\nabla\\cdot\(\\rho\\,v\)\(usingΔ​ρ=∇⋅\(ρ​∇log⁡ρ\)\\Delta\\rho=\\nabla\\cdot\(\\rho\\nabla\\log\\rho\)\):

∂tρt=−∇⋅\(ρt​\[𝐟−g22​∇log⁡ρt\]⏟=vt\)\\partial\_\{t\}\\rho\_\{t\}=\-\\nabla\\cdot\\\!\\left\(\\rho\_\{t\}\\underbrace\{\\left\[\\mathbf\{f\}\-\\tfrac\{g^\{2\}\}\{2\}\\nabla\\log\\rho\_\{t\}\\right\]\}\_\{=\\,v\_\{t\}\}\\right\)so the velocity field isvt=𝐟−g22​∇log⁡ρtv\_\{t\}=\\mathbf\{f\}\-\\frac\{g^\{2\}\}\{2\}\\nabla\\log\\rho\_\{t\}\. For VP\-SDE \(𝐟=−β2​𝐱\\mathbf\{f\}=\-\\frac\{\\beta\}\{2\}\\mathbf\{x\},g2=βg^\{2\}=\\beta\), this givesvt=−β2​𝐱−β2​∇log⁡ρt=−β2​∇\(\|x\|22\+log⁡ρt\)v\_\{t\}=\-\\frac\{\\beta\}\{2\}\\mathbf\{x\}\-\\frac\{\\beta\}\{2\}\\nabla\\log\\rho\_\{t\}=\-\\frac\{\\beta\}\{2\}\\nabla\\left\(\\frac\{\|x\|^\{2\}\}\{2\}\+\\log\\rho\_\{t\}\\right\)\.

Comparing with the Wasserstein gradientgradW​ℱ=∇V\+∇log⁡ρ\\mathrm\{grad\}\_\{W\}\\mathcal\{F\}=\\nabla V\+\\nabla\\log\\rhoderived in the main text \(takingV=12​\|x\|2V=\\frac\{1\}\{2\}\|x\|^\{2\}\), we seevt=−β2​gradW​ℱv\_\{t\}=\-\\frac\{\\beta\}\{2\}\\,\\mathrm\{grad\}\_\{W\}\\mathcal\{F\}—the velocity field is the*negative*gradient, up to the positive time\-scaling factorβ/2\\beta/2\.

Section 5 will prove:vtv\_\{t\}is precisely thenegative gradientofKL​\(ρt∥π\)\\mathrm\{KL\}\(\\rho\_\{t\}\\\|\\pi\)\(π=𝒩​\(0,I\)∝e−\|x\|2/2\\pi=\\mathcal\{N\}\(0,I\)\\propto e^\{\-\|x\|^\{2\}/2\}\) under the Wasserstein metric \(up to the positive factorβ/2\\beta/2, which is merely a time reparametrization\)\. That is, the forward process of VP\-SDE==thegradient flowof free energy in probability space\.

Section 6 will prove: each DDPM denoising step==one step of the JKO scheme—implicit gradient descent in probability space\.

In short, the derivation above uses Fokker–Planck only to extract the score\. The equation carries considerably more geometric structure, and recovering it is the purpose of this article\.

### Scenario 2: How you encounter the continuity equation in Flow Matching

Open Lipman et al\. \(2023\)“Flow Matching for Generative Modeling”, and the story is entirely different—no SDEs, no noise, everything is deterministic\.

Thecore objectsare a probability density pathptp\_\{t\}and a velocity fieldvtv\_\{t\}, satisfying:

∂tpt\+∇⋅\(pt​vt\)=0\(continuity equation\)\\partial\_\{t\}p\_\{t\}\+\\nabla\\cdot\(p\_\{t\}\\,v\_\{t\}\)=0\\qquad\\text\{\(continuity equation\)\}\(6\)
Where does this equation come from?If particles follow the ODEd​Xtd​t=vt​\(Xt\)\\frac\{dX\_\{t\}\}\{dt\}=v\_\{t\}\(X\_\{t\}\)with initial conditionX0∼p0X\_\{0\}\\sim p\_\{0\}, what equation does the densityptp\_\{t\}ofXtX\_\{t\}satisfy? For any test functionφ\\varphi:

dd​t​∫φ​pt​𝑑x=dd​t​𝔼​\[φ​\(Xt\)\]=𝔼​\[∇φ​\(Xt\)⋅vt​\(Xt\)\]=∫∇φ⋅vt​pt​d​x\\frac\{d\}\{dt\}\\int\\varphi\\,p\_\{t\}\\,dx=\\frac\{d\}\{dt\}\\mathbb\{E\}\[\\varphi\(X\_\{t\}\)\]=\\mathbb\{E\}\[\\nabla\\varphi\(X\_\{t\}\)\\cdot v\_\{t\}\(X\_\{t\}\)\]=\\int\\nabla\\varphi\\cdot v\_\{t\}\\,p\_\{t\}\\,dxIntegrating the right side by parts:=−∫φ​∇⋅\(vt​pt\)​𝑑x=\-\\int\\varphi\\,\\nabla\\cdot\(v\_\{t\}\\,p\_\{t\}\)\\,dx\. Comparing with the left side=∫φ​∂tpt​d​x=\\int\\varphi\\,\\partial\_\{t\}p\_\{t\}\\,dx, we get∂tpt=−∇⋅\(pt​vt\)\\partial\_\{t\}p\_\{t\}=\-\\nabla\\cdot\(p\_\{t\}v\_\{t\}\)\.

Intuitive meaning: the change in density==the negative of the net outflux\. Ifvtv\_\{t\}“carries mass away” fromxx\(∇⋅\(pt​vt\)\>0\\nabla\\cdot\(p\_\{t\}v\_\{t\}\)\>0\), then the density atxxdecreases\.

Thetraining objectiveis to fit a known target velocity fieldut​\(x\|x1\)u\_\{t\}\(x\|x\_\{1\}\)with a neural networkvθ​\(x,t\)v\_\{\\theta\}\(x,t\):

ℒCFM\(θ\)=𝔼t,x1∼q,x∼pt​\(x\|x1\)∥vθ\(x,t\)−ut\(x\|x1\)∥2\\mathcal\{L\}\_\{\\text\{CFM\}\}\(\\theta\)=\\mathbb\{E\}\_\{t,\\,x\_\{1\}\\sim q,\\,x\\sim p\_\{t\}\(x\|x\_\{1\}\)\}\\\|v\_\{\\theta\}\(x,t\)\-u\_\{t\}\(x\|x\_\{1\}\)\\\|^\{2\}
For the OT path \(the simplest choice\), the conditional probability path and conditional velocity field are:

pt​\(x\|x1\)=𝒩​\(x\|t​x1,\(1−t\)2​I\),ut​\(x\|x1\)=x1−x1−tp\_\{t\}\(x\|x\_\{1\}\)=\\mathcal\{N\}\\bigl\(x\\;\\big\|\\;t\\,x\_\{1\},\\;\(1\-t\)^\{2\}I\\bigr\),\\qquad u\_\{t\}\(x\|x\_\{1\}\)=\\frac\{x\_\{1\}\-x\}\{1\-t\}This is a “straight line” from𝒩​\(0,I\)\\mathcal\{N\}\(0,I\)toδx1\\delta\_\{x\_\{1\}\}\. Each particle travels at constant velocityx1−x01\\frac\{x\_\{1\}\-x\_\{0\}\}\{1\}from its starting pointx0x\_\{0\}to its endpointx1x\_\{1\}\.

For generation, one only needs to solve an ODE:d​xd​t=vθ​\(x,t\)\\frac\{dx\}\{dt\}=v\_\{\\theta\}\(x,t\), integrating fromx0∼𝒩​\(0,I\)x\_\{0\}\\sim\\mathcal\{N\}\(0,I\)tot=1t=1\.

Summary: in the context of Flow Matching, the continuity equation is the constraint linking “velocity field→\\todensity path\.”The paper uses it to prove: ifvθv\_\{\\theta\}accurately learns the conditional velocity field, then the marginal density pathptp\_\{t\}it generates is the desired one\. One question is left unaddressed, however: how does this equation relate to Fokker–Planck? Are the two different expressions of the same object?

### Scenario 3: Probability Flow ODE—where the two equations meet

A central observation of Song et al\. \(2021\) connects the two pictures\. They proved:

Key InsightForanyforward SDEd​𝐱=𝐟​d​t\+g​d​𝐰d\\mathbf\{x\}=\\mathbf\{f\}\\,dt\+g\\,d\\mathbf\{w\}, there exists adeterministic ODE\(probability flow ODE\):d​𝐱d​t=𝐟​\(𝐱,t\)⏟drift−12​g​\(t\)2​∇log⁡ρt​\(𝐱\)⏟score\\frac\{d\\mathbf\{x\}\}\{dt\}=\\underbrace\{\\mathbf\{f\}\(\\mathbf\{x\},t\)\}\_\{\\text\{drift\}\}\-\\frac\{1\}\{2\}g\(t\)^\{2\}\\underbrace\{\\nabla\\log\\rho\_\{t\}\(\\mathbf\{x\}\)\}\_\{\\text\{score\}\}\(7\)such that the marginal distributionρt\\rho\_\{t\}of particles isexactly the same as that of the SDE\.

Rather than state the result, let us derive it\.

Derivation: reading off the probability flow ODE from Fokker–Planck

Starting point: we know thatρt\\rho\_\{t\}satisfies the Fokker\-Planck equation

∂tρt=−∇⋅\(𝐟​ρt\)\+g22​Δ​ρt\\partial\_\{t\}\\rho\_\{t\}=\-\\nabla\\cdot\(\\mathbf\{f\}\\,\\rho\_\{t\}\)\+\\frac\{g^\{2\}\}\{2\}\\Delta\\rho\_\{t\}Now ask:can we find a velocity fieldvtv\_\{t\}such thatρt\\rho\_\{t\}also satisfies the continuity equation∂tρt=−∇⋅\(ρt​vt\)\\partial\_\{t\}\\rho\_\{t\}=\-\\nabla\\cdot\(\\rho\_\{t\}\\,v\_\{t\}\)?

If so, there exists a deterministic ODEx˙=vt​\(x\)\\dot\{x\}=v\_\{t\}\(x\)that produces the same density evolution—no noise needed\.

Key trick: “disguise” the diffusion term of Fokker\-Planck as a transport term\. Note the identity:

Δ​ρ=∇⋅\(∇ρ\)=∇⋅\(ρ⋅∇ρρ\)=∇⋅\(ρ​∇log⁡ρ\)\\Delta\\rho=\\nabla\\cdot\(\\nabla\\rho\)=\\nabla\\cdot\\\!\\left\(\\rho\\cdot\\frac\{\\nabla\\rho\}\{\\rho\}\\right\)=\\nabla\\cdot\(\\rho\\,\\nabla\\log\\rho\)Substituting into Fokker\-Planck:

∂tρt\\displaystyle\\partial\_\{t\}\\rho\_\{t\}=−∇⋅\(𝐟​ρt\)\+g22​∇⋅\(ρt​∇log⁡ρt\)\\displaystyle=\-\\nabla\\cdot\(\\mathbf\{f\}\\,\\rho\_\{t\}\)\+\\frac\{g^\{2\}\}\{2\}\\nabla\\cdot\(\\rho\_\{t\}\\,\\nabla\\log\\rho\_\{t\}\)=−∇⋅\(𝐟​ρt−g22​ρt​∇log⁡ρt\)\\displaystyle=\-\\nabla\\cdot\\\!\\left\(\\mathbf\{f\}\\,\\rho\_\{t\}\-\\frac\{g^\{2\}\}\{2\}\\rho\_\{t\}\\,\\nabla\\log\\rho\_\{t\}\\right\)=−∇⋅\(ρt​\[𝐟−g22​∇log⁡ρt\]⏟=⁣:vt\)\\displaystyle=\-\\nabla\\cdot\\\!\\left\(\\rho\_\{t\}\\underbrace\{\\left\[\\mathbf\{f\}\-\\frac\{g^\{2\}\}\{2\}\\nabla\\log\\rho\_\{t\}\\right\]\}\_\{=:\\,v\_\{t\}\}\\right\)
The velocity field can be read off directly:vt=𝐟−g22​∇log⁡ρtv\_\{t\}=\\mathbf\{f\}\-\\frac\{g^\{2\}\}\{2\}\\nabla\\log\\rho\_\{t\}ensures∂tρt\+∇⋅\(ρt​vt\)=0\\partial\_\{t\}\\rho\_\{t\}\+\\nabla\\cdot\(\\rho\_\{t\}v\_\{t\}\)=0\. The corresponding ODE is \([7](https://arxiv.org/html/2606.24157#Sx2.E7)\)\.

Verification \(from the other direction\):substitutevt=𝐟−g22​∇log⁡ρtv\_\{t\}=\\mathbf\{f\}\-\\frac\{g^\{2\}\}\{2\}\\nabla\\log\\rho\_\{t\}into the continuity equation and expand:

−∇⋅\(ρt​vt\)\\displaystyle\-\\nabla\\cdot\(\\rho\_\{t\}v\_\{t\}\)=−∇⋅\(ρt​𝐟−g22​ρt​∇log⁡ρt\)\\displaystyle=\-\\nabla\\cdot\\\!\\left\(\\rho\_\{t\}\\mathbf\{f\}\-\\frac\{g^\{2\}\}\{2\}\\rho\_\{t\}\\nabla\\log\\rho\_\{t\}\\right\)=−∇⋅\(ρt​𝐟\)\+g22​∇⋅\(ρt⋅∇ρtρt\)\\displaystyle=\-\\nabla\\cdot\(\\rho\_\{t\}\\mathbf\{f\}\)\+\\frac\{g^\{2\}\}\{2\}\\nabla\\cdot\\\!\\left\(\\rho\_\{t\}\\cdot\\frac\{\\nabla\\rho\_\{t\}\}\{\\rho\_\{t\}\}\\right\)=−∇⋅\(ρt​𝐟\)\+g22​∇⋅\(∇ρt\)\\displaystyle=\-\\nabla\\cdot\(\\rho\_\{t\}\\mathbf\{f\}\)\+\\frac\{g^\{2\}\}\{2\}\\nabla\\cdot\(\\nabla\\rho\_\{t\}\)=−∇⋅\(ρt​𝐟\)\+g22​Δ​ρt✓\\displaystyle=\-\\nabla\\cdot\(\\rho\_\{t\}\\mathbf\{f\}\)\+\\frac\{g^\{2\}\}\{2\}\\Delta\\rho\_\{t\}\\quad\\checkmarkThis is exactly Fokker\-Planck\. The two equations describethe same density trajectoryρt\\rho\_\{t\}\.

The point deserves emphasis:

- •Fokker\-Planck is an equation “with diffusion”—it hasΔ​ρ\\Delta\\rho\(second\-order term\)
- •The continuity equation is a “pure transport” equation—only∇⋅\(ρ​v\)\\nabla\\cdot\(\\rho v\)\(first\-order term\)
- •Yet they describethe same density trajectoryρt\\rho\_\{t\}

What distinguishes them is the velocity field\. The effective velocityv=𝐟−g22​∇log⁡ρv=\\mathbf\{f\}\-\\frac\{g^\{2\}\}\{2\}\\nabla\\log\\rhoof the transport form absorbs the diffusion into the velocity: the term−g22​∇log⁡ρ=−g22​∇ρρ\-\\frac\{g^\{2\}\}\{2\}\\nabla\\log\\rho=\-\\frac\{g^\{2\}\}\{2\}\\frac\{\\nabla\\rho\}\{\\rho\}convertsg22​Δ​ρ\\frac\{g^\{2\}\}\{2\}\\Delta\\rhoexactly into−∇⋅\(ρ⋅\(−g22​∇log⁡ρ\)\)\-\\nabla\\cdot\(\\rho\\cdot\(\-\\frac\{g^\{2\}\}\{2\}\\nabla\\log\\rho\)\)\.

In Flow Matching language:the probability flow ODE of a diffusion model is itself a flow, with velocity fieldvt=f−g22​scorev\_\{t\}=f\-\\frac\{g^\{2\}\}\{2\}\\,\\text\{score\}\. Diffusion and Flow Matching are thus two perspectives on a single object\.

A note on EDM\.Karras et al\. \(2022\), in“Elucidating the Design Space of Diffusion\-Based Generative Models,”push this ODE\-centric view furthest\. They write the forward process simply as the data convolved with Gaussian noise of scaleσ\\sigma—marginalsp​\(x;σ\)p\(x;\\sigma\)—and center generation on the probability flow ODE

d​xd​t=−σ˙​\(t\)​σ​\(t\)​∇xlog⁡p​\(x;σ​\(t\)\)\.\\frac\{dx\}\{dt\}=\-\\dot\{\\sigma\}\(t\)\\,\\sigma\(t\)\\,\\nabla\_\{x\}\\log p\\bigl\(x;\\sigma\(t\)\\bigr\)\.VP\- and VE\-SDE then become particular choices among a small set of*orthogonal*design axes—the noise scheduleσ​\(t\)\\sigma\(t\), an overall scalings​\(t\)s\(t\), the network preconditioning, the training loss weighting, and the ODE solver \(they adopt a second\-order Heun integrator\)—each of which can be tuned in isolation\. In the language of this article, these axes are reparametrizations of*time and scale*along one and the same density trajectory: the underlying geometric object—the Fokker–Planck flow, equivalently the Wasserstein gradient flow of the free energy—is invariant under them\. EDM is thus practical evidence for the article’s thesis that what matters is the trajectory and its score, not the particular SDE used to present it\.

A note on energy\-based formulations\.A complementary line of recent work makes the*equilibrium*side of this geometry explicit\.*Energy Matching*\(Balcerak et al\., 2025\) replaces the time\-dependent score or velocity network by a single*time\-independent*scalar potentialVθV\_\{\\theta\}: far from the data it transports samples by−∇Vθ\-\\nabla V\_\{\\theta\}\(an optimal\-transport flow\), and near the data it relaxes to the Boltzmann equilibriumρ∝e−Vθ/ε\\rho\\propto e^\{\-V\_\{\\theta\}/\\varepsilon\}\.*Equilibrium Matching*\(Wang et al\., 2025\) likewise discards time conditioning, learning a single time\-invariant gradient field of an implicit energy landscape whose stationary points are the data; generation then becomes gradient descent on that landscape, with step size and compute chosen at inference\. In the language of this article both are direct parametrizations of the*free\-energy / equilibrium*structure: the equilibriumρ∝e−V\\rho\\propto e^\{\-V\}is the minimizer of the free energy, and the relaxation toward it is the Wasserstein gradient flow \(Section 6 derives Energy Matching from the JKO scheme\)\. Where EDM emphasizes the transport, probability\-flow\-ODE face of the geometry, these emphasize its gradient\-flow, Gibbs face—two faces of one object\.

### How deep do these connections run?

Most expositions stop here: Fokker–Planck and the continuity equation are “equivalent,” and score and velocity differ by a drift\. A few further questions, however, repay attention\.

1. 1\.Why can Fokker\-Planck be written in the form of a continuity equation? This is not merely an algebraic coincidence\.∂tρ\+∇⋅\(ρ​v\)=0\\partial\_\{t\}\\rho\+\\nabla\\cdot\(\\rho v\)=0means the evolution of density can be understood as probability mass “flowing” under avelocity field\. But in what “space” does this flow take place? The answer: it takes place in the space of probability distributions𝒫2​\(ℝd\)\\mathcal\{P\}\_\{2\}\(\\mathbb\{R\}^\{d\}\), and the “velocity field”vvis precisely atangent vectorin this space\.
2. 2\.Why does VP\-SDE ultimately converge to𝒩​\(0,I\)\\mathcal\{N\}\(0,I\), and not some other distribution? For VP\-SDE,𝐟=−β2​𝐱\\mathbf\{f\}=\-\\frac\{\\beta\}\{2\}\\mathbf\{x\}implies potentialV​\(x\)=12​\|x\|2V\(x\)=\\frac\{1\}\{2\}\|x\|^\{2\}, so the Gibbs distribution isπ∝e−V=e−\|x\|2/2=𝒩​\(0,I\)\\pi\\propto e^\{\-V\}=e^\{\-\|x\|^\{2\}/2\}=\\mathcal\{N\}\(0,I\)\. The Fokker\-Planck equation drivesρt\\rho\_\{t\}to converge toπ\\pi—but this is no accident: it isminimizingKL​\(ρt∥π\)\\mathrm\{KL\}\(\\rho\_\{t\}\\\|\\pi\)\. Proof:computedd​t​KL​\(ρt∥π\)\\frac\{d\}\{dt\}\\mathrm\{KL\}\(\\rho\_\{t\}\\\|\\pi\)\. Letℱ​\(ρ\)=KL​\(ρ∥π\)=∫ρ​log⁡ρπ​d​x\\mathcal\{F\}\(\\rho\)=\\mathrm\{KL\}\(\\rho\\\|\\pi\)=\\int\\rho\\log\\frac\{\\rho\}\{\\pi\}\\,dx\. dd​t​ℱ​\(ρt\)\\displaystyle\\frac\{d\}\{dt\}\\mathcal\{F\}\(\\rho\_\{t\}\)=∫\(log⁡ρtπ\+1\)​∂tρt​d​x=∫log⁡ρtπ​∂tρt​d​x\(since∫∂tρt=0\)\\displaystyle=\\int\\left\(\\log\\frac\{\\rho\_\{t\}\}\{\\pi\}\+1\\right\)\\partial\_\{t\}\\rho\_\{t\}\\,dx=\\int\\log\\frac\{\\rho\_\{t\}\}\{\\pi\}\\,\\partial\_\{t\}\\rho\_\{t\}\\,dx\\quad\\text\{\(since $\\int\\partial\_\{t\}\\rho\_\{t\}=0$\)\}Substituting∂tρt=∇⋅\(ρt​∇log⁡ρtπ\)\\partial\_\{t\}\\rho\_\{t\}=\\nabla\\cdot\(\\rho\_\{t\}\\nabla\\log\\frac\{\\rho\_\{t\}\}\{\\pi\}\)\(Fokker\-Planck rewritten with velocity fieldv=−∇log⁡ρπv=\-\\nabla\\log\\frac\{\\rho\}\{\\pi\}\): =∫log⁡ρtπ⋅∇⋅\(ρt​∇log⁡ρtπ\)​𝑑x\\displaystyle=\\int\\log\\frac\{\\rho\_\{t\}\}\{\\pi\}\\cdot\\nabla\\cdot\\\!\\left\(\\rho\_\{t\}\\nabla\\log\\frac\{\\rho\_\{t\}\}\{\\pi\}\\right\)dx=−∫ρt​\|∇log⁡ρtπ\|2​𝑑x\(integration by parts\)\\displaystyle=\-\\int\\rho\_\{t\}\\left\|\\nabla\\log\\frac\{\\rho\_\{t\}\}\{\\pi\}\\right\|^\{2\}dx\\quad\\text\{\(integration by parts\)\}≤0\\displaystyle\\leq 0Equality holds if and only if∇log⁡ρtπ=0\\nabla\\log\\frac\{\\rho\_\{t\}\}\{\\pi\}=0, i\.e\.,ρt=π\\rho\_\{t\}=\\pi\. ThereforeKL​\(ρt∥π\)\\mathrm\{KL\}\(\\rho\_\{t\}\\\|\\pi\)isstrictly monotonically decreasinguntilρt=π\\rho\_\{t\}=\\pi\. This is the hallmark of “gradient descent”: the objective function decreases monotonically along the trajectory, at a rate proportional to the “squared norm of the gradient”∫ρ​\|∇log⁡ρπ\|2\\int\\rho\|\\nabla\\log\\frac\{\\rho\}\{\\pi\}\|^\{2\}\(this is called the relative Fisher information\)\.
3. 3\.Why is the OT path in Flow Matching “optimal”? Lipman et al\. found that models trained with the conditional OT pathxt=\(1−t\)​x0\+t​x1x\_\{t\}=\(1\-t\)x\_\{0\}\+t\\,x\_\{1\}\(straight lines\) perform better and sample faster than models trained with diffusion paths \(curved, roundabout\)\. Why should straight lines be preferable? A short computation settles it\. For the conditional OT path, particles move fromx0x\_\{0\}tox1x\_\{1\}at constant speed: Xt=\(1−t\)​x0\+t​x1,X˙t=x1−x0\(constant velocity\)X\_\{t\}=\(1\-t\)\\,x\_\{0\}\+t\\,x\_\{1\},\\qquad\\dot\{X\}\_\{t\}=x\_\{1\}\-x\_\{0\}\\quad\\text\{\(constant velocity\)\}“Kinetic energy integral” \(measuring how “costly” the path is\): ∫01\|x1−x0\|2​𝑑t=\|x1−x0\|2\\int\_\{0\}^\{1\}\|x\_\{1\}\-x\_\{0\}\|^\{2\}\\,dt=\|x\_\{1\}\-x\_\{0\}\|^\{2\}Averaging over all particles:𝔼​\[\|x1−x0\|2\]=W22​\(p0,p1\)\\mathbb\{E\}\[\|x\_\{1\}\-x\_\{0\}\|^\{2\}\]=W\_\{2\}^\{2\}\(p\_\{0\},p\_\{1\}\)\(for the optimal pairing\)\. For diffusion paths \(such as VP\-SDE\), particle trajectories are curved spirals \(first diffusing outward, then pulled back\)\. The path length for the same pair\(x0,x1\)\(x\_\{0\},x\_\{1\}\)is longer, and the kinetic energy integral is larger: ∫01\|X˙tdiffusion\|2​𝑑t\>\|x1−x0\|2\\int\_\{0\}^\{1\}\|\\dot\{X\}\_\{t\}^\{\\text\{diffusion\}\}\|^\{2\}\\,dt\>\|x\_\{1\}\-x\_\{0\}\|^\{2\} The Benamou\-Brenier formula states: W22​\(ρ0,ρ1\)=inf\(ρt,vt\):∂tρt\+∇⋅\(ρt​vt\)=0ρ\|t=0=ρ0,ρ\|t=1=ρ1∫01∫\|vt​\(x\)\|2​ρt​\(x\)​𝑑x​𝑑tW\_\{2\}^\{2\}\(\\rho\_\{0\},\\rho\_\{1\}\)=\\inf\_\{\\begin\{subarray\}\{c\}\(\\rho\_\{t\},v\_\{t\}\):\\\\ \\partial\_\{t\}\\rho\_\{t\}\+\\nabla\\cdot\(\\rho\_\{t\}v\_\{t\}\)=0\\\\ \\rho\|\_\{t=0\}=\\rho\_\{0\},\\;\\rho\|\_\{t=1\}=\\rho\_\{1\}\\end\{subarray\}\}\\int\_\{0\}^\{1\}\\\!\\\!\\int\|v\_\{t\}\(x\)\|^\{2\}\\rho\_\{t\}\(x\)\\,dx\\,dtThe right side is the “total kinetic energy” over all paths satisfying the continuity equation—the OT path \(straight\-line motion\) achieves the minimumW22W\_\{2\}^\{2\}\. Diffusion paths are not minimal→\\tovelocity fields are larger and more complex→\\toharder for networks to learn\. Conclusion: OT paths are geodesics \(“straight lines”\) in Wasserstein space, while diffusion paths “take detours\.”This is the mathematical root of why Flow Matching with OT paths is more efficient than Diffusion\.
4. 4\.What is each DDPM denoising step “optimizing”? DDPM starts fromxT∼𝒩​\(0,I\)x\_\{T\}\\sim\\mathcal\{N\}\(0,I\)and progressively denoises to obtainx0∼pdatax\_\{0\}\\sim p\_\{\\text\{data\}\}\. Each step takes the form: xt−1=1αt​\(xt−1−αt1−α¯t​ϵθ​\(xt,t\)\)\+σt​zx\_\{t\-1\}=\\frac\{1\}\{\\sqrt\{\\alpha\_\{t\}\}\}\\left\(x\_\{t\}\-\\frac\{1\-\\alpha\_\{t\}\}\{\\sqrt\{1\-\\bar\{\\alpha\}\_\{t\}\}\}\\epsilon\_\{\\theta\}\(x\_\{t\},t\)\\right\)\+\\sigma\_\{t\}\\,zThis formula appears to simply “subtract a bit of predicted noise\.” But what is its deeper meaning? Recall our gradient flow: the Fokker\-Planck velocity field isv=−∇log⁡ρ−∇Vv=\-\\nabla\\log\\rho\-\\nabla V\. The DDPM predictionϵθ\\epsilon\_\{\\theta\}satisfiesϵθ≈−1−α¯t​∇log⁡ρt\\epsilon\_\{\\theta\}\\approx\-\\sqrt\{1\-\\bar\{\\alpha\}\_\{t\}\}\\,\\nabla\\log\\rho\_\{t\}\(i\.e\., the negatively normalized score\)\. Substituting “subtractingϵ\\epsilon” amounts to “taking a step in the score direction”—a discretized gradient flow\. More precisely: in probability space, this step is equivalent to solving theJKO scheme\(implicit Euler\): ρk\+1=arg⁡minρ⁡\{ℱ​\(ρ\)⏟free energy \(wants to decrease\)\+12​τ​W22​\(ρ,ρk\)⏟don’t stray too far from the previous step\}\\rho\_\{k\+1\}=\\arg\\min\_\{\\rho\}\\left\\\{\\underbrace\{\\mathcal\{F\}\(\\rho\)\}\_\{\\text\{free energy \(wants to decrease\)\}\}\+\\underbrace\{\\frac\{1\}\{2\\tau\}W\_\{2\}^\{2\}\(\\rho,\\rho\_\{k\}\)\}\_\{\\text\{don't stray too far from the previous step\}\}\\right\\\}whereℱ​\(ρ\)=∫ρ​log⁡ρ\+∫V​ρ\\mathcal\{F\}\(\\rho\)=\\int\\rho\\log\\rho\+\\int V\\rhois the free energy\. TheJKO optimality condition\(derived in detail in Section 6\) states that the minimizer satisfies: ∇log⁡ρk\+1\+∇V\+x−Tk​\(x\)τ=0\\nabla\\log\\rho\_\{k\+1\}\+\\nabla V\+\\frac\{x\-T\_\{k\}\(x\)\}\{\\tau\}=0whereTkT\_\{k\}is the optimal transport map\. Rearranging givesx−Tk​\(x\)τ=−∇log⁡ρk\+1−∇V\\frac\{x\-T\_\{k\}\(x\)\}\{\\tau\}=\-\\nabla\\log\\rho\_\{k\+1\}\-\\nabla V, whose right\-hand side is precisely the negative velocity field of Fokker–Planck\. So each DDPM denoising step is one step of implicit gradient descent in probability space—which accounts for its stability \(implicit methods are unconditionally stable\) and its convergence to the data distribution\.
5. 5\.Why can Energy Matching \(2025\) use a single scalar field for both transport and equilibrium? Balcerak et al\. directly use the JKO scheme as a generative framework: ρt\+Δ​t=arg⁡minρ⁡\{W22​\(ρ,ρt\)2​Δ​t\+∫Vθ​𝑑ρ\+ε​\(t\)​∫ρ​log⁡ρ​d​x\}\\rho\_\{t\+\\Delta t\}=\\arg\\min\_\{\\rho\}\\left\\\{\\frac\{W\_\{2\}^\{2\}\(\\rho,\\rho\_\{t\}\)\}\{2\\Delta t\}\+\\int V\_\{\\theta\}\\,d\\rho\+\\varepsilon\(t\)\\int\\rho\\log\\rho\\,dx\\right\\\}The first\-order optimality condition yieldsx˙=−∇Vθ\\dot\{x\}=\-\\nabla V\_\{\\theta\}\(transport\) whenε=0\\varepsilon=0, andρ∝e−Vθ/ε\\rho\\propto e^\{\-V\_\{\\theta\}/\\varepsilon\}\(Boltzmann distribution\) at equilibrium\. A single scalar fieldVθV\_\{\\theta\}simultaneously encodes the “path” and the “endpoint”—because the JKO scheme naturally unifies both\.

### This article’s mission

The answers to all the questions above point to the same mathematical structure:

The space of probability distributions𝒫2​\(ℝd\)\\mathcal\{P\}\_\{2\}\(\\mathbb\{R\}^\{d\}\)forms a \(formal\)Riemannian manifoldunder theWasserstein distance\.Thecontinuity equationis the “equation of motion” on this manifold \(how tangent vectors induce density changes\),TheFokker\-Planck equationis thegradient flowof the free energyKL​\(ρ∥π\)\\mathrm\{KL\}\(\\rho\\\|\\pi\)on this manifold,TheJKO schemeis the implicit Euler discretization of this gradient flow\.

Translated into language you are familiar with:

These are not analogies—they are precise mathematical theorems\. The following six sections will build up this structure brick by brick:

- •Section 1 defines the Wasserstein distance \(“the metric of the space”\)
- •Section 2 derives the continuity equation \(“the equation of motion”\)
- •Section 3 provesW2=W\_\{2\}=geodesic distance \(Benamou\-Brenier, “why OT paths are optimal”\)
- •Section 4 derives Fokker\-Planck \(“from SDE to PDE”\)
- •Section 5 proves Fokker\-Planck==gradient flow \(“what is being optimized”\)
- •Section 6 derives the JKO scheme \(“how to discretize,” “each denoising step==one optimization step”\)

Prerequisites:If you have read either Song et al\. \(2021\) or Lipman et al\. \(2023\), you already have sufficient background\. We assume familiarity with multivariable calculus, linear algebra, and basic probability theory; prerequisites in measure theory, differential geometry, etc\. are developed from scratch in the appendices\.

### Related perspectives

This article gives a unified, end\-to\-end account of diffusion and flow matching from a single geometric object: the Wasserstein manifold of probability measures\. On it the two are not the same thing but complementary principles—diffusion is the gradient flow of the free energy \(an initial\-value problem\), flow matching the geodesic, i\.e\. the minimal\-action path \(a boundary\-value problem\)\. Where most treatments proceed through stochastic processes and score functions, here one geometric idea is carried, without gaps, from the definition of the Wasserstein metric to the algorithms used in practice\. Several lines of work reach this same picture from other directions; we place the article among them below\.

Unifying frameworks for flows and diffusions\.The cleanest statement that diffusion and flow\-based models are facets of one object is the*stochastic interpolants*framework of Albergo, Boffi & Vanden\-Eijnden \(2025\): a single interpolant between two densities induces both a transport \(continuity\) equation and a family of forward/backward Fokker–Planck equations with tunable noise, hence both deterministic \(ODE\) and stochastic \(SDE\) generators\. Lipman et al\. \(2024\), in the*Flow Matching Guide and Code*, give a comprehensive exposition presenting diffusion paths as special cases of Flow Matching\. Both unify the two through the SDE/ODE and velocity\-field language; this article instead routes the unification through Wasserstein geometry—gradient flow versus geodesic\.

Diffusion as a Wasserstein gradient flow\.Closest to the central identity of Section 5 is the work of Vuong, McCann, Santos & Lin \(2025\), who argue that diffusion training is better read as flow matching to the velocity field of a Wasserstein gradient flow than as score learning, and give numerical evidence that the learned field is generally not conservative \(hence not a genuine score\)\.

JKO as a generative algorithm\.A line of work builds generative models directly on the JKO scheme of Section 6: Mokrov et al\. \(2021\) scale Wasserstein gradient flows using input\-convex networks, and Xu, Cheng & Xie \(2023, JKO\-iFlow\) stack residual blocks that each realize one JKO step\. Energy Matching \(Balcerak et al\., 2025\) and Equilibrium Matching \(Wang et al\., 2025\) belong to the same family\.

A geometric caveat\.The two principles should not be over\-identified\. Lavenant & Santambrogio \(2022\) prove that the flow map of the Fokker–Planck equation is*not*the optimal\-transport map: the diffusion trajectory is in general not theW2W\_\{2\}geodesic between its endpoints\. This is precisely the distinction drawn here—diffusion descends the free energy \(an initial\-value gradient flow, along a curved path\), whereas OT Flow Matching follows the geodesic \(a boundary\-value problem\)—and it is why the two reach the same endpoints by different paths\.

Complementary expositions\.For tutorials that unify Flow Matching and diffusion through the SDE/score lens rather than optimal transport, see the lecture notes of Holderrieth & Erives \(2025, MIT 6\.S184\)\.

## 1Wasserstein Distance: The Cost of Moving Dirt

### 1\.1Why do we need a new distance?

Suppose you have two piles of sand \(probability distributions\) and want to measure “how different” they are\.

Consider two probability distributionsμ\\muandν\\nuonℝd\\mathbb\{R\}^\{d\}\. Classical divergences like KL or total variation have a fundamental limitation: they ignore the*geometry*of the underlying space\. Let us first recall their definitions:

###### Definition 1\.1\(KL divergence\)\.

Letμ\\muandν\\nuhave density functionsp​\(x\)p\(x\)andq​\(x\)q\(x\), respectively\. TheKL divergence\(Kullback–Leibler divergence; Kullback & Leibler, 1951\) is defined as:

KL​\(μ∥ν\)=∫ℝdp​\(x\)​log⁡p​\(x\)q​\(x\)​d​x\\mathrm\{KL\}\(\\mu\\\|\\nu\)=\\int\_\{\\mathbb\{R\}^\{d\}\}p\(x\)\\log\\frac\{p\(x\)\}\{q\(x\)\}\\,dx\(8\)If there existsxxsuch thatp​\(x\)\>0p\(x\)\>0butq​\(x\)=0q\(x\)=0\(i\.e\., the support ofμ\\muis not contained in the support ofν\\nu\), thenKL​\(μ∥ν\)=\+∞\\mathrm\{KL\}\(\\mu\\\|\\nu\)=\+\\infty\.

Intuition: KL measures “how much information is lost when approximatingμ\\mubyν\\nu”\. It is not symmetric:KL​\(μ∥ν\)≠KL​\(ν∥μ\)\\mathrm\{KL\}\(\\mu\\\|\\nu\)\\neq\\mathrm\{KL\}\(\\nu\\\|\\mu\)\.

###### Definition 1\.2\(Total Variation distance\)\.

Thetotal variation distanceis defined as:

TV​\(μ,ν\)=supA⊆ℝd\|μ​\(A\)−ν​\(A\)\|=12​∫ℝd\|p​\(x\)−q​\(x\)\|​𝑑x\\mathrm\{TV\}\(\\mu,\\nu\)=\\sup\_\{A\\subseteq\\mathbb\{R\}^\{d\}\}\\left\|\\mu\(A\)\-\\nu\(A\)\\right\|=\\frac\{1\}\{2\}\\int\_\{\\mathbb\{R\}^\{d\}\}\|p\(x\)\-q\(x\)\|\\,dx\(9\)The first expression is the most general definition \(taking the supremum over all measurable setsAA\); the second is an equivalent form when both measures have densities\.

Intuition: TV finds the “eventAAthat most effectively distinguishesμ\\mufromν\\nu”—for this event, the difference in probability assigned by the two measures is maximized\. TV∈\[0,1\]\\in\[0,1\]\.

The essential limitation of both metrics: each is blind to the geometry of the underlying space\.

###### Example 1\.3\.

Letμ=δx\\mu=\\delta\_\{x\}andν=δy\\nu=\\delta\_\{y\}be two Dirac masses\. Then:

- •KL​\(μ∥ν\)=\+∞\\mathrm\{KL\}\(\\mu\\\|\\nu\)=\+\\inftyfor anyx≠yx\\neq y\(they have disjoint support\)\.
- •TV​\(μ,ν\)=1\\mathrm\{TV\}\(\\mu,\\nu\)=1for anyx≠yx\\neq y, whether\|x−y\|=0\.001\|x\-y\|=0\.001or\|x−y\|=1000\|x\-y\|=1000\.

Neither captures thatδx\\delta\_\{x\}andδy\\delta\_\{y\}should be “close” whenx≈yx\\approx y\.

In contrast, for Wasserstein:W2​\(δx,δy\)=\|x−y\|W\_\{2\}\(\\delta\_\{x\},\\delta\_\{y\}\)=\|x\-y\|\. Asxxapproachesyy, the distance tends to zero—this is exactly the behavior we expect\.

IntuitionImagine you are a laborer with a pile of sand \(shaped likeμ\\mu\), and you need to reshape it into a different configuration \(ν\\nu\)\. The Wasserstein distance measures precisely this:the minimum total transportation cost required to accomplish this task\. Sand that is closer costs less to move, so this metric naturally respects the geometry of the underlying space\.

### 1\.2The Kantorovich formulation

The relaxation below is due to Kantorovich \(1942\); the resulting metric is named after Vaserstein \(1969\)\.

###### Definition 1\.4\(2\-Wasserstein distance\)\.

Let𝒫2​\(ℝd\)\\mathcal\{P\}\_\{2\}\(\\mathbb\{R\}^\{d\}\)denote the space of probability measures with finite second moment\. Forμ,ν∈𝒫2​\(ℝd\)\\mu,\\nu\\in\\mathcal\{P\}\_\{2\}\(\\mathbb\{R\}^\{d\}\):

W2​\(μ,ν\):=\(infγ∈Γ​\(μ,ν\)∫ℝd×ℝd\|x−y\|2​𝑑γ​\(x,y\)\)1/2W\_\{2\}\(\\mu,\\nu\):=\\left\(\\inf\_\{\\gamma\\in\\Gamma\(\\mu,\\nu\)\}\\int\_\{\\mathbb\{R\}^\{d\}\\times\\mathbb\{R\}^\{d\}\}\|x\-y\|^\{2\}\\,d\\gamma\(x,y\)\\right\)^\{1/2\}\(10\)whereΓ​\(μ,ν\)\\Gamma\(\\mu,\\nu\)is the set ofcouplings:

Γ​\(μ,ν\):=\{γ∈𝒫​\(ℝd×ℝd\)\|γ​\(A×ℝd\)=μ​\(A\)​∀Aγ​\(ℝd×B\)=ν​\(B\)​∀B\}\\Gamma\(\\mu,\\nu\):=\\left\\\{\\gamma\\in\\mathcal\{P\}\(\\mathbb\{R\}^\{d\}\\times\\mathbb\{R\}^\{d\}\)\\;\\middle\|\\;\\begin\{aligned\} &\\gamma\(A\\times\\mathbb\{R\}^\{d\}\)=\\mu\(A\)\\;\\;\\forall A\\\\ &\\gamma\(\\mathbb\{R\}^\{d\}\\times B\)=\\nu\(B\)\\;\\;\\forall B\\end\{aligned\}\\right\\\}

IntuitionUnpacking the formula term by term:W2​\(μ,ν\)=\(infγ∫\|x−y\|2​𝑑γ​\(x,y\)\)1/2W\_\{2\}\(\\mu,\\nu\)=\\Big\(\\inf\_\{\\gamma\}\\int\|x\-y\|^\{2\}\\,d\\gamma\(x,y\)\\Big\)^\{1/2\}•γ​\(x,y\)\\gamma\(x,y\): a transport plan\. The “density” ofγ\\gammanear the point\(x,y\)\(x,y\)represents “how much mass is moved from locationxxto locationyy”\.•\|x−y\|2\|x\-y\|^\{2\}: the unit cost of moving mass fromxxtoyy\(the squared distance\)\. Why squared? \(1\) It is mathematically more convenient \(related to inner products and energy\); \(2\) it penalizes long\-distance transport more heavily \(moving 10 meters costs 100 times as much as moving 1 meter, not 10 times\)\.•∫\|x−y\|2​𝑑γ\\int\|x\-y\|^\{2\}\\,d\\gamma: thetotal transport costunder planγ\\gamma=∑\\sum\(each unit of mass×\\timesits squared transport distance\)\.•infγ\\inf\_\{\\gamma\}: find the plan with theminimum costamong all feasible plans\.•\(⋅\)1/2\(\\cdot\)^\{1/2\}: taking the square root givesW2W\_\{2\}the dimension of a “distance” \(same order as\|x−y\|\|x\-y\|\)\.Intuition for the marginal conditions:•γ​\(A×ℝd\)=μ​\(A\)\\gamma\(A\\times\\mathbb\{R\}^\{d\}\)=\\mu\(A\): the total mass transported out of regionAAequals the total mass ofμ\\muonAA\. In other words, all the dirt inAAmust be moved out—none may be left behind\.•γ​\(ℝd×B\)=ν​\(B\)\\gamma\(\\mathbb\{R\}^\{d\}\\times B\)=\\nu\(B\): the total mass arriving at regionBBequals the total mass ofν\\nuonBB\. In other words, the demand atBBmust be met exactly—no more, no less\.An assignment analogy:think of a dispatch problem—μ\\muis the distribution of available vehicles,ν\\nuthe distribution of requested destinations,γ\\gammathe assignment “which vehicle serves which destination,” and\|x−y\|2\|x\-y\|^\{2\}the cost of a trip\. ThenW2W\_\{2\}is the minimum total cost of the optimal assignment \(after taking the square root\)\.

### 1\.3The Monge formulation and Brenier’s theorem

Historical background:In 1781, the French mathematician Gaspard Monge posed a practical problem: how to transport a pile of earth \(excavated from a quarry\) to designated construction sites at minimum cost\. His idea was very natural:assign each grain of earth a definite destination\.

###### Definition 1\.5\(Monge’s problem\)\.

Find a mapT:ℝd→ℝdT:\\mathbb\{R\}^\{d\}\\to\\mathbb\{R\}^\{d\}\(a “transport map”\) that moves the mass at each locationxxto locationT​\(x\)T\(x\), subject to:

1. 1\.Pushforward constraint:T\#​μ=νT\_\{\\\#\}\\mu=\\nu, meaningν​\(B\)=μ​\(T−1​\(B\)\)\\nu\(B\)=\\mu\(T^\{\-1\}\(B\)\)for all BorelBB\.
2. 2\.Minimize total cost: infT:T\#​μ=ν∫ℝd\|x−T​\(x\)\|2​𝑑μ​\(x\)\\inf\_\{T:\\,T\_\{\\\#\}\\mu=\\nu\}\\int\_\{\\mathbb\{R\}^\{d\}\}\|x\-T\(x\)\|^\{2\}\\,d\\mu\(x\)\(11\)

IntuitionUnpacking the Monge problem term by term:Meaning of the mapTT:TTis an “instruction table”—for each locationxxin the source distributionμ\\mu,T​\(x\)T\(x\)tells you “where the dirt atxxshould be sent”\.Meaning of the constraintT\#​μ=νT\_\{\\\#\}\\mu=\\nu:After the transport is complete, the new distribution of dirt is exactlyν\\nu\.•In plain language: if the dirt at every locationxxis moved according toTT, the resulting shape must beν\\nu\.•In formulas:ν​\(B\)=μ​\(T−1​\(B\)\)=μ​\(\{x:T​\(x\)∈B\}\)\\nu\(B\)=\\mu\(T^\{\-1\}\(B\)\)=\\mu\(\\\{x:T\(x\)\\in B\\\}\)= “the total amount of dirt transported into regionBBequals the amount thatν\\nurequires inBB”\.Meaning of the cost∫\|x−T​\(x\)\|2​𝑑μ​\(x\)\\int\|x\-T\(x\)\|^\{2\}\\,d\\mu\(x\):•\|x−T​\(x\)\|2\|x\-T\(x\)\|^\{2\}: the dirt at locationxxis moved toT​\(x\)T\(x\); this is the squared transport distance\.•d​μ​\(x\)d\\mu\(x\): how much dirt is at locationxx\(weighted byμ\\mu\)\.•The integral: total transport cost of all dirt =∑\\sum\(amount of each portion of dirt×\\timessquared transport distance\)\.One\-dimensional example:Letμ=Uniform​\[0,1\]\\mu=\\text\{Uniform\}\[0,1\]andν=Uniform​\[1,2\]\\nu=\\text\{Uniform\}\[1,2\]\. The optimal map is clearlyT​\(x\)=x\+1T\(x\)=x\+1\(shift every grain of sand one unit to the right\)\. Cost =∫01\|x−\(x\+1\)\|2​𝑑x=∫011​𝑑x=1\\int\_\{0\}^\{1\}\|x\-\(x\+1\)\|^\{2\}\\,dx=\\int\_\{0\}^\{1\}1\\,dx=1\.Another example:μ=Uniform​\[0,2\]\\mu=\\text\{Uniform\}\[0,2\],ν=Uniform​\[0,1\]\\nu=\\text\{Uniform\}\[0,1\]\(compressing a wide distribution into a narrow one\)\. The optimal map isT​\(x\)=x/2T\(x\)=x/2\. Cost =∫02\|x−x/2\|2⋅12​𝑑x=∫02x24⋅12​𝑑x=18⋅x33\|02=13\\int\_\{0\}^\{2\}\|x\-x/2\|^\{2\}\\cdot\\frac\{1\}\{2\}\\,dx=\\int\_\{0\}^\{2\}\\frac\{x^\{2\}\}\{4\}\\cdot\\frac\{1\}\{2\}\\,dx=\\frac\{1\}\{8\}\\cdot\\frac\{x^\{3\}\}\{3\}\\Big\|\_\{0\}^\{2\}=\\frac\{1\}\{3\}\.

The principal difficulty with Monge’s formulation:

The mapTTrequires thatallmass at eachxxgoes to the same place—splitting is not allowed\. This leads to two problems:

Problem 1: the solution may not exist\.Letμ=δ0\\mu=\\delta\_\{0\}\(all mass at the origin\),ν=12​δ−1\+12​δ1\\nu=\\frac\{1\}\{2\}\\delta\_\{\-1\}\+\\frac\{1\}\{2\}\\delta\_\{1\}\(half at−1\-1, half at11\)\. No mapTTsatisfyingT\#​δ0=νT\_\{\\\#\}\\delta\_\{0\}=\\nuexists—sinceT​\(0\)T\(0\)can only be a single point, it cannot simultaneously be−1\-1and11\. But the Kantorovich couplingγ=12​δ\(0,−1\)\+12​δ\(0,1\)\\gamma=\\frac\{1\}\{2\}\\delta\_\{\(0,\-1\)\}\+\\frac\{1\}\{2\}\\delta\_\{\(0,1\)\}is feasible \(splitting the mass at the origin in half\)\.

Problem 2: the optimization problem is non\-convex\.The constraintT\#​μ=νT\_\{\\\#\}\\mu=\\nuis highly nonlinear\. If bothT1T\_\{1\}andT2T\_\{2\}satisfyT1​\#​μ=T2​\#​μ=νT\_\{1\\\#\}\\mu=T\_\{2\\\#\}\\mu=\\nu, their averageT1\+T22\\frac\{T\_\{1\}\+T\_\{2\}\}\{2\}typically doesnotsatisfy the constraint\. This makes direct optimization very difficult\.

Kantorovich’s relaxation:replace the mapTTwith a joint distributionγ\\gamma, turning the non\-convex constraint into alinear constraint\(the marginal conditions are linear inγ\\gamma\), and the non\-convex problem into a linear program—which always has a solution and can be computed efficiently\.

The relationship between Monge and Kantorovich:

If the optimal mapTTexists, then the corresponding Kantorovich coupling isγ=\(id,T\)\#​μ\\gamma=\(\\mathrm\{id\},T\)\_\{\\\#\}\\mu, i\.e\., the joint distribution is concentrated on the graph of the mapTT:\{\(x,T​\(x\)\):x∈ℝd\}\\\{\(x,T\(x\)\):x\\in\\mathbb\{R\}^\{d\}\\\}\. Conversely, if the optimal couplingγ\\gammahappens to be concentrated on the graph of some map, then that map is the solution to the Monge problem\.

###### Theorem 1\.6\(Brenier, 1991\)\.

Ifμ\\muis absolutely continuous with respect to the Lebesgue measure, then:

1. 1\.The optimal transport mapTTexists and is uniqueμ\\mu\-a\.e\.
2. 2\.T=∇ϕT=\\nabla\\phifor some convex functionϕ:ℝd→ℝ\\phi:\\mathbb\{R\}^\{d\}\\to\\mathbb\{R\}\.

Consequently:

W22​\(μ,ν\)=∫ℝd\|x−∇ϕ​\(x\)\|2​𝑑μ​\(x\)W\_\{2\}^\{2\}\(\\mu,\\nu\)=\\int\_\{\\mathbb\{R\}^\{d\}\}\|x\-\\nabla\\phi\(x\)\|^\{2\}\\,d\\mu\(x\)\(12\)

IntuitionWhat does Brenier’s theorem say? Three levels of understanding:Level 1: Optimal transport is deterministic\.Although Kantorovich allows “splitting the mass atxxamong multiple destinationsyy” \(probabilistic transport\), whenμ\\muis absolutely continuous, the optimal planrequires no splitting at all—the mass at eachxxis moved entirely to a single definite locationT​\(x\)T\(x\)\. This greatly simplifies the problem: instead of optimizing over a joint distributionγ\\gamma\(an infinite\-dimensional object\), we optimize over a mapTT\(relatively simpler\)\.Level 2: The optimal map does not cross\.T=∇ϕT=\\nabla\\phiwithϕ\\phiconvex means thatTTis “monotone” \(the high\-dimensional analogue of monotonicity = gradient of a convex function\)\. Intuition: if two transport paths cross \(x1x\_\{1\}’s dirt goes toy2y\_\{2\},x2x\_\{2\}’s dirt goes toy1y\_\{1\}, and the paths cross\), then we can “uncross” them \(x1→y1x\_\{1\}\\to y\_\{1\},x2→y2x\_\{2\}\\to y\_\{2\}\), reducing the total cost\. Therefore the optimal plan must be uncrossed, and “uncrossed”⇔\\iffthe map is the gradient of a convex function\.Sub\-optimal \(crossed\):x1x\_\{1\}x2x\_\{2\}y2y\_\{2\}y1y\_\{1\}Optimal \(uncrossed\):x1x\_\{1\}x2x\_\{2\}y1y\_\{1\}y2y\_\{2\}Cost of the crossed plan:\|x1−y2\|2\+\|x2−y1\|2\|x\_\{1\}\-y\_\{2\}\|^\{2\}\+\|x\_\{2\}\-y\_\{1\}\|^\{2\}\. Cost of the uncrossed plan:\|x1−y1\|2\+\|x2−y2\|2\|x\_\{1\}\-y\_\{1\}\|^\{2\}\+\|x\_\{2\}\-y\_\{2\}\|^\{2\}\. The difference is2​\(x1−x2\)⋅\(y2−y1\)2\(x\_\{1\}\-x\_\{2\}\)\\cdot\(y\_\{2\}\-y\_\{1\}\)\(obtained by expanding the squares\)\. When the paths cross,\(x1−x2\)\(x\_\{1\}\-x\_\{2\}\)and\(y1−y2\)\(y\_\{1\}\-y\_\{2\}\)point in opposite directions, so this difference is≥0\\geq 0—crossing is always more expensive\.Level 3: The geometric meaning of convexity ofϕ\\phi\.∇ϕ\\nabla\\phiis the gradient of a convex function, so it is a “spreading” map—it does not send two distant points to the same neighborhood\. More precisely: the gradient of a convex function satisfies the monotonicity condition\(∇ϕ​\(x1\)−∇ϕ​\(x2\)\)⋅\(x1−x2\)≥0\(\\nabla\\phi\(x\_\{1\}\)\-\\nabla\\phi\(x\_\{2\}\)\)\\cdot\(x\_\{1\}\-x\_\{2\}\)\\geq 0This guarantees that the map “preserves order”—points on the left are sent to the left, points on the right to the right, with no reversals \(crossings\)\.

A common point of confusion: Brenier’s theorem fails for discrete distributions\.

Consider the following counterexample:

μ=0\.6​δa\+0\.4​δb,ν=0\.5​δc\+0\.5​δd\\mu=0\.6\\,\\delta\_\{a\}\+0\.4\\,\\delta\_\{b\},\\qquad\\nu=0\.5\\,\\delta\_\{c\}\+0\.5\\,\\delta\_\{d\}
If a mapTTsatisfyingT\#​μ=νT\_\{\\\#\}\\mu=\\nuexists, thenTTcan only sendaato one ofccordd, and likewise forbb\. But no assignment works:

- •T​\(a\)=c,T​\(b\)=dT\(a\)=c,\\,T\(b\)=d: thenν​\(\{c\}\)=μ​\(T−1​\(\{c\}\)\)=μ​\(\{a\}\)=0\.6≠0\.5\\nu\(\\\{c\\\}\)=\\mu\(T^\{\-1\}\(\\\{c\\\}\)\)=\\mu\(\\\{a\\\}\)=0\.6\\neq 0\.5✗
- •T​\(a\)=d,T​\(b\)=cT\(a\)=d,\\,T\(b\)=c: thenν​\(\{d\}\)=μ​\(\{a\}\)=0\.6≠0\.5\\nu\(\\\{d\\\}\)=\\mu\(\\\{a\\\}\)=0\.6\\neq 0\.5✗
- •T​\(a\)=c,T​\(b\)=cT\(a\)=c,\\,T\(b\)=c: thenν​\(\{c\}\)=0\.6\+0\.4=1≠0\.5\\nu\(\\\{c\\\}\)=0\.6\+0\.4=1\\neq 0\.5✗
- •All other cases fail similarly\.

Conclusion: No transport map fromμ\\mutoν\\nuexists\.

An optimalcoupling, however, does exist\. The optimal plan is: split 0\.5 fromaatocc, send the remaining 0\.1 todd; send all 0\.4 frombbtodd:

γ=0\.5​δ\(a,c\)\+0\.1​δ\(a,d\)\+0\.4​δ\(b,d\)\\gamma=0\.5\\,\\delta\_\{\(a,c\)\}\+0\.1\\,\\delta\_\{\(a,d\)\}\+0\.4\\,\\delta\_\{\(b,d\)\}Verification of marginals: the total mass ofγ\\gammaonaais0\.5\+0\.1=0\.6=μ​\(\{a\}\)0\.5\+0\.1=0\.6=\\mu\(\\\{a\\\}\)✓; the total mass ofγ\\gammaonccis0\.5=ν​\(\{c\}\)0\.5=\\nu\(\\\{c\\\}\)✓; and so on\.

This is why Brenier’s theorem requires the condition “μ\\muis absolutely continuous”:

- •Discrete distributions\(as above\): a finite amount of mass is concentrated at a single point, and transport maynecessarily require splitting, so a mapTTis insufficient—one must use a couplingγ\\gamma\.
- •Absolutely continuous distributions\(with densityρ\\rho\): each point carries only an “infinitesimal” massρ​\(x\)​d​x\\rho\(x\)dxand never needs to be split—even ifν\\nuis discrete, no source point has “a bulk of mass that must be allocated to different destinations” \(since each point ofμ\\muhas zero mass\)\. Therefore a deterministic mapTTsuffices\. Note: onlyμ\\muneeds to be absolutely continuous; there isnorestriction onν\\nu\.

In one sentence:Absolutely continuous⇒\\Rightarrowmass is “continuously spread”⇒\\Rightarrowno splitting needed⇒\\RightarrowmapTTexists⇒\\RightarrowBrenier applies\.

Discrete/singular distribution⇒\\Rightarrowthere are “bulk” masses⇒\\Rightarrowsplitting may be necessary⇒\\Rightarrowonly a couplingγ\\gammacan describe the solution⇒\\RightarrowBrenier does not apply\.

### 1\.4A concrete example: Gaussians

###### Example 1\.7\.

Forμ=𝒩​\(m1,Σ1\)\\mu=\\mathcal\{N\}\(m\_\{1\},\\Sigma\_\{1\}\)andν=𝒩​\(m2,Σ2\)\\nu=\\mathcal\{N\}\(m\_\{2\},\\Sigma\_\{2\}\):

W22​\(μ,ν\)=\|m1−m2\|2\+tr​\(Σ1\+Σ2−2​\(Σ11/2​Σ2​Σ11/2\)1/2\)W\_\{2\}^\{2\}\(\\mu,\\nu\)=\|m\_\{1\}\-m\_\{2\}\|^\{2\}\+\\mathrm\{tr\}\\\!\\left\(\\Sigma\_\{1\}\+\\Sigma\_\{2\}\-2\\bigl\(\\Sigma\_\{1\}^\{1/2\}\\Sigma\_\{2\}\\Sigma\_\{1\}^\{1/2\}\\bigr\)^\{1/2\}\\right\)\(13\)Special case: ifΣ1=Σ2=σ2​I\\Sigma\_\{1\}=\\Sigma\_\{2\}=\\sigma^\{2\}I, thenW2​\(μ,ν\)=\|m1−m2\|W\_\{2\}\(\\mu,\\nu\)=\|m\_\{1\}\-m\_\{2\}\|\.

Key steps of the derivation:For Gaussian distributions, the optimal transport map is affine:T​\(x\)=A​x\+bT\(x\)=Ax\+b\. By Brenier’s theorem,T=∇ϕT=\\nabla\\phi\(ϕ\\phiconvex\) requiresAAto be asymmetric positive definitematrix\. The pushforward conditionT\#​μ=νT\_\{\\\#\}\\mu=\\nugivesb=m2−A​m1b=m\_\{2\}\-Am\_\{1\}andA​Σ1​AT=Σ2A\\Sigma\_\{1\}A^\{T\}=\\Sigma\_\{2\}\. Combined withA=AT≻0A=A^\{T\}\\succ 0, the unique solution isA=Σ1−1/2​\(Σ11/2​Σ2​Σ11/2\)1/2​Σ1−1/2A=\\Sigma\_\{1\}^\{\-1/2\}\(\\Sigma\_\{1\}^\{1/2\}\\Sigma\_\{2\}\\Sigma\_\{1\}^\{1/2\}\)^\{1/2\}\\Sigma\_\{1\}^\{\-1/2\}\. The cost∫\|x−T​x\|2​𝑑μ=\|m1−m2\|2\+tr​\(\(I−A\)​Σ1​\(I−A\)T\)\\int\|x\-Tx\|^\{2\}\\,d\\mu=\|m\_\{1\}\-m\_\{2\}\|^\{2\}\+\\mathrm\{tr\}\\bigl\(\(I\-A\)\\Sigma\_\{1\}\(I\-A\)^\{T\}\\bigr\), which simplifies to the stated formula usingA2​Σ1=Σ2A^\{2\}\\Sigma\_\{1\}=\\Sigma\_\{2\}\(by the symmetry ofAA\)\.

### 1\.5Key properties

The space\(𝒫2​\(ℝd\),W2\)\(\\mathcal\{P\}\_\{2\}\(\\mathbb\{R\}^\{d\}\),W\_\{2\}\)is:

1. 1\.Acomplete separable metric space\(Polish space\)\.
2. 2\.It metrizesweak convergence \+ convergence of second moments\.
3. 3\.Most profoundly: it admits aformal Riemannian structure\(Otto, 2001\)\.

We now explain each of these three properties in detail\.

#### Complete \(completeness\)

###### Definition 1\.8\(Completeness\)\.

A metric space\(X,d\)\(X,d\)iscompleteif every Cauchy sequence in it converges\.

Cauchy sequence:\{xn\}\\\{x\_\{n\}\\\}satisfies that for anyϵ\>0\\epsilon\>0, there existsNNsuch thatd​\(xm,xn\)<ϵd\(x\_\{m\},x\_\{n\}\)<\\epsilonwheneverm,n\>Nm,n\>N\(the points get arbitrarily close to each other\)\.

Completeness means: as long as points get closer and closer together, they must converge to some limit within the space—they cannot “escape outside the space”\.

Examples:

- •ℝ\\mathbb\{R\}is complete: every Cauchy sequence converges to some real number\.
- •ℚ\\mathbb\{Q\}\(the rationals\) is not complete:1,1\.4,1\.41,1\.414,…1,1\.4,1\.41,1\.414,\\ldotsis a Cauchy sequence inℚ\\mathbb\{Q\}, but its limit2∉ℚ\\sqrt\{2\}\\notin\\mathbb\{Q\}\.
- •The open interval\(0,1\)\(0,1\)is not complete:1n\\frac\{1\}\{n\}is a Cauchy sequence, but its limit0∉\(0,1\)0\\notin\(0,1\)\.

Why this matters:Completeness guarantees that when we take limits \(such as theτ→0\\tau\\to 0limit in the JKO scheme\), the limiting object is still a legitimate probability measure\. In an incomplete space, “taking limits” may lead outside the space, making rigorous analysis much harder\.

#### Separable \(separability\)

###### Definition 1\.9\(Separability\)\.

A metric space\(X,d\)\(X,d\)isseparableif it has acountable dense subset— i\.e\., there exists a countable set\{x1,x2,…\}⊂X\\\{x\_\{1\},x\_\{2\},\\ldots\\\}\\subset Xsuch that every point inXXcan be approximated arbitrarily well by elements of this set\.

Examples:

- •ℝ\\mathbb\{R\}is separable: the rationalsℚ\\mathbb\{Q\}are countable and dense inℝ\\mathbb\{R\}\.
- •ℝd\\mathbb\{R\}^\{d\}is separable: points with all rational coordinatesℚd\\mathbb\{Q\}^\{d\}are dense\.
- •𝒫2​\(ℝd\)\\mathcal\{P\}\_\{2\}\(\\mathbb\{R\}^\{d\}\)is separable: finite mixtures of Dirac masses∑i=1nwi​δxi\\sum\_\{i=1\}^\{n\}w\_\{i\}\\delta\_\{x\_\{i\}\}\(wi∈ℚw\_\{i\}\\in\\mathbb\{Q\},xi∈ℚdx\_\{i\}\\in\\mathbb\{Q\}^\{d\}\) can approximate any probability measure\.

Why this matters:Separability ensures that the space is “not too large”—we can approximate everything using countably many “basic elements”\. This is a necessary condition for many theorems in topology and functional analysis \(e\.g\., Prokhorov’s theorem, weak compactness\)\.

#### Weak convergence

###### Definition 1\.10\(Weak convergence of measures\)\.

A sequence of probability measuresμn\\mu\_\{n\}converges weaklytoμ\\mu\(writtenμn⇀μ\\mu\_\{n\}\\rightharpoonup\\mu\) if for all bounded continuous functionsff:

∫f​𝑑μn→∫f​𝑑μas​n→∞\\int f\\,d\\mu\_\{n\}\\to\\int f\\,d\\mu\\quad\\text\{as \}n\\to\\infty

Intuition:Weak convergence means “the shapes of the distributions become increasingly similar”—for any “smooth observable”ff, the expectation underμn\\mu\_\{n\},𝔼μn​\[f\]\\mathbb\{E\}\_\{\\mu\_\{n\}\}\[f\], approaches the expectation underμ\\mu,𝔼μ​\[f\]\\mathbb\{E\}\_\{\\mu\}\[f\]\.

Examples:

- •μn=𝒩​\(0,1/n\)\\mu\_\{n\}=\\mathcal\{N\}\(0,1/n\)\(Gaussians with vanishing variance\) converges weakly toδ0\\delta\_\{0\}\. Intuition: the bell curve becomes narrower and narrower, eventually collapsing to a point\.
- •μn=Uniform​\[−1/n,1/n\]\\mu\_\{n\}=\\text\{Uniform\}\[\-1/n,1/n\]converges weakly toδ0\\delta\_\{0\}\.
- •μn=1n​∑i=1nδi/n\\mu\_\{n\}=\\frac\{1\}\{n\}\\sum\_\{i=1\}^\{n\}\\delta\_\{i/n\}converges weakly toUniform​\[0,1\]\\text\{Uniform\}\[0,1\]\. Intuition: the uniformly scattered points become increasingly dense, approaching the continuous uniform distribution\.

Weak convergence vs\. Wasserstein convergence:

W2​\(μn,μ\)→0W\_\{2\}\(\\mu\_\{n\},\\mu\)\\to 0⇔\\iffμn⇀μ\\mu\_\{n\}\\rightharpoonup\\muandthe second moments converge \(∫\|x\|2​𝑑μn→∫\|x\|2​𝑑μ\\int\|x\|^\{2\}\\,d\\mu\_\{n\}\\to\\int\|x\|^\{2\}\\,d\\mu\)\.

Pure weak convergence does not control the “tails”—for example,μn=\(1−1n\)​δ0\+1n​δn\\mu\_\{n\}=\(1\-\\frac\{1\}\{n\}\)\\delta\_\{0\}\+\\frac\{1\}\{n\}\\delta\_\{n\}converges weakly toδ0\\delta\_\{0\}, butW2​\(μn,δ0\)=1n⋅n2=n→∞W\_\{2\}\(\\mu\_\{n\},\\delta\_\{0\}\)=\\sqrt\{\\frac\{1\}\{n\}\\cdot n^\{2\}\}=\\sqrt\{n\}\\to\\infty\. Wasserstein convergence is stronger because it additionally requires that mass cannot “escape to infinity”\.

#### Formal Riemannian structure

We now use the language above to describe the Wasserstein space:

###### Definition 1\.11\(Riemannian manifold — informal\)\.

ARiemannian manifold\(M,g\)\(M,g\)is a smooth manifoldMMequipped with aninner productgpg\_\{p\}on the tangent spaceTp​MT\_\{p\}Mat each pointpp\. This inner product allows us to:

- •Measure the “length” of tangent vectors:‖v‖p=gp​\(v,v\)\\\|v\\\|\_\{p\}=\\sqrt\{g\_\{p\}\(v,v\)\}
- •Compute the “angle” between two tangent vectors:cos⁡θ=gp​\(u,v\)‖u‖p​‖v‖p\\cos\\theta=\\frac\{g\_\{p\}\(u,v\)\}\{\\\|u\\\|\_\{p\}\\\|v\\\|\_\{p\}\}
- •Define “shortest paths” \(geodesics\) and “distance”:d​\(p,q\)=infγ∫01‖γ˙​\(t\)‖γ​\(t\)​𝑑td\(p,q\)=\\inf\_\{\\gamma\}\\int\_\{0\}^\{1\}\\\|\\dot\{\\gamma\}\(t\)\\\|\_\{\\gamma\(t\)\}\\,dt
- •Define the “gradient” of a function:gp​\(grad​f,v\)=D​fp​\[v\]g\_\{p\}\(\\text\{grad\}\\,f,v\)=Df\_\{p\}\[v\]

Simple examples:

- •ℝn\\mathbb\{R\}^\{n\}with the standard inner productgp​\(u,v\)=u⋅vg\_\{p\}\(u,v\)=u\\cdot v: the simplest Riemannian manifold\. The tangent space isℝn\\mathbb\{R\}^\{n\}itself, and geodesics are straight lines\.
- •The sphereS2S^\{2\}with the induced metric: the tangent space at a point is the tangent plane to the sphere there, and geodesics are great circle arcs\.

Otto’s discovery: the probability space\(𝒫2​\(ℝd\),W2\)\(\\mathcal\{P\}\_\{2\}\(\\mathbb\{R\}^\{d\}\),W\_\{2\}\)also “looks like” a Riemannian manifold\.

Why “formal”?

Because𝒫2​\(ℝd\)\\mathcal\{P\}\_\{2\}\(\\mathbb\{R\}^\{d\}\)is aninfinite\-dimensionalspace and is not, strictly speaking, a finite\-dimensional smooth manifold\. Otto’s “Riemannian structure” faces several technical difficulties in a rigorous mathematical sense:

- •The tangent spaceTρ​𝒫2T\_\{\\rho\}\\mathcal\{P\}\_\{2\}is not a standard Banach/Hilbert space \(ρ\\rhoappears as a weight in the inner product\)\.
- •There is no uniform smooth structure \(the tangent spaces at different “points”ρ\\rhocannot be naturally “glued together”\)\.
- •Operations like the exponential map are only well\-defined in special cases\.

However, this “formal” Riemannian perspective iscomputationallyentirely correct—the PDEs, gradient formulas, and the JKO scheme derived from it are all rigorously valid\. Later, Ambrosio–Gigli–Savaré \(2008\) made these conclusions rigorous using the theory of “gradient flows in metric spaces”, without relying on manifold structure\. Thus, although Otto calculus is “formal”, the conclusions it yields can all be rigorously proved\.

Key Insight\(𝒫2​\(ℝd\),W2\)\(\\mathcal\{P\}\_\{2\}\(\\mathbb\{R\}^\{d\}\),W\_\{2\}\)is not merely a metric space—it is \(formally\) an infinite\-dimensional Riemannian manifold\. This means we can speak of “tangent vectors”, “inner products”, “gradients”, and “gradient flows” in the space of probability measures\. This is the foundation of the entire story\.

## 2The Continuity Equation: Conservation of Mass

Where we are:Section 1 defined the Wasserstein distance—a “static” concept \(comparing the distance between two fixed distributions\)\. But to discuss gradient flows, we need to describe how distributionsevolve over time\. The continuity equation is precisely the tool for this: it tells us “if mass flows according to a velocity fieldvv, how does the density evolve?” The Benamou–Brenier formula, Fokker–Planck equation, and JKO scheme that follow are all built upon this equation\.

Core idea:the continuity equation is thelocal differential form of mass conservation: the density inside any infinitesimal region changes only through the net flux across its boundary\.

### 2\.1Physical picture

Imagine ink diffusing in a glass of water\. At each instanttt, the ink concentration isρt​\(x\)\\rho\_\{t\}\(x\), and the average velocity of ink particles at positionxxisvt​\(x\)v\_\{t\}\(x\)\. The most fundamental physical law is:mass is neither created nor destroyed\.

### 2\.2Derivation from first principles

Take any fixed regionΩ⊂ℝd\\Omega\\subset\\mathbb\{R\}^\{d\}\. The total mass inside is:

M​\(t\)=∫Ωρt​\(x\)​𝑑xM\(t\)=\\int\_\{\\Omega\}\\rho\_\{t\}\(x\)\\,dx
The rate of change of mass equals the net flux through the boundary∂Ω\\partial\\Omega:

dd​t​∫Ωρt​\(x\)​𝑑x=−∫∂Ωρt​\(x\)​vt​\(x\)⋅n^​\(x\)​𝑑S​\(x\)\\frac\{d\}\{dt\}\\int\_\{\\Omega\}\\rho\_\{t\}\(x\)\\,dx=\-\\int\_\{\\partial\\Omega\}\\rho\_\{t\}\(x\)\\,v\_\{t\}\(x\)\\cdot\\hat\{n\}\(x\)\\,dS\(x\)\(14\)wheren^\\hat\{n\}is the outward unit normal \(the minus sign: outflow*decreases*mass inside\)\.

IntuitionWhere does the integral∫∂Ωρ​v⋅n^​𝑑S\\int\_\{\\partial\\Omega\}\\rho\\,v\\cdot\\hat\{n\}\\,dScome from?Consider a small surface elementd​SdSon the boundary∂Ω\\partial\\Omega, with outward unit normaln^\\hat\{n\}\. Near this surface element:1\.Themass densityisρ​\(x\)\\rho\(x\)\(how much mass per unit volume\)\.2\.Thevelocityisv​\(x\)v\(x\)\(the direction and speed of mass at that point\)\.3\.Thenormal componentv⋅n^v\\cdot\\hat\{n\}: the projection of velocity onto the normal direction of the surface element\.•v⋅n^\>0v\\cdot\\hat\{n\}\>0: velocity has an outward component—mass isflowing out\.•v⋅n^<0v\\cdot\\hat\{n\}<0: velocity has an inward component—mass isflowing in\.•v⋅n^=0v\\cdot\\hat\{n\}=0: velocity is parallel to the surface element—massmoves along the boundarywithout entering or exiting\.Mass passing through this surface element per unit time= density×\\timesnormal velocity×\\timesarea:d​Φ=ρ​\(x\)⋅\(v​\(x\)⋅n^​\(x\)\)⋅d​Sd\\Phi=\\rho\(x\)\\cdot\(v\(x\)\\cdot\\hat\{n\}\(x\)\)\\cdot dSWhy “density×\\timesnormal velocity”?Picture a checkpoint of aread​SdSacross a road \(the surface element\)\. The number of vehicles crossing per unit time equals the traffic density times the velocity component normal to the checkpoint\. When traffic approaches at an angle, only the normal componentv⋅n^v\\cdot\\hat\{n\}contributes to throughput; the tangential component moves vehicles along the checkpoint without crossing it\.Integrating over the entire boundary gives thetotal net outflux:Net outflux=∫∂Ωρ​v⋅n^​𝑑S\\text\{Net outflux\}=\\int\_\{\\partial\\Omega\}\\rho\\,v\\cdot\\hat\{n\}\\,dSBy conservation of mass: the rate of decrease of mass insideΩ\\Omega= net outflux, hence the minus sign\.

Step\-by\-step derivation:

Step 1\.Apply the divergence theorem to the right side:

−∫∂Ωρt​vt⋅n^​𝑑S=−∫Ω∇⋅\(ρt​vt\)​𝑑x\-\\int\_\{\\partial\\Omega\}\\rho\_\{t\}v\_\{t\}\\cdot\\hat\{n\}\\,dS=\-\\int\_\{\\Omega\}\\nabla\\cdot\(\\rho\_\{t\}v\_\{t\}\)\\,dx
Step 2\.Movedd​t\\frac\{d\}\{dt\}inside the integral on the left \(justified sinceΩ\\Omegais fixed\):

∫Ω∂tρt​\(x\)​d​x=−∫Ω∇⋅\(ρt​vt\)​\(x\)​𝑑x\\int\_\{\\Omega\}\\partial\_\{t\}\\rho\_\{t\}\(x\)\\,dx=\-\\int\_\{\\Omega\}\\nabla\\cdot\(\\rho\_\{t\}v\_\{t\}\)\(x\)\\,dx
Step 3\.Since this holds for allΩ\\Omegaand the integrand is continuous, it must vanish pointwise \(du Bois\-Reymond lemma, Appendix[C](https://arxiv.org/html/2606.24157#A3)\):

∂tρt\+∇⋅\(ρt​vt\)=0\\boxed\{\\partial\_\{t\}\\rho\_\{t\}\+\\nabla\\cdot\(\\rho\_\{t\}\\,v\_\{t\}\)=0\}\(15\)
This is thecontinuity equation\.

IntuitionExpanding the divergence operator:∇⋅\(ρ​v\)=v⋅∇ρ\+ρ​∇⋅v\\nabla\\cdot\(\\rho v\)=v\\cdot\\nabla\\rho\+\\rho\\,\\nabla\\cdot v\. So the continuity equation can also be written as:∂tρ\+v⋅∇ρ=−ρ​∇⋅v\\partial\_\{t\}\\rho\+v\\cdot\\nabla\\rho=\-\\rho\\,\\nabla\\cdot vThe left\-hand sideD​ρD​t=∂tρ\+v⋅∇ρ\\frac\{D\\rho\}\{Dt\}=\\partial\_\{t\}\\rho\+v\\cdot\\nabla\\rhois the rate of change of density following the fluid \(the material derivative\), and the right\-hand side−ρ​∇⋅v\-\\rho\\nabla\\cdot vrepresents density changes due to compression/expansion of the fluid\. If the fluid is incompressible \(∇⋅v=0\\nabla\\cdot v=0\), then the density is constant along streamlines\.

### 2\.3Weak formulation

For non\-smoothρt\\rho\_\{t\}\(or whenρt\\rho\_\{t\}is a measure, not a function\), we use the weak form\. Multiply \([15](https://arxiv.org/html/2606.24157#S2.E15)\) by a test functionφ∈Cc∞​\(ℝd×\[0,T\]\)\\varphi\\in C\_\{c\}^\{\\infty\}\(\\mathbb\{R\}^\{d\}\\times\[0,T\]\)and integrate by parts:

∫0T∫ℝd\[∂tφ​\(x,t\)\+vt​\(x\)⋅∇xφ​\(x,t\)\]​𝑑μt​\(x\)​𝑑t=0\\int\_\{0\}^\{T\}\\\!\\int\_\{\\mathbb\{R\}^\{d\}\}\\left\[\\partial\_\{t\}\\varphi\(x,t\)\+v\_\{t\}\(x\)\\cdot\\nabla\_\{x\}\\varphi\(x,t\)\\right\]d\\mu\_\{t\}\(x\)\\,dt=0\(16\)
Derivation details:Starting from∂tρ\+∇⋅\(ρ​v\)=0\\partial\_\{t\}\\rho\+\\nabla\\cdot\(\\rho v\)=0, multiply both sides byφ\\varphiand integrate:

∫0T∫φ​∂tρ​d​x​d​t\+∫0T∫φ​∇⋅\(ρ​v\)​𝑑x​𝑑t=0\\int\_\{0\}^\{T\}\\\!\\int\\varphi\\,\\partial\_\{t\}\\rho\\,dx\\,dt\+\\int\_\{0\}^\{T\}\\\!\\int\\varphi\\,\\nabla\\cdot\(\\rho v\)\\,dx\\,dt=0Integration by parts inttfor the first term:∫φ​∂tρ​d​t=−∫∂tφ​ρ​d​t\\int\\varphi\\,\\partial\_\{t\}\\rho\\,dt=\-\\int\\partial\_\{t\}\\varphi\\,\\rho\\,dt\(boundary terms vanish sinceφ\\varphihas compact support\)\. Integration by parts inxxfor the second term:∫φ​∇⋅\(ρ​v\)​𝑑x=−∫∇φ⋅\(ρ​v\)​𝑑x\\int\\varphi\\,\\nabla\\cdot\(\\rho v\)\\,dx=\-\\int\\nabla\\varphi\\cdot\(\\rho v\)\\,dx\. Combining:−∫0T∫\[∂tφ\+v⋅∇φ\]​ρ​𝑑x​𝑑t=0\-\\int\_\{0\}^\{T\}\\\!\\int\[\\partial\_\{t\}\\varphi\+v\\cdot\\nabla\\varphi\]\\rho\\,dx\\,dt=0; negating both sides gives the result\.

### 2\.4Connection to particle ODEs

###### Proposition 2\.1\.

If particles follow the ODEX˙t=vt​\(Xt\)\\dot\{X\}\_\{t\}=v\_\{t\}\(X\_\{t\}\)with initial distributionX0∼μ0X\_\{0\}\\sim\\mu\_\{0\}, thenμt:=\(Xt\)\#​μ0\\mu\_\{t\}:=\(X\_\{t\}\)\_\{\\\#\}\\mu\_\{0\}satisfies the continuity equation with velocity fieldvtv\_\{t\}\.

###### Proof\.

For anyφ∈Cc∞​\(ℝd\)\\varphi\\in C\_\{c\}^\{\\infty\}\(\\mathbb\{R\}^\{d\}\):

dd​t​∫φ​𝑑μt\\displaystyle\\frac\{d\}\{dt\}\\int\\varphi\\,d\\mu\_\{t\}=dd​t​∫φ​\(Xt​\(x\)\)​𝑑μ0​\(x\)\\displaystyle=\\frac\{d\}\{dt\}\\int\\varphi\(X\_\{t\}\(x\)\)\\,d\\mu\_\{0\}\(x\)\(17\)=∫∇φ​\(Xt​\(x\)\)⋅X˙t​\(x\)​𝑑μ0​\(x\)\\displaystyle=\\int\\nabla\\varphi\(X\_\{t\}\(x\)\)\\cdot\\dot\{X\}\_\{t\}\(x\)\\,d\\mu\_\{0\}\(x\)\(18\)=∫∇φ​\(Xt​\(x\)\)⋅vt​\(Xt​\(x\)\)​𝑑μ0​\(x\)\\displaystyle=\\int\\nabla\\varphi\(X\_\{t\}\(x\)\)\\cdot v\_\{t\}\(X\_\{t\}\(x\)\)\\,d\\mu\_\{0\}\(x\)\(19\)=∫∇φ​\(y\)⋅vt​\(y\)​𝑑μt​\(y\)\\displaystyle=\\int\\nabla\\varphi\(y\)\\cdot v\_\{t\}\(y\)\\,d\\mu\_\{t\}\(y\)\(20\)Step \([17](https://arxiv.org/html/2606.24157#S2.E17)\): definition of pushforward\. Step \([18](https://arxiv.org/html/2606.24157#S2.E18)\): chain rule\. Step \([19](https://arxiv.org/html/2606.24157#S2.E19)\): substitute ODE\. Step \([20](https://arxiv.org/html/2606.24157#S2.E20)\): change of variablesy=Xt​\(x\)y=X\_\{t\}\(x\), back to pushforward\.

This is exactly the weak form of∂tμt\+∇⋅\(μt​vt\)=0\\partial\_\{t\}\\mu\_\{t\}\+\\nabla\\cdot\(\\mu\_\{t\}v\_\{t\}\)=0\. ∎

## 3The Riemannian Structure of Wasserstein Space

Where we are:Section 1 gave us the Wasserstein distance, and Section 2 gave us the continuity equation\. Now we unify the two within a single geometric framework:the Wasserstein distance = the geodesic distance in probability space, and the velocity field described by the continuity equation is precisely the “tangent vector”\. This Riemannian structure is the foundation for everything that follows \(gradients, gradient flows, the JKO scheme\)\.

The prerequisite differential geometry \(manifolds, tangent spaces, cotangent spaces, dual spaces, musical isomorphisms, Riemannian metrics, gradients, gradient flows\) can be found inAppendix[B](https://arxiv.org/html/2606.24157#A2)\. If you are already familiar with these concepts, you may proceed directly\.

### 3\.1The Benamou–Brenier formula

So far, the Wasserstein distance has been a “static” optimization problem \(finding an optimal coupling\)\. Benamou and Brenier discovered an equivalent “dynamic” formulation—this bridge directly connects optimal transport with the continuity equation\.

###### Theorem 3\.1\(Benamou–Brenier, 2000\)\.

W22​\(μ0,μ1\)=inf\(ρt,vt\)\{∫01∫ℝd\|vt​\(x\)\|2​ρt​\(x\)​𝑑x​𝑑t\}W\_\{2\}^\{2\}\(\\mu\_\{0\},\\mu\_\{1\}\)=\\inf\_\{\(\\rho\_\{t\},v\_\{t\}\)\}\\left\\\{\\int\_\{0\}^\{1\}\\\!\\int\_\{\\mathbb\{R\}^\{d\}\}\|v\_\{t\}\(x\)\|^\{2\}\\,\\rho\_\{t\}\(x\)\\,dx\\,dt\\right\\\}\(21\)where the infimum is over all pairs\(ρt,vt\)t∈\[0,1\]\(\\rho\_\{t\},v\_\{t\}\)\_\{t\\in\[0,1\]\}withvt∈L2​\(ρt;ℝd\)v\_\{t\}\\in L^\{2\}\(\\rho\_\{t\};\\mathbb\{R\}^\{d\}\), satisfying:

- •Continuity equation \(in the weak sense\):∂tρt\+∇⋅\(ρt​vt\)=0\\partial\_\{t\}\\rho\_\{t\}\+\\nabla\\cdot\(\\rho\_\{t\}v\_\{t\}\)=0
- •Boundary conditions:ρ0=μ0,ρ1=μ1\\rho\_\{0\}=\\mu\_\{0\},\\;\\rho\_\{1\}=\\mu\_\{1\}

IntuitionAnalogy:In Euclidean space, the distance between two pointsa,b∈ℝna,b\\in\\mathbb\{R\}^\{n\}can be written as:\|a−b\|2=infγ:\[0,1\]→ℝn\{∫01\|γ˙\(t\)\|2dt\|γ\(0\)=a,γ\(1\)=b\}\|a\-b\|^\{2\}=\\inf\_\{\\gamma:\[0,1\]\\to\\mathbb\{R\}^\{n\}\}\\left\\\{\\int\_\{0\}^\{1\}\|\\dot\{\\gamma\}\(t\)\|^\{2\}\\,dt\\;\\middle\|\\;\\gamma\(0\)=a,\\,\\gamma\(1\)=b\\right\\\}\(The optimal path is a straight line with constant velocityγ˙=b−a\\dot\{\\gamma\}=b\-a, and cost=\|b−a\|2=\|b\-a\|^\{2\}\.\)The Benamou–Brenier formula embodies exactly the same idea, except:•“Points” become probability distributionsρt\\rho\_\{t\}\.•“Paths” become continuous evolutions of distributions \(satisfying mass conservation\)\.•“Squared velocity” becomes∫\|vt\|2​ρt​𝑑x\\int\|v\_\{t\}\|^\{2\}\\rho\_\{t\}\\,dx\(kinetic energy weighted byρt\\rho\_\{t\}\)\.Therefore, the Wasserstein distance is precisely thegeodesic length in probability space\.

Why∫\|v\|2​ρ​𝑑x\\int\|v\|^\{2\}\\rho\\,dxand not∫\|v\|2​𝑑x\\int\|v\|^\{2\}\\,dx?

Consider an infinitesimal fluid element with massd​m=ρ​d​xdm=\\rho\\,dx; its kinetic energy is12​\|v\|2​d​m=12​\|v\|2​ρ​d​x\\frac\{1\}\{2\}\|v\|^\{2\}\\,dm=\\frac\{1\}\{2\}\|v\|^\{2\}\\rho\\,dx\. Thus∫\|v\|2​ρ​𝑑x\\int\|v\|^\{2\}\\rho\\,dxis the total kinetic energy \(up to the factor12\\frac\{1\}\{2\}\)\. This weighting is crucial: if the density in some region is zero \(no mass\), the velocity there “costs nothing”\.

IntuitionWhy can “total kinetic energy” describe the Wasserstein distance? The deeper reason:Question 1: Why are the static and dynamic formulations equivalent?The static Kantorovich formulation says:W22=minγ​∫\|x−y\|2​𝑑γW\_\{2\}^\{2\}=\\min\_\{\\gamma\}\\int\|x\-y\|^\{2\}\\,d\\gamma—find the optimal “who goes where”\.The dynamic Benamou–Brenier formulation says:W22=min\(ρ,v\)​∫01∫\|v\|2​ρ​𝑑x​𝑑tW\_\{2\}^\{2\}=\\min\_\{\(\\rho,v\)\}\\int\_\{0\}^\{1\}\\\!\\int\|v\|^\{2\}\\rho\\,dx\\,dt—find the optimal “how to get there”\.Why do these two entirely different formulations give the same answer?The key insight:when the constraint enforces “constant\-speed straight\-line motion”, the two are naturally equivalent\.Suppose the optimal transport mapTTsendsxxtoT​\(x\)T\(x\)\. If each mass particle travels at constant speed in a straight line fromxxtoT​\(x\)T\(x\), then at timettit is located atxt=\(1−t\)​x\+t​T​\(x\)x\_\{t\}=\(1\-t\)x\+tT\(x\), with constant velocityx˙t=T​\(x\)−x\\dot\{x\}\_\{t\}=T\(x\)\-x\.The time integral of kinetic energy along this path:∫01\|T​\(x\)−x\|2​𝑑t=\|T​\(x\)−x\|2\\int\_\{0\}^\{1\}\|T\(x\)\-x\|^\{2\}\\,dt=\|T\(x\)\-x\|^\{2\}Integrating over all mass:∫\|T​\(x\)−x\|2​𝑑μ​\(x\)=W22​\(μ,ν\)\\int\|T\(x\)\-x\|^\{2\}\\,d\\mu\(x\)=W\_\{2\}^\{2\}\(\\mu,\\nu\)\.So constant\-speed straight\-line motion gives kinetic energy = static transport cost\. Any other path \(non\-straight or non\-constant\-speed\) only increases the kinetic energy \(by the Cauchy–Schwarz inequality\)\.Question 2: Why use kinetic energy \(velocity2\) rather than velocity itself?If we use∫\|v\|​ρ​𝑑x\\int\|v\|\\rho\\,dx\(first power of velocity\), we obtain theW1W\_\{1\}distance \(1\-Wasserstein\)\. Using∫\|v\|2​ρ​𝑑x\\int\|v\|^\{2\}\\rho\\,dx\(squared velocity\) givesW22W\_\{2\}^\{2\}\. Why preferW2W\_\{2\}?1\.Physical naturalness: minimum kinetic energy⇔\\iffgeodesic \(shortest path\)\.This requires explanation\. In Riemannian geometry, the “distance” between two points is defined as:d​\(a,b\)2=infγ∫01\|γ˙​\(t\)\|2​𝑑t,γ​\(0\)=a,γ​\(1\)=bd\(a,b\)^\{2\}=\\inf\_\{\\gamma\}\\int\_\{0\}^\{1\}\|\\dot\{\\gamma\}\(t\)\|^\{2\}\\,dt,\\quad\\gamma\(0\)=a,\\;\\gamma\(1\)=bThe optimalγ\\gammais called ageodesic—it is the “straight line” \(shortest path\) in the given geometry\.In physics, in the absence of external forces \(a free particle\), the equation of motion isγ¨=0\\ddot\{\\gamma\}=0\(constant\-speed straight\-line motion\)\. This is equivalent to minimizing the actionS=∫0112​\|γ˙\|2​𝑑tS=\\int\_\{0\}^\{1\}\\frac\{1\}\{2\}\|\\dot\{\\gamma\}\|^\{2\}\\,dt\(the time integral of kinetic energy\)\. Thereforefree particle trajectory = geodesic = minimum kinetic energy path\.On curved spaces \(e\.g\., the sphere\), geodesics are no longer “straight lines” but great circle arcs—yet they are still the paths that minimize∫\|γ˙\|2​𝑑t\\int\|\\dot\{\\gamma\}\|^\{2\}\\,dt\.What Benamou–Brenier does is a perfect analogy:∫01∫\|v\|2​ρ​𝑑x​𝑑t\\int\_\{0\}^\{1\}\\\!\\int\|v\|^\{2\}\\rho\\,dx\\,dtis the “kinetic energy integral in probability space”, and the pathρt\\rho\_\{t\}that minimizes it is the geodesic in Wasserstein space\.2\.Mathematical convenience:\|v\|2\|v\|^\{2\}is smooth and strictly convex, makingW22W\_\{2\}^\{2\}differentiable with respect toρ\\rho, whereas\|v\|\|v\|is non\-smooth at zero, andW1W\_\{1\}has worse mathematical properties \(non\-differentiable, no Riemannian structure\)\.3\.Riemannian structure:Only when using the squared velocity \(rather than the absolute value\) as the “metric” does the probability space acquire an inner product structure—an inner product requires‖v‖2\\\|v\\\|^\{2\}to be a quadratic function ofvv\. If one uses‖v‖=∫\|v\|​ρ​𝑑x\\\|v\\\|=\\int\|v\|\\rho\\,dx, that is merely a norm \(Finsler geometry\) without an inner product, and “gradients” cannot be defined\.Question 3: What is the deeper intuition behind the dynamic formulation?Imagine you are a fluid dispatcher\. You must deform a body of fluid from the shapeμ0\\mu\_\{0\}into the shapeμ1\\mu\_\{1\}\. The rules are:•Fluid cannot be created or destroyed \(continuity equation∂tρ\+∇⋅\(ρ​v\)=0\\partial\_\{t\}\\rho\+\\nabla\\cdot\(\\rho v\)=0\)\.•You must minimize the total “fuel” consumed—each fluid element’s fuel consumption is proportional to velocity×2\{\}^\{2\}\\timesmass×\\timestime\.This is precisely the Benamou–Brenier formula\.What is the optimal strategy?Have each mass particle travel in astraight line at constant speedfrom start to finish\.•Why constant speed? Because∫01\|v​\(t\)\|2​𝑑t≥\(∫01\|v​\(t\)\|​𝑑t\)2\\int\_\{0\}^\{1\}\|v\(t\)\|^\{2\}\\,dt\\geq\(\\int\_\{0\}^\{1\}\|v\(t\)\|\\,dt\)^\{2\}\(Jensen’s inequality\), with equality if and only if\|v\|\|v\|is constant\. Going slow then fast, or fast then slow, always wastes more energy than traveling at constant speed\.•Why straight lines? Because detours = longer path = higher average speed required = greater kinetic energy\.Thereforethe Wasserstein geodesic is “all mass particles traveling in straight lines at constant speed”, and the total energy is∫\|x−T​\(x\)\|2​𝑑μ=W22\\int\|x\-T\(x\)\|^\{2\}\\,d\\mu=W\_\{2\}^\{2\}\.

### 3\.2Why continuity equation \+ minimum kinetic energy gives the OT path

Let us make the previous statement precise\. The continuity equation alone is not an optimization principle:

∂tρt\+∇⋅\(ρt​vt\)=0\\partial\_\{t\}\\rho\_\{t\}\+\\nabla\\cdot\(\\rho\_\{t\}v\_\{t\}\)=0only says that probability mass is transported without being created or destroyed\. The optimization enters through the kinetic action

𝒜​\(ρ,v\)=∫01∫ℝd\|vt​\(x\)\|2​ρt​\(x\)​𝑑x​𝑑t\.\\mathcal\{A\}\(\\rho,v\)=\\int\_\{0\}^\{1\}\\\!\\int\_\{\\mathbb\{R\}^\{d\}\}\|v\_\{t\}\(x\)\|^\{2\}\\rho\_\{t\}\(x\)\\,dx\\,dt\.The Benamou–Brenier problem is:

infρt,vt𝒜​\(ρ,v\)\\inf\_\{\\rho\_\{t\},v\_\{t\}\}\\mathcal\{A\}\(\\rho,v\)subject to

∂tρt\+∇⋅\(ρt​vt\)=0,ρ0=μ0,ρ1=μ1\.\\partial\_\{t\}\\rho\_\{t\}\+\\nabla\\cdot\(\\rho\_\{t\}v\_\{t\}\)=0,\\qquad\\rho\_\{0\}=\\mu\_\{0\},\\quad\\rho\_\{1\}=\\mu\_\{1\}\.In words: among all mass\-preserving ways to deformμ0\\mu\_\{0\}intoμ1\\mu\_\{1\}, choose the one with the smallest total kinetic energy\.

IntuitionThe two halves of the argument\.First, every admissible dynamic path contains a static transport plan\.If particles move according toX˙t=vt​\(Xt\),\\dot\{X\}\_\{t\}=v\_\{t\}\(X\_\{t\}\),then the continuity equation says thatρt=Law​\(Xt\)\\rho\_\{t\}=\\mathrm\{Law\}\(X\_\{t\}\)\. Each particle has an initial pointX0=xX\_\{0\}=xand a final pointX1=yX\_\{1\}=y, so the whole flow induces a couplingγ=Law​\(X0,X1\)∈Π​\(μ0,μ1\)\.\\gamma=\\mathrm\{Law\}\(X\_\{0\},X\_\{1\}\)\\in\\Pi\(\\mu\_\{0\},\\mu\_\{1\}\)\.Thus any dynamic path tells us “who goes where\.”Second, the kinetic energy of any particle path dominates its squared displacement\.For one trajectoryXtX\_\{t\}, Cauchy–Schwarz gives∫01\|X˙t\|2​𝑑t≥\|∫01X˙t​𝑑t\|2=\|X1−X0\|2\.\\int\_\{0\}^\{1\}\|\\dot\{X\}\_\{t\}\|^\{2\}\\,dt\\geq\\left\|\\int\_\{0\}^\{1\}\\dot\{X\}\_\{t\}\\,dt\\right\|^\{2\}=\|X\_\{1\}\-X\_\{0\}\|^\{2\}\.Equality holds exactly when the particle moves along a straight line at constant speed\. Averaging over all particles,∫01∫\|vt​\(x\)\|2​ρt​\(x\)​𝑑x​𝑑t≥∫\|y−x\|2​𝑑γ​\(x,y\)\.\\int\_\{0\}^\{1\}\\\!\\int\|v\_\{t\}\(x\)\|^\{2\}\\rho\_\{t\}\(x\)\\,dx\\,dt\\geq\\int\|y\-x\|^\{2\}\\,d\\gamma\(x,y\)\.SinceW22​\(μ0,μ1\)W\_\{2\}^\{2\}\(\\mu\_\{0\},\\mu\_\{1\}\)is the minimum of the right\-hand side over all couplingsγ\\gamma, every admissible dynamic path satisfies𝒜​\(ρ,v\)≥W22​\(μ0,μ1\)\.\\mathcal\{A\}\(\\rho,v\)\\geq W\_\{2\}^\{2\}\(\\mu\_\{0\},\\mu\_\{1\}\)\.No mass\-preserving path can spend less kinetic energy than the optimal transport cost\.

Conversely, take an optimal couplingγ⋆∈Π​\(μ0,μ1\)\\gamma^\{\\star\}\\in\\Pi\(\\mu\_\{0\},\\mu\_\{1\}\)\. Move each pair\(x,y\)\(x,y\)along the constant\-speed straight line

Xt=\(1−t\)​x\+t​y,X˙t=y−x\.X\_\{t\}=\(1\-t\)x\+ty,\\qquad\\dot\{X\}\_\{t\}=y\-x\.Define

ρt=\(\(1−t\)​x\+t​y\)\#​γ⋆\.\\rho\_\{t\}=\\bigl\(\(1\-t\)x\+ty\\bigr\)\_\{\\\#\}\\gamma^\{\\star\}\.This path satisfies the continuity equation because it is generated by moving particles\. Its kinetic action is

∫01∫\|y−x\|2​𝑑γ⋆​\(x,y\)​𝑑t=∫\|y−x\|2​𝑑γ⋆​\(x,y\)=W22​\(μ0,μ1\)\.\\int\_\{0\}^\{1\}\\\!\\int\|y\-x\|^\{2\}\\,d\\gamma^\{\\star\}\(x,y\)\\,dt=\\int\|y\-x\|^\{2\}\\,d\\gamma^\{\\star\}\(x,y\)=W\_\{2\}^\{2\}\(\\mu\_\{0\},\\mu\_\{1\}\)\.So the lower bound is achieved, and we obtain

inf∂tρ\+∇⋅\(ρ​v\)=0ρ0=μ0,ρ1=μ1∫01∫\|vt\|2​ρt​𝑑x​𝑑t=W22​\(μ0,μ1\)\\boxed\{\\inf\_\{\\begin\{subarray\}\{c\}\\partial\_\{t\}\\rho\+\\nabla\\cdot\(\\rho v\)=0\\\\ \\rho\_\{0\}=\\mu\_\{0\},\\;\\rho\_\{1\}=\\mu\_\{1\}\\end\{subarray\}\}\\int\_\{0\}^\{1\}\\\!\\int\|v\_\{t\}\|^\{2\}\\rho\_\{t\}\\,dx\\,dt=W\_\{2\}^\{2\}\(\\mu\_\{0\},\\mu\_\{1\}\)\}This is the Benamou–Brenier formula\.

Ifμ0\\mu\_\{0\}is sufficiently regular, the optimal coupling is induced by a Brenier mapTT, soy=T​\(x\)y=T\(x\)and the minimizer becomes

ρt=\(\(1−t\)​id\+t​T\)\#​μ0\.\\rho\_\{t\}=\\bigl\(\(1\-t\)\\mathrm\{id\}\+tT\\bigr\)\_\{\\\#\}\\mu\_\{0\}\.This is theMcCann displacement interpolation: every infinitesimal piece of mass travels along a straight line at constant speed according to the optimal transport map\. This is why the OT path is the “straight line” in Wasserstein space\.

Key InsightThe continuity equation provides the*constraint*: admissible mass\-preserving paths\. The kinetic energy provides the*criterion*: choose the least costly path\. Together they recover the optimal transport path, i\.e\. the Wasserstein geodesic\.

### 3\.3Otto’s Riemannian interpretation

Otto \(2001\) proposed viewing\(𝒫2​\(ℝd\),W2\)\(\\mathcal\{P\}\_\{2\}\(\\mathbb\{R\}^\{d\}\),W\_\{2\}\)as an infinite\-dimensional Riemannian manifold:

###### Definition 3\.2\(Tangent space atρ\\rho\)\.

The tangent space at a pointρ∈𝒫2​\(ℝd\)\\rho\\in\\mathcal\{P\}\_\{2\}\(\\mathbb\{R\}^\{d\}\)is:

Tρ​𝒫2:=\{∇ϕ:ϕ∈Cc∞​\(ℝd\)\}¯L2​\(ρ;ℝd\)T\_\{\\rho\}\\mathcal\{P\}\_\{2\}:=\\overline\{\\\{\\nabla\\phi:\\phi\\in C\_\{c\}^\{\\infty\}\(\\mathbb\{R\}^\{d\}\)\\\}\}^\{L^\{2\}\(\\rho;\\mathbb\{R\}^\{d\}\)\}\(22\)i\.e\., theL2​\(ρ;ℝd\)L^\{2\}\(\\rho;\\mathbb\{R\}^\{d\}\)\-closure of gradient vector fields \(hereL2​\(ρ;ℝd\)L^\{2\}\(\\rho;\\mathbb\{R\}^\{d\}\)denotes the space of vector fieldsv:ℝd→ℝdv:\\mathbb\{R\}^\{d\}\\to\\mathbb\{R\}^\{d\}with∫\|v\|2​ρ​𝑑x<∞\\int\|v\|^\{2\}\\rho\\,dx<\\infty\)\.

Why are tangent vectors gradient fields?

A “tangent vector” describes the instantaneous direction of change ofρ\\rho\. By the continuity equation, the evolution ofρ\\rhois determined by the velocity fieldvv:ρ˙=−∇⋅\(ρ​v\)\\dot\{\\rho\}=\-\\nabla\\cdot\(\\rho v\)\. However, the sameρ˙\\dot\{\\rho\}may correspond to multiple differentvv\(since∇⋅\(ρ​v\)=∇⋅\(ρ​\(v\+w\)\)\\nabla\\cdot\(\\rho v\)=\\nabla\\cdot\(\\rho\(v\+w\)\)whenever∇⋅\(ρ​w\)=0\\nabla\\cdot\(\\rho w\)=0\)\.

To remove this ambiguity, we choose thevvwithminimumL2​\(ρ\)L^\{2\}\(\\rho\)norm\. For givenρ\>0\\rho\>0, any vector fieldvvcan be uniquely decomposed asv=∇ϕ\+wv=\\nabla\\phi\+w, where∇⋅\(ρ​w\)=0\\nabla\\cdot\(\\rho w\)=0\. These two components are orthogonal in theL2​\(ρ\)L^\{2\}\(\\rho\)inner product \(verification:∫ρ​∇ϕ⋅w​d​x=−∫ϕ​∇⋅\(ρ​w\)​𝑑x=0\\int\\rho\\,\\nabla\\phi\\cdot w\\,dx=\-\\int\\phi\\,\\nabla\\cdot\(\\rho w\)\\,dx=0\)\. Therefore‖v‖L2​\(ρ\)2=‖∇ϕ‖L2​\(ρ\)2\+‖w‖L2​\(ρ\)2\\\|v\\\|\_\{L^\{2\}\(\\rho\)\}^\{2\}=\\\|\\nabla\\phi\\\|\_\{L^\{2\}\(\\rho\)\}^\{2\}\+\\\|w\\\|\_\{L^\{2\}\(\\rho\)\}^\{2\}, and the minimum\-normvvmust be a gradient fieldv=∇ϕv=\\nabla\\phi\.

###### Definition 3\.3\(Riemannian metric\)\.

For tangent vectorsξ=∇ϕ\\xi=\\nabla\\phi,η=∇ψ∈Tρ​𝒫2\\eta=\\nabla\\psi\\in T\_\{\\rho\}\\mathcal\{P\}\_\{2\}:

⟨ξ,η⟩ρ:=∫ℝd∇ϕ​\(x\)⋅∇ψ​\(x\)​ρ​\(x\)​𝑑x\\langle\\xi,\\eta\\rangle\_\{\\rho\}:=\\int\_\{\\mathbb\{R\}^\{d\}\}\\nabla\\phi\(x\)\\cdot\\nabla\\psi\(x\)\\,\\rho\(x\)\\,dx\(23\)

This is the inner product version of the instantaneous cost∫\|v\|2​ρ​𝑑x=‖∇ϕ‖ρ2\\int\|v\|^\{2\}\\rho\\,dx=\\\|\\nabla\\phi\\\|\_\{\\rho\}^\{2\}appearing in the Benamou–Brenier formula\.

### 3\.4Geodesics = displacement interpolation

The geodesic \(shortest path\) betweenμ0\\mu\_\{0\}andμ1\\mu\_\{1\}in\(𝒫2,W2\)\(\\mathcal\{P\}\_\{2\},W\_\{2\}\)is theMcCann displacement interpolation:

ρt=\(\(1−t\)​id\+t​T\)\#​μ0,t∈\[0,1\]\\rho\_\{t\}=\\bigl\(\(1\-t\)\\mathrm\{id\}\+t\\,T\\bigr\)\_\{\\\#\}\\mu\_\{0\},\\quad t\\in\[0,1\]\(24\)whereT=∇ϕT=\\nabla\\phiis the optimal transport map fromμ0\\mu\_\{0\}toμ1\\mu\_\{1\}\.

Intuition: each mass particle travels in a straight line from its initial positionxxto its target positionT​\(x\)T\(x\), and at timettit is located at\(1−t\)​x\+t​T​\(x\)\(1\-t\)x\+tT\(x\)\.

Summary of progress so far:We have established the geometric structure of probability space\(𝒫2,W2\)\(\\mathcal\{P\}\_\{2\},W\_\{2\}\)—it is a Riemannian manifold whose geodesics are given by the Benamou–Brenier formula\.Next step: we will derive a physical equation \(Fokker–Planck\), and then in Section 5 prove that it is precisely thegradient flowin this geometry\.

## 4The Fokker–Planck Equation

Where we are:The first three sections established the geometric framework—the Wasserstein distance \(Section 1\), the continuity equation \(Section 2\), and the Riemannian structure \(Section 3\)\. But these are all descriptions of the “space” itself, without a concrete physical equation\. This section derives the Fokker–Planck equation: a PDE describing “how the density evolves when particles undergo Brownian motion in a potential field”\. The crucial final step is to rewrite it in the form of a continuity equation∂tρ\+∇⋅\(ρ​v\)=0\\partial\_\{t\}\\rho\+\\nabla\\cdot\(\\rho v\)=0, thereby identifying the velocity fieldvv\. The next section \(Section 5\) will prove that thisvvis precisely the negative Wasserstein gradient of the free energy\.

### 4\.1Overview

We begin with a microscopic stochastic process—particles undergoing Brownian motion in a potential field—and ask how the macroscopicprobability densityof the particles evolves\. The answer is the Fokker–Planck equation, a partial differential equation governing the density over time\.

The goals of this section:

1. 1\.Explain the physical meaning of SDEs \(stochastic differential equations\)\.
2. 2\.Derive the Fokker–Planck equation from SDEs \(with detailed explanation of each step\)\.
3. 3\.Rewrite the Fokker–Planck equation in the form of a continuity equation \(preparing the connection to gradient flows\)\.

### 4\.2Motivation: the physical setup

Setting:Imagine a pollen grain floating in water \(Brownian motion\)\. At the same time, the bottom of the pool has a bowl shape—the pollen tends toward the bottom due to gravity\. The particle’s motion is driven by two forces:

1. 1\.Deterministic force−∇V​\(x\)\-\\nabla V\(x\): the negative gradient of the potentialV​\(x\)V\(x\), pushing the particle toward the location of lowest potential energy \(the bottom of the bowl\)
2. 2\.Random force\(noise\): random collisions from water molecules, causing the particle to move irregularly

###### Definition 4\.1\(The SDE\)\.

The particle’s positionXt∈ℝdX\_\{t\}\\in\\mathbb\{R\}^\{d\}satisfies the following stochastic differential equation \(SDE\):

d​Xt=−∇V​\(Xt\)​d​t⏟drift: deterministic force\+2​d​Bt⏟diffusion: random noisedX\_\{t\}=\\underbrace\{\-\\nabla V\(X\_\{t\}\)\\,dt\}\_\{\\text\{drift: deterministic force\}\}\+\\underbrace\{\\sqrt\{2\}\\,dB\_\{t\}\}\_\{\\text\{diffusion: random noise\}\}\(25\)whereBtB\_\{t\}is standarddd\-dimensional Brownian motion\.

IntuitionTerm\-by\-term explanation of the SDE:d​XtdX\_\{t\}: the displacement of the particle during the time interval\[t,t\+d​t\]\[t,t\+dt\]\(how far it moves in an infinitesimal time step\)\.−∇V​\(Xt\)​d​t\-\\nabla V\(X\_\{t\}\)\\,dt\(drift term\):•V​\(x\)V\(x\)is the potential energy function—for example,V​\(x\)=12​\|x\|2V\(x\)=\\frac\{1\}\{2\}\|x\|^\{2\}represents a bowl\-shaped potential \(low at center, high at edges\)•−∇V\-\\nabla Vis the direction of steepest descent of the potential \(like a ball rolling down the bowl wall toward the bottom\)•The factord​tdtarises because this is displacement produced by a deterministic velocity: displacement==velocity×\\timestime2​d​Bt\\sqrt\{2\}\\,dB\_\{t\}\(diffusion term\):•d​BtdB\_\{t\}is the increment of Brownian motion—a random vector with zero mean and varianced​tdt•Intuition: at each time stepd​tdt, the particle receives a random kick in a random direction with magnitude∼d​t\\sim\\sqrt\{dt\}•2\\sqrt\{2\}is the diffusion strength \(temperature\)—the larger the coefficient, the stronger the randomnessCompetition between the two terms:•The drift term seeks to “concentrate” particles near the minimum ofVV•The diffusion term seeks to “spread” particles everywhere•The eventual balance \(stationary distribution\) is determined by the ratio of their strengths

Central question:If we simultaneously release alarge numberof particles \(each independently following the SDE above\), with initial distributionρ0​\(x\)\\rho\_\{0\}\(x\), what PDE does the probability densityρt​\(x\)\\rho\_\{t\}\(x\)at timettsatisfy?

This is precisely the question that the Fokker\-Planck equation answers—going from “individual particle’s random trajectory” to “macroscopic density evolution of many particles\.”

### 4\.3Prerequisite: Itô’s formula

Deriving the Fokker\-Planck equation requires one tool: Itô’s formula \(Itô, 1951\)\. It is the “chain rule” of stochastic calculus—but with one extra term compared to the ordinary chain rule\.

###### Theorem 4\.2\(Itô’s formula — informal version\)\.

LetXtX\_\{t\}satisfyd​Xt=b​d​t\+σ​d​BtdX\_\{t\}=b\\,dt\+\\sigma\\,dB\_\{t\}, and letffbe a twice\-differentiable function\. Then:

d​f​\(Xt\)=∇f​\(Xt\)⋅d​Xt⏟ordinary chain rule\+12​∑i,j∂2f∂xi​∂xj​\(Xt\)⋅\(σ​σT\)i​j​d​t⏟Itô correction termdf\(X\_\{t\}\)=\\underbrace\{\\nabla f\(X\_\{t\}\)\\cdot dX\_\{t\}\}\_\{\\text\{ordinary chain rule\}\}\+\\underbrace\{\\frac\{1\}\{2\}\\sum\_\{i,j\}\\frac\{\\partial^\{2\}f\}\{\\partial x\_\{i\}\\partial x\_\{j\}\}\(X\_\{t\}\)\\cdot\(\\sigma\\sigma^\{T\}\)\_\{ij\}\\,dt\}\_\{\\text\{It\\^\{o\} correction term\}\}\(26\)

Why is there an extra term?

In ordinary calculus,d​f=f′​\(x\)​d​xdf=f^\{\\prime\}\(x\)\\,dx, and we ignore the\(d​x\)2\(dx\)^\{2\}term \(because\(d​x\)2\(dx\)^\{2\}is of higher order thand​xdx\)\.

But in stochastic calculus, the “size” of the Brownian motion incrementd​BtdB\_\{t\}isd​t\\sqrt\{dt\}rather thand​tdt\. This means\(d​Bt\)2≈d​t\(dB\_\{t\}\)^\{2\}\\approx dtis*not*a higher\-order infinitesimal; it is of the same order asd​tdt\. Therefore, the second\-order term in the Taylor expansion cannot be discarded and must be retained\. This is the Itô correction\.

Specifically:

- •Ordinary calculus:\(d​x\)2=0\(dx\)^\{2\}=0\(higher\-order infinitesimal, discarded\)
- •Stochastic calculus:\(d​Bt\)2=d​t\(dB\_\{t\}\)^\{2\}=dt\(same order asd​tdt,cannotbe discarded\)
- •Cross term:d​Bt⋅d​t=0dB\_\{t\}\\cdot dt=0\(d​t⋅d​t=d​t3/2\\sqrt\{dt\}\\cdot dt=dt^\{3/2\}, higher\-order, discarded\)

### 4\.4Derivation of Fokker–Planck: complete details

We now derive the equation step by step\. The strategy is:

1. 1\.For an arbitrary “test function”φ\\varphi, compute the rate of change of𝔼​\[φ​\(Xt\)\]\\mathbb\{E\}\[\\varphi\(X\_\{t\}\)\]
2. 2\.Using𝔼​\[φ​\(Xt\)\]=∫φ​ρt​𝑑x\\mathbb\{E\}\[\\varphi\(X\_\{t\}\)\]=\\int\\varphi\\,\\rho\_\{t\}\\,dx, translate the result into an equation forρt\\rho\_\{t\}

Step 1: Apply Itô’s formula\.

Take an arbitrary smooth functionφ∈Cc∞​\(ℝd\)\\varphi\\in C\_\{c\}^\{\\infty\}\(\\mathbb\{R\}^\{d\}\)\(a “probe”—we will use it to “probe” changes in the density\)\.

For our SDE \([25](https://arxiv.org/html/2606.24157#S4.E25)\),b=−∇Vb=\-\\nabla V,σ=2​I\\sigma=\\sqrt\{2\}\\,I, soσ​σT=2​I\\sigma\\sigma^\{T\}=2I\. Applying Itô’s formula:

d​φ​\(Xt\)\\displaystyle d\\varphi\(X\_\{t\}\)=∇φ​\(Xt\)⋅d​Xt\+12​∑i,j∂2φ∂xi​∂xj​\(Xt\)⋅2​δi​j​d​t\\displaystyle=\\nabla\\varphi\(X\_\{t\}\)\\cdot dX\_\{t\}\+\\frac\{1\}\{2\}\\sum\_\{i,j\}\\frac\{\\partial^\{2\}\\varphi\}\{\\partial x\_\{i\}\\partial x\_\{j\}\}\(X\_\{t\}\)\\cdot 2\\delta\_\{ij\}\\,dt=∇φ​\(Xt\)⋅d​Xt\+Δ​φ​\(Xt\)​d​t\\displaystyle=\\nabla\\varphi\(X\_\{t\}\)\\cdot dX\_\{t\}\+\\Delta\\varphi\(X\_\{t\}\)\\,dt\(27\)
Explanation of each term:

- •∇φ⋅d​Xt\\nabla\\varphi\\cdot dX\_\{t\}: the “ordinary” chain rule part\. The change inφ≈\\varphi\\approxgradient×\\timesdisplacement\.
- •Δ​φ​d​t\\Delta\\varphi\\,dt: the Itô correction\.Δ​φ=∑i∂2φ∂xi2\\Delta\\varphi=\\sum\_\{i\}\\frac\{\\partial^\{2\}\\varphi\}\{\\partial x\_\{i\}^\{2\}\}is the Laplacian \(the sum of second derivatives ofφ\\varphiin all directions\)\. It arises from\(d​Bt\)2=d​t\(dB\_\{t\}\)^\{2\}=dt\.
- •2​δi​j2\\delta\_\{ij\}comes from\(σ​σT\)i​j=\(2​I\)​\(2​I\)i​jT=2​δi​j\(\\sigma\\sigma^\{T\}\)\_\{ij\}=\(\\sqrt\{2\}I\)\(\\sqrt\{2\}I\)^\{T\}\_\{ij\}=2\\delta\_\{ij\}\.
- •12×2=1\\frac\{1\}\{2\}\\times 2=1, so the correction term simplifies toΔ​φ​d​t\\Delta\\varphi\\,dt\.

Step 2: Expand∇φ⋅d​Xt\\nabla\\varphi\\cdot dX\_\{t\}\.

Substitutingd​Xt=−∇V​d​t\+2​d​BtdX\_\{t\}=\-\\nabla V\\,dt\+\\sqrt\{2\}\\,dB\_\{t\}:

d​φ​\(Xt\)\\displaystyle d\\varphi\(X\_\{t\}\)=∇φ​\(Xt\)⋅\(−∇V​\(Xt\)​d​t\+2​d​Bt\)\+Δ​φ​\(Xt\)​d​t\\displaystyle=\\nabla\\varphi\(X\_\{t\}\)\\cdot\\bigl\(\-\\nabla V\(X\_\{t\}\)\\,dt\+\\sqrt\{2\}\\,dB\_\{t\}\\bigr\)\+\\Delta\\varphi\(X\_\{t\}\)\\,dt=\[−∇φ⋅∇V\+Δ​φ\]​\(Xt\)​d​t\+2​∇φ​\(Xt\)⋅d​Bt\\displaystyle=\\bigl\[\-\\nabla\\varphi\\cdot\\nabla V\+\\Delta\\varphi\\bigr\]\(X\_\{t\}\)\\,dt\+\\sqrt\{2\}\\,\\nabla\\varphi\(X\_\{t\}\)\\cdot dB\_\{t\}\(28\)
Rearranging:d​φ​\(Xt\)=d\\varphi\(X\_\{t\}\)=\(deterministic part\)d​tdt\+ \(stochastic part\)d​BtdB\_\{t\}\.

Step 3: Take expectations\.

Take expectations of both sides\. The key fact:the expectation of an Itô integral is zero\.

Why is𝔼​\[∫∇φ⋅d​Bt\]=0\\mathbb\{E\}\[\\int\\nabla\\varphi\\cdot dB\_\{t\}\]=0? Because the Itô integral is amartingale—intuitively,d​BtdB\_\{t\}has zero expectation \(random noise does not favor any direction on average\), so its contribution to the expected value ofφ\\varphiiszero on average\.

Therefore:

dd​t​𝔼​\[φ​\(Xt\)\]=𝔼​\[−∇φ​\(Xt\)⋅∇V​\(Xt\)\+Δ​φ​\(Xt\)\]\\frac\{d\}\{dt\}\\mathbb\{E\}\[\\varphi\(X\_\{t\}\)\]=\\mathbb\{E\}\\bigl\[\-\\nabla\\varphi\(X\_\{t\}\)\\cdot\\nabla V\(X\_\{t\}\)\+\\Delta\\varphi\(X\_\{t\}\)\\bigr\]\(29\)
Interpretation:The rate of change of𝔼​\[φ​\(Xt\)\]\\mathbb\{E\}\[\\varphi\(X\_\{t\}\)\]= “contribution of deterministic drift toφ\\varphi” \+ “contribution of diffusion toφ\\varphi\.” The noise term vanishes under expectation\.

Step 4: Rewrite using densityρt\\rho\_\{t\}\.

IfXtX\_\{t\}has probability densityρt​\(x\)\\rho\_\{t\}\(x\), then𝔼​\[g​\(Xt\)\]=∫g​\(x\)​ρt​\(x\)​𝑑x\\mathbb\{E\}\[g\(X\_\{t\}\)\]=\\int g\(x\)\\rho\_\{t\}\(x\)\\,dx\(expectation = integral against the density\)\.

Substituting:

dd​t​∫ℝdφ​\(x\)​ρt​\(x\)​𝑑x=∫ℝd\[−∇φ​\(x\)⋅∇V​\(x\)\+Δ​φ​\(x\)\]​ρt​\(x\)​𝑑x\\frac\{d\}\{dt\}\\int\_\{\\mathbb\{R\}^\{d\}\}\\varphi\(x\)\\,\\rho\_\{t\}\(x\)\\,dx=\\int\_\{\\mathbb\{R\}^\{d\}\}\\bigl\[\-\\nabla\\varphi\(x\)\\cdot\\nabla V\(x\)\+\\Delta\\varphi\(x\)\\bigr\]\\rho\_\{t\}\(x\)\\,dx\(30\)
The left side=∫φ​∂tρt​d​x=\\int\\varphi\\,\\partial\_\{t\}\\rho\_\{t\}\\,dx\(sinceφ\\varphidoes not depend ontt, the derivative acts onρt\\rho\_\{t\}\)\.

Now the problem becomes: the right side contains∇φ\\nabla\\varphiandΔ​φ\\Delta\\varphi—but the equation we want should be aboutρt\\rho\_\{t\}, not aboutφ\\varphi\. The solution:integration by parts, transferring the derivatives fromφ\\varphiontoρt\\rho\_\{t\}\.

Step 5: Integration by parts — the key step\.

We use integration by parts \(see Appendix[C](https://arxiv.org/html/2606.24157#A3)for details\) to transfer the derivatives fromφ\\varphiontoρt\\rho\_\{t\}\. The key formulas are:

∫∇φ⋅𝐅​d​x=−∫φ​\(∇⋅𝐅\)​𝑑x,∫\(Δ​φ\)​g​𝑑x=∫φ​\(Δ​g\)​𝑑x\\int\\nabla\\varphi\\cdot\\mathbf\{F\}\\,dx=\-\\int\\varphi\\,\(\\nabla\\cdot\\mathbf\{F\}\)\\,dx,\\qquad\\int\(\\Delta\\varphi\)\\,g\\,dx=\\int\\varphi\\,\(\\Delta g\)\\,dx
Term 1:−∫∇φ⋅∇V​ρt​d​x\-\\int\\nabla\\varphi\\cdot\\nabla V\\,\\rho\_\{t\}\\,dx

Treatingρt​∇V\\rho\_\{t\}\\nabla Vas the vector field𝐅=ρt​∇V\\mathbf\{F\}=\\rho\_\{t\}\\nabla Vand applying integration by parts:

−∫∇φ⋅\(ρt​∇V\)​𝑑x\\displaystyle\-\\int\\nabla\\varphi\\cdot\(\\rho\_\{t\}\\nabla V\)\\,dx=\+∫φ​∇⋅\(ρt​∇V\)​𝑑x\\displaystyle=\+\\int\\varphi\\,\\nabla\\cdot\(\\rho\_\{t\}\\nabla V\)\\,dx\(31\)
Term 2:∫Δ​φ​ρt​𝑑x\\int\\Delta\\varphi\\,\\rho\_\{t\}\\,dx

This requires two integrations by parts\. The first:

∫\(Δ​φ\)​ρt​𝑑x\\displaystyle\\int\(\\Delta\\varphi\)\\,\\rho\_\{t\}\\,dx=∫∇⋅\(∇φ\)​ρt​𝑑x=−∫∇φ⋅∇ρt​d​x\\displaystyle=\\int\\nabla\\cdot\(\\nabla\\varphi\)\\,\\rho\_\{t\}\\,dx=\-\\int\\nabla\\varphi\\cdot\\nabla\\rho\_\{t\}\\,dx\(32\)The second \(applying integration by parts once more to−∫∇φ⋅∇ρt​d​x\-\\int\\nabla\\varphi\\cdot\\nabla\\rho\_\{t\}\\,dx\):

−∫∇φ⋅∇ρt​d​x\\displaystyle\-\\int\\nabla\\varphi\\cdot\\nabla\\rho\_\{t\}\\,dx=\+∫φ​∇⋅\(∇ρt\)​𝑑x=\+∫φ​Δ​ρt​𝑑x\\displaystyle=\+\\int\\varphi\\,\\nabla\\cdot\(\\nabla\\rho\_\{t\}\)\\,dx=\+\\int\\varphi\\,\\Delta\\rho\_\{t\}\\,dx\(33\)
Summary:Two integrations by parts convertΔ​φ\\Delta\\varphiintoΔ​ρt\\Delta\\rho\_\{t\}—“the second\-order derivative has been transferred fromφ\\varphitoρt\\rho\_\{t\}”\.

Step 6: Combine and conclude\.

Substituting the results of Step 5 back into Step 4:

∫φ​∂tρt​d​x=∫φ​\[∇⋅\(ρt​∇V\)\+Δ​ρt\]​𝑑x\\int\\varphi\\,\\partial\_\{t\}\\rho\_\{t\}\\,dx=\\int\\varphi\\,\\bigl\[\\nabla\\cdot\(\\rho\_\{t\}\\nabla V\)\+\\Delta\\rho\_\{t\}\\bigr\]\\,dx
Sinceφ\\varphiis anarbitrarysmooth function, if∫φ⋅\(A−B\)​𝑑x=0\\int\\varphi\\cdot\(A\-B\)\\,dx=0holds for allφ\\varphi, thenA=BA=B\(the “fundamental lemma” / du Bois\-Reymond lemma\)\. Therefore:

∂tρt=∇⋅\(ρt​∇V\)\+Δ​ρt\\boxed\{\\partial\_\{t\}\\rho\_\{t\}=\\nabla\\cdot\(\\rho\_\{t\}\\nabla V\)\+\\Delta\\rho\_\{t\}\}\(34\)
This is theFokker–Planck equation\(also known as the Kolmogorov forward equation\)\.

IntuitionWhat do the two terms in the Fokker–Planck equation mean?∂tρ=∇⋅\(ρ​∇V\)⏟\(a\) drift term\+Δ​ρ⏟\(b\) diffusion term\\partial\_\{t\}\\rho=\\underbrace\{\\nabla\\cdot\(\\rho\\nabla V\)\}\_\{\\text\{\(a\) drift term\}\}\+\\underbrace\{\\Delta\\rho\}\_\{\\text\{\(b\) diffusion term\}\}\(a\)∇⋅\(ρ​∇V\)\\nabla\\cdot\(\\rho\\nabla V\): the effect driven by the external potential\. Let us unpack this layer by layer:Layer 1: What is∇V\\nabla V?•V​\(x\)V\(x\)is the potential energy function\.∇V​\(x\)\\nabla V\(x\)is a vector pointing in the direction ofsteepest increaseofVVatxx\.•For example, ifV​\(x\)=12​\|x\|2V\(x\)=\\frac\{1\}\{2\}\|x\|^\{2\}\(bowl\-shaped\), then∇V​\(x\)=x\\nabla V\(x\)=x, pointing from the bottom toward the rim \(outward\)\.Layer 2: What isρ​∇V\\rho\\nabla V?•This is a vector field \(with direction and magnitude at each point\), representing the “mass flux”\.•Imagine each particle being pushed by∇V\\nabla V\. Thenρ​\(x\)​∇V​\(x\)\\rho\(x\)\\nabla V\(x\)= at positionxx,mass per unit volume×\\timesthe direction and speed each mass element is pushed\.•It is density times velocity, whose physical meaning is “mass passing through a unit area per unit time”—the same asρ​v\\rho vin the continuity equation\.•Note the direction:ρ​∇V\\rho\\nabla Vpoints towardincreasingpotential energy \(“uphill”\)\. So in the Fokker–Planck equation∂tρ=\+∇⋅\(ρ​∇V\)\\partial\_\{t\}\\rho=\+\\nabla\\cdot\(\\rho\\nabla V\), density is actually “flowing away” from regions of high potential—we explain why below\.Layer 3: What is∇⋅\(ρ​∇V\)\\nabla\\cdot\(\\rho\\nabla V\)?•The divergence∇⋅𝐅\\nabla\\cdot\\mathbf\{F\}measures the “degree of outward flux” of a vector field𝐅\\mathbf\{F\}at a point:–∇⋅𝐅​\(x\)\>0\\nabla\\cdot\\mathbf\{F\}\(x\)\>0:xxis a “source”—flux diverges outward fromxx, matter is leaving–∇⋅𝐅​\(x\)<0\\nabla\\cdot\\mathbf\{F\}\(x\)<0:xxis a “sink”—flux converges towardxx, matter is accumulating•So∇⋅\(ρ​∇V\)​\(x\)\>0\\nabla\\cdot\(\\rho\\nabla V\)\(x\)\>0means: atxx, the fluxρ​∇V\\rho\\nabla Visdiverging outward\.Layer 4: Why does∂tρ=\+∇⋅\(ρ​∇V\)\\partial\_\{t\}\\rho=\+\\nabla\\cdot\(\\rho\\nabla V\)imply that density accumulates at low potential?Reading the sign:At first glance, the positive sign in\+∇⋅\(ρ​∇V\)\+\\nabla\\cdot\(\\rho\\nabla V\)seems to suggest density is “moving uphill”\. But in fact this positive sign corresponds precisely to “particles moving downhill”\. Here is why:The continuity equation reads∂tρ\+∇⋅\(ρ​v\)=0\\partial\_\{t\}\\rho\+\\nabla\\cdot\(\\rho\\,v\)=0, and the particle velocity isv=−∇Vv=\-\\nabla V\(in the direction of decreasing potential\)\. Substituting:∂tρ=−∇⋅\(ρ​v\)=−∇⋅\(ρ⋅\(−∇V\)\)=\+∇⋅\(ρ​∇V\)\\partial\_\{t\}\\rho=\-\\nabla\\cdot\(\\rho\\,v\)=\-\\nabla\\cdot\\bigl\(\\rho\\cdot\(\-\\nabla V\)\\bigr\)=\+\\nabla\\cdot\(\\rho\\nabla V\)So the physical meaning of\+∇⋅\(ρ​∇V\)\+\\nabla\\cdot\(\\rho\\nabla V\)is “the density change caused by particles moving downhill with velocity−∇V\-\\nabla V”\.One\-dimensional verification:LetV​\(x\)=12​x2V\(x\)=\\frac\{1\}\{2\}x^\{2\},∇V=x\\nabla V=x,Δ​V=1\\Delta V=1\. Expanding the divergence \(product rule\):∇⋅\(ρ​∇V\)=∇ρ⋅∇V\+ρ​Δ​V=x​ρ′\+ρ\\nabla\\cdot\(\\rho\\nabla V\)=\\nabla\\rho\\cdot\\nabla V\+\\rho\\,\\Delta V=x\\rho^\{\\prime\}\+\\rhoAnalyzing the behavior at different locations:•Atx=0x=0\(bottom of bowl\):∇⋅\(ρ​∇V\)=ρ​\(0\)\>0\\nabla\\cdot\(\\rho\\nabla V\)=\\rho\(0\)\>0, so∂tρ​\(0\)\>0\\partial\_\{t\}\\rho\(0\)\>0—density at the bottom isincreasing•At largexx\(bowl rim\):ρ′<0\\rho^\{\\prime\}<0and\|x\|\|x\|is large, sox​ρ′x\\rho^\{\\prime\}is a large negative number, giving∂tρ<0\\partial\_\{t\}\\rho<0overall—density at the rim isdecreasingConclusion: density flows from the rim to the bottom—particles are indeedaccumulatingat the low\-potential region, consistent with the intuition fromv=−∇Vv=\-\\nabla V\.\(b\)Δ​ρ\\Delta\\rho: diffusion \(the “heat equation” effect\)\.•Δ​ρ=∇2ρ\\Delta\\rho=\\nabla^\{2\}\\rhois the Laplacian of the density•Sign of the Laplacian: ifρ​\(x\)\\rho\(x\)islargerthan the average of its neighbors \(a density peak\), thenΔ​ρ<0\\Delta\\rho<0and the densitydecreases•Conversely, ifρ​\(x\)\\rho\(x\)issmallerthan its neighbors \(a density trough\), thenΔ​ρ\>0\\Delta\\rho\>0and the densityincreases•The effect: smoothing out density inhomogeneities—high\-density regions spread out, low\-density regions fill in \(just like a drop of ink naturally diffusing in water\)Competition between the two terms:•The drift term tries to concentrate particles at the minimum ofVV\(density becomes sharper\)•The diffusion term tries to spread particles uniformly \(density becomes flatter\)•The final equilibrium = stationary distributionρ∞∝e−V\\rho\_\{\\infty\}\\propto e^\{\-V\}

### 4\.5Rewriting as a continuity equation

This step is the key link between Fokker–Planck and Wasserstein gradient flows\. We rewrite the Fokker–Planck equation in the form∂tρ\+∇⋅\(ρ​v\)=0\\partial\_\{t\}\\rho\+\\nabla\\cdot\(\\rho v\)=0, thereby identifying the “velocity fieldvv”\.

Step 1:Note thatΔ​ρ=∇⋅\(∇ρ\)\\Delta\\rho=\\nabla\\cdot\(\\nabla\\rho\)\(Laplacian = divergence of gradient\):

∂tρt\\displaystyle\\partial\_\{t\}\\rho\_\{t\}=∇⋅\(ρt​∇V\)\+∇⋅\(∇ρt\)\\displaystyle=\\nabla\\cdot\(\\rho\_\{t\}\\nabla V\)\+\\nabla\\cdot\(\\nabla\\rho\_\{t\}\)=∇⋅\(ρt​∇V\+∇ρt\)\\displaystyle=\\nabla\\cdot\\bigl\(\\rho\_\{t\}\\nabla V\+\\nabla\\rho\_\{t\}\\bigr\)\(35\)
\(Combining both terms into a single divergence—since “∇⋅A\+∇⋅B=∇⋅\(A\+B\)\\nabla\\cdot A\+\\nabla\\cdot B=\\nabla\\cdot\(A\+B\)”\.\)Step 2:Use the key identity∇ρ=ρ​∇log⁡ρ\\nabla\\rho=\\rho\\,\\nabla\\log\\rho:

Deriving this identity:By the chain rule,∇log⁡ρ=1ρ​∇ρ\\nabla\\log\\rho=\\frac\{1\}\{\\rho\}\\nabla\\rho, soρ​∇log⁡ρ=∇ρ\\rho\\,\\nabla\\log\\rho=\\nabla\\rho\. ✓

Substituting:

∂tρt\\displaystyle\\partial\_\{t\}\\rho\_\{t\}=∇⋅\(ρt​∇V\+ρt​∇log⁡ρt\)\\displaystyle=\\nabla\\cdot\\bigl\(\\rho\_\{t\}\\nabla V\+\\rho\_\{t\}\\nabla\\log\\rho\_\{t\}\\bigr\)=∇⋅\(ρt​\(∇V\+∇log⁡ρt\)⏟=−vt\)\\displaystyle=\\nabla\\cdot\\bigl\(\\rho\_\{t\}\\,\\underbrace\{\(\\nabla V\+\\nabla\\log\\rho\_\{t\}\)\}\_\{=\-v\_\{t\}\}\\bigr\)\(36\)
Step 3:Compare with the continuity equation∂tρ\+∇⋅\(ρ​v\)=0\\partial\_\{t\}\\rho\+\\nabla\\cdot\(\\rho\\,v\)=0, i\.e\.,∂tρ=−∇⋅\(ρ​v\)\\partial\_\{t\}\\rho=\-\\nabla\\cdot\(\\rho\\,v\):

vt=−∇V−∇log⁡ρt\\boxed\{v\_\{t\}=\-\\nabla V\-\\nabla\\log\\rho\_\{t\}\}\(37\)
IntuitionThe two terms in the velocity field—the physics of two “forces”:\(1\)−∇V\-\\nabla V: external force \(deterministic\)•This is the “push” from the potential—driving particles toward the minimum ofVV•Analogy: a ball rolling in a bowl, always pushed toward the bottom•This forcedoes not depend on the densityρ\\rho—it is the same regardless of how many other particles are nearby\(2\)−∇log⁡ρt=−∇ρtρt\-\\nabla\\log\\rho\_\{t\}=\-\\frac\{\\nabla\\rho\_\{t\}\}\{\\rho\_\{t\}\}: entropic/diffusive force \(generated by the density itself\)•This is a force “from crowded to sparse”—∇ρ\\nabla\\rhopoints toward increasing density; the negative sign reverses it to point towarddecreasingdensity•Division byρ\\rhois normalization: it is therelativedensity difference, not the absolute one, that drives diffusion•Analogy: in a subway car, people naturally “flow” from crowded cars to empty ones•This forcedepends on the density itself—the more non\-uniform the density, the stronger the force•Physically called “osmotic pressure” or “entropic force”—no one is actually pushing; it is the macroscopic manifestation of statistical effectsKey differences:

### 4\.6Stationary distribution

When the system reaches equilibrium \(∂tρ=0\\partial\_\{t\}\\rho=0\), the density no longer changes\. For our gradient flow, the velocityv=−∇\(log⁡ρ\+V\)v=\-\\nabla\(\\log\\rho\+V\)is a gradient field, and∇⋅\(ρ​v\)=0\\nabla\\cdot\(\\rho v\)=0impliesv=0v=0\(settingψ=log⁡ρ\+V\\psi=\\log\\rho\+Vand integrating by parts gives∫ρ​\|v\|2​𝑑x=∫ρ​\|∇ψ\|2​𝑑x=−∫ψ​∇⋅\(ρ​∇ψ\)​𝑑x=∫ψ​∇⋅\(ρ​v\)​𝑑x=0\\int\\rho\|v\|^\{2\}\\,dx=\\int\\rho\|\\nabla\\psi\|^\{2\}\\,dx=\-\\int\\psi\\,\\nabla\\cdot\(\\rho\\nabla\\psi\)\\,dx=\\int\\psi\\,\\nabla\\cdot\(\\rho v\)\\,dx=0\)\. So the equilibrium condition is equivalent tov=0v=0:

v=0\\displaystyle v=0⟹∇V\+∇log⁡ρ∞=0\\displaystyle\\implies\\nabla V\+\\nabla\\log\\rho\_\{\\infty\}=0⟹∇log⁡ρ∞=−∇V\\displaystyle\\implies\\nabla\\log\\rho\_\{\\infty\}=\-\\nabla V⟹log⁡ρ∞=−V\+C\(integrate both sides\)\\displaystyle\\implies\\log\\rho\_\{\\infty\}=\-V\+C\\quad\\text\{\(integrate both sides\)\}⟹ρ∞=eC⋅e−V∝e−V\\displaystyle\\implies\\rho\_\{\\infty\}=e^\{C\}\\cdot e^\{\-V\}\\propto e^\{\-V\}\(38\)
The normalization constant:eC=1/Ze^\{C\}=1/Z, whereZ=∫e−V​\(x\)​𝑑xZ=\\int e^\{\-V\(x\)\}\\,dx\(the partition function\)\. Therefore:

ρ∞​\(x\)=e−V​\(x\)∫e−V​\(y\)​𝑑y\\rho\_\{\\infty\}\(x\)=\\frac\{e^\{\-V\(x\)\}\}\{\\int e^\{\-V\(y\)\}\\,dy\}
This is theGibbs–Boltzmann distribution\.

IntuitionIntuition for the equilibrium:ρ∞∝e−V\\rho\_\{\\infty\}\\propto e^\{\-V\}means particles tend to accumulate at low\-potential regions, but do not concentrate entirely at a single point \(diffusion prevents this\)\.Verification:Atρ∞∝e−V\\rho\_\{\\infty\}\\propto e^\{\-V\}, the potential force−∇V\-\\nabla V\(toward the bottom\) and the diffusive force−∇log⁡ρ∞=\+∇V\-\\nabla\\log\\rho\_\{\\infty\}=\+\\nabla V\(away from the bottom\) exactly cancel, giving zero net velocity\.This is theGibbs–Boltzmann distribution\. For the role of temperature \(concentration at low temperature vs\. uniformity at high temperature\) and its wide applications in statistical mechanics, machine learning, and Bayesian inference, see Appendix[D](https://arxiv.org/html/2606.24157#A4)\.

### 4\.7Connection to modern generative models

The Fokker–Planck velocity fieldvt=−∇V−∇log⁡ρtv\_\{t\}=\-\\nabla V\-\\nabla\\log\\rho\_\{t\}directly reveals the relationship between score and velocity\. For the more general SDEd​Xt=f​\(Xt,t\)​d​t\+g​\(t\)​d​BtdX\_\{t\}=f\(X\_\{t\},t\)\\,dt\+g\(t\)\\,dB\_\{t\}, the Fokker–Planck equation∂tρ=−∇⋅\(ρ​f\)\+12​g2​Δ​ρ\\partial\_\{t\}\\rho=\-\\nabla\\cdot\(\\rho f\)\+\\frac\{1\}\{2\}g^\{2\}\\Delta\\rhocan be rewritten as the continuity equation∂tρ\+∇⋅\(ρ​vt\)=0\\partial\_\{t\}\\rho\+\\nabla\\cdot\(\\rho\\,v\_\{t\}\)=0, where the probability flow ODE velocity field is:

vt=f​\(x,t\)⏟drift−12​g​\(t\)2​∇log⁡ρt⏟score\\boxed\{v\_\{t\}=\\underbrace\{f\(x,t\)\}\_\{\\text\{drift\}\}\-\\frac\{1\}\{2\}g\(t\)^\{2\}\\underbrace\{\\nabla\\log\\rho\_\{t\}\}\_\{\\text\{score\}\}\}\(Verification: for our settingf=−∇Vf=\-\\nabla V,g=2g=\\sqrt\{2\}, we getvt=−∇V−∇log⁡ρtv\_\{t\}=\-\\nabla V\-\\nabla\\log\\rho\_\{t\}✓\.\) Here, the scorest:=∇log⁡ρts\_\{t\}:=\\nabla\\log\\rho\_\{t\}is what diffusion models learn, while the velocityvtv\_\{t\}is what flow matching learns\. They differ only by a known drift term—they are two equivalent descriptions of the same probability evolution\.

For the geometric meaning of the score \(==the Wasserstein gradient of negative entropy\), detailed formulas, and comparisons among major models \(DDPM, NCSN, VE\-SDE, VP\-SDE, Flow Matching\), see Appendix[F](https://arxiv.org/html/2606.24157#A6)\.

### 4\.8Conceptual clarification: what needs an energy, and what does not?

The continuity equation is more general than the Fokker–Planck equation\. It is akinematicstatement:

∂tρt\+∇⋅\(ρt​vt\)=0\.\\partial\_\{t\}\\rho\_\{t\}\+\\nabla\\cdot\(\\rho\_\{t\}v\_\{t\}\)=0\.It only says: if probability mass moves with velocity fieldvtv\_\{t\}, then the density changes by mass conservation\. No potentialVV, no free energyℱ\\mathcal\{F\}, and no entropy term is needed to write this equation\.

However, the continuity equation does not decide the motion by itself\. One must still provide the velocity fieldvtv\_\{t\}\. Different choices ofvtv\_\{t\}produce different dynamics:

continuity equation\+chosen velocity field⟹a probability path\.\\text\{continuity equation\}\+\\text\{chosen velocity field\}\\quad\\Longrightarrow\\quad\\text\{a probability path\}\.
The Fokker–Planck equation is a special case where the velocity field is not arbitrary\. It is generated by a free energy:

vt=−∇δ​ℱδ​ρ​\(ρt\)\.v\_\{t\}=\-\\nabla\\frac\{\\delta\\mathcal\{F\}\}\{\\delta\\rho\}\(\\rho\_\{t\}\)\.For

ℱ​\(ρ\)=∫V​ρ​𝑑x\+∫ρ​log⁡ρ​d​x,\\mathcal\{F\}\(\\rho\)=\\int V\\rho\\,dx\+\\int\\rho\\log\\rho\\,dx,we get

vt=−∇V−∇log⁡ρt,v\_\{t\}=\-\\nabla V\-\\nabla\\log\\rho\_\{t\},and substituting this into the continuity equation gives exactly

∂tρt=∇⋅\(ρt​∇V\)\+Δ​ρt\.\\partial\_\{t\}\\rho\_\{t\}=\\nabla\\cdot\(\\rho\_\{t\}\\nabla V\)\+\\Delta\\rho\_\{t\}\.
Key InsightThe continuity equation defines “how mass can flow\.” The Fokker–Planck equation additionally specifies “why it flows that way”: it follows the negative Wasserstein gradient of a free energy\.

### 4\.9Where is the modeling choice in diffusion?

The hand\-crafted part of diffusion models is not the Itô rule itself\. Once Brownian motion is chosen, the identity\(d​Bt\)2=d​t\(dB\_\{t\}\)^\{2\}=dtis a mathematical consequence of its quadratic variation\. The modeling choice is the SDE:

d​Xt=f​\(Xt,t\)​d​t\+g​\(t\)​d​Bt\.dX\_\{t\}=f\(X\_\{t\},t\)\\,dt\+g\(t\)\\,dB\_\{t\}\.Choosing the driftff, the diffusion strengthgg, and the noise schedule \(for example theβ​\(t\)\\beta\(t\)schedule in VP\-SDE\) specifies the forward noising process\. The Fokker–Planck equation is then the deterministic density evolution implied by that SDE\.

Flow Matching removes this SDE modeling layer\. It does not first choose Brownian noise and then derive a Fokker–Planck equation\. Instead, it directly learns a velocity fieldvtv\_\{t\}for a chosen probability path, often an OT\-like path from noise to data\.

What we now know:The Fokker\-Planck equation describes the evolution of particle density and has the unique stationary distributionπ∝e−V\\pi\\propto e^\{\-V\}\.Central question:Why mustρt\\rho\_\{t\}converge toπ\\pi? What objective is it “optimizing”? The answer is in the next section\.

## 5Free Energy and the Wasserstein Gradient

Where we are:Section 4 derived the Fokker\-Planck equation and wrote it as∂tρ=−∇⋅\(ρ​v\)\\partial\_\{t\}\\rho=\-\\nabla\\cdot\(\\rho\\,v\)withv=−∇log⁡ρ−∇Vv=\-\\nabla\\log\\rho\-\\nabla V\. Section 3 showed that the space of probability distributions carries a Riemannian structure\. This section brings the two together: we define the free energyℱ​\(ρ\)=KL​\(ρ∥π\)\\mathcal\{F\}\(\\rho\)=\\mathrm\{KL\}\(\\rho\\\|\\pi\), compute its Wasserstein gradientgradW​ℱ=∇log⁡ρ\+∇V\\mathrm\{grad\}\_\{W\}\\mathcal\{F\}=\\nabla\\log\\rho\+\\nabla V, and thereby provev=−gradW​ℱv=\-\\mathrm\{grad\}\_\{W\}\\mathcal\{F\}—the Fokker\-Planck equation is the gradient flow of the free energy in probability space\. This is the central theorem of the entire article\.

### 5\.1The free energy functional

We now address a central question: What is the Fokker\-Planck equation “optimizing”?

The answer: it is minimizing thefree energy—the KL divergence between the distributionρ\\rhoand the equilibriumπ∝e−V\\pi\\propto e^\{\-V\}\. But before writing down the formula, let us first build physical intuition for what free energy is\.

IntuitionFree energy as the balance of two competing tendenciesA system in statistical mechanics \(say, a confined gas\) is governed by two opposing tendencies\.Tendency one: lower the energy\.•Each particle is driven toward the location of lowest potentialV​\(x\)V\(x\)\(the bottom of the well\)•Acting alone, this collapses all mass onto the global minimum ofVV, formingδxmin\\delta\_\{x\_\{\\min\}\}•It is the dominant tendency at low temperature, where the system settles toward its ground stateTendency two: raise the entropy\.•Thermal agitation drives the particles to spread out•Acting alone, this produces a uniform distribution, independent of the potential•It is the dominant tendency at high temperatureFree energy is the quantitative trade\-off between the two:Free energy=Potential energy \(favors concentration\)−Temperature×Entropy \(favors dispersion\)\\text\{Free energy\}=\\text\{Potential energy \(favors concentration\)\}\-\\text\{Temperature\}\\times\\text\{Entropy \(favors dispersion\)\}Minimizing free energy selects the optimal balance between concentration and dispersion\.The role of temperature:•Low temperature: the energy term dominates; mass concentrates near the minimum ofVV•High temperature: the entropy term dominates; the distribution approaches uniform•Temperature=1=1\(our setting\): the two terms carry equal weight

###### Definition 5\.1\(Free energy\)\.

Defineπ​\(x\)=1Z​e−V​\(x\)\\pi\(x\)=\\frac\{1\}\{Z\}e^\{\-V\(x\)\}whereZ=∫e−V​𝑑xZ=\\int e^\{\-V\}dx\. The free energy is:

ℱ​\(ρ\)\\displaystyle\\mathcal\{F\}\(\\rho\):=KL​\(ρ∥π\)=∫ℝdρ​log⁡ρπ​d​x\\displaystyle:=\\mathrm\{KL\}\(\\rho\\\|\\pi\)=\\int\_\{\\mathbb\{R\}^\{d\}\}\\rho\\log\\frac\{\\rho\}\{\\pi\}\\,dx=∫ρ​log⁡ρ​d​x\+∫V​ρ​𝑑x\+log⁡Z\\displaystyle=\\int\\rho\\log\\rho\\,dx\+\\int V\\rho\\,dx\+\\log Z\(39\)

Term\-by\-term breakdown of the formula:

Step one: Why doesKL​\(ρ∥π\)\\mathrm\{KL\}\(\\rho\\\|\\pi\)expand as above?

KL​\(ρ∥π\)\\displaystyle\\mathrm\{KL\}\(\\rho\\\|\\pi\)=∫ρ​log⁡ρπ​d​x=∫ρ​\(log⁡ρ−log⁡π\)​𝑑x\\displaystyle=\\int\\rho\\log\\frac\{\\rho\}\{\\pi\}\\,dx=\\int\\rho\\,\(\\log\\rho\-\\log\\pi\)\\,dx=∫ρ​log⁡ρ​d​x−∫ρ​log⁡π​d​x\\displaystyle=\\int\\rho\\log\\rho\\,dx\-\\int\\rho\\log\\pi\\,dxSinceπ=1Z​e−V\\pi=\\frac\{1\}\{Z\}e^\{\-V\}, we havelog⁡π=−V−log⁡Z\\log\\pi=\-V\-\\log Z\. Substituting:

=∫ρ​log⁡ρ​d​x−∫ρ​\(−V−log⁡Z\)​𝑑x\\displaystyle=\\int\\rho\\log\\rho\\,dx\-\\int\\rho\(\-V\-\\log Z\)\\,dx=∫ρ​log⁡ρ​d​x\+∫V​ρ​𝑑x\+log⁡Z​∫ρ​𝑑x⏟=1\\displaystyle=\\int\\rho\\log\\rho\\,dx\+\\int V\\rho\\,dx\+\\log Z\\underbrace\{\\int\\rho\\,dx\}\_\{=1\}
Ignoring the constantlog⁡Z\\log Z\(which doesn’t depend onρ\\rho\), the free energy decomposes as:

ℱ​\(ρ\)=∫ρ​log⁡ρ​d​x⏟negative entropy:−S​\(ρ\)\+∫V​ρ​𝑑x⏟potential energy:​𝔼ρ​\[V\]\\mathcal\{F\}\(\\rho\)=\\underbrace\{\\int\\rho\\log\\rho\\,dx\}\_\{\\text\{negative entropy: \}\-S\(\\rho\)\}\+\\underbrace\{\\int V\\,\\rho\\,dx\}\_\{\\text\{potential energy: \}\\mathbb\{E\}\_\{\\rho\}\[V\]\}\(40\)
More generally, if the diffusion strength isε\>0\\varepsilon\>0, the free energy is

ℱε​\(ρ\)=∫V​ρ​𝑑x\+ε​∫ρ​log⁡ρ​d​x\.\\mathcal\{F\}\_\{\\varepsilon\}\(\\rho\)=\\int V\\rho\\,dx\+\\varepsilon\\int\\rho\\log\\rho\\,dx\.\(41\)Then

δ​ℱεδ​ρ=V\+ε​\(log⁡ρ\+1\),v=−∇δ​ℱεδ​ρ=−∇V−ε​∇log⁡ρ\.\\frac\{\\delta\\mathcal\{F\}\_\{\\varepsilon\}\}\{\\delta\\rho\}=V\+\\varepsilon\(\\log\\rho\+1\),\\qquad v=\-\\nabla\\frac\{\\delta\\mathcal\{F\}\_\{\\varepsilon\}\}\{\\delta\\rho\}=\-\\nabla V\-\\varepsilon\\nabla\\log\\rho\.This is the clean way to “add the diffusion force”: we do*not*add it into the external potentialV​\(x\)V\(x\)\. Instead, we add the negative entropy termε​∫ρ​log⁡ρ\\varepsilon\\int\\rho\\log\\rhoto the free energy\. The resulting force−ε​∇log⁡ρ\-\\varepsilon\\nabla\\log\\rhodepends on the current density itself, so it is a collective entropic force rather than a fixed external potential\.

IntuitionWhat does each term mean?\(1\)∫ρ​log⁡ρ​d​x\\int\\rho\\log\\rho\\,dx: negative entropy•The entropyS​\(ρ\)=−∫ρ​log⁡ρ​d​xS\(\\rho\)=\-\\int\\rho\\log\\rho\\,dxmeasures the “degree of disorder / spread” of the distribution•The more uniform the distribution, the higher the entropy \(maximum entropy = uniform distribution\)•The more concentrated the distribution \(e\.g\.,ρ≈δx\\rho\\approx\\delta\_\{x\}\), the lower the entropy \(extreme case: Dirac entropy=−∞=\-\\infty\)•The free energy contains\+∫ρ​log⁡ρ=−S\+\\int\\rho\\log\\rho=\-S\(negative entropy\), so minimizing free energy is equivalent tomaximizing entropy\(subject to the energy constraint\)•Intuition: this term “penalizes concentration” and encourages the distribution to spread out\(2\)∫V​ρ​𝑑x=𝔼ρ​\[V\]\\int V\\rho\\,dx=\\mathbb\{E\}\_\{\\rho\}\[V\]: mean potential energy•This is the expected value of the potentialV​\(x\)V\(x\)under the distributionρ\\rho•Ifρ\\rhoconcentrates whereVVis low,𝔼ρ​\[V\]\\mathbb\{E\}\_\{\\rho\}\[V\]is small•Ifρ\\rhohas mass whereVVis high,𝔼ρ​\[V\]\\mathbb\{E\}\_\{\\rho\}\[V\]is large•Minimizing this term⇒\\Rightarrowmakingρ\\rhoconcentrate near the minimum ofVV•Intuition: this term “rewards concentration” in low\-energy regionsThe tension between the two terms:The minimum of the free energyis the optimal compromise between the two—which turns out to be the Gibbs distributionρ∗=1Z​e−V\\rho^\{\*\}=\\frac\{1\}\{Z\}e^\{\-V\}\.

Verification: Why is the minimum ofℱ\\mathcal\{F\}attained atρ∗=π\\rho^\{\*\}=\\pi?

Becauseℱ​\(ρ\)=KL​\(ρ∥π\)\\mathcal\{F\}\(\\rho\)=\\mathrm\{KL\}\(\\rho\\\|\\pi\), and the KL divergence satisfies:

- •KL​\(ρ∥π\)≥0\\mathrm\{KL\}\(\\rho\\\|\\pi\)\\geq 0\(always non\-negative—Gibbs’ inequality\)
- •KL​\(ρ∥π\)=0⇔ρ=π\\mathrm\{KL\}\(\\rho\\\|\\pi\)=0\\iff\\rho=\\pi\(equals zero if and only if the two distributions are identical\)

So the global minimum ofℱ\\mathcal\{F\}is0, attained atρ=π=1Z​e−V\\rho=\\pi=\\frac\{1\}\{Z\}e^\{\-V\}\.

In one sentence:The Fokker\-Planck equation describes the process of the system continuously decreasingKL​\(ρt∥π\)\\mathrm\{KL\}\(\\rho\_\{t\}\\\|\\pi\)—starting from any initial distributionρ0\\rho\_\{0\},ρt\\rho\_\{t\}gradually approachesπ\\pi, with the free energy monotonically decreasing to zero\.

What is the justification for this claim? We give a rigorous proof below\.

Theorem:Letρt\\rho\_\{t\}be a solution of the Fokker\-Planck equation andπ∝e−V\\pi\\propto e^\{\-V\}\. Then:

dd​t​KL​\(ρt∥π\)=−∫ρt​\|∇log⁡ρtπ\|2​𝑑x≤0\\frac\{d\}\{dt\}\\mathrm\{KL\}\(\\rho\_\{t\}\\\|\\pi\)=\-\\int\\rho\_\{t\}\\left\|\\nabla\\log\\frac\{\\rho\_\{t\}\}\{\\pi\}\\right\|^\{2\}\\,dx\\leq 0That is, the free energy isstrictly monotonically decreasing\(unlessρt=π\\rho\_\{t\}=\\pi\)\.

Complete derivation:

Step 1: Write outdd​t​KL​\(ρt∥π\)\\frac\{d\}\{dt\}\\mathrm\{KL\}\(\\rho\_\{t\}\\\|\\pi\)\.

π\\pidoes not depend on time, so:

dd​t​KL​\(ρt∥π\)\\displaystyle\\frac\{d\}\{dt\}\\mathrm\{KL\}\(\\rho\_\{t\}\\\|\\pi\)=dd​t​∫ρt​log⁡ρtπ​d​x\\displaystyle=\\frac\{d\}\{dt\}\\int\\rho\_\{t\}\\log\\frac\{\\rho\_\{t\}\}\{\\pi\}\\,dx=∫∂tρt⋅log⁡ρtπ​d​x\+∫ρt⋅∂tρtρt​𝑑x\\displaystyle=\\int\\partial\_\{t\}\\rho\_\{t\}\\cdot\\log\\frac\{\\rho\_\{t\}\}\{\\pi\}\\,dx\+\\int\\rho\_\{t\}\\cdot\\frac\{\\partial\_\{t\}\\rho\_\{t\}\}\{\\rho\_\{t\}\}\\,dx=∫∂tρt⋅log⁡ρtπ​d​x\+∫∂tρt​d​x\\displaystyle=\\int\\partial\_\{t\}\\rho\_\{t\}\\cdot\\log\\frac\{\\rho\_\{t\}\}\{\\pi\}\\,dx\+\\int\\partial\_\{t\}\\rho\_\{t\}\\,dx=∫∂tρt⋅log⁡ρtπ​d​x\+0\\displaystyle=\\int\\partial\_\{t\}\\rho\_\{t\}\\cdot\\log\\frac\{\\rho\_\{t\}\}\{\\pi\}\\,dx\+0\(42\)
\(The last step uses∫∂tρt​d​x=dd​t​∫ρt​𝑑x=dd​t​1=0\\int\\partial\_\{t\}\\rho\_\{t\}\\,dx=\\frac\{d\}\{dt\}\\int\\rho\_\{t\}\\,dx=\\frac\{d\}\{dt\}1=0, since total mass is conserved\.\)

Step 2: Substitute the Fokker–Planck equation\.

The Fokker\-Planck equation can be written as∂tρt=∇⋅\(ρt​∇log⁡ρtπ\)\\partial\_\{t\}\\rho\_\{t\}=\\nabla\\cdot\\bigl\(\\rho\_\{t\}\\nabla\\log\\frac\{\\rho\_\{t\}\}\{\\pi\}\\bigr\)\.

Why? Becauselog⁡ρπ=log⁡ρ−log⁡π=log⁡ρ\+V\+log⁡Z\\log\\frac\{\\rho\}\{\\pi\}=\\log\\rho\-\\log\\pi=\\log\\rho\+V\+\\log Z, so∇log⁡ρπ=∇log⁡ρ\+∇V\\nabla\\log\\frac\{\\rho\}\{\\pi\}=\\nabla\\log\\rho\+\\nabla V\. And we previously derived∂tρ=∇⋅\(ρ​\(∇log⁡ρ\+∇V\)\)\\partial\_\{t\}\\rho=\\nabla\\cdot\(\\rho\(\\nabla\\log\\rho\+\\nabla V\)\)—which is precisely∇⋅\(ρ​∇log⁡ρπ\)\\nabla\\cdot\(\\rho\\nabla\\log\\frac\{\\rho\}\{\\pi\}\)\.

Substituting into Step 1:

dd​t​KL​\(ρt∥π\)=∫∇⋅\(ρt​∇log⁡ρtπ\)⋅log⁡ρtπ​d​x\\frac\{d\}\{dt\}\\mathrm\{KL\}\(\\rho\_\{t\}\\\|\\pi\)=\\int\\nabla\\cdot\\\!\\left\(\\rho\_\{t\}\\nabla\\log\\frac\{\\rho\_\{t\}\}\{\\pi\}\\right\)\\cdot\\log\\frac\{\\rho\_\{t\}\}\{\\pi\}\\,dx\(43\)
Step 3: Integration by parts\.

Apply integration by parts to \([43](https://arxiv.org/html/2606.24157#S5.E43)\)\. Let𝐅=ρt​∇log⁡ρtπ\\mathbf\{F\}=\\rho\_\{t\}\\nabla\\log\\frac\{\\rho\_\{t\}\}\{\\pi\}andg=log⁡ρtπg=\\log\\frac\{\\rho\_\{t\}\}\{\\pi\}:

∫\(∇⋅𝐅\)​g​𝑑x=−∫𝐅⋅∇g​d​x\\int\(\\nabla\\cdot\\mathbf\{F\}\)\\,g\\,dx=\-\\int\\mathbf\{F\}\\cdot\\nabla g\\,dx\(The boundary terms vanish\.\) Substituting𝐅\\mathbf\{F\}and∇g\\nabla g:

dd​t​KL​\(ρt∥π\)\\displaystyle\\frac\{d\}\{dt\}\\mathrm\{KL\}\(\\rho\_\{t\}\\\|\\pi\)=−∫ρt​∇log⁡ρtπ⋅∇log⁡ρtπ​d​x\\displaystyle=\-\\int\\rho\_\{t\}\\nabla\\log\\frac\{\\rho\_\{t\}\}\{\\pi\}\\cdot\\nabla\\log\\frac\{\\rho\_\{t\}\}\{\\pi\}\\,dx=−∫ρt​\|∇log⁡ρtπ\|2​𝑑x\\displaystyle=\-\\int\\rho\_\{t\}\\left\|\\nabla\\log\\frac\{\\rho\_\{t\}\}\{\\pi\}\\right\|^\{2\}\\,dx\(44\)
Step 4: Conclude\.

The right side=−∫ρt​\|∇log⁡ρtπ\|2​𝑑x=\-\\int\\rho\_\{t\}\|\\nabla\\log\\frac\{\\rho\_\{t\}\}\{\\pi\}\|^\{2\}\\,dx\. Sinceρt≥0\\rho\_\{t\}\\geq 0and\|⋅\|2≥0\|\\cdot\|^\{2\}\\geq 0, the entire integral is≥0\\geq 0, so with the negative sign it is≤0\\leq 0\.

dd​t​KL​\(ρt∥π\)=−∫ρt​\|∇log⁡ρtπ\|2​𝑑x≤0\\boxed\{\\frac\{d\}\{dt\}\\mathrm\{KL\}\(\\rho\_\{t\}\\\|\\pi\)=\-\\int\\rho\_\{t\}\\left\|\\nabla\\log\\frac\{\\rho\_\{t\}\}\{\\pi\}\\right\|^\{2\}\\,dx\\leq 0\}\(45\)
Equality holds if and only if∇log⁡ρtπ=0\\nabla\\log\\frac\{\\rho\_\{t\}\}\{\\pi\}=0everywhere, i\.e\.,ρt=c⋅π\\rho\_\{t\}=c\\cdot\\pi\. By the normalization condition∫ρt=1=∫π\\int\\rho\_\{t\}=1=\\int\\pi, we getc=1c=1, i\.e\.,ρt=π\\rho\_\{t\}=\\pi\.

Summary:

- •KL​\(ρt∥π\)\\mathrm\{KL\}\(\\rho\_\{t\}\\\|\\pi\)ismonotonically decreasingalong the Fokker\-Planck equation
- •The rate of decrease=∫ρt​\|∇log⁡ρtπ\|2​𝑑x=\\int\\rho\_\{t\}\|\\nabla\\log\\frac\{\\rho\_\{t\}\}\{\\pi\}\|^\{2\}\\,dx, a quantity called therelative Fisher information
- •The rate of decrease is zero only whenρt=π\\rho\_\{t\}=\\pi—once equilibrium is reached, there is no further change
- •Physical intuition: the system cannot spontaneously move away from equilibrium \(a mathematical version of the second law of thermodynamics\)

For a complete introduction to Fisher information \(the parametric version in statistics, the distributional version in information theory, how they are unified, and the definition of relative Fisher information and information inequalities\), see Appendix[E](https://arxiv.org/html/2606.24157#A5)\.

IntuitionWhy is the result a negative “squared term”? This is no coincidence\.This structure is ubiquitous in mathematics\. Consider a gradient flow in Euclidean space,x˙=−∇f​\(x\)\\dot\{x\}=\-\\nabla f\(x\):dd​t​f​\(xt\)=∇f​\(xt\)⋅x˙t=∇f⋅\(−∇f\)=−\|∇f\|2≤0\\frac\{d\}\{dt\}f\(x\_\{t\}\)=\\nabla f\(x\_\{t\}\)\\cdot\\dot\{x\}\_\{t\}=\\nabla f\\cdot\(\-\\nabla f\)=\-\|\\nabla f\|^\{2\}\\leq 0The rate of decrease offf= the squared norm of the gradient \(always non\-positive\)\.Fokker\-Planck is a perfect analogue: it is the gradient flow ofℱ=KL\(⋅∥π\)\\mathcal\{F\}=\\mathrm\{KL\}\(\\cdot\\\|\\pi\)in Wasserstein space\. The “Wasserstein gradient”=∇log⁡ρπ=\\nabla\\log\\frac\{\\rho\}\{\\pi\}, so the rate of decrease=−‖gradient‖ρ2=−∫ρ​\|∇log⁡ρπ\|2​𝑑x=\-\\\|\\text\{gradient\}\\\|\_\{\\rho\}^\{2\}=\-\\int\\rho\|\\nabla\\log\\frac\{\\rho\}\{\\pi\}\|^\{2\}\\,dx\.The structure is identical, lifted from finite to infinite dimensions\.

A concrete example:

LetV​\(x\)=12​x2V\(x\)=\\frac\{1\}\{2\}x^\{2\}\(one\-dimensional quadratic potential\), soπ=𝒩​\(0,1\)\\pi=\\mathcal\{N\}\(0,1\)\.

- •Initialρ0=δ5\\rho\_\{0\}=\\delta\_\{5\}\(all particles atx=5x=5\): - –Potential energy∫V​ρ0=V​\(5\)=12\.5\\int V\\rho\_\{0\}=V\(5\)=12\.5\(very high—the particles are on the bowl rim\) - –Negative entropy=−∞=\-\\infty\(completely concentrated at a single point\) - –Fokker\-Planck will causeρt\\rho\_\{t\}to slide toward the bottom while simultaneously spreading out, eventually becoming𝒩​\(0,1\)\\mathcal\{N\}\(0,1\)
- •Initialρ0=Uniform​\[−10,10\]\\rho\_\{0\}=\\text\{Uniform\}\[\-10,10\]\(uniformly distributed over a large range\): - –Potential energy∫V​ρ0=120​∫−1010x22​𝑑x=1006≈16\.7\\int V\\rho\_\{0\}=\\frac\{1\}\{20\}\\int\_\{\-10\}^\{10\}\\frac\{x^\{2\}\}\{2\}dx=\\frac\{100\}\{6\}\\approx 16\.7\(very high—too much mass on the bowl rim\) - –Negative entropy=log⁡\(20\)≈3=\\log\(20\)\\approx 3\(very spread out, high entropy, low negative entropy\) - –Fokker\-Planck will cause the mass at the edges to contract toward the center, eventually also becoming𝒩​\(0,1\)\\mathcal\{N\}\(0,1\)
- •Regardless of the startingρ0\\rho\_\{0\}, the endpoint is alwaysπ=𝒩​\(0,1\)\\pi=\\mathcal\{N\}\(0,1\)—the unique minimizer of the free energy\.

Free energy in the real world—milk diffusing in coffee \(V≈0V\\approx 0, purely entropy\-driven\); atmospheric pressureρ​\(h\)∝e−m​g​h/k​T\\rho\(h\)\\propto e^\{\-mgh/kT\}\(gravity vs\. thermal motion\); protein folding \(conformational energy vs\. conformational entropy\); in Bayesian inference, SGLD \(Welling & Teh, 2011\) is a discrete version of Fokker\-Planck\.

Key InsightFree energy vs\. ordinary potential: the essential distinctionFree energy==mean potential energy\+\+negative entropy\. The extra “entropy” term is a purely collective effect—a single particle has no entropy to speak of\.

### 5\.2First variation \(functional derivative\)

Goal:Computeδ​ℱδ​ρ\\frac\{\\delta\\mathcal\{F\}\}\{\\delta\\rho\}—that is, howℱ\\mathcal\{F\}responds when the density changes fromρ\\rhotoρ\+ϵ​δ​ρ\\rho\+\\epsilon\\,\\delta\\rho\.

ℱ​\(ρ\+ϵ​δ​ρ\)\\displaystyle\\mathcal\{F\}\(\\rho\+\\epsilon\\,\\delta\\rho\)=∫\(ρ\+ϵ​δ​ρ\)​log⁡\(ρ\+ϵ​δ​ρ\)​𝑑x\+∫V​\(ρ\+ϵ​δ​ρ\)​𝑑x\\displaystyle=\\int\(\\rho\+\\epsilon\\,\\delta\\rho\)\\log\(\\rho\+\\epsilon\\,\\delta\\rho\)\\,dx\+\\int V\(\\rho\+\\epsilon\\,\\delta\\rho\)\\,dx≈∫ρ​log⁡ρ​d​x\+ϵ​∫δ​ρ​\(log⁡ρ\+1\)​𝑑x\+∫V​ρ​𝑑x\+ϵ​∫V​δ​ρ​𝑑x\\displaystyle\\approx\\int\\rho\\log\\rho\\,dx\+\\epsilon\\int\\delta\\rho\\,\(\\log\\rho\+1\)\\,dx\+\\int V\\rho\\,dx\+\\epsilon\\int V\\,\\delta\\rho\\,dx
\(Expanding\(ρ\+ϵ​δ​ρ\)​log⁡\(ρ\+ϵ​δ​ρ\)≈ρ​log⁡ρ\+ϵ​δ​ρ​\(log⁡ρ\+1\)\(\\rho\+\\epsilon\\delta\\rho\)\\log\(\\rho\+\\epsilon\\delta\\rho\)\\approx\\rho\\log\\rho\+\\epsilon\\delta\\rho\(\\log\\rho\+1\), which is a first\-order Taylor expansion\.\)

Therefore:

δ​ℱδ​ρ=log⁡ρ\+1\+V\\frac\{\\delta\\mathcal\{F\}\}\{\\delta\\rho\}=\\log\\rho\+1\+V\(46\)

### 5\.3Computing the Wasserstein gradient

The core computation:We now compute the gradient ofℱ\\mathcal\{F\}with respect to the Wasserstein metric\.

Recall first: Why does the Euclidean gradient∇f\\nabla fsatisfy⟨∇f,h⟩=D​f​\[h\]\\langle\\nabla f,h\\rangle=Df\[h\]?

This is in fact thedefinitionof the gradient\. Specifically:

Forf:ℝn→ℝf:\\mathbb\{R\}^\{n\}\\to\\mathbb\{R\}, the directional derivative \(the Fréchet derivative acting on the directionhh\) is:

D​f​\(x\)​\[h\]=limϵ→0f​\(x\+ϵ​h\)−f​\(x\)ϵ=∑i∂f∂xi​hiDf\(x\)\[h\]=\\lim\_\{\\epsilon\\to 0\}\\frac\{f\(x\+\\epsilon h\)\-f\(x\)\}\{\\epsilon\}=\\sum\_\{i\}\\frac\{\\partial f\}\{\\partial x\_\{i\}\}h\_\{i\}The right side can be written as an inner product:D​f​\(x\)​\[h\]=\(∂f∂x1,…,∂f∂xn\)⋅\(h1,…,hn\)=⟨∇f,h⟩Df\(x\)\[h\]=\\left\(\\frac\{\\partial f\}\{\\partial x\_\{1\}\},\\ldots,\\frac\{\\partial f\}\{\\partial x\_\{n\}\}\\right\)\\cdot\(h\_\{1\},\\ldots,h\_\{n\}\)=\\langle\\nabla f,h\\rangle\.

In other words:D​f​\(x\)Df\(x\)is alinear functional\(takes a directionhhas input and outputs a real number\), and∇f\\nabla fis thevectorthat “represents” this linear functional via the inner product—the Riesz representative\.

Why is an inner product needed to define the gradient?BecauseD​fDfby itself is only a “cotangent vector” \(linear functional\); to “raise” it to a “tangent vector” \(vector\), one needs a metric\. Under the standard Euclidean metric⟨⋅,⋅⟩\\langle\\cdot,\\cdot\\rangle,∇f\\nabla fis simply the vector of partial derivatives\. But under a different metricgg, the gradient changes—gradg​f=g−1​∇f\\mathrm\{grad\}\_\{g\}f=g^\{\-1\}\\nabla f\(see the differential geometry appendix\)\.

Core idea:Gradient==using the metric to convert a “differential” \(linear functional\) into a “direction” \(vector\)\. Different metric→\\todifferent gradient→\\todifferent gradient flow\.

In Wasserstein space, by complete analogy,gradW​ℱ​\(ρ\)\\mathrm\{grad\}\_\{W\}\\mathcal\{F\}\(\\rho\)is a vector fieldξ\\xisatisfying:

⟨ξ,η⟩ρ=D​ℱ​\(ρ\)​\[ρ˙\]for all tangent vectors​η\\langle\\xi,\\eta\\rangle\_\{\\rho\}=D\\mathcal\{F\}\(\\rho\)\[\\dot\{\\rho\}\]\\quad\\text\{for all tangent vectors\}\\;\\etawhereρ˙=−∇⋅\(ρ​η\)\\dot\{\\rho\}=\-\\nabla\\cdot\(\\rho\\eta\)is the density change induced byη\\eta\.

Detailed explanation of each concept:

\(i\) What is a “tangent vector” in Wasserstein space?

Recall: in finite dimensions, a tangent vectorv∈Tp​Mv\\in T\_\{p\}Mat a pointppon a manifoldMMdescribes the “instantaneous velocity starting frompp\.”

In Wasserstein space, a “point” is a probability densityρ\\rho\. A “motion” starting fromρ\\rhois a time\-varying densityρt\\rho\_\{t\}\. The continuity equation∂tρt\+∇⋅\(ρt​vt\)=0\\partial\_\{t\}\\rho\_\{t\}\+\\nabla\\cdot\(\\rho\_\{t\}\\,v\_\{t\}\)=0tells us: any change in density is driven by avelocity fieldvtv\_\{t\}\.

Therefore, a tangent vector atρ\\rhoin Wasserstein space is avelocity fieldη​\(x\)\\eta\(x\)\. But not an arbitrary vector field—Otto showed that the tangent space is spanned bygradient fieldsη=∇ψ\\eta=\\nabla\\psi\(i\.e\., the velocity field must be the gradient of some scalar potential—“irrotational flow”\)\.

\(ii\) Why isρ˙=−∇⋅\(ρ​η\)\\dot\{\\rho\}=\-\\nabla\\cdot\(\\rho\\eta\)? \(What does “density perturbation induced by a tangent vector” mean?\)

In Wasserstein space: “point”==densityρ​\(x\)\\rho\(x\), “tangent vector”==velocity fieldη​\(x\)\\eta\(x\)\.

Question: If all mass flows according to the velocity fieldη\\etafor an instant, how does the density change?

The answer is a rewriting of the continuity equation∂tρ\+∇⋅\(ρ​v\)=0\\partial\_\{t\}\\rho\+\\nabla\\cdot\(\\rho\\,v\)=0:

ρ˙:=∂tρ\|t=0=−∇⋅\(ρ​η\)\\dot\{\\rho\}:=\\partial\_\{t\}\\rho\\big\|\_\{t=0\}=\-\\nabla\\cdot\(\\rho\\,\\eta\)Intuition:ρ​η\\rho\\etais the mass flux \(density×\\timesvelocity\),∇⋅\(ρ​η\)\\nabla\\cdot\(\\rho\\eta\)is the net outflow, so the density change=−=\-net outflow \(however much flows out is subtracted\)\.

This is the “density perturbation induced byη\\eta”—translating from “how mass flows” to “how density changes\.”

One\-dimensional example:

Letρ​\(x\)=1\\rho\(x\)=1\(uniform distribution on\[0,1\]\[0,1\]\)\.

- •Takeη​\(x\)=x\\eta\(x\)=x\(the further right, the faster the flow\): ρ˙=−∂∂x​\(ρ⋅x\)=−∂∂x​\(x\)=−1\\dot\{\\rho\}=\-\\frac\{\\partial\}\{\\partial x\}\(\\rho\\cdot x\)=\-\\frac\{\\partial\}\{\\partial x\}\(x\)=\-1Density decreases everywhere—mass accelerates to the right, and the mass that flows away from the left is not replenished\.
- •Takeη​\(x\)=1\\eta\(x\)=1\(all mass translates to the right at the same speed\): ρ˙=−∂∂x​\(ρ⋅1\)=−ρ′​\(x\)=0\\dot\{\\rho\}=\-\\frac\{\\partial\}\{\\partial x\}\(\\rho\\cdot 1\)=\-\\rho^\{\\prime\}\(x\)=0The uniform distribution shifts as a whole; its shape is unchanged,ρ˙=0\\dot\{\\rho\}=0\.
- •Takeη​\(x\)=−x\\eta\(x\)=\-x\(contraction toward the origin\): ρ˙=−∂∂x​\(ρ⋅\(−x\)\)=−∂∂x​\(−x\)=1\\dot\{\\rho\}=\-\\frac\{\\partial\}\{\\partial x\}\(\\rho\\cdot\(\-x\)\)=\-\\frac\{\\partial\}\{\\partial x\}\(\-x\)=1Density increases everywhere—mass gathers from both sides toward the center\.

Finite\-dimensional analogy:Given a velocityvv\(tangent vector\), the rate of change of position isx˙=v\\dot\{x\}=v\. Here, “position” becomes “density,” “velocity” becomes “velocity field,” and “x˙=v\\dot\{x\}=v” becomes “ρ˙=−∇⋅\(ρ​η\)\\dot\{\\rho\}=\-\\nabla\\cdot\(\\rho\\eta\)” \(the continuity equation\)\.

\(iii\) Why isD​ℱ​\(ρ\)​\[ρ˙\]=∫δ​ℱδ​ρ⋅ρ˙​𝑑xD\\mathcal\{F\}\(\\rho\)\[\\dot\{\\rho\}\]=\\int\\frac\{\\delta\\mathcal\{F\}\}\{\\delta\\rho\}\\cdot\\dot\{\\rho\}\\,dx?

This is thedefinitionof the functional derivative\. Recall:

dd​ϵ\|ϵ=0​ℱ​\(ρ\+ϵ​h\)=∫δ​ℱδ​ρ​\(x\)​h​\(x\)​𝑑x\\frac\{d\}\{d\\epsilon\}\\bigg\|\_\{\\epsilon=0\}\\mathcal\{F\}\(\\rho\+\\epsilon\\,h\)=\\int\\frac\{\\delta\\mathcal\{F\}\}\{\\delta\\rho\}\(x\)\\,h\(x\)\\,dxHereh=ρ˙h=\\dot\{\\rho\}is the perturbation direction of the density\. The left side is the directional derivative ofℱ\\mathcal\{F\}in directionρ˙\\dot\{\\rho\}, namelyD​ℱ​\(ρ\)​\[ρ˙\]D\\mathcal\{F\}\(\\rho\)\[\\dot\{\\rho\}\]; the right side is theL2L^\{2\}inner product of the functional derivativeδ​ℱδ​ρ\\frac\{\\delta\\mathcal\{F\}\}\{\\delta\\rho\}with the perturbationρ˙\\dot\{\\rho\}\.

This is the exact analogue of the finite\-dimensional identityD​f​\(x\)​\[h\]=⟨∇f,h⟩=∑i∂f∂xi​hiDf\(x\)\[h\]=\\langle\\nabla f,h\\rangle=\\sum\_\{i\}\\frac\{\\partial f\}\{\\partial x\_\{i\}\}h\_\{i\}\.

Step 1: Directional derivative along a Wasserstein tangent vector\.

Take a tangent vectorη=∇ψ\\eta=\\nabla\\psi\(the gradient of some scalar functionψ\\psi\)\. The density perturbation it induces is:

ρ˙=−∇⋅\(ρ​∇ψ\)\\dot\{\\rho\}=\-\\nabla\\cdot\(\\rho\\,\\nabla\\psi\)
The directional derivative ofℱ\\mathcal\{F\}in this direction is \(chain: first expand using the definition of functional derivative, then substitute the expression forρ˙\\dot\{\\rho\}\):

D​ℱ​\(ρ\)​\[ρ˙\]\\displaystyle D\\mathcal\{F\}\(\\rho\)\[\\dot\{\\rho\}\]=∫δ​ℱδ​ρ​\(x\)⋅ρ˙​\(x\)​𝑑x\\displaystyle=\\int\\frac\{\\delta\\mathcal\{F\}\}\{\\delta\\rho\}\(x\)\\cdot\\dot\{\\rho\}\(x\)\\,dx\(definition of functional derivative\)=∫\(log⁡ρ\+1\+V\)⏟δ​ℱδ​ρ​from previous section⋅\(−∇⋅\(ρ​∇ψ\)\)⏟expression for​ρ˙​𝑑x\\displaystyle=\\int\\underbrace\{\(\\log\\rho\+1\+V\)\}\_\{\\frac\{\\delta\\mathcal\{F\}\}\{\\delta\\rho\}\\text\{ from previous section\}\}\\cdot\\underbrace\{\\bigl\(\-\\nabla\\cdot\(\\rho\\nabla\\psi\)\\bigr\)\}\_\{\\text\{expression for \}\\dot\{\\rho\}\}\\,dx\(47\)
The task now: use integration by parts to simplify the above into the form∫ρ​\(something\)⋅∇ψ​d​x\\int\\rho\\,\(\\text\{something\}\)\\cdot\\nabla\\psi\\,dx, then compare with the Wasserstein inner product⟨ξ,η⟩ρ=∫ρ​ξ⋅η​𝑑x=∫ρ​ξ⋅∇ψ​d​x\\langle\\xi,\\eta\\rangle\_\{\\rho\}=\\int\\rho\\,\\xi\\cdot\\eta\\,dx=\\int\\rho\\,\\xi\\cdot\\nabla\\psi\\,dxto read off whatξ\\xiis\.

Step 2: Integration by parts\.

The integral obtained from Step 1 is:

D​ℱ​\(ρ\)​\[ρ˙\]=∫\(log⁡ρ\+1\+V\)⏟=⁣:f​\(x\)⋅\(−∇⋅\(ρ​∇ψ\)\)⏟=⁣:−∇⋅𝐅​𝑑xD\\mathcal\{F\}\(\\rho\)\[\\dot\{\\rho\}\]=\\int\\underbrace\{\(\\log\\rho\+1\+V\)\}\_\{=:f\(x\)\}\\cdot\\underbrace\{\\bigl\(\-\\nabla\\cdot\(\\rho\\nabla\\psi\)\\bigr\)\}\_\{=:\-\\nabla\\cdot\\mathbf\{F\}\}\\,dxwhere we denotef:=log⁡ρ\+1\+Vf:=\\log\\rho\+1\+V\(a scalar function\) and𝐅:=ρ​∇ψ\\mathbf\{F\}:=\\rho\\nabla\\psi\(a vector field\)\.

Integration by parts formula\(see Appendix[C](https://arxiv.org/html/2606.24157#A3)\): for compactly supported functions,

∫f​\(∇⋅𝐅\)​𝑑x=−∫∇f⋅𝐅​d​x\\int f\\,\(\\nabla\\cdot\\mathbf\{F\}\)\\,dx=\-\\int\\nabla f\\cdot\\mathbf\{F\}\\,dxThis formula follows from the product rule∇⋅\(f​𝐅\)=∇f⋅𝐅\+f​∇⋅𝐅\\nabla\\cdot\(f\\mathbf\{F\}\)=\\nabla f\\cdot\\mathbf\{F\}\+f\\,\\nabla\\cdot\\mathbf\{F\}, integrating both sides, with the left side∫∇⋅\(f​𝐅\)​𝑑x=0\\int\\nabla\\cdot\(f\\mathbf\{F\}\)\\,dx=0\(divergence theorem\+\+vanishing boundary terms\), and rearranging\.

Applying to our integral\(note the negative sign−∇⋅𝐅\-\\nabla\\cdot\\mathbf\{F\}\):

∫f⋅\(−∇⋅𝐅\)​𝑑x\\displaystyle\\int f\\cdot\(\-\\nabla\\cdot\\mathbf\{F\}\)\\,dx=−∫f​\(∇⋅𝐅\)​𝑑x\\displaystyle=\-\\int f\\,\(\\nabla\\cdot\\mathbf\{F\}\)\\,dx=−\(−∫∇f⋅𝐅​d​x\)\(by the integration by parts formula\)\\displaystyle=\-\\left\(\-\\int\\nabla f\\cdot\\mathbf\{F\}\\,dx\\right\)\\qquad\\text\{\(by the integration by parts formula\)\}=∫∇f⋅𝐅​d​x\\displaystyle=\\int\\nabla f\\cdot\\mathbf\{F\}\\,dx\(48\)
Substituting backf=log⁡ρ\+1\+Vf=\\log\\rho\+1\+V,𝐅=ρ​∇ψ\\mathbf\{F\}=\\rho\\nabla\\psi:

=∫∇\(log⁡ρ\+1\+V\)⋅\(ρ​∇ψ\)​𝑑x\\displaystyle=\\int\\nabla\(\\log\\rho\+1\+V\)\\cdot\(\\rho\\nabla\\psi\)\\,dx=∫ρ​∇\(log⁡ρ\+V\)⋅∇ψ​d​x\\displaystyle=\\int\\rho\\,\\nabla\(\\log\\rho\+V\)\\cdot\\nabla\\psi\\,dxThe last step uses∇1=0\\nabla 1=0\(the gradient of a constant is zero\), so∇\(log⁡ρ\+1\+V\)=∇\(log⁡ρ\+V\)\\nabla\(\\log\\rho\+1\+V\)=\\nabla\(\\log\\rho\+V\)\.

Summary of Step 2: Integration by parts transfers the “∇⋅\\nabla\\cdot” from𝐅=ρ​∇ψ\\mathbf\{F\}=\\rho\\nabla\\psiinto a “∇\\nabla” acting onf=log⁡ρ\+1\+Vf=\\log\\rho\+1\+V\. Result: the integral becomes∫ρ​∇\(log⁡ρ\+V\)⋅∇ψ​d​x\\int\\rho\\,\\nabla\(\\log\\rho\+V\)\\cdot\\nabla\\psi\\,dx, already of the form∫ρ​\(vector field\)⋅∇ψ​d​x\\int\\rho\\,\(\\text\{vector field\}\)\\cdot\\nabla\\psi\\,dx\.

Step 3: Identify the gradient\.

We needξ\\xisuch that⟨ξ,∇ψ⟩ρ=D​ℱ​\(ρ\)​\[ρ˙\]\\langle\\xi,\\nabla\\psi\\rangle\_\{\\rho\}=D\\mathcal\{F\}\(\\rho\)\[\\dot\{\\rho\}\]for allψ\\psi:

∫ρ​ξ⋅∇ψ​d​x=∫ρ​∇\(log⁡ρ\+V\)⋅∇ψ​d​x∀ψ\\int\\rho\\,\\xi\\cdot\\nabla\\psi\\,dx=\\int\\rho\\,\\nabla\(\\log\\rho\+V\)\\cdot\\nabla\\psi\\,dx\\quad\\forall\\psi
This forces:

gradW​ℱ​\(ρ\)=∇\(log⁡ρ\+V\)=∇log⁡ρ\+∇V\\boxed\{\\mathrm\{grad\}\_\{W\}\\mathcal\{F\}\(\\rho\)=\\nabla\(\\log\\rho\+V\)=\\nabla\\log\\rho\+\\nabla V\}\(49\)

### 5\.4The main result: Fokker–Planck as a Wasserstein gradient flow

The definition of the Wasserstein gradient flow: “the curve of steepest descent ofℱ\\mathcal\{F\}in probability space\.”

The Wasserstein gradient flow ofℱ\\mathcal\{F\}is defined by:

∂tρt\+∇⋅\(ρt​vt\)=0,vt=−gradW​ℱ​\(ρt\)\\partial\_\{t\}\\rho\_\{t\}\+\\nabla\\cdot\\bigl\(\\rho\_\{t\}\\,v\_\{t\}\\bigr\)=0,\\quad v\_\{t\}=\-\\mathrm\{grad\}\_\{W\}\\mathcal\{F\}\(\\rho\_\{t\}\)\(50\)
Substituting \([49](https://arxiv.org/html/2606.24157#S5.E49)\):

vt\\displaystyle v\_\{t\}=−\(∇log⁡ρt\+∇V\)\\displaystyle=\-\(\\nabla\\log\\rho\_\{t\}\+\\nabla V\)∂tρt\\displaystyle\\partial\_\{t\}\\rho\_\{t\}=−∇⋅\(ρt​vt\)=∇⋅\(ρt​\(∇log⁡ρt\+∇V\)\)\\displaystyle=\-\\nabla\\cdot\(\\rho\_\{t\}\\,v\_\{t\}\)=\\nabla\\cdot\\bigl\(\\rho\_\{t\}\(\\nabla\\log\\rho\_\{t\}\+\\nabla V\)\\bigr\)=∇⋅\(∇ρt\+ρt​∇V\)\\displaystyle=\\nabla\\cdot\(\\nabla\\rho\_\{t\}\+\\rho\_\{t\}\\nabla V\)=Δ​ρt\+∇⋅\(ρt​∇V\)\\displaystyle=\\Delta\\rho\_\{t\}\+\\nabla\\cdot\(\\rho\_\{t\}\\nabla V\)\(51\)
Key InsightThe Fokker\-Planck equation is the gradient flow of the free energyℱ​\(ρ\)=KL​\(ρ∥π\)\\mathcal\{F\}\(\\rho\)=\\mathrm\{KL\}\(\\rho\\\|\\pi\)in Wasserstein space\.Just asx˙=−∇f​\(x\)\\dot\{x\}=\-\\nabla f\(x\)in Euclidean space movesxxin the direction of steepest descent offf, the Fokker\-Planck equation makesρt\\rho\_\{t\}“flow” in probability space along the direction of steepest descent ofKL\(⋅∥π\)\\mathrm\{KL\}\(\\cdot\\\|\\pi\), until reaching the minimizerρ∞=π\\rho\_\{\\infty\}=\\pi\.This is the central contribution of Jordan, Kinderlehrer, and Otto \(1998\)\.

### 5\.5End\-of\-section comparison: Euclidean vs\. Wasserstein gradient\-flow logic

The whole section is the Wasserstein analogue of the most familiar Euclidean story:

function⟶gradient⟶gradient flow⟶energy decreases\.\\text\{function\}\\;\\longrightarrow\\;\\text\{gradient\}\\;\\longrightarrow\\;\\text\{gradient flow\}\\;\\longrightarrow\\;\\text\{energy decreases\}\.The only difference is that the “point” is now a probability density and the “velocity” is represented by a vector field through the continuity equation\.

Key InsightThe proof strategy is exactly the same as in Euclidean space\. First identify the metric, then compute the gradient of the objective under that metric, then write the negative\-gradient dynamics\. In Wasserstein space, this negative\-gradient dynamics is not an ODE for points but a continuity equation for densities; for the free energy above, it becomes the Fokker–Planck equation\.

The next natural question:Since Fokker\-Planck is a gradient flow, can we—just as inℝn\\mathbb\{R\}^\{n\}—use discrete optimization to “approximate” this continuous descent process? This is the JKO scheme\.

## 6The JKO Scheme

Where we are:Section 5 proved that the Fokker\-Planck equation is the gradient flow of the free energyℱ\\mathcal\{F\}in Wasserstein space\. This section completes the final step:discretizingthis continuous process into a step\-by\-step optimization problem\.

### 6\.1Implicit Euler method: from gradient flow to optimization problem

Consider the gradient flowx˙=−∇f​\(x\)\\dot\{x\}=\-\\nabla f\(x\)in Euclidean space\. The implicit Euler method discretizes it using a backward difference:

xk\+1−xkτ=−∇f​\(xk\+1\)\\frac\{x\_\{k\+1\}\-x\_\{k\}\}\{\\tau\}=\-\\nabla f\(x\_\{k\+1\}\)Rearranging gives∇f​\(xk\+1\)\+1τ​\(xk\+1−xk\)=0\\nabla f\(x\_\{k\+1\}\)\+\\frac\{1\}\{\\tau\}\(x\_\{k\+1\}\-x\_\{k\}\)=0\. The left side is precisely∇x\[f​\(x\)\+12​τ​\|x−xk\|2\]\|xk\+1\\nabla\_\{x\}\\bigl\[f\(x\)\+\\frac\{1\}\{2\\tau\}\|x\-x\_\{k\}\|^\{2\}\\bigr\]\\big\|\_\{x\_\{k\+1\}\}, so implicit Euler is equivalent to:

xk\+1=arg⁡minx⁡\{f​\(x\)\+12​τ​\|x−xk\|2\}x\_\{k\+1\}=\\arg\\min\_\{x\}\\left\\\{f\(x\)\+\\frac\{1\}\{2\\tau\}\|x\-x\_\{k\}\|^\{2\}\\right\\\}\(52\)
Intuition: Each step finds the optimal trade\-off between decreasingffand not straying too far fromxkx\_\{k\}\. The step sizeτ\\taucontrols the step length: smallτ\\taulets the penalty term dominate, giving small but stable steps; largeτ\\taulets the objective term dominate, giving large steps that may overshoot\.

Why is the stationary point a minimum?The Hessian ofg​\(x\)=f​\(x\)\+12​τ​\|x−xk\|2g\(x\)=f\(x\)\+\\frac\{1\}\{2\\tau\}\|x\-x\_\{k\}\|^\{2\}is∇2f\+1τ​I\\nabla^\{2\}f\+\\frac\{1\}\{\\tau\}I\. Whenτ\\tauis sufficiently small \(1/τ1/\\tauexceeds the most negative eigenvalue of∇2f\\nabla^\{2\}f\),ggis strictly convex and the stationary point is the unique global minimum\.

### 6\.2JKO scheme: Implicit Euler in Wasserstein space

The core idea of Jordan–Kinderlehrer–Otto \(1998\): replace every Euclidean concept in implicit Euler with its Wasserstein\-space counterpart\.

###### Definition 6\.1\(JKO scheme\)\.

Given an initial distributionρ0τ=ρ0\\rho\_\{0\}^\{\\tau\}=\\rho\_\{0\}and step sizeτ\>0\\tau\>0, the JKO iteration is:

ρk\+1τ=arg⁡minρ∈𝒫2​\(ℝd\)⁡\{ℱ​\(ρ\)\+12​τ​W22​\(ρ,ρkτ\)\}\\boxed\{\\rho\_\{k\+1\}^\{\\tau\}=\\arg\\min\_\{\\rho\\in\\mathcal\{P\}\_\{2\}\(\\mathbb\{R\}^\{d\}\)\}\\left\\\{\\mathcal\{F\}\(\\rho\)\+\\frac\{1\}\{2\\tau\}W\_\{2\}^\{2\}\(\\rho,\\rho\_\{k\}^\{\\tau\}\)\\right\\\}\}\(53\)whereℱ​\(ρ\)=∫ρ​log⁡ρ​d​x\+∫V​ρ​𝑑x\\mathcal\{F\}\(\\rho\)=\\int\\rho\\log\\rho\\,dx\+\\int V\\rho\\,dxis the free energy functional\.

Key InsightOrientation convention for JKO\.In the JKO stepρk\+1=arg⁡minρ⁡\{ℱ​\(ρ\)\+12​τ​W22​\(ρ,ρk\)\},\\rho\_\{k\+1\}=\\arg\\min\_\{\\rho\}\\left\\\{\\mathcal\{F\}\(\\rho\)\+\\frac\{1\}\{2\\tau\}W\_\{2\}^\{2\}\(\\rho,\\rho\_\{k\}\)\\right\\\},ρk\\rho\_\{k\}is the old distribution andρk\+1\\rho\_\{k\+1\}is the new distribution\. We use the optimal mapTk:ρk\+1→ρk,\(Tk\)\#​ρk\+1=ρk\.T\_\{k\}:\\rho\_\{k\+1\}\\to\\rho\_\{k\},\\qquad\(T\_\{k\}\)\_\{\\\#\}\\rho\_\{k\+1\}=\\rho\_\{k\}\.Thus, forx∼ρk\+1x\\sim\\rho\_\{k\+1\}, the pointTk​\(x\)T\_\{k\}\(x\)is its previous location inρk\\rho\_\{k\}, andx−Tk​\(x\)x\-T\_\{k\}\(x\)is the displacement from old to new\. Therefore the discrete velocity isx−Tk​\(x\)τ\.\\frac\{x\-T\_\{k\}\(x\)\}\{\\tau\}\.The JKO optimality condition is consequentlyx−Tk​\(x\)τ=−∇δ​ℱδ​ρ​\(x\)\\boxed\{\\frac\{x\-T\_\{k\}\(x\)\}\{\\tau\}=\-\\nabla\\frac\{\\delta\\mathcal\{F\}\}\{\\delta\\rho\}\(x\)\}with the force evaluated at thenewpointxxand thenewdensityρk\+1\\rho\_\{k\+1\}\.

Complete analogy dictionary:

Why must the objective be the free energyℱ\\mathcal\{F\}?This is not a “choice\.” Section 5 already showed that the velocity field of the Fokker\-Planck equation isv=−∇\(log⁡ρ\+V\)v=\-\\nabla\(\\log\\rho\+V\), while the Wasserstein gradient flow satisfiesv=−gradW​ℱ=−∇δ​ℱδ​ρv=\-\\mathrm\{grad\}\_\{W\}\\mathcal\{F\}=\-\\nabla\\frac\{\\delta\\mathcal\{F\}\}\{\\delta\\rho\}\. Comparing the two immediately givesδ​ℱδ​ρ=log⁡ρ\+V\+C\\frac\{\\delta\\mathcal\{F\}\}\{\\delta\\rho\}=\\log\\rho\+V\+C, and the unique functional satisfying this isℱ​\(ρ\)=∫ρ​log⁡ρ\+∫V​ρ\\mathcal\{F\}\(\\rho\)=\\int\\rho\\log\\rho\+\\int V\\rho\. The Fokker\-Planck equationdeterminesthe JKO objective\.

### 6\.3Deriving JKO from the gradient flow

Starting point:The Wasserstein gradient flow \(conclusion of Section 5\):

vt=−gradW​ℱ​\(ρt\)=−∇\(log⁡ρt\+V\)v\_\{t\}=\-\\mathrm\{grad\}\_\{W\}\\mathcal\{F\}\(\\rho\_\{t\}\)=\-\\nabla\(\\log\\rho\_\{t\}\+V\)
Step 1: Backward difference\.LetTkT\_\{k\}be the optimal transport map fromρk\+1\\rho\_\{k\+1\}toρk\\rho\_\{k\}\(\(Tk\)\#​ρk\+1=ρk\(T\_\{k\}\)\_\{\\\#\}\\rho\_\{k\+1\}=\\rho\_\{k\}\)\. Thenx−Tk​\(x\)x\-T\_\{k\}\(x\)is the displacement from the previous step to the current step, and the approximate velocity isv≈id−Tkτv\\approx\\frac\{\\mathrm\{id\}\-T\_\{k\}\}\{\\tau\}\. Substituting into the gradient flow equation:

id−Tkτ=−gradW​ℱ​\(ρk\+1\)=−∇log⁡ρk\+1−∇V\\frac\{\\mathrm\{id\}\-T\_\{k\}\}\{\\tau\}=\-\\mathrm\{grad\}\_\{W\}\\mathcal\{F\}\(\\rho\_\{k\+1\}\)=\-\\nabla\\log\\rho\_\{k\+1\}\-\\nabla V\(54\)The right\-hand side is evaluated at thenewdistributionρk\+1\\rho\_\{k\+1\}—this is precisely the meaning of “implicit\.”

Step 2: Recognize the optimality condition of the optimization problem\.Rearranging \([54](https://arxiv.org/html/2606.24157#S6.E54)\) as

gradW​ℱ​\(ρk\+1\)\+id−Tkτ=0\\mathrm\{grad\}\_\{W\}\\mathcal\{F\}\(\\rho\_\{k\+1\}\)\+\\frac\{\\mathrm\{id\}\-T\_\{k\}\}\{\\tau\}=0andid−Tkτ\\frac\{\\mathrm\{id\}\-T\_\{k\}\}\{\\tau\}is precisely the Wasserstein gradient of12​τ​W22​\(ρ,ρk\)\\frac\{1\}\{2\\tau\}W\_\{2\}^\{2\}\(\\rho,\\rho\_\{k\}\)\(proved in the next subsection\), so the above is equivalent to

gradW​\[ℱ​\(ρ\)\+12​τ​W22​\(ρ,ρk\)\]\|ρk\+1=0\\mathrm\{grad\}\_\{W\}\\left\[\\mathcal\{F\}\(\\rho\)\+\\frac\{1\}\{2\\tau\}W\_\{2\}^\{2\}\(\\rho,\\rho\_\{k\}\)\\right\]\\bigg\|\_\{\\rho\_\{k\+1\}\}=0The critical point==the minimizer \(guaranteed by displacement convexity of theW22W\_\{2\}^\{2\}term\), yielding the JKO formula \([53](https://arxiv.org/html/2606.24157#S6.E53)\)\.

### 6\.4First\-order optimality conditions of JKO

The argmin formula of JKO specifies “what to do at each step,” but it must be expanded into a computable form\. This subsection derives thefirst\-order conditionssatisfied by the JKO minimizer, which are the key to subsequently proving convergence to the Fokker\-Planck equation\.

#### Step 1: Variational condition

ρk\+1\\rho\_\{k\+1\}is the solution to the following problem:

minρ:∫ρ=1,ρ≥0⁡\{ℱ​\(ρ\)\+12​τ​W22​\(ρ,ρk\)\}\\min\_\{\\rho:\\,\\int\\rho=1,\\,\\rho\\geq 0\}\\left\\\{\\mathcal\{F\}\(\\rho\)\+\\frac\{1\}\{2\\tau\}W\_\{2\}^\{2\}\(\\rho,\\rho\_\{k\}\)\\right\\\}For admissible perturbationsδ​ρ\\delta\\rho\(satisfying∫δ​ρ=0\\int\\delta\\rho=0to preserve mass normalization\), the necessary condition for a minimum is:

∫\[δ​ℱδ​ρ\|ρk\+1\+12​τ​δ​W22​\(⋅,ρk\)δ​ρ\|ρk\+1\]​δ​ρ​𝑑x=0,∀δ​ρ:∫δ​ρ=0\\int\\left\[\\frac\{\\delta\\mathcal\{F\}\}\{\\delta\\rho\}\\bigg\|\_\{\\rho\_\{k\+1\}\}\+\\frac\{1\}\{2\\tau\}\\frac\{\\delta W\_\{2\}^\{2\}\(\\cdot,\\rho\_\{k\}\)\}\{\\delta\\rho\}\\bigg\|\_\{\\rho\_\{k\+1\}\}\\right\]\\delta\\rho\\,dx=0,\\quad\\forall\\,\\delta\\rho:\\;\\int\\delta\\rho=0By the constrained du Bois\-Reymond lemma \(Appendix[C](https://arxiv.org/html/2606.24157#A3)\), the function inside the brackets must be a constantCC\(the Lagrange multiplier corresponding to the constraint∫ρ=1\\int\\rho=1\):

δ​ℱδ​ρ\|ρk\+1\+12​τ​δ​W22​\(⋅,ρk\)δ​ρ\|ρk\+1=C\\frac\{\\delta\\mathcal\{F\}\}\{\\delta\\rho\}\\bigg\|\_\{\\rho\_\{k\+1\}\}\+\\frac\{1\}\{2\\tau\}\\frac\{\\delta W\_\{2\}^\{2\}\(\\cdot,\\rho\_\{k\}\)\}\{\\delta\\rho\}\\bigg\|\_\{\\rho\_\{k\+1\}\}=C\(55\)

#### Step 2: Functional derivative ofW22W\_\{2\}^\{2\}

LetTkT\_\{k\}be the optimal transport map fromρk\+1\\rho\_\{k\+1\}toρk\\rho\_\{k\}\. The following is a classical result \(Santambrogio,Optimal Transport for Applied Mathematicians, Prop\. 7\.17\):

12​δ​W22​\(⋅,ρk\)δ​ρ\|ρk\+1​\(x\)=12​\|x−Tk​\(x\)\|2\\frac\{1\}\{2\}\\frac\{\\delta W\_\{2\}^\{2\}\(\\cdot,\\rho\_\{k\}\)\}\{\\delta\\rho\}\\bigg\|\_\{\\rho\_\{k\+1\}\}\(x\)=\\frac\{1\}\{2\}\|x\-T\_\{k\}\(x\)\|^\{2\}\(56\)
Intuition: Adding a small amount of mass atxxmeans that this mass also needs to be transported toρk\\rho\_\{k\}, at a cost of\|x−Tk​\(x\)\|2\|x\-T\_\{k\}\(x\)\|^\{2\}\.

Derivation: LetTϵT\_\{\\epsilon\}be the optimal map fromρϵ:=ρk\+1\+ϵ​δ​ρ\\rho\_\{\\epsilon\}:=\\rho\_\{k\+1\}\+\\epsilon\\,\\delta\\rhotoρk\\rho\_\{k\}\. By the optimality condition, the variation ofTϵT\_\{\\epsilon\}is of higher order inϵ\\epsilon\(the “envelope theorem”\), so

dd​ϵ\|0​12​∫\|x−Tϵ​\(x\)\|2​ρϵ​𝑑x=12​∫\|x−Tk​\(x\)\|2​δ​ρ​𝑑x\\left\.\\frac\{d\}\{d\\epsilon\}\\right\|\_\{0\}\\frac\{1\}\{2\}\\int\|x\-T\_\{\\epsilon\}\(x\)\|^\{2\}\\rho\_\{\\epsilon\}\\,dx=\\frac\{1\}\{2\}\\int\|x\-T\_\{k\}\(x\)\|^\{2\}\\,\\delta\\rho\\,dxComparing with the definition of the functional derivative∫δ​\(⋯\)δ​ρ​δ​ρ​𝑑x\\int\\frac\{\\delta\(\\cdots\)\}\{\\delta\\rho\}\\,\\delta\\rho\\,dx, we read off \([56](https://arxiv.org/html/2606.24157#S6.E56)\)\.

#### Step 3: Substitute and take the spatial gradient

Substitutingδ​ℱδ​ρ=log⁡ρk\+1\+1\+V\\frac\{\\delta\\mathcal\{F\}\}\{\\delta\\rho\}=\\log\\rho\_\{k\+1\}\+1\+Vand \([56](https://arxiv.org/html/2606.24157#S6.E56)\) into \([55](https://arxiv.org/html/2606.24157#S6.E55)\):

log⁡ρk\+1​\(x\)\+V​\(x\)\+12​τ​\|x−Tk​\(x\)\|2=C′\\log\\rho\_\{k\+1\}\(x\)\+V\(x\)\+\\frac\{1\}\{2\\tau\}\|x\-T\_\{k\}\(x\)\|^\{2\}=C^\{\\prime\}Taking the spatial gradient∇x\\nabla\_\{x\}of both sides \(the constantC′C^\{\\prime\}vanishes\):

∇log⁡ρk\+1\+∇V\+1τ​\(x−Tk​\(x\)\)=0\\nabla\\log\\rho\_\{k\+1\}\+\\nabla V\+\\frac\{1\}\{\\tau\}\(x\-T\_\{k\}\(x\)\)=0\(57\)\(The last term is the Wasserstein gradient of12​τ​W22​\(⋅,ρk\)\\frac\{1\}\{2\\tau\}W\_\{2\}^\{2\}\(\\cdot,\\rho\_\{k\}\)atρk\+1\\rho\_\{k\+1\}\. By a standard result of optimal transport, the first variationδδ​ρ​\[12​W22​\(⋅,ρk\)\]\|ρk\+1\\frac\{\\delta\}\{\\delta\\rho\}\\big\[\\frac\{1\}\{2\}W\_\{2\}^\{2\}\(\\cdot,\\rho\_\{k\}\)\\big\]\\big\|\_\{\\rho\_\{k\+1\}\}is the Kantorovich potentialφk\\varphi\_\{k\}, whose spatial gradient is*exactly*the displacement∇φk​\(x\)=x−Tk​\(x\)\\nabla\\varphi\_\{k\}\(x\)=x\-T\_\{k\}\(x\)\(by definition of the Brenier map,Tk​\(x\)=x−∇φk​\(x\)T\_\{k\}\(x\)=x\-\\nabla\\varphi\_\{k\}\(x\)\)\. Hence the gradient isx−Tk​\(x\)τ\\frac\{x\-T\_\{k\}\(x\)\}\{\\tau\}, with*no*∇Tk\\nabla T\_\{k\}correction term; see Santambrogio Prop\. 7\.17 / AGS Ch\. 10 for the rigorous statement\.\)

#### Step 4: Recognize the Fokker\-Planck velocity field

Solving for the discrete velocity from \([57](https://arxiv.org/html/2606.24157#S6.E57)\):

x−Tk​\(x\)τ=−∇log⁡ρk\+1​\(x\)−∇V​\(x\)\\frac\{x\-T\_\{k\}\(x\)\}\{\\tau\}=\-\\nabla\\log\\rho\_\{k\+1\}\(x\)\-\\nabla V\(x\)\(58\)The left\-hand side is the approximate velocity fromρk\\rho\_\{k\}toρk\+1\\rho\_\{k\+1\}\(TkT\_\{k\}pushesρk\+1\\rho\_\{k\+1\}back toρk\\rho\_\{k\}, sox−Tk​\(x\)x\-T\_\{k\}\(x\)is the displacement\)\. The right\-hand side is precisely the Fokker\-Planck velocity fieldv=−∇log⁡ρ−∇Vv=\-\\nabla\\log\\rho\-\\nabla V\.

Therefore:the discrete velocity field produced by each JKO step is exactly the Fokker\-Planck velocity field at that time instant\. Asτ→0\\tau\\to 0, the discrete JKO solution converges to the continuous\-time solution of the Fokker\-Planck equation\.

### 6\.5Rigorous convergence theorem

###### Theorem 6\.2\(JKO, 1998; Ambrosio–Gigli–Savaré, 2008\)\.

AssumeVVisλ\\lambda\-convex \(∇2V≥λ​I\\nabla^\{2\}V\\geq\\lambda I,λ∈ℝ\\lambda\\in\\mathbb\{R\}\) and satisfies a quadratic growth condition\. Define the piecewise constant interpolation:

ρτ​\(t\):=ρkτfor​t∈\[k​τ,\(k\+1\)​τ\)\\rho^\{\\tau\}\(t\):=\\rho\_\{k\}^\{\\tau\}\\quad\\text\{for \}t\\in\[k\\tau,\(k\+1\)\\tau\)Then asτ→0\\tau\\to 0,ρτ​\(t\)\\rho^\{\\tau\}\(t\)converges in theW2W\_\{2\}metric to the unique solution of the Fokker\-Planck equation\.

Four key steps of the proof:

1. 1\.Existence: Each JKO subproblem has a solution—theW22W\_\{2\}^\{2\}term provides compactness, andℱ\\mathcal\{F\}is lower semicontinuous in theW2W\_\{2\}topology\.
2. 2\.Energy decay:ℱ​\(ρk\+1τ\)≤ℱ​\(ρkτ\)\\mathcal\{F\}\(\\rho\_\{k\+1\}^\{\\tau\}\)\\leq\\mathcal\{F\}\(\\rho\_\{k\}^\{\\tau\}\), sinceρk\\rho\_\{k\}itself is a feasible solution \(with costℱ​\(ρk\)\+0\\mathcal\{F\}\(\\rho\_\{k\}\)\+0\)\.
3. 3\.Uniform estimates: Energy decay\+\+control of the Fisher information yields compactness of\{ρkτ\}\\\{\\rho\_\{k\}^\{\\tau\}\\\}, enabling passage to the limitτ→0\\tau\\to 0\.
4. 4\.Identification of the limit: Via the optimality condition \([57](https://arxiv.org/html/2606.24157#S6.E57)\), the discrete velocity field converges to−∇log⁡ρ−∇V\-\\nabla\\log\\rho\-\\nabla Vasτ→0\\tau\\to 0, recovering the Fokker\-Planck equation\.

### 6\.6Structural advantages of JKO

1. \(a\)Positivity and mass conservation are automatic: Each step optimizes over𝒫2​\(ℝd\)\\mathcal\{P\}\_\{2\}\(\\mathbb\{R\}^\{d\}\), soρk\+1≥0\\rho\_\{k\+1\}\\geq 0and∫ρk\+1=1\\int\\rho\_\{k\+1\}=1are built\-in constraints\.
2. \(b\)Unconditional stability: No CFL condition is needed; arbitrarily largeτ\\taustill produces a meaningful \(though coarse\) approximation\.
3. \(c\)Built\-in energy dissipation:ℱ​\(ρk\+1\)≤ℱ​\(ρk\)\\mathcal\{F\}\(\\rho\_\{k\+1\}\)\\leq\\mathcal\{F\}\(\\rho\_\{k\}\)without any additional verification\.
4. \(d\)Modularity: Changingℱ\\mathcal\{F\}handles different PDEs:

### 6\.7Example: Heat equation

Takeℱ​\(ρ\)=∫ρ​log⁡ρ​d​x\\mathcal\{F\}\(\\rho\)=\\int\\rho\\log\\rho\\,dxandV=0V=0\. Thenδ​ℱδ​ρ=log⁡ρ\+1\\frac\{\\delta\\mathcal\{F\}\}\{\\delta\\rho\}=\\log\\rho\+1, and the Wasserstein gradient=∇log⁡ρ=\\nabla\\log\\rho\. The gradient flow velocity isv=−∇log⁡ρv=\-\\nabla\\log\\rho; substituting into the continuity equation:

∂tρ=−∇⋅\(ρ​v\)=∇⋅\(ρ​∇log⁡ρ\)=∇⋅\(∇ρ\)=Δ​ρ\\partial\_\{t\}\\rho=\-\\nabla\\cdot\(\\rho v\)=\\nabla\\cdot\(\\rho\\nabla\\log\\rho\)=\\nabla\\cdot\(\\nabla\\rho\)=\\Delta\\rhoThe heat equation describes “diffusion purely due to crowding”—the system only seeks to maximize entropy, with no external force\.

Each JKO step becomes:ρk\+1=arg⁡minρ⁡\{∫ρ​log⁡ρ​d​x\+12​τ​W22​\(ρ,ρk\)\}\\rho\_\{k\+1\}=\\arg\\min\_\{\\rho\}\\bigl\\\{\\int\\rho\\log\\rho\\,dx\+\\frac\{1\}\{2\\tau\}W\_\{2\}^\{2\}\(\\rho,\\rho\_\{k\}\)\\bigr\\\}\. This provides an “optimization perspective” on the heat equation: the progressive spreading of the Gaussian heat kernel==stepwise entropy maximization \(subject to transport cost constraints\)\.

### 6\.8Application: Energy Matching \(Balcerak et al\., 2025\)

The JKO scheme is not only a theoretical device; it is used directly in the design of generative models\.

Key idea: Energy Matching uses a singletime\-independentscalar potentialVθ​\(x\)V\_\{\\theta\}\(x\)as the only learnable parameter, achieving both transport and equilibrium through the JKO scheme:

ρt\+Δ​t=arg⁡minρ⁡\{W22​\(ρ,ρt\)2​Δ​t\+∫Vθ​𝑑ρ\+ε​\(t\)​∫ρ​log⁡ρ​d​x\}\\rho\_\{t\+\\Delta t\}=\\arg\\min\_\{\\rho\}\\left\\\{\\frac\{W\_\{2\}^\{2\}\(\\rho,\\rho\_\{t\}\)\}\{2\\Delta t\}\+\\int V\_\{\\theta\}\\,d\\rho\+\\varepsilon\(t\)\\int\\rho\\log\\rho\\,dx\\right\\\}\(59\)Compared to the standard JKO, the innovation lies in allowing thetemperatureε​\(t\)\\varepsilon\(t\)to vary with time\.

First\-order optimality condition\(directly from the derivation in Section 6\.4\):

x−yΔ​t\+∇Vθ​\(x\)\+ε​\(t\)​∇log⁡ρt\+Δ​t​\(x\)=0\\frac\{x\-y\}\{\\Delta t\}\+\\nabla V\_\{\\theta\}\(x\)\+\\varepsilon\(t\)\\nabla\\log\\rho\_\{t\+\\Delta t\}\(x\)=0wherey=T​\(x\)y=T\(x\)is the point inρt\\rho\_\{t\}corresponding toxx\. This condition yields strikingly different behavior in two limiting regimes:

Phase 1 \(ε=0\\varepsilon=0, far from data\): The condition reduces tox−yΔ​t=−∇Vθ​\(x\)\\frac\{x\-y\}\{\\Delta t\}=\-\\nabla V\_\{\\theta\}\(x\)\. Particles undergo deterministic transport with velocity−∇Vθ\-\\nabla V\_\{\\theta\}—a pure OT flow, similar to Flow Matching\.

Phase 2 \(ε=εmax\\varepsilon=\\varepsilon\_\{\\max\}, near equilibrium\): Samples barely move \(x≈yx\\approx y\), and the condition becomes∇Vθ\+εmax​∇log⁡ρeq=0\\nabla V\_\{\\theta\}\+\\varepsilon\_\{\\max\}\\nabla\\log\\rho\_\{\\text\{eq\}\}=0, yielding the equilibrium distributionρeq∝e−Vθ/εmax\\rho\_\{\\text\{eq\}\}\\propto e^\{\-V\_\{\\theta\}/\\varepsilon\_\{\\max\}\}—a Boltzmann distribution\.VθV\_\{\\theta\}directly encodes the log\-likelihood of the data\.

Key InsightOneVθ​\(x\)V\_\{\\theta\}\(x\), two roles: Far from data,−∇Vθ\-\\nabla V\_\{\\theta\}serves as the transport velocity \(Flow Matching\); near data,Vθ/εV\_\{\\theta\}/\\varepsilonserves as the energy function \(EBM\)\. There is no need to learn a time\-dependentsθ​\(x,t\)s\_\{\\theta\}\(x,t\)orvθ​\(x,t\)v\_\{\\theta\}\(x,t\)\.

Correspondence with concepts in this article:

### 6\.9Unified derivation from JKO to mainstream generative algorithms

Instead of going through the continuous\-time Fokker–Planck equation, we can read JKO directly as aniterative generative algorithm\. The basic loop is:

JKO generative template1\.Choose a free energyℱk​\(ρ\)=∫Vk​ρ​𝑑x\+εk​∫ρ​log⁡ρ​d​x\.\\mathcal\{F\}\_\{k\}\(\\rho\)=\\int V\_\{k\}\\rho\\,dx\+\\varepsilon\_\{k\}\\int\\rho\\log\\rho\\,dx\.2\.Initialize from a simple distribution, e\.g\.ρ0=𝒩​\(0,I\)\\rho\_\{0\}=\\mathcal\{N\}\(0,I\)\.3\.Fork=0,…,K−1k=0,\\ldots,K\-1, compute the proximal Wasserstein stepρk\+1=arg⁡minρ⁡\{ℱk​\(ρ\)\+12​τk​W22​\(ρ,ρk\)\}\.\\rho\_\{k\+1\}=\\arg\\min\_\{\\rho\}\\left\\\{\\mathcal\{F\}\_\{k\}\(\\rho\)\+\\frac\{1\}\{2\\tau\_\{k\}\}W\_\{2\}^\{2\}\(\\rho,\\rho\_\{k\}\)\\right\\\}\.4\.Implement this density step on particles using a learned score, energy, or velocity model\.

For one step, letTkT\_\{k\}be the optimal map from the new densityρk\+1\\rho\_\{k\+1\}back to the old densityρk\\rho\_\{k\}:

Tk:ρk\+1→ρk,\(Tk\)\#​ρk\+1=ρk\.T\_\{k\}:\\rho\_\{k\+1\}\\to\\rho\_\{k\},\\qquad\(T\_\{k\}\)\_\{\\\#\}\\rho\_\{k\+1\}=\\rho\_\{k\}\.Then the JKO optimality condition is

x−Tk​\(x\)τk=−∇Vk​\(x\)−εk​∇log⁡ρk\+1​\(x\),x∼ρk\+1\.\\frac\{x\-T\_\{k\}\(x\)\}\{\\tau\_\{k\}\}=\-\\nabla V\_\{k\}\(x\)\-\\varepsilon\_\{k\}\\nabla\\log\\rho\_\{k\+1\}\(x\),\\qquad x\\sim\\rho\_\{k\+1\}\.\(60\)Equivalently,

x=Tk​\(x\)\+τk​\[−∇Vk​\(x\)−εk​∇log⁡ρk\+1​\(x\)\]\.x=T\_\{k\}\(x\)\+\\tau\_\{k\}\\bigl\[\-\\nabla V\_\{k\}\(x\)\-\\varepsilon\_\{k\}\\nabla\\log\\rho\_\{k\+1\}\(x\)\\bigr\]\.This is already an algorithmic update rule: the new point equals the old point plus a force evaluated at the new point \(implicit Euler\)\. Mainstream generative algorithms differ in how they chooseℱk\\mathcal\{F\}\_\{k\}, how they schedule\(τk,εk\)\(\\tau\_\{k\},\\varepsilon\_\{k\}\), and how they approximate the unknown score or energy in \([60](https://arxiv.org/html/2606.24157#S6.E60)\)\.

Indexing note for diffusion papers: DDPM papers often index sampling backward, fromxtx\_\{t\}toxt−1x\_\{t\-1\}\. In this subsection, we use the JKO convention consistently:kkis old andk\+1k\+1is new\. To compare with DDPM notation, read “k\+1k\+1” here as the next denoising state, even if the diffusion paper labels that state by a smaller physical time index\.

#### 1\. DDPM: stochastic JKO with Gaussian reference potential

JKO choice: use the Gaussian reference potentialVk​\(x\)=12​\|x\|2V\_\{k\}\(x\)=\\frac\{1\}\{2\}\|x\|^\{2\}and retain the entropy term\. One step is

ρk\+1=arg⁡minρ⁡\{∫12​\|x\|2​ρ​𝑑x\+εk​∫ρ​log⁡ρ​d​x\+12​τk​W22​\(ρ,ρk\)\}\.\\rho\_\{k\+1\}=\\arg\\min\_\{\\rho\}\\left\\\{\\int\\frac\{1\}\{2\}\|x\|^\{2\}\\rho\\,dx\+\\varepsilon\_\{k\}\\int\\rho\\log\\rho\\,dx\+\\frac\{1\}\{2\\tau\_\{k\}\}W\_\{2\}^\{2\}\(\\rho,\\rho\_\{k\}\)\\right\\\}\.The optimality condition is

x−Tk​\(x\)τk=−x−εk​∇log⁡ρk\+1​\(x\),x∼ρk\+1\.\\frac\{x\-T\_\{k\}\(x\)\}\{\\tau\_\{k\}\}=\-x\-\\varepsilon\_\{k\}\\nabla\\log\\rho\_\{k\+1\}\(x\),\\qquad x\\sim\\rho\_\{k\+1\}\.Let a particle in the old distribution be denoted byxk∼ρkx\_\{k\}\\sim\\rho\_\{k\}, and suppose its paired new position isxk\+1∼ρk\+1x\_\{k\+1\}\\sim\\rho\_\{k\+1\}, soTk​\(xk\+1\)=xkT\_\{k\}\(x\_\{k\+1\}\)=x\_\{k\}\. Substitutingx=xk\+1x=x\_\{k\+1\}into the optimality condition gives

xk\+1−xkτk=−xk\+1−εk​∇log⁡ρk\+1​\(xk\+1\)\.\\frac\{x\_\{k\+1\}\-x\_\{k\}\}\{\\tau\_\{k\}\}=\-x\_\{k\+1\}\-\\varepsilon\_\{k\}\\nabla\\log\\rho\_\{k\+1\}\(x\_\{k\+1\}\)\.Therefore the exact JKO particle equation is

xk\+1=xk\+τk​\[−xk\+1−εk​∇log⁡ρk\+1​\(xk\+1\)\]\.x\_\{k\+1\}=x\_\{k\}\+\\tau\_\{k\}\\left\[\-x\_\{k\+1\}\-\\varepsilon\_\{k\}\\nabla\\log\\rho\_\{k\+1\}\(x\_\{k\+1\}\)\\right\]\.\(61\)If we include finite\-temperature Langevin sampling inside the same proximal step, we add the Gaussian fluctuation:

xk\+1≈xk\+τk​\[−xk\+1−εk​∇log⁡ρk\+1​\(xk\+1\)\]\+2​εk​τk​z,z∼𝒩​\(0,I\)\.x\_\{k\+1\}\\approx x\_\{k\}\+\\tau\_\{k\}\\left\[\-x\_\{k\+1\}\-\\varepsilon\_\{k\}\\nabla\\log\\rho\_\{k\+1\}\(x\_\{k\+1\}\)\\right\]\+\\sqrt\{2\\varepsilon\_\{k\}\\tau\_\{k\}\}\\,z,\\qquad z\\sim\\mathcal\{N\}\(0,I\)\.\(62\)
This is the exact place where JKO is implicit: both the drift−xk\+1\-x\_\{k\+1\}and the score∇log⁡ρk\+1​\(xk\+1\)\\nabla\\log\\rho\_\{k\+1\}\(x\_\{k\+1\}\)are evaluated at the unknown new point\.

Let us first see where thexk\+1x\_\{k\+1\}on the right\-hand side goes\. Ignoring the Gaussian fluctuation for one line, \([61](https://arxiv.org/html/2606.24157#S6.E61)\) gives

xk\+1=xk−τk​xk\+1−τk​εk​∇log⁡ρk\+1​\(xk\+1\)\.x\_\{k\+1\}=x\_\{k\}\-\\tau\_\{k\}x\_\{k\+1\}\-\\tau\_\{k\}\\varepsilon\_\{k\}\\nabla\\log\\rho\_\{k\+1\}\(x\_\{k\+1\}\)\.Move the linear drift term to the left:

\(1\+τk\)​xk\+1=xk−τk​εk​∇log⁡ρk\+1​\(xk\+1\)\.\(1\+\\tau\_\{k\}\)x\_\{k\+1\}=x\_\{k\}\-\\tau\_\{k\}\\varepsilon\_\{k\}\\nabla\\log\\rho\_\{k\+1\}\(x\_\{k\+1\}\)\.Therefore

xk\+1=11\+τk​xk−τk​εk1\+τk​∇log⁡ρk\+1​\(xk\+1\)\.x\_\{k\+1\}=\\frac\{1\}\{1\+\\tau\_\{k\}\}x\_\{k\}\-\\frac\{\\tau\_\{k\}\\varepsilon\_\{k\}\}\{1\+\\tau\_\{k\}\}\\nabla\\log\\rho\_\{k\+1\}\(x\_\{k\+1\}\)\.\(63\)So thexk\+1x\_\{k\+1\}on the right has not disappeared: it has been absorbed into the prefactor11\+τk\\frac\{1\}\{1\+\\tau\_\{k\}\}after solving the implicit linear drift term\.

Reverse/generative implementation\.DDPM generation runs a reverse denoising update\. In the samek→k\+1k\\to k\+1notation, write the explicit score\-form update as

xk\+1=1αk​xk\+γk​sθ​\(xk,k\)\+σk​z\.x\_\{k\+1\}=\\frac\{1\}\{\\sqrt\{\\alpha\_\{k\}\}\}x\_\{k\}\+\\gamma\_\{k\}\\,s\_\{\\theta\}\(x\_\{k\},k\)\+\\sigma\_\{k\}z\.\(64\)Hereαk∈\(0,1\)\\alpha\_\{k\}\\in\(0,1\)is the usual DDPM schedule coefficient, andγk\\gamma\_\{k\}is the score coefficient\. To match the standard DDPM formula, the relation is

γk=1−αkαk\.\\boxed\{\\gamma\_\{k\}=\\frac\{1\-\\alpha\_\{k\}\}\{\\sqrt\{\\alpha\_\{k\}\}\}\.\}\(65\)
Important: what is the relation between the two scores?The implicit JKO equation contains

∇log⁡ρk\+1​\(xk\+1\),\\nabla\\log\\rho\_\{k\+1\}\(x\_\{k\+1\}\),where both the distribution and the point are the*new*unknowns\. The explicit DDPM\-style update uses

sθ​\(xk,k\)≈∇log⁡ρk​\(xk\),s\_\{\\theta\}\(x\_\{k\},k\)\\approx\\nabla\\log\\rho\_\{k\}\(x\_\{k\}\),the score at the currently available point and current marginal\. These two quantities are not exactly equal:

∇log⁡ρk\+1​\(xk\+1\)≠∇log⁡ρk​\(xk\)in general\.\\nabla\\log\\rho\_\{k\+1\}\(x\_\{k\+1\}\)\\neq\\nabla\\log\\rho\_\{k\}\(x\_\{k\}\)\\quad\\text\{in general\.\}The connection is anexplicit time\-discretization approximation: for small steps and a smooth density path,

ρk\+1≈ρk,xk\+1≈xk,\\rho\_\{k\+1\}\\approx\\rho\_\{k\},\\qquad x\_\{k\+1\}\\approx x\_\{k\},so one replaces the implicit new\-point score by the available old\-point score:

∇log⁡ρk\+1​\(xk\+1\)≈∇log⁡ρk​\(xk\)≈sθ​\(xk,k\)\.\\nabla\\log\\rho\_\{k\+1\}\(x\_\{k\+1\}\)\\approx\\nabla\\log\\rho\_\{k\}\(x\_\{k\}\)\\approx s\_\{\\theta\}\(x\_\{k\},k\)\.This is exactly the same conceptual move as replacing implicit Euler by an explicit Euler\-like implementation\. Thus DDPM’s practical update should not be read as an exact solution of the JKO subproblem; it is a score\-trained, explicit particle approximation to that implicit proximal step\.

Now introduce the score estimator and the noise\-prediction parameterization:

sθ​\(xk,k\)≈∇log⁡ρk​\(xk\),ϵθ​\(xk,k\)≈−1−α¯k​sθ​\(xk,k\)\.s\_\{\\theta\}\(x\_\{k\},k\)\\approx\\nabla\\log\\rho\_\{k\}\(x\_\{k\}\),\\qquad\\epsilon\_\{\\theta\}\(x\_\{k\},k\)\\approx\-\\sqrt\{1\-\\bar\{\\alpha\}\_\{k\}\}\\,s\_\{\\theta\}\(x\_\{k\},k\)\.Equivalently,

sθ​\(xk,k\)=−ϵθ​\(xk,k\)1−α¯k\.s\_\{\\theta\}\(x\_\{k\},k\)=\-\\frac\{\\epsilon\_\{\\theta\}\(x\_\{k\},k\)\}\{\\sqrt\{1\-\\bar\{\\alpha\}\_\{k\}\}\}\.Substitute this into \([64](https://arxiv.org/html/2606.24157#S6.E64)\):

xk\+1=1αk​xk−γk1−α¯k​ϵθ​\(xk,k\)\+σk​z\.x\_\{k\+1\}=\\frac\{1\}\{\\sqrt\{\\alpha\_\{k\}\}\}x\_\{k\}\-\\frac\{\\gamma\_\{k\}\}\{\\sqrt\{1\-\\bar\{\\alpha\}\_\{k\}\}\}\\epsilon\_\{\\theta\}\(x\_\{k\},k\)\+\\sigma\_\{k\}z\.Using \([65](https://arxiv.org/html/2606.24157#S6.E65)\),

γk1−α¯k=1−αkαk​1−α¯k\.\\frac\{\\gamma\_\{k\}\}\{\\sqrt\{1\-\\bar\{\\alpha\}\_\{k\}\}\}=\\frac\{1\-\\alpha\_\{k\}\}\{\\sqrt\{\\alpha\_\{k\}\}\\sqrt\{1\-\\bar\{\\alpha\}\_\{k\}\}\}\.Therefore

xk\+1=1αk​xk−1−αkαk​1−α¯k​ϵθ​\(xk,k\)\+σk​z\.x\_\{k\+1\}=\\frac\{1\}\{\\sqrt\{\\alpha\_\{k\}\}\}x\_\{k\}\-\\frac\{1\-\\alpha\_\{k\}\}\{\\sqrt\{\\alpha\_\{k\}\}\\sqrt\{1\-\\bar\{\\alpha\}\_\{k\}\}\}\\epsilon\_\{\\theta\}\(x\_\{k\},k\)\+\\sigma\_\{k\}z\.Factoring out1αk\\frac\{1\}\{\\sqrt\{\\alpha\_\{k\}\}\}gives the standard DDPM update:

xk\+1=1αk​\(xk−1−αk1−α¯k​ϵθ​\(xk,k\)\)\+σk​z\.x\_\{k\+1\}=\\frac\{1\}\{\\sqrt\{\\alpha\_\{k\}\}\}\\left\(x\_\{k\}\-\\frac\{1\-\\alpha\_\{k\}\}\{\\sqrt\{1\-\\bar\{\\alpha\}\_\{k\}\}\}\\epsilon\_\{\\theta\}\(x\_\{k\},k\)\\right\)\+\\sigma\_\{k\}z\.\(66\)Thus \([66](https://arxiv.org/html/2606.24157#S6.E66)\) should be read as anexplicit reverse\-time score implementation/approximationof the implicit JKO particle equation \([62](https://arxiv.org/html/2606.24157#S6.E62)\)\. The implicit new\-point score∇log⁡ρk\+1​\(xk\+1\)\\nabla\\log\\rho\_\{k\+1\}\(x\_\{k\+1\}\)is replaced by the learned score at the currently available pointxkx\_\{k\}, and the coefficient relationγk=\(1−αk\)/αk\\gamma\_\{k\}=\(1\-\\alpha\_\{k\}\)/\\sqrt\{\\alpha\_\{k\}\}converts the score\-form update into the standardϵθ\\epsilon\_\{\\theta\}\-form update\. Thus DDPM is a stochastic particle implementation of repeated JKO proximal steps: the score implements the entropy force, the Gaussian drift implements the reference potential, and the randomzzimplements finite\-temperature Langevin sampling inside each proximal step\.

#### 2\. DDIM: deterministic JKO transport using the same learned score

DDIM \(Song, Meng & Ermon, 2021\) uses the same learned score but removes the fresh stochastic sampling noise\. In JKO language, it keeps the deterministic transport part of the proximal map:

xk\+1≈xk\+τk​\[−∇Vk​\(xk\+1\)−εk​sθ​\(xk\+1,k\+1\)\]\.x\_\{k\+1\}\\approx x\_\{k\}\+\\tau\_\{k\}\[\-\\nabla V\_\{k\}\(x\_\{k\+1\}\)\-\\varepsilon\_\{k\}s\_\{\\theta\}\(x\_\{k\+1\},k\+1\)\]\.In the DDPM parameterization, this gives the deterministic DDIM\-style update

xnext=α¯next​x^0​\(xt,t\)\+1−α¯next​ϵθ​\(xt,t\),x\_\{\\mathrm\{next\}\}=\\sqrt\{\\bar\{\\alpha\}\_\{\\mathrm\{next\}\}\}\\,\\hat\{x\}\_\{0\}\(x\_\{t\},t\)\+\\sqrt\{1\-\\bar\{\\alpha\}\_\{\\mathrm\{next\}\}\}\\,\\epsilon\_\{\\theta\}\(x\_\{t\},t\),wherex^0\\hat\{x\}\_\{0\}is the usual estimate of the clean sample from\(xt,ϵθ\)\(x\_\{t\},\\epsilon\_\{\\theta\}\)\. So DDIM is not a different free energy; it is a deterministic implementation of the same learned JKO displacement field\. This explains why DDIM is deterministic and can use larger step sizes: it follows a smoother transport map rather than repeatedly injecting fresh noise\.

#### 3\. NCSN/SMLD: pure\-entropy JKO at multiple scales

JKO choice: setV=0V=0and use only entropy,

ℱk​\(ρ\)=εk​∫ρ​log⁡ρ​d​x\.\\mathcal\{F\}\_\{k\}\(\\rho\)=\\varepsilon\_\{k\}\\int\\rho\\log\\rho\\,dx\.The proximal step is

ρk\+1=arg⁡minρ⁡\{εk​∫ρ​log⁡ρ​d​x\+12​τk​W22​\(ρ,ρk\)\},\\rho\_\{k\+1\}=\\arg\\min\_\{\\rho\}\\left\\\{\\varepsilon\_\{k\}\\int\\rho\\log\\rho\\,dx\+\\frac\{1\}\{2\\tau\_\{k\}\}W\_\{2\}^\{2\}\(\\rho,\\rho\_\{k\}\)\\right\\\},with optimality condition

x−Tk​\(x\)τk=−εk​∇log⁡ρk\+1​\(x\),x∼ρk\+1\.\\frac\{x\-T\_\{k\}\(x\)\}\{\\tau\_\{k\}\}=\-\\varepsilon\_\{k\}\\nabla\\log\\rho\_\{k\+1\}\(x\),\\qquad x\\sim\\rho\_\{k\+1\}\.NCSN/SMLD learns the score at a sequence of noise scalesσ1\>⋯\>σN\\sigma\_\{1\}\>\\cdots\>\\sigma\_\{N\}\. Replacing∇log⁡ρk\+1\\nabla\\log\\rho\_\{k\+1\}bysθ​\(x,σk\+1\)s\_\{\\theta\}\(x,\\sigma\_\{k\+1\}\)gives the annealed Langevin/JKO particle update

xk\+1=xk\+ηk​sθ​\(xk\+1,σk\+1\)\+2​ηk​z,x\_\{k\+1\}=x\_\{k\}\+\\eta\_\{k\}s\_\{\\theta\}\(x\_\{k\+1\},\\sigma\_\{k\+1\}\)\+\\sqrt\{2\\eta\_\{k\}\}\\,z,Largeσk\\sigma\_\{k\}corresponds to coarse entropy\-driven moves; smallσk\\sigma\_\{k\}gives fine denoising near the data manifold\.

#### 4\. Energy Matching: learning the JKO energy directly

Energy Matching keeps the JKO iteration explicit:

ρk\+1=arg⁡minρ⁡\{∫Vθ​𝑑ρ\+εk​∫ρ​log⁡ρ​d​x\+12​τk​W22​\(ρ,ρk\)\}\.\\rho\_\{k\+1\}=\\arg\\min\_\{\\rho\}\\left\\\{\\int V\_\{\\theta\}\\,d\\rho\+\\varepsilon\_\{k\}\\int\\rho\\log\\rho\\,dx\+\\frac\{1\}\{2\\tau\_\{k\}\}W\_\{2\}^\{2\}\(\\rho,\\rho\_\{k\}\)\\right\\\}\.Unlike DDPM/NCSN, which learn the score∇log⁡ρk\+1\\nabla\\log\\rho\_\{k\+1\}, Energy Matching learns the scalar energyVθV\_\{\\theta\}itself\. The sameVθV\_\{\\theta\}plays two roles: whenεk\\varepsilon\_\{k\}is small,−∇Vθ\-\\nabla V\_\{\\theta\}acts as the transport velocity; near equilibrium,VθV\_\{\\theta\}defines the Boltzmann densityρ∝e−Vθ/εk\\rho\\propto e^\{\-V\_\{\\theta\}/\\varepsilon\_\{k\}\}\.

#### 5\. Flow Matching: Benamou–Brenier OT paths

Choice: Do not define the target distribution through a potential energyVVor a free energyℱ\\mathcal\{F\}\. Instead, prescribe a source distributionρ0\\rho\_\{0\}\(noise\), a target distributionρ1\\rho\_\{1\}\(data\), and learn a velocity field that transports one into the other\.

Important distinction: Flow Matching is not literally the JKO scheme withℱ=0\\mathcal\{F\}=0\. If we putℱ=0\\mathcal\{F\}=0into the standard JKO update

ρk\+1=arg⁡minρ⁡\{ℱ​\(ρ\)\+12​τ​W22​\(ρ,ρk\)\},\\rho\_\{k\+1\}=\\arg\\min\_\{\\rho\}\\left\\\{\\mathcal\{F\}\(\\rho\)\+\\frac\{1\}\{2\\tau\}W\_\{2\}^\{2\}\(\\rho,\\rho\_\{k\}\)\\right\\\},then the minimizer is simplyρk\+1=ρk\\rho\_\{k\+1\}=\\rho\_\{k\}—nothing moves\. A Wasserstein geodesic appears only after we also prescribe the endpointρ1\\rho\_\{1\}\. Thus OT Flow Matching is better viewed through theBenamou–Brenier dynamic OT formulation:

W22​\(ρ0,ρ1\)=inf∂tρt\+∇⋅\(ρt​vt\)=0ρ0,ρ1​fixed∫01∫\|vt​\(x\)\|2​ρt​\(x\)​𝑑x​𝑑t\.W\_\{2\}^\{2\}\(\\rho\_\{0\},\\rho\_\{1\}\)=\\inf\_\{\\begin\{subarray\}\{c\}\\partial\_\{t\}\\rho\_\{t\}\+\\nabla\\cdot\(\\rho\_\{t\}v\_\{t\}\)=0\\\\ \\rho\_\{0\},\\rho\_\{1\}\\ \\text\{fixed\}\\end\{subarray\}\}\\int\_\{0\}^\{1\}\\\!\\int\|v\_\{t\}\(x\)\|^\{2\}\\rho\_\{t\}\(x\)\\,dx\\,dt\.
The geodesic fromρ0\\rho\_\{0\}toρ1\\rho\_\{1\}in Wasserstein space is precisely theMcCann displacement interpolation:

ρt=\[\(1−t\)​id\+t​T\]\#​ρ0,T=∇φ​\(Brenier map\)\\rho\_\{t\}=\[\(1\-t\)\\mathrm\{id\}\+t\\,T\]\_\{\\\#\}\\rho\_\{0\},\\quad T=\\nabla\\varphi\\;\\text\{\(Brenier map\)\}Each particle moves at constant speed along the straight lineXt=\(1−t\)​x0\+t​T​\(x0\)X\_\{t\}=\(1\-t\)x\_\{0\}\+t\\,T\(x\_\{0\}\), with velocityvt​\(Xt\)=T​\(x0\)−x0v\_\{t\}\(X\_\{t\}\)=T\(x\_\{0\}\)\-x\_\{0\}\.

From dynamic OT to Flow Matching: In practice,TTis unknown, so Flow Matching approximates it usingconditional optimal transport\(mini\-batch OT\):

vt​\(x\)≈x1−x01,x0∼ρ0,x1∼ρ1,\(x0,x1\)​paired by OTv\_\{t\}\(x\)\\approx\\frac\{x\_\{1\}\-x\_\{0\}\}\{1\},\\quad x\_\{0\}\\sim\\rho\_\{0\},\\;x\_\{1\}\\sim\\rho\_\{1\},\\;\(x\_\{0\},x\_\{1\}\)\\text\{ paired by OT\}A neural networkvθ​\(x,t\)v\_\{\\theta\}\(x,t\)is trained to fit this velocity field:

ℒFM=𝔼t,\(x0,x1\)​\[‖vθ​\(xt,t\)−\(x1−x0\)‖2\],xt=\(1−t\)​x0\+t​x1\\boxed\{\\mathcal\{L\}\_\{\\text\{FM\}\}=\\mathbb\{E\}\_\{t,\(x\_\{0\},x\_\{1\}\)\}\\left\[\\\|v\_\{\\theta\}\(x\_\{t\},t\)\-\(x\_\{1\}\-x\_\{0\}\)\\\|^\{2\}\\right\],\\quad x\_\{t\}=\(1\-t\)x\_\{0\}\+t\\,x\_\{1\}\}
Variational perspective: OT Flow Matching learns the velocity field of the Benamou–Brenier minimizer\. There is no free energy whose gradient is being followed, and there is no equilibrium distributionπ∝e−V\\pi\\propto e^\{\-V\}\. The target information enters through the boundary conditionρ1=pdata\\rho\_\{1\}=p\_\{\\text\{data\}\}, not through a potentialVV\.

Same endpoint, different path: In the idealized limit of infinite model capacity, perfect training, and exact numerical integration, both diffusion models and Flow Matching can transport a simple base distribution to the data distribution\. The endpoint can be the same:

ρ1=pdata\.\\rho\_\{1\}=p\_\{\\text\{data\}\}\.But the paths are different\. Diffusion constructs a stochastic noising/denoising process and follows the free\-energy/Fokker–Planck geometry; OT Flow Matching directly learns a deterministic transport path, often close to a Wasserstein geodesic\. Thus Flow Matching removes some SDE machinery that is unnecessary if the only goal is to learn a transport map, while diffusion retains useful extra structure: scores, reverse SDEs, Langevin sampling, likelihood tools, and inverse\-problem machinery\.

Can we recover a potential from a learned Flow Matching velocity?Suppose a Flow Matching model gives a velocity fieldvt​\(x\)v\_\{t\}\(x\)and the induced density pathpt​\(x\)p\_\{t\}\(x\)\. If we try to reinterpret it as a free\-energy gradient flow with entropy coefficientε\\varepsilon, we would need

vt​\(x\)=−∇Vt​\(x\)−ε​∇log⁡pt​\(x\)\.v\_\{t\}\(x\)=\-\\nabla V\_\{t\}\(x\)\-\\varepsilon\\nabla\\log p\_\{t\}\(x\)\.Formally this means

∇Vt​\(x\)=−vt​\(x\)−ε​∇log⁡pt​\(x\)\.\\nabla V\_\{t\}\(x\)=\-v\_\{t\}\(x\)\-\\varepsilon\\nabla\\log p\_\{t\}\(x\)\.This defines a scalar potentialVtV\_\{t\}only if the right\-hand side is a conservative vector field:

∇×\(−vt−ε​∇log⁡pt\)=0\.\\nabla\\times\\bigl\(\-v\_\{t\}\-\\varepsilon\\nabla\\log p\_\{t\}\\bigr\)=0\.Since∇log⁡pt\\nabla\\log p\_\{t\}is already a gradient field, the obstruction is the rotational component ofvtv\_\{t\}\. A general neural velocity field need not be curl\-free, so a scalar potential may not exist\. Even when it exists, it is usually atime\-dependent effective potentialVt​\(x\)V\_\{t\}\(x\), not a fixed equilibrium potentialV​\(x\)V\(x\)withπ∝e−V\\pi\\propto e^\{\-V\}\.

#### Unified summary

Key insight: JKO and Benamou–Brenier are two complementary variational principles on Wasserstein space:

1. 1\.JKOis an initial\-value problem: givenρk\\rho\_\{k\}, descend a free energyℱ\\mathcal\{F\}while staying close inW2W\_\{2\}\.
2. 2\.Benamou–Brenieris a boundary\-value problem: given bothρ0\\rho\_\{0\}andρ1\\rho\_\{1\}, minimize kinetic energy among all mass\-preserving paths\.
3. 3\.Diffusion modelslive naturally in the JKO/free\-energy picture;OT Flow Matchinglives naturally in the Benamou–Brenier/geodesic picture\.
4. 4\.Probability flow ODEsare the deterministic continuity\-equation representation of diffusion: the entropy\-driven diffusion term is absorbed into the velocity field through the score\.

## 7The Big Picture

Let us review the logical chain of the entire story:

Summary1\.Wasserstein distanceW2W\_\{2\}gives𝒫2​\(ℝd\)\\mathcal\{P\}\_\{2\}\(\\mathbb\{R\}^\{d\}\)the structure of a metric space — in fact, a \(formal\) Riemannian manifold\.2\.Thecontinuity equation∂tρ\+∇⋅\(ρ​v\)=0\\partial\_\{t\}\\rho\+\\nabla\\cdot\(\\rho v\)=0describes how probability mass flows\. It is the equation of motion in Wasserstein space\. It does not require an energy functional; it only requires a velocity field\.3\.TheBenamou–Brenier formulareveals thatW2W\_\{2\}= geodesic distance in this Riemannian structure, with kinetic energy∫\|v\|2​ρ​𝑑x\\int\|v\|^\{2\}\\rho\\,dxas the metric\.4\.TheFokker–Planck equationis the Wasserstein gradient flow of the free energyℱ​\(ρ\)=KL​\(ρ∥π\)\\mathcal\{F\}\(\\rho\)=\\mathrm\{KL\}\(\\rho\\\|\\pi\):“​ρ˙=−gradW​ℱ​\(ρ\)​”\\text\{\`\`\}\\dot\{\\rho\}=\-\\mathrm\{grad\}\_\{W\}\\mathcal\{F\}\(\\rho\)\\text\{''\}This is the infinite\-dimensional analog ofx˙=−∇f​\(x\)\\dot\{x\}=\-\\nabla f\(x\)\. In this special case, the velocity is not arbitrary: it is determined by the free energy\.5\.TheJKO schemediscretizes this gradient flow:ρk\+1=arg⁡minρ⁡\{ℱ​\(ρ\)\+W22​\(ρ,ρk\)2​τ\}\\rho\_\{k\+1\}=\\arg\\min\_\{\\rho\}\\left\\\{\\mathcal\{F\}\(\\rho\)\+\\frac\{W\_\{2\}^\{2\}\(\\rho,\\rho\_\{k\}\)\}\{2\\tau\}\\right\\\}Asτ→0\\tau\\to 0, it recovers the continuous\-time Fokker–Planck equation\.

One\-sentence summary:

Wasserstein distance endows probability space with geometry;the continuity equation describes mass\-conserving flow;the Fokker–Planck equation is the free\-energy\-driven special case;Flow Matching learns a transport velocity, often the OT geodesic velocity;the JKO scheme discretizes this continuous “descent” into step\-by\-step optimization problems\.

Wasserstein geometryW2W\_\{2\}\+ continuity equationFree\-energygradient flowBenamou–Brenierminimum actionFokker–PlanckJKOimplicit EulerOT Flow Matchinggeodesic pathτ→0\\tau\\to 0
## References and Further Reading

### Optimal transport and gradient flows

1. 1\.Monge, G\. \(1781\)\.Mémoire sur la théorie des déblais et des remblais\.Histoire de l’Académie Royale des Sciences de Paris, 666–704\.
2. 2\.Kantorovich, L\. V\. \(1942\)\.On the translocation of masses\.Dokl\. Akad\. Nauk SSSR37, 199–201\.
3. 3\.Brenier, Y\. \(1991\)\.Polar factorization and monotone rearrangement of vector\-valued functions\.Comm\. Pure Appl\. Math\.44\(4\), 375–417\.
4. 4\.Vaserstein, L\. N\. \(1969\)\.Markov processes over denumerable products of spaces, describing large systems of automata\.Problemy Peredachi Informatsii5\(3\), 64–72\. \[The source of the name “Wasserstein distance\.”\]
5. 5\.McCann, R\. J\. \(1997\)\.A convexity principle for interacting gases\.Adv\. Math\.128\(1\), 153–179\. \[Displacement interpolation\.\]
6. 6\.Benamou, J\.\-D\. & Brenier, Y\. \(2000\)\.A computational fluid mechanics solution to the Monge–Kantorovich mass transfer problem\.Numer\. Math\.84\(3\), 375–393\.
7. 7\.Jordan, R\., Kinderlehrer, D\., & Otto, F\. \(1998\)\.The variational formulation of the Fokker–Planck equation\.SIAM J\. Math\. Anal\.29\(1\), 1–17\.
8. 8\.Otto, F\. \(2001\)\.The geometry of dissipative evolution equations: the porous medium equation\.Comm\. PDE26\(1\-2\), 101–174\.
9. 9\.Ambrosio, L\., Gigli, N\., & Savaré, G\. \(2008\)\.Gradient Flows in Metric Spaces and in the Space of Probability Measures\.Birkhäuser\. \[The definitive reference\.\]
10. 10\.Villani, C\. \(2003\)\.Topics in Optimal Transportation\.AMS\. \[Excellent introduction\.\]
11. 11\.Santambrogio, F\. \(2015\)\.Optimal Transport for Applied Mathematicians\.Birkhäuser\. \[Very readable\.\]
12. 12\.Peyré, G\. & Cuturi, M\. \(2019\)\.Computational Optimal Transport\.Found\. Trends ML11\(5\-6\), 355–607\. \[Computational perspective\.\]
13. 13\.Lavenant, H\. & Santambrogio, F\. \(2022\)\.The flow map of the Fokker–Planck equation does not provide optimal transport\.Appl\. Math\. Lett\.133, 108225\.
14. 14\.Mokrov, P\., Korotin, A\., Li, L\., Genevay, A\., Solomon, J\., & Burnaev, E\. \(2021\)\.Large\-scale Wasserstein gradient flows\.NeurIPS\.
15. 15\.Xu, C\., Cheng, X\., & Xie, Y\. \(2023\)\.Normalizing flow neural networks by JKO scheme\.NeurIPS\. \[JKO\-iFlow\.\]

### Generative models

1. 1\.Hyvärinen, A\. \(2005\)\.Estimation of non\-normalized statistical models by score matching\.J\. Mach\. Learn\. Res\.6, 695–709\.
2. 2\.Vincent, P\. \(2011\)\.A connection between score matching and denoising autoencoders\.Neural Comput\.23\(7\), 1661–1674\.
3. 3\.Welling, M\. & Teh, Y\. W\. \(2011\)\.Bayesian learning via stochastic gradient Langevin dynamics\.ICML, 681–688\.
4. 4\.Song, Y\. & Ermon, S\. \(2019\)\.Generative modeling by estimating gradients of the data distribution\.NeurIPS\. \[NCSN/SMLD\.\]
5. 5\.Ho, J\., Jain, A\., & Abbeel, P\. \(2020\)\.Denoising diffusion probabilistic models\.NeurIPS\. \[DDPM\.\]
6. 6\.Song, J\., Meng, C\., & Ermon, S\. \(2021\)\.Denoising diffusion implicit models\.ICLR\. \[DDIM\.\]
7. 7\.Song, Y\., Sohl\-Dickstein, J\., Kingma, D\. P\., Kumar, A\., Ermon, S\., & Poole, B\. \(2021\)\.Score\-based generative modeling through stochastic differential equations\.ICLR\.
8. 8\.Karras, T\., Aittala, M\., Aila, T\., & Laine, S\. \(2022\)\.Elucidating the design space of diffusion\-based generative models\.NeurIPS\. \[EDM\.\]
9. 9\.Lipman, Y\., Chen, R\. T\. Q\., Ben\-Hamu, H\., Nickel, M\., & Le, M\. \(2023\)\.Flow matching for generative modeling\.ICLR\.
10. 10\.Balcerak, M\., et al\. \(2025\)\.Energy Matching: unifying flow matching and energy\-based models for generative modeling\.arXiv:2504\.10612 \(NeurIPS 2025\)\.
11. 11\.Wang, R\., et al\. \(2025\)\.Equilibrium Matching: generative modeling with implicit energy\-based models\.arXiv:2510\.02300\.
12. 12\.Albergo, M\. S\., Boffi, N\. M\., & Vanden\-Eijnden, E\. \(2025\)\.Stochastic interpolants: a unifying framework for flows and diffusions\.J\. Mach\. Learn\. Res\.26\(209\), 1–80\.
13. 13\.Lipman, Y\., et al\. \(2024\)\.Flow Matching Guide and Code\.arXiv:2412\.06264\.
14. 14\.Vuong, A\. B\., McCann, M\. T\., Santos, J\. E\., & Lin, Y\. T\. \(2025\)\.Are we really learning the score function? Reinterpreting diffusion models through Wasserstein gradient flow matching\.CIKM\. \(arXiv:2509\.00336\)
15. 15\.Holderrieth, P\. & Erives, E\. \(2025\)\.An Introduction to Flow Matching and Diffusion Models\.MIT 6\.S184 lecture notes\.

### Classical foundations

1. 1\.Fisher, R\. A\. \(1925\)\.Theory of statistical estimation\.Proc\. Cambridge Philos\. Soc\.22, 700–725\.
2. 2\.Langevin, P\. \(1908\)\.Sur la théorie du mouvement brownien\.C\. R\. Acad\. Sci\. Paris146, 530–533\.
3. 3\.Uhlenbeck, G\. E\. & Ornstein, L\. S\. \(1930\)\.On the theory of the Brownian motion\.Phys\. Rev\.36\(5\), 823–841\.
4. 4\.Kolmogorov, A\. N\. \(1931\)\.Über die analytischen Methoden in der Wahrscheinlichkeitsrechnung\.Math\. Ann\.104, 415–458\. \[Kolmogorov forward/backward equations\.\]
5. 5\.Itô, K\. \(1951\)\.On stochastic differential equations\.Mem\. Amer\. Math\. Soc\.4, 1–51\.
6. 6\.Kullback, S\. & Leibler, R\. A\. \(1951\)\.On information and sufficiency\.Ann\. Math\. Statist\.22\(1\), 79–86\.
7. 7\.McKean, H\. P\. \(1966\)\.A class of Markov processes associated with nonlinear parabolic equations\.Proc\. Natl\. Acad\. Sci\. USA56\(6\), 1907–1911\. \[McKean–Vlasov\.\]
8. 8\.Anderson, B\. D\. O\. \(1982\)\.Reverse\-time diffusion equation models\.Stochastic Process\. Appl\.12\(3\), 313–326\.

## Appendix AMeasures and Couplings

Before introducing the Wasserstein distance, we need to clarify two foundational concepts:measuresandcouplings\. If you are already familiar with measure theory, you may skip this section\.

### A\.1What is a measure?

The most intuitive understanding:A measure is a way of “assigning size to sets\.”

There are many notions of “size” that we use in everyday life:

- •Thelengthof a line segment:\[0,3\]\[0,3\]has length 3
- •Theareaof a region: a circle of radiusrrhas areaπ​r2\\pi r^\{2\}
- •Thevolumeof a solid: a cube has volumel3l^\{3\}

These are all special cases of the Lebesgue measure\. But measures go far beyond these—they provide a unified framework for expressing various ways of “distributing mass\.”

###### Definition A\.1\(Measure — informal\)\.

A measureμ\\muonℝd\\mathbb\{R\}^\{d\}is a function that assigns a non\-negative numberμ​\(A\)≥0\\mu\(A\)\\geq 0to each \(measurable\) subsetA⊆ℝdA\\subseteq\\mathbb\{R\}^\{d\}, satisfying:

1. 1\.μ​\(∅\)=0\\mu\(\\emptyset\)=0\(empty set has zero measure\)
2. 2\.Countable additivity:IfA1,A2,…A\_\{1\},A\_\{2\},\\ldotsare disjoint, thenμ​\(⋃i=1∞Ai\)=∑i=1∞μ​\(Ai\)\\mu\\\!\\left\(\\bigcup\_\{i=1\}^\{\\infty\}A\_\{i\}\\right\)=\\sum\_\{i=1\}^\{\\infty\}\\mu\(A\_\{i\}\)

Aprobability measureadditionally satisfiesμ​\(ℝd\)=1\\mu\(\\mathbb\{R\}^\{d\}\)=1\(total mass is 1\)\.

Three measures you must know:

1. \(1\)Lebesgue measureλ\\lambda: This is simply “volume\.”λ​\(\[0,2\]×\[0,3\]\)=6\\lambda\(\[0,2\]\\times\[0,3\]\)=6\. It uniformly assigns “mass” to every region of space\.
2. \(2\)Dirac measureδx\\delta\_\{x\}: Concentrates all mass at a single pointxx: δx​\(A\)=\{1if​x∈A0if​x∉A\\delta\_\{x\}\(A\)=\\begin\{cases\}1&\\text\{if \}x\\in A\\\\ 0&\\text\{if \}x\\notin A\\end\{cases\}Think of it as a single grain of sand placed at positionxx\.
3. \(3\)Absolutely continuous measure: A measure with a density functionρ​\(x\)\\rho\(x\):μ​\(A\)=∫Aρ​\(x\)​𝑑x\\mu\(A\)=\\int\_\{A\}\\rho\(x\)\\,dx\. For example, the normal distribution𝒩​\(0,1\)\\mathcal\{N\}\(0,1\):μ​\(A\)=∫A12​π​e−x2/2​𝑑x\\mu\(A\)=\\int\_\{A\}\\frac\{1\}\{\\sqrt\{2\\pi\}\}e^\{\-x^\{2\}/2\}\\,dx\. Intuition:ρ​\(x\)\\rho\(x\)describes the “concentration” of mass in space—whereρ​\(x\)\\rho\(x\)is large the mass is dense, whereρ​\(x\)\\rho\(x\)is small the mass is sparse\.

IntuitionMeasures vs\. density functions:Beginners often ask: “What is the difference between a measure and a probability density function \(PDF\)?”The answer:a PDF is one way of representing a measure, but not every measure has a PDF\.•δx\\delta\_\{x\}\(a single grain of sand\) has no PDF—you cannot write down a functionρ​\(x\)\\rho\(x\)such that∫Aρ​𝑑x=δx​\(A\)\\int\_\{A\}\\rho\\,dx=\\delta\_\{x\}\(A\)•The normal distribution does have a PDF:ρ​\(x\)=12​π​e−x2/2\\rho\(x\)=\\frac\{1\}\{\\sqrt\{2\\pi\}\}e^\{\-x^\{2\}/2\}•The mixture12​δ0\+12​𝒩​\(1,1\)\\frac\{1\}\{2\}\\delta\_\{0\}\+\\frac\{1\}\{2\}\\mathcal\{N\}\(1,1\)also has no PDF \(in the classical sense\)Working with measures instead of PDFs allows us to handle all these cases in a unified framework\.

### A\.2Absolutely continuous: what does it really mean?

In optimal transport, you will repeatedly encounter the condition: “μ\\muis absolutely continuous with respect to the Lebesgue measure\.” What does this mean? Why is it important?

###### Definition A\.2\(Absolute continuity of measures\)\.

A measureμ\\muisabsolutely continuouswith respect to another measureλ\\lambda\(writtenμ≪λ\\mu\\ll\\lambda\) if:

λ​\(A\)=0⟹μ​\(A\)=0for all measurable sets​A\\lambda\(A\)=0\\implies\\mu\(A\)=0\\quad\\text\{for all measurable sets \}AIn words: every set that is “invisible” toλ\\lambdais also “invisible” toμ\\mu\.

In plain language:If a set has “zero volume” \(λ​\(A\)=0\\lambda\(A\)=0\), thenμ\\mualso places no mass on it \(μ​\(A\)=0\\mu\(A\)=0\)\.

IntuitionEveryday analogy: Spreading butter vs\. placing pebblesImagine distributing something on a slice of bread \(ℝd\\mathbb\{R\}^\{d\}space\)\. There are two ways:•Spreading butter\(absolutely continuous\): Butter is spread uniformly or non\-uniformly on the bread surface\. If you cut along a line of zero area, there will be no butter on it—because butter is “continuously spread,” and no area means no butter\.•Placing pebbles\(not absolutely continuous\): You place several pebbles on the bread\. The pebbles occupy “zero area” \(points have no area\), yet they have finite mass\. This violates absolute continuity—there exists a set of zero area \(the points where the pebbles sit\), yet it carries positive mass\.The essence of absolute continuity is: mass is spread like butter across space, not piled up like pebbles at isolated points\.

A finer understanding: What are “sets of zero volume”?

In spaces of different dimensions, “zero volume” means different things:

- •Inℝ1\\mathbb\{R\}^\{1\}\(one\-dimensional/line\): Points, countable sets \(e\.g\., the rationalsℚ\\mathbb\{Q\}\), and the Cantor set are all null sets
- •Inℝ2\\mathbb\{R\}^\{2\}\(two\-dimensional/plane\): Points, line segments, and curves are all null sets \(they have no “area”\)
- •Inℝ3\\mathbb\{R\}^\{3\}\(three\-dimensional/space\): Points, line segments, and surfaces are all null sets \(they have no “volume”\)

General rule: Inℝd\\mathbb\{R\}^\{d\}, subsets of dimension<d<dare typically null\.

Key corollary:An absolutely continuous measureμ\\muonℝd\\mathbb\{R\}^\{d\}satisfiesμ​\(\{x\}\)=0\\mu\(\\\{x\\\}\)=0for all individual pointsxx\. That is,the mass at any single point is zero—there are no “atoms\.” Mass must be “spread out” over regions with positive volume\.

###### Theorem A\.3\(Radon–Nikodym\)\.

μ≪λ\\mu\\ll\\lambdaif and only if there exists a non\-negative measurable functionρ:ℝd→\[0,∞\)\\rho:\\mathbb\{R\}^\{d\}\\to\[0,\\infty\)such that:

μ​\(A\)=∫Aρ​\(x\)​𝑑xfor all measurable​A\\mu\(A\)=\\int\_\{A\}\\rho\(x\)\\,dx\\quad\\text\{for all measurable \}A\(67\)The functionρ\\rhois called theRadon–Nikodym derivative\(ordensity\) ofμ\\muw\.r\.t\.λ\\lambda, writtenρ=d​μd​λ\\rho=\\frac\{d\\mu\}\{d\\lambda\}\.

Significance of the Radon–Nikodym theorem:

“μ\\muis absolutely continuous”⇔\\iff“μ\\muhas a density functionρ​\(x\)\\rho\(x\)”

These two statements arecompletely equivalent\! So whenever you see “μ\\muis absolutely continuous with respect to Lebesgue measure,” you can mentally translate it to “μ\\muhas a PDFρ​\(x\)\\rho\(x\)\.”

Intuition for the densityρ​\(x\)\\rho\(x\):ρ​\(x\)\\rho\(x\)describes how much mass is contained per unit volume nearxx:

μ​\(small ballBϵ​\(x\)centered atx\)≈ρ​\(x\)⋅Vol​\(Bϵ​\(x\)\)\\mu\(\\text\{small ball $B\_\{\\epsilon\}\(x\)$ centered at $x$\}\)\\approx\\rho\(x\)\\cdot\\text\{Vol\}\(B\_\{\\epsilon\}\(x\)\)Whereρ​\(x\)\\rho\(x\)is large the mass is dense, whereρ​\(x\)\\rho\(x\)is small the mass is sparse, and whereρ​\(x\)=0\\rho\(x\)=0there is no mass at all\. The key point is:no matter how largeρ​\(x\)\\rho\(x\)is, the mass at a single point\{x\}\\\{x\\\}is always zero\(becauseVol​\(\{x\}\)=0\\text\{Vol\}\(\\\{x\\\}\)=0\)\.

Caution: “Absolutely continuous” can be confused with another concept

In mathematics, “absolutely continuous” hastwo distinct usages:

1. 1\.Absolute continuity of measures\(what we are discussing\):μ≪λ\\mu\\ll\\lambda, meaning one measure does not concentrate on null sets of another\.
2. 2\.Absolute continuity of functions: A functionf:\[a,b\]→ℝf:\[a,b\]\\to\\mathbb\{R\}is absolutely continuous, roughly meaning “the variation offfcan be controlled by integration”—weaker than Lipschitz, stronger than uniform continuity\.

The two are deeply connected \(ffis absolutely continuous⇔\\ifff′f^\{\\prime\}exists a\.e\. andf​\(x\)=f​\(a\)\+∫axf′​\(t\)​𝑑tf\(x\)=f\(a\)\+\\int\_\{a\}^\{x\}f^\{\\prime\}\(t\)dt, which is essentially saying the measure induced byffis absolutely continuous with respect to Lebesgue measure\), but in the context of optimal transport we always mean the first\.

Why does optimal transport care about absolute continuity?

Recall the Brenier theorem: Ifμ\\muis absolutely continuous \(has a density\), then the optimal transport mapTTexists and is unique\. Why is this condition needed? The intuition is as follows:

Core idea: “Infinitesimal mass does not need to be split”

A mapTTrequires thatallmass at positionxxgoes to the same destinationT​\(x\)T\(x\)\.

- •Ifμ\\muhas an “atom” \(finite massμ​\(\{x0\}\)\>0\\mu\(\\\{x\_\{0\}\\\}\)\>0at some pointx0x\_\{0\}\), andν\\nurequires this mass to be distributed to multiple locations, then a mapTTcannot achieve this— sinceT​\(x0\)T\(x\_\{0\}\)can only be a single point\.
- •But ifμ\\muis absolutely continuous, then every pointxxcarries only “infinitesimal” massρ​\(x\)​d​x\\rho\(x\)dx\. This infinitesimal mass going entirely to one destinationT​\(x\)T\(x\)is perfectly fine— since each portion of mass from the source is infinitesimal, there is no need to “split” it for the mapTTto transport it to a unique destination\. The targetν\\nuis formed by “accumulating” infinitesimal masses sent from many differentxx\. Note:ν\\nuitselfneed notbe absolutely continuous—the Brenier theorem only requiresμ\\muto be absolutely continuous\.

Analogy: Distributing fruit

- •10 apples \(discrete/atomic\) must be divided among 3 people, each requiring 3\.33—cannot divide into integers, must cut \(split\)\.
- •10 liters of juice \(continuous/absolutely continuous\) must be divided among 3 people, each requiring 3\.33 liters—easy to divide, liquid can be poured to arbitrary precision\.

An absolutely continuous measure is like a liquid: it can be “poured” by the mapTTto any destination without needing to be “cut\.”

###### Example A\.4\.

Classification exercises:

1. 1\.μ=Uniform​\[0,1\]\\mu=\\text\{Uniform\}\[0,1\]: absolutely continuous ✓ \(densityρ​\(x\)=𝟏\[0,1\]​\(x\)\\rho\(x\)=\\mathbf\{1\}\_\{\[0,1\]\}\(x\)\)
2. 2\.μ=𝒩​\(0,1\)\\mu=\\mathcal\{N\}\(0,1\): absolutely continuous ✓ \(densityρ​\(x\)=12​π​e−x2/2\\rho\(x\)=\\frac\{1\}\{\\sqrt\{2\\pi\}\}e^\{\-x^\{2\}/2\}\)
3. 3\.μ=δ0\\mu=\\delta\_\{0\}: not absolutely continuous×\\times\(mass concentrated on the null set\{0\}\\\{0\\\}\)
4. 4\.μ=12​δ0\+12​Uniform​\[0,1\]\\mu=\\frac\{1\}\{2\}\\delta\_\{0\}\+\\frac\{1\}\{2\}\\text\{Uniform\}\[0,1\]: not absolutely continuous×\\times\(has an atom at\{0\}\\\{0\\\}\)
5. 5\.The measure onℝ2\\mathbb\{R\}^\{2\}uniformly distributed on the segment\{\(x,0\):x∈\[0,1\]\}\\\{\(x,0\):x\\in\[0,1\]\\\}: not absolutely continuous×\\times \(The segment has Lebesgue measure zero inℝ2\\mathbb\{R\}^\{2\}, yet this measure places all its mass there\)
6. 6\.μ\\muwith densityρ​\(x\)=1\|x\|1/2​𝟏\[−1,1\]​\(x\)\\rho\(x\)=\\frac\{1\}\{\|x\|^\{1/2\}\}\\mathbf\{1\}\_\{\[\-1,1\]\}\(x\)\(density tends to∞\\inftyatx=0x=0\): absolutely continuous ✓ \(Althoughρ​\(0\)=∞\\rho\(0\)=\\infty, we haveμ​\(\{0\}\)=∫\{0\}ρ​𝑑x=0\\mu\(\\\{0\\\}\)=\\int\_\{\\\{0\\\}\}\\rho\\,dx=0—the density can “blow up” at certain points, as long as it remains integrable\)

One subtle point: the density can be∞\\inftyyet the measure is still absolutely continuous

Beginners often have a misconception: “the densityρ​\(x\)\\rho\(x\)being large at some point” means “there is a large concentration of mass at that point\.” This iswrong\!

ρ​\(x\)=100\\rho\(x\)=100means there is a lot of mass perunit volumenearxx, but the mass at the single point\{x\}\\\{x\\\}is still zero\. Think of a glass of concentrated sugar water—the concentration in some region is very high, but if you take a single drop \(volume→0\\to 0\), the amount of sugar also tends to zero\.

What actually causes a measure to be “not absolutely continuous” is anatom: a finite amount of mass concentrated on a set of zero volume\. This cannot be represented within the density function framework—because∫\{x\}ρ​\(x\)​𝑑x\\int\_\{\\\{x\\\}\}\\rho\(x\)\\,dxis always zero, no matter how largeρ​\(x\)\\rho\(x\)is\.

### A\.3Integration against a measure

When we write∫f​\(x\)​𝑑μ​\(x\)\\int f\(x\)\\,d\\mu\(x\), it means “weighted summation offfusingμ\\mu”:

- •Ifμ\\muhas densityρ\\rho:∫f​𝑑μ=∫f​\(x\)​ρ​\(x\)​𝑑x\\int f\\,d\\mu=\\int f\(x\)\\rho\(x\)\\,dx\(ordinary integral\)
- •Ifμ=δa\\mu=\\delta\_\{a\}:∫f​𝑑μ=f​\(a\)\\int f\\,d\\mu=f\(a\)\(point evaluation\)
- •Ifμ=∑iwi​δxi\\mu=\\sum\_\{i\}w\_\{i\}\\delta\_\{x\_\{i\}\}\(discrete measure\):∫f​𝑑μ=∑iwi​f​\(xi\)\\int f\\,d\\mu=\\sum\_\{i\}w\_\{i\}f\(x\_\{i\}\)\(weighted sum\)

So∫f​𝑑μ\\int f\\,d\\muis an extremely flexible notation—it unifies “integration,” “evaluation,” and “weighted summation\.”

### A\.4Pushforward measure

###### Definition A\.5\(Pushforward\)\.

Given a measureμ\\muonℝd\\mathbb\{R\}^\{d\}and a mapT:ℝd→ℝdT:\\mathbb\{R\}^\{d\}\\to\\mathbb\{R\}^\{d\}, thepushforwardT\#​μT\_\{\\\#\}\\muis a measure onℝd\\mathbb\{R\}^\{d\}defined by:

\(T\#​μ\)​\(A\):=μ​\(T−1​\(A\)\)=μ​\(\{x:T​\(x\)∈A\}\)\(T\_\{\\\#\}\\mu\)\(A\):=\\mu\(T^\{\-1\}\(A\)\)=\\mu\(\\\{x:T\(x\)\\in A\\\}\)

Intuition:Imagine you have a bag of marbles \(distributed asμ\\mu\), and you apply the transformationTTto every marble—moving the marble at positionxxtoT​\(x\)T\(x\)\. The new distribution of marbles after the transformation isT\#​μT\_\{\\\#\}\\mu\.

Equivalent formulation\(using integrals\): For any functionff,

∫f​\(y\)​d​\(T\#​μ\)​\(y\)=∫f​\(T​\(x\)\)​𝑑μ​\(x\)\\int f\(y\)\\,d\(T\_\{\\\#\}\\mu\)\(y\)=\\int f\(T\(x\)\)\\,d\\mu\(x\)Left side: integrate against the new distribution\. Right side: integrate against the old distribution with a change of variables\.

Examples:

- •μ=δ0\\mu=\\delta\_\{0\},T​\(x\)=x\+1T\(x\)=x\+1\. ThenT\#​μ=δ1T\_\{\\\#\}\\mu=\\delta\_\{1\}\(point moves from 0 to 1\)\.
- •μ=Uniform​\[0,1\]\\mu=\\text\{Uniform\}\[0,1\],T​\(x\)=2​xT\(x\)=2x\. ThenT\#​μ=Uniform​\[0,2\]T\_\{\\\#\}\\mu=\\text\{Uniform\}\[0,2\]\(stretching\)\.
- •μ=𝒩​\(0,1\)\\mu=\\mathcal\{N\}\(0,1\),T​\(x\)=σ​x\+mT\(x\)=\\sigma x\+m\. ThenT\#​μ=𝒩​\(m,σ2\)T\_\{\\\#\}\\mu=\\mathcal\{N\}\(m,\\sigma^\{2\}\)\.

### A\.5Coupling: the key concept for optimal transport

We now arrive at the most important concept of this section—coupling\.

###### Definition A\.6\(Coupling\)\.

Given two probability measuresμ\\muon𝒳\\mathcal\{X\}andν\\nuon𝒴\\mathcal\{Y\}, acouplingof\(μ,ν\)\(\\mu,\\nu\)is a joint probability measureγ\\gammaon𝒳×𝒴\\mathcal\{X\}\\times\\mathcal\{Y\}whose marginals areμ\\muandν\\nu:

First marginal:γ​\(A×𝒴\)=μ​\(A\)∀A\\displaystyle\\quad\\gamma\(A\\times\\mathcal\{Y\}\)=\\mu\(A\)\\quad\\forall A\(68\)Second marginal:γ​\(𝒳×B\)=ν​\(B\)∀B\\displaystyle\\quad\\gamma\(\\mathcal\{X\}\\times B\)=\\nu\(B\)\\quad\\forall B\(69\)

What does this mean? Let us understand it on three levels:

Level one: The discrete case \(most intuitive\)

Supposeμ\\muandν\\nuare both discrete distributions:

- •μ\\mu: Factory A produces 3 tons, Factory B produces 2 tons \(total 5 tons\)
- •ν\\nu: City 1 demands 1 ton, City 2 demands 4 tons \(total 5 tons\)

A couplingγ\\gammais atransport plan table:

City 1City 2Row sum=μFactory AγA​1γA​23Factory BγB​1γB​22Col sum=ν145\\begin\{array\}\[\]\{c\|cc\|c\}&\\text\{City 1\}&\\text\{City 2\}&\\text\{Row sum\}=\\mu\\\\ \\hline\\cr\\text\{Factory A\}&\\gamma\_\{A1\}&\\gamma\_\{A2\}&3\\\\ \\text\{Factory B\}&\\gamma\_\{B1\}&\\gamma\_\{B2\}&2\\\\ \\hline\\cr\\text\{Col sum\}=\\nu&1&4&5\\end\{array\}
The constraints are:

- •Each row sum = the factory’s total production \(marginal isμ\\mu\)
- •Each column sum = the city’s total demand \(marginal isν\\nu\)

For example, one feasible plan isγA​1=1,γA​2=2,γB​1=0,γB​2=2\\gamma\_\{A1\}=1,\\gamma\_\{A2\}=2,\\gamma\_\{B1\}=0,\\gamma\_\{B2\}=2\. Another feasible plan isγA​1=0,γA​2=3,γB​1=1,γB​2=1\\gamma\_\{A1\}=0,\\gamma\_\{A2\}=3,\\gamma\_\{B1\}=1,\\gamma\_\{B2\}=1\. There are many plans satisfying the constraints—optimal transport finds the one with minimum cost\.

Level two: The random variable perspective

If\(X,Y\)\(X,Y\)is a joint random variable withX∼μX\\sim\\muandY∼νY\\sim\\nu, then the joint distribution of\(X,Y\)\(X,Y\)is a coupling ofμ\\muandν\\nu\.

Themarginal conditionsmean: regardless of what valueYYtakes, the \(marginal\) distribution ofXXisμ\\mu; and vice versa\.

Note: Given marginal distributionsμ\\muandν\\nu, the joint distribution isnot unique—it also depends on the correlation betweenXXandYY\.

- •Independent coupling:γ=μ⊗ν\\gamma=\\mu\\otimes\\nu\(XXandYYare independent\)
- •Perfect correlation:γ=\(id,T\)\#​μ\\gamma=\(\\text\{id\},T\)\_\{\\\#\}\\mu\(Y=T​\(X\)Y=T\(X\)is deterministically determined byXX\)
- •Various intermediate cases

The set of couplingsΓ​\(μ,ν\)\\Gamma\(\\mu,\\nu\)encompasses all these possibilities\.

Level three: The transport plan perspective

Returning to the earth\-moving metaphor\.γ​\(x,y\)\\gamma\(x,y\)represents “the amount of earth moved from positionxxto positionyy\.”

The marginal conditions guarantee:

- •The total amount moved fromxx= the amount of earth atμ​\(x\)\\mu\(x\)\(all earth atxxis moved away\)
- •The total amount arriving atyy= the amount needed atν​\(y\)\\nu\(y\)\(all demand atyyis satisfied\)

###### Example A\.7\(Couplings between two Dirac masses\)\.

Letμ=δa\\mu=\\delta\_\{a\},ν=δb\\nu=\\delta\_\{b\}\. The unique coupling isγ=δ\(a,b\)\\gamma=\\delta\_\{\(a,b\)\}—all mass goes fromaatobb, with no other choice\.

###### Example A\.8\(Couplings between uniform distributions\)\.

Letμ=ν=Uniform​\[0,1\]\\mu=\\nu=\\text\{Uniform\}\[0,1\]\. Some examples of couplings:

- •Independent coupling:γ=\\gamma=uniform distribution on\[0,1\]2\[0,1\]^\{2\}\(scatter points uniformly in the square\)
- •Identity coupling:γ\\gammais concentrated on the diagonal\{\(x,x\):x∈\[0,1\]\}\\\{\(x,x\):x\\in\[0,1\]\\\}\(every pointxxstays put\)
- •Reflection coupling:γ\\gammais concentrated on the anti\-diagonal\{\(x,1−x\):x∈\[0,1\]\}\\\{\(x,1\-x\):x\\in\[0,1\]\\\}\(xxis moved to1−x1\-x\)

The “transport costs”∫\|x−y\|2​𝑑γ\\int\|x\-y\|^\{2\}\\,d\\gammafor these three plans are16\\frac\{1\}\{6\},0, and13\\frac\{1\}\{3\}, respectively\. The optimal one is the identity coupling \(cost 0\), which makes sense—if the two distributions are identical, the best “transport plan” is to move nothing\.

Understanding the “transport style” of the independent coupling:

The independent couplingγ=μ⊗ν\\gamma=\\mu\\otimes\\nuis the most “chaotic” transport plan\. It means:

Each grain of sand forgets where it came from and is randomly thrown to some location in the target distributionν\\nu\.

Specifically, whenμ=ν=Uniform​\[0,1\]\\mu=\\nu=\\text\{Uniform\}\[0,1\],γ\\gammais uniformly distributed on the square\[0,1\]2\[0,1\]^\{2\}\. This means:

- •Sand at positionx=0\.2x=0\.2might be moved toy=0\.1y=0\.1, or might be moved toy=0\.9y=0\.9— the destinationyyis uniformly distributed on\[0,1\]\[0,1\], completely independent of the originxx\.
- •Sand at positionx=0\.8x=0\.8behaves the same—the destination is uniformly random\.
- •There isno correlation whatsoeverbetween origin and destination \(statistical independence\)\.

Why is this a valid transport plan?Verify the marginal conditions:

γ​\(A×\[0,1\]\)\\displaystyle\\gamma\(A\\times\[0,1\]\)=Uniform\[0,1\]2​\(A×\[0,1\]\)=\|A\|=μ​\(A\)​✓\\displaystyle=\\text\{Uniform\}\_\{\[0,1\]^\{2\}\}\(A\\times\[0,1\]\)=\|A\|=\\mu\(A\)\\;\\checkmarkγ​\(\[0,1\]×B\)\\displaystyle\\gamma\(\[0,1\]\\times B\)=Uniform\[0,1\]2​\(\[0,1\]×B\)=\|B\|=ν​\(B\)​✓\\displaystyle=\\text\{Uniform\}\_\{\[0,1\]^\{2\}\}\(\[0,1\]\\times B\)=\|B\|=\\nu\(B\)\\;\\checkmarkSo the marginals are indeedμ\\muandν\\nu—the total amount moved from eachxxis correct, and the total arriving at eachyyis also correct\. The transport is simply extremely chaotic\.

Computing the cost∫\|x−y\|2​𝑑γ\\int\|x\-y\|^\{2\}\\,d\\gamma: Sincexxandyyareindependentunder the independent coupling,

∫\[0,1\]2\|x−y\|2​𝑑x​𝑑y\\displaystyle\\int\_\{\[0,1\]^\{2\}\}\|x\-y\|^\{2\}\\,dxdy=∫01∫01\(x−y\)2​𝑑x​𝑑y\\displaystyle=\\int\_\{0\}^\{1\}\\int\_\{0\}^\{1\}\(x\-y\)^\{2\}\\,dxdy=∫01∫01\(x2−2​x​y\+y2\)​𝑑x​𝑑y\\displaystyle=\\int\_\{0\}^\{1\}\\int\_\{0\}^\{1\}\(x^\{2\}\-2xy\+y^\{2\}\)\\,dxdy=𝔼​\[X2\]−2​𝔼​\[X\]​𝔼​\[Y\]\+𝔼​\[Y2\]\\displaystyle=\\mathbb\{E\}\[X^\{2\}\]\-2\\mathbb\{E\}\[X\]\\mathbb\{E\}\[Y\]\+\\mathbb\{E\}\[Y^\{2\}\]=13−2⋅12⋅12\+13=16\\displaystyle=\\frac\{1\}\{3\}\-2\\cdot\\frac\{1\}\{2\}\\cdot\\frac\{1\}\{2\}\+\\frac\{1\}\{3\}=\\frac\{1\}\{6\}\(where𝔼​\[X2\]=∫01x2​𝑑x=13\\mathbb\{E\}\[X^\{2\}\]=\\int\_\{0\}^\{1\}x^\{2\}dx=\\frac\{1\}\{3\},𝔼​\[X\]=12\\mathbb\{E\}\[X\]=\\frac\{1\}\{2\}\.\)

Comparing the “transport styles” of the three couplings:

In one sentence:The independent coupling has nonzero cost because it “forgets” the connection between origin and destination— even though nothing needs to move \(cost zero\), it insists on randomly reshuffling the sand, needlessly increasing the transport distance\. This is why the independent coupling is typically avery poortransport plan \(though still valid\)\.

Key InsightCoupling = encoding a transport plan via a joint distributionThis is the most crucial intuition for understanding optimal transport\. A joint distributionγ​\(x,y\)\\gamma\(x,y\)simultaneously tells you two things:1\.Who is paired with whom:The larger the “density” ofγ\\gammaat\(x,y\)\(x,y\), the more mass is transported fromxxtoyy2\.The pairing satisfies supply\-demand balance:The marginal conditions ensure every source point’s mass is fully allocated, and every target point’s demand is fully metDifferent joint distributions = different transport strategies:•γ\\gammaconcentrated on the diagonalx=yx=y⇒\\Rightarrow“stay in place”•γ\\gammaconcentrated on some curvey=T​\(x\)y=T\(x\)⇒\\Rightarrow“deterministic transport” \(Monge map\)•γ\\gammaspread over the entirexx\-yyplane⇒\\Rightarrow“random/chaotic transport” \(mass from eachxxis dispersed to multipleyy\)A joint distribution is anextremely flexiblerepresentation—it can express deterministic maps \(γ\\gammadegenerates onto the graph of the map\), as well as “one\-to\-many” mass splitting \(γ\\gammaspreads in theyydirection\)\. This is why the Kantorovich formulation is more powerful than the Monge formulation\.

Key InsightThe essence of coupling:Given two distributionsμ\\muandν\\nu, a couplingγ\\gammadescribes one specific plan for “who goes where\.” The same pair\(μ,ν\)\(\\mu,\\nu\)admits infinitely many couplings \(transport plans\), each with different cost\.Optimal transportfinds the coupling with minimum cost\. TheWasserstein distanceis that minimum cost \(square root\)\.

Why use couplings instead of mapsTTdirectly?

Monge’s original idea was to find a mapTTsuch thatT\#​μ=νT\_\{\\\#\}\\mu=\\nu\. But a map requires “one\-to\-one”—all mass atxxmust go to a single destination\. Ifμ=δ0\\mu=\\delta\_\{0\}andν=12​δ−1\+12​δ1\\nu=\\frac\{1\}\{2\}\\delta\_\{\-1\}\+\\frac\{1\}\{2\}\\delta\_\{1\}\(one point splits into two\), no such map exists\!

Kantorovich’s coupling relaxes this restriction: the mass at a pointxxcan besplitamong multiple destinations\. This ensures: \(1\) the optimization problemalwayshas a solution; \(2\) the problem becomeslinear programming, which is easier to handle\.

## Appendix BDifferential Geometry Prerequisites

IntuitionWhy are these concepts needed?Later we will say “Wasserstein space is a Riemannian manifold” and “Fokker–Planck is a gradient flow\.” The core tools underlying these statements aremanifolds,tangent spaces,Riemannian metrics, andgradients\.If you are unfamiliar with these concepts, the following builds them from scratch using only intuition and examples\.

\(1\) Manifold = a space that locally looks likeℝn\\mathbb\{R\}^\{n\}

- •Intuition: A manifold is a “possibly curved space,” but if you zoom into a small neighborhood, it looks just like ordinaryℝn\\mathbb\{R\}^\{n\}\.
- •Example 1: The surface of the EarthS2S^\{2\}—globally a sphere \(curved\), but locally described by latitude and longitude \(like a patch ofℝ2\\mathbb\{R\}^\{2\}\)\.
- •Example 2:ℝn\\mathbb\{R\}^\{n\}itself—the simplest manifold, globally flat\.
- •Example 3: The probability simplex\{\(p1,…,pn\):pi≥0,∑pi=1\}\\\{\(p\_\{1\},\\ldots,p\_\{n\}\):p\_\{i\}\\geq 0,\\sum p\_\{i\}=1\\\}—an\(n−1\)\(n\-1\)\-dimensional hyperplane inℝn\\mathbb\{R\}^\{n\}\.

\(2\) Tangent space = all possible “instantaneous velocity directions” at a point

- •Intuition: You stand at a pointppon the manifold; all directionsyou can currently moveform a vector space—this isTp​MT\_\{p\}M\.
- •Example 1 \(a pointppinℝn\\mathbb\{R\}^\{n\}\): The tangent spaceTp​ℝn=ℝnT\_\{p\}\\mathbb\{R\}^\{n\}=\\mathbb\{R\}^\{n\}\. You can move in any direction without restriction\.
- •Example 2 \(the north poleNNonS2S^\{2\}\): The tangent spaceTN​S2T\_\{N\}S^\{2\}is the “tangent plane” at the north pole—anℝ2\\mathbb\{R\}^\{2\}plane tangent to the sphere\. You cannot “tunnel into the sphere” or “fly off the surface”—you can only move along the surface\.

ppv1v\_\{1\}v2v\_\{2\}Tangent planeTp​MT\_\{p\}MTangent spaceTp​MT\_\{p\}Mat a pointppon the sphereM=S2M=S^\{2\}v1,v2v\_\{1\},v\_\{2\}are tangent vectors—directions one can move along the sphere

More rigorous definitions \(two equivalent approaches\):

- •Curve definition:A tangent vectorv∈Tp​Mv\\in T\_\{p\}Mis the velocityγ˙​\(0\)\\dot\{\\gamma\}\(0\)att=0t=0of some curveγ​\(t\)\\gamma\(t\)passing throughpp\(γ​\(0\)=p\\gamma\(0\)=p\)\. Intuition: An ant stands at pointppand begins walking along some pathγ\\gamma—the instantaneous velocity at the moment of departure is a tangent vector\.
- •Directional derivative definition:A tangent vectorvvis an operator that “takes directional derivatives of functions”—v​\(f\):=dd​t\|t=0​f​\(γ​\(t\)\)v\(f\):=\\left\.\\frac\{d\}\{dt\}\\right\|\_\{t=0\}f\(\\gamma\(t\)\)\. Intuition: A tangent vectorvvtells you “how much any functionffchanges if you take a small step in the directionvv\.”

\(3\) Key properties of tangent vectors

- •The tangent spaceTp​MT\_\{p\}Mis avector space\(with addition and scalar multiplication\), of dimension=dimM=\\dim M\.
- •Tangent spaces at different points aredifferent:Tp​M≠Tq​MT\_\{p\}M\\neq T\_\{q\}M\(one cannot directly compare tangent vectors at different points—this is why one needs “connections” or “parallel transport,” but we do not need this here\)\.
- •The velocity of a curveγ​\(t\)\\gamma\(t\)satisfiesγ˙​\(t\)∈Tγ​\(t\)​M\\dot\{\\gamma\}\(t\)\\in T\_\{\\gamma\(t\)\}M—at each moment the velocity lives in the tangent spaceat the point where the curve currently is\.

\(4\) Dual space and cotangent space

First, the dual space—a purely linear\-algebraic concept:

###### Definition B\.1\(Dual space\)\.

LetVVbe a vector space overℝ\\mathbb\{R\}\. Thedual spaceV∗V^\{\*\}is defined as:

V∗:=\{ω:V→ℝ∣ω​is a linear map\}V^\{\*\}:=\\\{\\omega:V\\to\\mathbb\{R\}\\mid\\omega\\text\{~is a linear map\}\\\}Elements ofV∗V^\{\*\}are calledlinear functionalsonVVorcovectors\.

Concrete example \(V=ℝnV=\\mathbb\{R\}^\{n\}\):

Elements ofV=ℝnV=\\mathbb\{R\}^\{n\}are column vectorsv=\(v1⋮vn\)v=\\begin\{pmatrix\}v^\{1\}\\\\ \\vdots\\\\ v^\{n\}\\end\{pmatrix\}\.

What are elements ofV∗V^\{\*\}? They are linear maps that “eat a column vector and output a number\.” The most general form is:ω​\(v\)=a1​v1\+a2​v2\+⋯\+an​vn\\omega\(v\)=a\_\{1\}v^\{1\}\+a\_\{2\}v^\{2\}\+\\cdots\+a\_\{n\}v^\{n\}— i\.e\., taking the inner product ofvvwith fixed coefficients\(a1,…,an\)\(a\_\{1\},\\ldots,a\_\{n\}\)\.

In matrix form:ω=\(a1,a2,…,an\)\\omega=\(a\_\{1\},a\_\{2\},\\ldots,a\_\{n\}\)is arow vector, and the pairing is row×\\timescolumn:

ω​\(v\)=\(a1⋯an\)​\(v1⋮vn\)=∑i=1nai​vi\\omega\(v\)=\\begin\{pmatrix\}a\_\{1\}&\\cdots&a\_\{n\}\\end\{pmatrix\}\\begin\{pmatrix\}v^\{1\}\\\\ \\vdots\\\\ v^\{n\}\\end\{pmatrix\}=\\sum\_\{i=1\}^\{n\}a\_\{i\}v^\{i\}
So\(ℝn\)∗\(\\mathbb\{R\}^\{n\}\)^\{\*\}is the space of all row vectors—alsonn\-dimensional, but “living in a different place\.”

Key distinction: Pairing vs\. inner product

- •Pairingω​\(v\)=∑ai​vi\\omega\(v\)=\\sum a\_\{i\}v^\{i\}: A row vectorω∈V∗\\omega\\in V^\{\*\}acts on a column vectorv∈Vv\\in V\. Thisdoes not require an inner product—it is a purely algebraic operation\.
- •Inner product⟨u,v⟩=∑gi​j​ui​vj\\langle u,v\\rangle=\\sum g\_\{ij\}u^\{i\}v^\{j\}: An operation between two vectorsu,v∈Vu,v\\in Vof the same type\. Thisrequires an additional choiceof metricgg\.

Pairing is “row×\\timescolumn” \(a natural operation between different types\); inner product is “column×\\timescolumn” \(between the same type, requiring the introduction of a metric matrixggto “flip” one of them\)\.

What is an isomorphism?

###### Definition B\.2\(Vector space isomorphism\)\.

Anisomorphismbetween two vector spacesVVandWWis alinear bijectionφ:V→W\\varphi:V\\to W\(linear \+ injective \+ surjective\)\.

If such aφ\\varphiexists, we sayVVandWWareisomorphic, writtenV≅WV\\cong W\.

Intuition:V≅WV\\cong Wmeans the two spaces have “exactly the same structure”—you can losslessly “translate” elements of one space into elements of the other viaφ\\varphi, preserving all linear relationships\.

Simple examples:

- •ℝ2≅ℝ2\\mathbb\{R\}^\{2\}\\cong\\mathbb\{R\}^\{2\}—the identity mapφ​\(v\)=v\\varphi\(v\)=vis an isomorphism \(trivial\)\.
- •ℝ2≅ℝ2\\mathbb\{R\}^\{2\}\\cong\\mathbb\{R\}^\{2\}—a90∘90^\{\\circ\}rotation is also an isomorphism \(a different “translation”\)\.
- •ℝ2\\mathbb\{R\}^\{2\}and\{\\\{all2×12\\times 1matrices\}\\\}are isomorphic—column vectors↔\\leftrightarrowmatrices, just different notation\.
- •ℝ2\\mathbb\{R\}^\{2\}andℝ3\\mathbb\{R\}^\{3\}arenotisomorphic—different dimensions make a linear bijection impossible\.

Key concept: “Natural” isomorphisms vs\. “choice\-dependent” isomorphisms

- •Canonical isomorphism: An isomorphism that can be written down without making any choices\. Example:V≅V∗∗V\\cong V^\{\*\*\}\(double dual\) has a canonical isomorphismv↦evvv\\mapsto\\mathrm\{ev\}\_\{v\}, whereevv​\(ω\):=ω​\(v\)\\mathrm\{ev\}\_\{v\}\(\\omega\):=\\omega\(v\)\. The definition of this map involves no “choices\.”
- •Non\-canonical isomorphism: An isomorphism that requires additional choices to specify\. Example:V≅V∗V\\cong V^\{\*\}—one must choose a basis or an inner product to write down a specific map\. Different choices⇒\\Rightarrowdifferent isomorphisms\.

Where exactly is the difference betweenVVandV∗V^\{\*\}?

As abstract vector spaces,V≅V∗≅ℝnV\\cong V^\{\*\}\\cong\\mathbb\{R\}^\{n\}\(same dimension⇒\\Rightarrowan isomorphism exists\)\. But the key point is:

- •There is nocanonicalisomorphismV→∼V∗V\\xrightarrow\{\\sim\}V^\{\*\}—you cannot turn a column vector into a row vector without making a choice\.
- •Choosing an inner productgggives aspecificisomorphism:v↦g​\(v,⋅\)v\\mapsto g\(v,\\cdot\), i\.e\.,vi↦gi​j​vjv^\{i\}\\mapsto g\_\{ij\}v^\{j\}\(column to row, multiply bygg\)\.
- •Choosing a different inner productg~\\tilde\{g\}yields adifferentisomorphism\.

“Isn’t the transpose a canonical isomorphism?”—A common misconception

You might think:v=\(v1v2\)↦vT=\(v1,v2\)v=\\begin\{pmatrix\}v^\{1\}\\\\ v^\{2\}\\end\{pmatrix\}\\mapsto v^\{T\}=\(v^\{1\},v^\{2\}\), isn’t this a “choice\-free” column→\\torow map?

Answer: This mapdepends on the basis you chose\. Specifically:

LetV=ℝ2V=\\mathbb\{R\}^\{2\}, basis\{e1,e2\}\\\{e\_\{1\},e\_\{2\}\\\}\. The vectorv=3​e1\+2​e2v=3e\_\{1\}\+2e\_\{2\}has coordinates\(32\)\\binom\{3\}\{2\}in this basis, and “transposing” gives the row vector\(3,2\)\(3,2\), corresponding to the linear functionalω1\\omega\_\{1\}withω1​\(e1\)=3\\omega\_\{1\}\(e\_\{1\}\)=3,ω1​\(e2\)=2\\omega\_\{1\}\(e\_\{2\}\)=2\.

Now change to a different basis\{e1′=e1\+e2,e2′=e2\}\\\{e\_\{1\}^\{\\prime\}=e\_\{1\}\+e\_\{2\},\\;e\_\{2\}^\{\\prime\}=e\_\{2\}\\\}\. The samevvhas coordinates\(3−1\)\\binom\{3\}\{\-1\}in the new basis \(sincev=3​e1′\+\(−1\)​e2′v=3e\_\{1\}^\{\\prime\}\+\(\-1\)e\_\{2\}^\{\\prime\}\), and “transposing” gives\(3,−1\)\(3,\-1\), corresponding to the linear functionalω2\\omega\_\{2\}withω2​\(e1′\)=3\\omega\_\{2\}\(e\_\{1\}^\{\\prime\}\)=3,ω2​\(e2′\)=−1\\omega\_\{2\}\(e\_\{2\}^\{\\prime\}\)=\-1\.

ω1≠ω2\\omega\_\{1\}\\neq\\omega\_\{2\}\!—the samevv, after changing basis, “transposing” yields a different covector\.

So “transpose” is not a map fromVVtoV∗V^\{\*\}—it is a map from “coordinates ofVVin some basis” to “coordinates ofV∗V^\{\*\}in the dual basis\.” The moment you say “take coordinates,” you have already made a choice \(choosing a basis\)\.

A more fundamental statement:The mathematical meaning of “natural” is anatural transformation in the sense of category theory—V→V∗V\\to V^\{\*\}is not natural because it does not commute with all linear mapsA:V→WA:V\\to W\. In contrast,V→V∗∗V\\to V^\{\*\*\}is natural becauseev\\mathrm\{ev\}commutes with anyAA:\(A∗∗∘evV\)​\(v\)=evW​\(A​v\)\(A^\{\*\*\}\\circ\\mathrm\{ev\}\_\{V\}\)\(v\)=\\mathrm\{ev\}\_\{W\}\(Av\)\.

This distinction is invisible inℝn\\mathbb\{R\}^\{n\}with the standard inner product \(sinceg=Ig=I, columns and rows “look the same”\)\. But onceg≠Ig\\neq I, or on a curved surface \(whereggvaries from point to point\), the distinction becomes crucial\.

Now back to manifolds: The cotangent spaceTp∗​MT\_\{p\}^\{\*\}M

###### Definition B\.3\(Cotangent space\)\.

Thecotangent spaceat the pointppis defined as the dual of the tangent space:

Tp∗​M:=\(Tp​M\)∗=\{ω:Tp​M→ℝ∣ω​is a linear map\}T\_\{p\}^\{\*\}M:=\(T\_\{p\}M\)^\{\*\}=\\\{\\omega:T\_\{p\}M\\to\\mathbb\{R\}\\mid\\omega\\text\{~is a linear map\}\\\}Elements ofTp∗​MT\_\{p\}^\{\*\}Mare calledcovectorsor1\-forms\(atpp\)\.

Intuition:A tangent vectorv∈Tp​Mv\\in T\_\{p\}Mis a “direction” \(an arrow\); a covectorω∈Tp∗​M\\omega\\in T\_\{p\}^\{\*\}Mis a “ruler for measuring directions”— it accepts a directionvvand outputs a numberω​\(v\)∈ℝ\\omega\(v\)\\in\\mathbb\{R\}\.

Coordinate representation:Let\(x1,…,xn\)\(x^\{1\},\\ldots,x^\{n\}\)be local coordinates\.

- •Natural basis ofTp​MT\_\{p\}M:\{∂1,…,∂n\}\\\{\\partial\_\{1\},\\ldots,\\partial\_\{n\}\\\}, where∂i:=∂∂xi\|p\\partial\_\{i\}:=\\frac\{\\partial\}\{\\partial x^\{i\}\}\\big\|\_\{p\}
- •Dual basis ofTp∗​MT\_\{p\}^\{\*\}M:\{d​x1,…,d​xn\}\\\{dx^\{1\},\\ldots,dx^\{n\}\\\}, defined byd​xi​\(∂j\)=δjidx^\{i\}\(\\partial\_\{j\}\)=\\delta^\{i\}\_\{j\}

Tangent vectors are written asv=vi​∂iv=v^\{i\}\\partial\_\{i\}\(upper indices, column vectors\); covectors are written asω=ωi​d​xi\\omega=\\omega\_\{i\}\\,dx^\{i\}\(lower indices, row vectors\)\.

Pairing:ω​\(v\)=ωi​vi\\omega\(v\)=\\omega\_\{i\}v^\{i\}\(row×\\timescolumn==scalar, no metric needed\)\.

The differentiald​fdfof a function is a covector:

###### Definition B\.4\(Differential of a function\)\.

Letf:M→ℝf:M\\to\\mathbb\{R\}be smooth\. Thedifferentialofffatpp,d​fp∈Tp∗​Mdf\_\{p\}\\in T\_\{p\}^\{\*\}M, is defined as:

d​fp​\(v\):=v​\(f\)=dd​t\|t=0​f​\(γ​\(t\)\),γ​\(0\)=p,γ˙​\(0\)=vdf\_\{p\}\(v\):=v\(f\)=\\left\.\\frac\{d\}\{dt\}\\right\|\_\{t=0\}f\(\\gamma\(t\)\),\\quad\\gamma\(0\)=p,\\;\\dot\{\\gamma\}\(0\)=v

In coordinates:d​fp=∂f∂xi\|p​d​xidf\_\{p\}=\\frac\{\\partial f\}\{\\partial x^\{i\}\}\\big\|\_\{p\}\\,dx^\{i\}\. This is a row vector\(∂f∂x1,…,∂f∂xn\)\\bigl\(\\frac\{\\partial f\}\{\\partial x^\{1\}\},\\ldots,\\frac\{\\partial f\}\{\\partial x^\{n\}\}\\bigr\)\.

Key point:d​fpdf\_\{p\}requires only the smooth structure, not a metric\.It is a “natural” object— given the functionffand the manifold structure,d​fdfis completely determined\.

Comparison ofTp​MT\_\{p\}MvsTp∗​MT\_\{p\}^\{\*\}M:

Why areTp​MT\_\{p\}MandTp∗​MT\_\{p\}^\{\*\}Mnot “the same thing”?

In finite dimensions,Tp​MT\_\{p\}MandTp∗​MT\_\{p\}^\{\*\}Mhave the same dimension and are isomorphic as abstract vector spaces\. But there isno canonical isomorphism—you cannot “turn” a tangent vectorv∈Tp​Mv\\in T\_\{p\}Minto a covectorω∈Tp∗​M\\omega\\in T\_\{p\}^\{\*\}Mwithout introducing additional structure\.

If one chooses an inner productgpg\_\{p\}, one obtains themusical isomorphism:

♭:Tp​M\\displaystyle\\flat:\\;T\_\{p\}M→Tp∗​M,v↦v♭:=gp​\(v,⋅\)\(“lower indices”: column→row\)\\displaystyle\\to T\_\{p\}^\{\*\}M,\\quad v\\mapsto v^\{\\flat\}:=g\_\{p\}\(v,\\,\\cdot\\,\)\\quad\\text\{\(\`\`lower indices'': column$\\to$row\)\}♯:Tp∗​M\\displaystyle\\sharp:\\;T\_\{p\}^\{\*\}M→Tp​M,ω↦ω♯\(“raise indices”: row→column\)\\displaystyle\\to T\_\{p\}M,\\quad\\omega\\mapsto\\omega^\{\\sharp\}\\quad\\text\{\(\`\`raise indices'': row$\\to$column\)\}In coordinates:\(v♭\)i=gi​j​vj\(v^\{\\flat\}\)\_\{i\}=g\_\{ij\}v^\{j\}\(multiply bygg\),\(ω♯\)i=gi​j​ωj\(\\omega^\{\\sharp\}\)^\{i\}=g^\{ij\}\\omega\_\{j\}\(multiply byg−1g^\{\-1\}\)\.

Concrete example:LetV=ℝ2V=\\mathbb\{R\}^\{2\}, metricg=\(2113\)g=\\begin\{pmatrix\}2&1\\\\ 1&3\\end\{pmatrix\}\.

Lowering indices♭\\flat\(tangent vector→\\tocovector\): Takev=\(12\)∈Tp​Mv=\\binom\{1\}\{2\}\\in T\_\{p\}M\.

The definition ofv♭v^\{\\flat\}is:v♭​\(w\):=g​\(v,w\)v^\{\\flat\}\(w\):=g\(v,w\)for allww\. In coordinates:

\(v♭\)i=gi​jvj:v♭=\(2113\)\(12\)=\(47\)T=4dx1\+7dx2\(v^\{\\flat\}\)\_\{i\}=g\_\{ij\}v^\{j\}:\\qquad v^\{\\flat\}=\\begin\{pmatrix\}2&1\\\\ 1&3\\end\{pmatrix\}\\begin\{pmatrix\}1\\\\ 2\\end\{pmatrix\}=\\begin\{pmatrix\}4\\\\ 7\\end\{pmatrix\}^\{T\}\\\!=\\;4\\,dx^\{1\}\+7\\,dx^\{2\}\(Note: the result is the row vector\(4,7\)\(4,7\), living inTp∗​MT\_\{p\}^\{\*\}M\.\)

Verification:v♭​\(w\)=4​w1\+7​w2v^\{\\flat\}\(w\)=4w^\{1\}\+7w^\{2\}\. Andg​\(v,w\)=\(1,2\)​\(2113\)​\(w1w2\)=4​w1\+7​w2g\(v,w\)=\(1,2\)\\begin\{pmatrix\}2&1\\\\ 1&3\\end\{pmatrix\}\\binom\{w^\{1\}\}\{w^\{2\}\}=4w^\{1\}\+7w^\{2\}\. ✓

Raising indices♯\\sharp\(cotangent vector→\\totangent vector\): Takeω=4​d​x1\+7​d​x2\\omega=4\\,dx^\{1\}\+7\\,dx^\{2\}, i\.e\.,\(ωi\)=\(4,7\)\(\\omega\_\{i\}\)=\(4,7\)\.

\(ω♯\)i=gi​jωj:ω♯=g−1\(47\)=15\(3−1−12\)\(47\)=15\(510\)=\(12\)\(\\omega^\{\\sharp\}\)^\{i\}=g^\{ij\}\\omega\_\{j\}:\\qquad\\omega^\{\\sharp\}=g^\{\-1\}\\begin\{pmatrix\}4\\\\ 7\\end\{pmatrix\}=\\frac\{1\}\{5\}\\begin\{pmatrix\}3&\-1\\\\ \-1&2\\end\{pmatrix\}\\begin\{pmatrix\}4\\\\ 7\\end\{pmatrix\}=\\frac\{1\}\{5\}\\begin\{pmatrix\}5\\\\ 10\\end\{pmatrix\}=\\begin\{pmatrix\}1\\\\ 2\\end\{pmatrix\}We recoverv=\(12\)v=\\binom\{1\}\{2\}\! This verifies that♯=♭−1\\sharp=\\flat^\{\-1\}\(raising and lowering are inverses of each other\)\.

Origin of the names:♭\\flat\(flat\) lowers “high” objects \(tangent vectors, superscriptviv^\{i\}\) to “low” ones \(cotangent vectors, subscriptωi\\omega\_\{i\}\);♯\\sharp\(sharp\) does the reverse\. The notation matches the musical symbols for lowering a note by a half step \(♭\\flat\) and raising it by a half step \(♯\\sharp\)\.

Intuitive meaning of the metricg=\(2113\)g=\\begin\{pmatrix\}2&1\\\\ 1&3\\end\{pmatrix\}:

The metricggtells you the “cost of moving in space\.” Specifically, the squared length of a tangent vectorvvis:

‖v‖g2=g​\(v,v\)=vT​g​v=2​\(v1\)2\+2​v1​v2\+3​\(v2\)2\\\|v\\\|\_\{g\}^\{2\}=g\(v,v\)=v^\{T\}g\\,v=2\(v^\{1\}\)^\{2\}\+2v^\{1\}v^\{2\}\+3\(v^\{2\}\)^\{2\}
- •g11=2g\_\{11\}=2: the cost of taking one step in thex1x^\{1\}direction is22\(more expensive than the standard metric’s11\)
- •g22=3g\_\{22\}=3: the cost of taking one step in thex2x^\{2\}direction is33\(even more expensive\)
- •g12=1≠0g\_\{12\}=1\\neq 0: thex1x^\{1\}andx2x^\{2\}directions arenot orthogonal—moving simultaneously in both directions incurs a “coupling cost”

Geometric picture: under the standard metricg=Ig=I, the set of all directions at “unit distance” forms a circle \(‖v‖=1\\\|v\\\|=1is a circle\)\. Under thisgg,‖v‖g=1\\\|v\\\|\_\{g\}=1becomes anellipse:

2​\(v1\)2\+2​v1​v2\+3​\(v2\)2=12\(v^\{1\}\)^\{2\}\+2v^\{1\}v^\{2\}\+3\(v^\{2\}\)^\{2\}=1The major axis of the ellipse points in the “cheap direction” \(low cost for the same coordinate distance\), and the minor axis points in the “expensive direction\.”

Physical analogy:Imagine walking on muddy ground\.g=Ig=Iis a flat dry road \(equal cost in all directions\)\.g≠Ig\\neq Iis uneven terrain—some directions are downhill \(cheap\), some are muddy uphill \(expensive\), and there may be slopes that cause “walking north” to inevitably “slide a bit east” \(off\-diagonal termsg12≠0g\_\{12\}\\neq 0\)\.

Key InsightGradient = raising the index of the differential:grad​f=\(d​f\)♯\\mathrm\{grad\}\\,f=\(df\)^\{\\sharp\}That is:grad​f\\mathrm\{grad\}\\,fis obtained by “lifting” the cotangent vectord​f∈Tp∗​Mdf\\in T\_\{p\}^\{\*\}MintoTp​MT\_\{p\}Mvia the metricgg\. In coordinates:\(grad​f\)i=gi​j​∂f∂xj\(\\mathrm\{grad\}\\,f\)^\{i\}=g^\{ij\}\\frac\{\\partial f\}\{\\partial x^\{j\}\}\.This is whythe samed​fdf, with differentgg, yields differentgrad​f\\mathrm\{grad\}\\,f\.d​fdfis a fixed row vector;grad​f=g−1⋅d​fT\\mathrm\{grad\}\\,f=g^\{\-1\}\\cdot df^\{T\}depends on whichg−1g^\{\-1\}is used to “flip” it\.

\(5\) Riemannian Metric

So far, we have the manifoldMM, the tangent spaceTp​MT\_\{p\}M, the cotangent spaceTp∗​MT\_\{p\}^\{\*\}M, and the musical isomorphisms\. But the musical isomorphisms require an inner productgg—this is the Riemannian metric\.

###### Definition B\.5\(Riemannian metric\)\.

ARiemannian metricggon a smooth manifoldMMassigns to each pointp∈Mp\\in Maninner productgpg\_\{p\}onTp​MT\_\{p\}M, varying smoothly withpp\.

That is,gp:Tp​M×Tp​M→ℝg\_\{p\}:T\_\{p\}M\\times T\_\{p\}M\\to\\mathbb\{R\}satisfies:

- •Bilinearity:gp​\(α​u\+β​v,w\)=α​gp​\(u,w\)\+β​gp​\(v,w\)g\_\{p\}\(\\alpha u\+\\beta v,w\)=\\alpha g\_\{p\}\(u,w\)\+\\beta g\_\{p\}\(v,w\)
- •Symmetry:gp​\(u,v\)=gp​\(v,u\)g\_\{p\}\(u,v\)=g\_\{p\}\(v,u\)
- •Positive definiteness:gp​\(v,v\)\>0g\_\{p\}\(v,v\)\>0\(whenv≠0v\\neq 0\)

In coordinates:gp​\(∂i,∂j\)=gi​j​\(p\)g\_\{p\}\(\\partial\_\{i\},\\partial\_\{j\}\)=g\_\{ij\}\(p\), soggis represented by a positive definite symmetric matrix function\[gi​j​\(x\)\]\[g\_\{ij\}\(x\)\]\.

What can a Riemannian metric do?

- •Length of a tangent vector:‖v‖g=gp​\(v,v\)\\\|v\\\|\_\{g\}=\\sqrt\{g\_\{p\}\(v,v\)\}
- •Angle between two tangent vectors:cos⁡θ=gp​\(u,v\)‖u‖g​‖v‖g\\cos\\theta=\\frac\{g\_\{p\}\(u,v\)\}\{\\\|u\\\|\_\{g\}\\,\\\|v\\\|\_\{g\}\}
- •Length of a curve:L​\(γ\)=∫01‖γ˙​\(t\)‖γ​\(t\)​𝑑tL\(\\gamma\)=\\int\_\{0\}^\{1\}\\\|\\dot\{\\gamma\}\(t\)\\\|\_\{\\gamma\(t\)\}\\,dt
- •Distance between two points:d​\(p,q\)=infγL​\(γ\)d\(p,q\)=\\inf\_\{\\gamma\}L\(\\gamma\)\(geodesic distance\)
- •Raising and lowering indices:♭:Tp​M→Tp∗​M\\flat:T\_\{p\}M\\to T\_\{p\}^\{\*\}M,♯:Tp∗​M→Tp​M\\sharp:T\_\{p\}^\{\*\}M\\to T\_\{p\}M
- •Gradient:grad​f=\(d​f\)♯\\mathrm\{grad\}\\,f=\(df\)^\{\\sharp\}

A manifold\(M,g\)\(M,g\)equipped with a metricggis called aRiemannian manifold\.

\(6\) Gradient =\(d​f\)♯\(df\)^\{\\sharp\}= raising the index of the differential

###### Definition B\.6\(Riemannian gradient\)\.

Let\(M,g\)\(M,g\)be a Riemannian manifold andf:M→ℝf:M\\to\\mathbb\{R\}smooth\. Thegradientofffatppis:

grad​f\|p:=\(d​fp\)♯∈Tp​M\\mathrm\{grad\}\\,f\\big\|\_\{p\}:=\(df\_\{p\}\)^\{\\sharp\}\\in T\_\{p\}MThat is: first take the differentiald​fp∈Tp∗​Mdf\_\{p\}\\in T\_\{p\}^\{\*\}M\(a cotangent vector\), then raise the index using the metric to obtain a tangent vector\.

Equivalent definition \(expanding the meaning of♯\\sharp\):grad​f\|p\\mathrm\{grad\}\\,f\\big\|\_\{p\}is the unique tangent vector satisfying

gp​\(grad​f,v\)=d​fp​\(v\)∀v∈Tp​Mg\_\{p\}\\bigl\(\\mathrm\{grad\}\\,f,\\;v\\bigr\)=df\_\{p\}\(v\)\\quad\\forall\\,v\\in T\_\{p\}M\(70\)

Coordinate formula:

\(grad​f\)i=∑jgi​j​∂f∂xj,i\.e\.,grad​f=g−1⋅\(column vector of partial derivatives\)\(\\mathrm\{grad\}\\,f\)^\{i\}=\\sum\_\{j\}g^\{ij\}\\,\\frac\{\\partial f\}\{\\partial x^\{j\}\},\\qquad\\text\{i\.e\.,\}\\quad\\mathrm\{grad\}\\,f=g^\{\-1\}\\cdot\(\\text\{column vector of partial derivatives\}\)
Concrete example: full computation of “gradient = raising the index of the differential”

LetM=ℝ2M=\\mathbb\{R\}^\{2\},f​\(x1,x2\)=2​x1\+3​x2f\(x^\{1\},x^\{2\}\)=2x^\{1\}\+3x^\{2\}, with metricg=\(1114\)g=\\begin\{pmatrix\}1&1\\\\ 1&4\\end\{pmatrix\}\.

Step 1\.Compute the differentiald​f∈Tp∗​Mdf\\in T\_\{p\}^\{\*\}M\(independent of the metric\):

d​f=2​d​x1\+3​d​x2,\(d​f\)i=\(2,3\)​—a row vectordf=2\\,dx^\{1\}\+3\\,dx^\{2\},\\qquad\(df\)\_\{i\}=\(2,3\)\\;\\text\{\-\-\-a row vector\}
Step 2\.Raise the index:grad​f=\(d​f\)♯=g−1⋅\(d​f\)T\\mathrm\{grad\}\\,f=\(df\)^\{\\sharp\}=g^\{\-1\}\\cdot\(df\)^\{T\}:

g−1=13​\(4−1−11\),grad​f=13​\(4−1−11\)​\(23\)=\(5/31/3\)g^\{\-1\}=\\frac\{1\}\{3\}\\begin\{pmatrix\}4&\-1\\\\ \-1&1\\end\{pmatrix\},\\qquad\\mathrm\{grad\}\\,f=\\frac\{1\}\{3\}\\begin\{pmatrix\}4&\-1\\\\ \-1&1\\end\{pmatrix\}\\begin\{pmatrix\}2\\\\ 3\\end\{pmatrix\}=\\begin\{pmatrix\}5/3\\\\ 1/3\\end\{pmatrix\}
Step 3\.Verifyg​\(grad​f,v\)=d​f​\(v\)g\(\\mathrm\{grad\}\\,f,v\)=df\(v\):

Take arbitraryv=\(v1,v2\)Tv=\(v^\{1\},v^\{2\}\)^\{T\}:

g​\(grad​f,v\)=\(5/31/3\)T​\(1114\)​\(v1v2\)=2​v1\+3​v2=d​f​\(v\)✓g\(\\mathrm\{grad\}\\,f,\\,v\)=\\begin\{pmatrix\}5/3\\\\ 1/3\\end\{pmatrix\}^\{T\}\\\!\\begin\{pmatrix\}1&1\\\\ 1&4\\end\{pmatrix\}\\begin\{pmatrix\}v^\{1\}\\\\ v^\{2\}\\end\{pmatrix\}=2v^\{1\}\+3v^\{2\}=df\(v\)\\quad\\checkmark
Comparison with the standard metricg=Ig=I:grad​f=I−1​\(23\)=\(23\)\\mathrm\{grad\}\\,f=I^\{\-1\}\\binom\{2\}\{3\}=\\binom\{2\}\{3\}—recovering the familiar∇f\\nabla ffrom calculus\.

Key InsightCore summary:•Standard metricg=Ig=I: the numerical values ofgrad​f\\mathrm\{grad\}\\,fequal\(d​f\)T\(df\)^\{T\}; they “look the same”•Non\-standard metric:grad​f=g−1⋅\(d​f\)T≠\(d​f\)T\\mathrm\{grad\}\\,f=g^\{\-1\}\\cdot\(df\)^\{T\}\\neq\(df\)^\{T\}•The Wasserstein space has a non\-standard metric—so the gradient≠\\neqthe functional derivative

Physical intuition:g−1g^\{\-1\}performs a “cost correction”—in directions that are expensive under the metricgg, it moves less; in cheap directions, it moves more\. The gradient does not point in the “direction of largest partial derivatives,” but rather in the “direction of steepest descent offfper unit cost\.”

Three equivalent geometric meanings of the gradient:

Let‖v‖g=1\\\|v\\\|\_\{g\}=1\(a unit tangent vector in the metricgg\)\. Then:

1. 1\.The direction ofgrad​f\\mathrm\{grad\}\\,fis the direction that maximizesd​f​\(v\)df\(v\)\(finding the steepest direction in thegg\-unit ball\)
2. 2\.The direction ofgrad​f\\mathrm\{grad\}\\,fis thegg\-normal direction to the level set\{f=c\}\\\{f=c\\\}
3. 3\.‖grad​f‖g=max‖v‖g=1⁡d​f​\(v\)\\\|\\mathrm\{grad\}\\,f\\\|\_\{g\}=\\max\_\{\\\|v\\\|\_\{g\}=1\}df\(v\)

IntuitionPreview: why are the Wasserstein gradient and theL2L^\{2\}gradient different?For the same functionalℱ​\(ρ\)\\mathcal\{F\}\(\\rho\)and the same functional derivativeδ​ℱδ​ρ\\frac\{\\delta\\mathcal\{F\}\}\{\\delta\\rho\}, but:•L2L^\{2\}metric⟨u,v⟩L2=∫u​v​𝑑x\\langle u,v\\rangle\_\{L^\{2\}\}=\\int u\\,v\\,dx⇒\\;\\Rightarrow\\;gradL2​ℱ=δ​ℱδ​ρ\\mathrm\{grad\}\_\{L^\{2\}\}\\mathcal\{F\}=\\frac\{\\delta\\mathcal\{F\}\}\{\\delta\\rho\}\(theg=Ig=Icase\)•Wasserstein metric⟨ξ,η⟩ρ=∫ρ​ξ⋅η​𝑑x\\langle\\xi,\\eta\\rangle\_\{\\rho\}=\\int\\rho\\,\\xi\\cdot\\eta\\,dx⇒\\;\\Rightarrow\\;gradW​ℱ=∇δ​ℱδ​ρ\\mathrm\{grad\}\_\{W\}\\mathcal\{F\}=\\nabla\\frac\{\\delta\\mathcal\{F\}\}\{\\delta\\rho\}\(g≠Ig\\neq I, with an extra∇\\nabla\)This will be derived in detail in Section 5\.

\(7\) Gradient Flow

###### Definition B\.7\(Gradient flow\)\.

Let\(M,g\)\(M,g\)be a Riemannian manifold andf:M→ℝf:M\\to\\mathbb\{R\}smooth\. Thegradient flowofffis a curveγ​\(t\)\\gamma\(t\)satisfying

γ˙​\(t\)=−grad​f\|γ​\(t\)\\dot\{\\gamma\}\(t\)=\-\\mathrm\{grad\}\\,f\\big\|\_\{\\gamma\(t\)\}That is: at each instant, move in the direction of steepest descent offf\.

- •Inℝn\\mathbb\{R\}^\{n\}\(standard metric\):x˙=−∇f​\(x\)\\dot\{x\}=\-\\nabla f\(x\)\(continuous\-time gradient descent\)
- •On a Riemannian manifold:γ˙=−grad​f=−\(d​f\)♯\\dot\{\\gamma\}=\-\\mathrm\{grad\}\\,f=\-\(df\)^\{\\sharp\}\(“rolling downhill” along the surface\)
- •In Wasserstein space: the gradient flow yields a PDE \(the Fokker–Planck equation\)

Key property of gradient flow:ffstrictly decreases along the flow—dd​t​f​\(γ​\(t\)\)=−‖grad​f‖g2≤0\\frac\{d\}\{dt\}f\(\\gamma\(t\)\)=\-\\\|\\mathrm\{grad\}\\,f\\\|\_\{g\}^\{2\}\\leq 0\.

FAQ 1: Isd​fdfa complete object on its own? Or does it need a direction to be defined?

d​fpdf\_\{p\}is itself afully defined object—it is an element ofTp∗​MT\_\{p\}^\{\*\}M\(a cotangent vector\) and does not need a “given direction” to exist\. It simultaneously encodes information aboutall directions: given anyv∈Tp​Mv\\in T\_\{p\}M, it outputsd​fp​\(v\)∈ℝdf\_\{p\}\(v\)\\in\\mathbb\{R\}\.

Analogy: inℝn\\mathbb\{R\}^\{n\},d​fpdf\_\{p\}corresponds to the row vector\(∂f∂x1,…,∂f∂xn\)\\bigl\(\\frac\{\\partial f\}\{\\partial x^\{1\}\},\\ldots,\\frac\{\\partial f\}\{\\partial x^\{n\}\}\\bigr\)\. This row vector is already completely determined; no direction needs to be specified\. Its “action” is multiplication with a column vectorvv:d​fp​\(v\)=∑i∂f∂xi​vidf\_\{p\}\(v\)=\\sum\_\{i\}\\frac\{\\partial f\}\{\\partial x^\{i\}\}v^\{i\}\.

FAQ 2: Isgrad​f\\mathrm\{grad\}\\,fa “direction”?

Yes\.grad​f\|p∈Tp​M\\mathrm\{grad\}\\,f\\big\|\_\{p\}\\in T\_\{p\}Mis a tangent vector—it lives in the tangent space, so it has both adirectionand amagnitude\.

Specifically:

- •d​fp∈Tp∗​Mdf\_\{p\}\\in T\_\{p\}^\{\*\}Mis a cotangent vector \(a “measuring instrument”\); it isnotitself a direction
- •grad​f\|p=\(d​fp\)♯∈Tp​M\\mathrm\{grad\}\\,f\\big\|\_\{p\}=\(df\_\{p\}\)^\{\\sharp\}\\in T\_\{p\}Mis a tangent vector \(an “arrow”\); itisa direction

The meaning of this direction: under the metricgg, it is the direction of fastest increase offf\(the direction ofvvthat maximizesd​f​\(v\)/‖v‖gdf\(v\)/\\\|v\\\|\_\{g\}\)\.

So in the gradient flowγ˙​\(t\)=−grad​f\\dot\{\\gamma\}\(t\)=\-\\mathrm\{grad\}\\,f, both sides of the equation are tangent vectors inTγ​\(t\)​MT\_\{\\gamma\(t\)\}M: “the velocity of the curve”==“the negative gradient direction”—two vectors in the same tangent space being equal\.

FAQ 3: Why doesdd​t​f​\(γ​\(t\)\)=−‖grad​f‖g2\\frac\{d\}\{dt\}f\(\\gamma\(t\)\)=\-\\\|\\mathrm\{grad\}\\,f\\\|\_\{g\}^\{2\}involve a square?

Derivation:

Step 1\.Chain rule \(for functions on manifolds\):

dd​t​f​\(γ​\(t\)\)=d​fγ​\(t\)​\(γ˙​\(t\)\)\\frac\{d\}\{dt\}f\(\\gamma\(t\)\)=df\_\{\\gamma\(t\)\}\\bigl\(\\dot\{\\gamma\}\(t\)\\bigr\)
Step 2\.Substitute the gradient flow definitionγ˙=−grad​f\\dot\{\\gamma\}=\-\\mathrm\{grad\}\\,f:

=d​f​\(−grad​f\)=−d​f​\(grad​f\)=df\\bigl\(\-\\mathrm\{grad\}\\,f\\bigr\)=\-df\\bigl\(\\mathrm\{grad\}\\,f\\bigr\)
Step 3\.Use the definition of the gradientg​\(grad​f,v\)=d​f​\(v\)g\(\\mathrm\{grad\}\\,f,\\,v\)=df\(v\), takingv=grad​fv=\\mathrm\{grad\}\\,f:

d​f​\(grad​f\)=g​\(grad​f,grad​f\)=‖grad​f‖g2df\(\\mathrm\{grad\}\\,f\)=g\(\\mathrm\{grad\}\\,f,\\,\\mathrm\{grad\}\\,f\)=\\\|\\mathrm\{grad\}\\,f\\\|\_\{g\}^\{2\}
Conclusion:

dd​t​f​\(γ​\(t\)\)=−‖grad​f‖g2≤0\\frac\{d\}\{dt\}f\(\\gamma\(t\)\)=\-\\\|\\mathrm\{grad\}\\,f\\\|\_\{g\}^\{2\}\\leq 0
The reason a square appears:the gradient plays two roles simultaneously—it is both the velocity \(viaγ˙=−grad​f\\dot\{\\gamma\}=\-\\mathrm\{grad\}\\,f\) and the “metric dual” for the directional derivative \(viad​f​\(v\)=g​\(grad​f,v\)df\(v\)=g\(\\mathrm\{grad\}\\,f,v\)\)\. Whenvvhappens to equalgrad​f\\mathrm\{grad\}\\,fitself,g​\(grad​f,grad​f\)=‖grad​f‖2g\(\\mathrm\{grad\}\\,f,\\mathrm\{grad\}\\,f\)=\\\|\\mathrm\{grad\}\\,f\\\|^\{2\}\.

Analogy inℝn\\mathbb\{R\}^\{n\}:dd​t​f​\(x​\(t\)\)=∇f⋅x˙=∇f⋅\(−∇f\)=−\|∇f\|2\\frac\{d\}\{dt\}f\(x\(t\)\)=\\nabla f\\cdot\\dot\{x\}=\\nabla f\\cdot\(\-\\nabla f\)=\-\|\\nabla f\|^\{2\}\.

\(8\) Summary—the logical chain of all concepts:

ManifoldM→smooth structuredefinesTp​M,Tp∗​M→choosegextra structure♯,♭→raise indexd​f↦\(d​f\)♯grad​f→negative direction−grad​fgradient flow\\boxed\{\\text\{Manifold $M$\}\\xrightarrow\[\\text\{smooth structure\}\]\{\\text\{defines\}\}T\_\{p\}M,\\,T\_\{p\}^\{\*\}M\\xrightarrow\[\\text\{choose $g$\}\]\{\\text\{extra structure\}\}\\sharp,\\flat\\xrightarrow\[\\text\{raise index\}\]\{df\\mapsto\(df\)^\{\\sharp\}\}\\mathrm\{grad\}\\,f\\xrightarrow\[\\text\{negative direction\}\]\{\-\\mathrm\{grad\}\\,f\}\\text\{gradient flow\}\}
Key insight:The gradient and gradient flow depend on the choice of metric\.Different metrics⇒\\Rightarrowdifferent♯\\sharp⇒\\Rightarrowdifferent gradients⇒\\Rightarrowdifferent gradient flow PDEs\.

This is why the same free energy functionalℱ\\mathcal\{F\}, under theL2L^\{2\}metric gives the heat equation, while under the Wasserstein metric gives the Fokker–Planck equation\.

## Appendix CAnalysis Tools: Integration by Parts and du Bois\-Reymond Lemma

This appendix collects the analytical tools used repeatedly throughout the text: the family of integration by parts formulas and the du Bois\-Reymond lemma\.

### C\.1Integration by Parts

The essence of integration by parts istransferring a derivative from one function to another, at the cost of introducing a minus sign \(and boundary terms\)\.

#### C\.1\.1One\-dimensional case

###### Theorem C\.1\(One\-dimensional integration by parts\)\.

Letf,g∈C1​\(\[a,b\]\)f,g\\in C^\{1\}\(\[a,b\]\)\. Then:

∫abf′​\(x\)​g​\(x\)​𝑑x=\[f​\(x\)​g​\(x\)\]ab−∫abf​\(x\)​g′​\(x\)​𝑑x\\int\_\{a\}^\{b\}f^\{\\prime\}\(x\)\\,g\(x\)\\,dx=\\big\[f\(x\)g\(x\)\\big\]\_\{a\}^\{b\}\-\\int\_\{a\}^\{b\}f\(x\)\\,g^\{\\prime\}\(x\)\\,dxIffforggvanishes at the endpoints \(e\.g\., has compact support\), then the boundary terms vanish:

∫abf′​g​𝑑x=−∫abf​g′​𝑑x\\int\_\{a\}^\{b\}f^\{\\prime\}\\,g\\,dx=\-\\int\_\{a\}^\{b\}f\\,g^\{\\prime\}\\,dx

Intuition:Integrate the product rule\(f​g\)′=f′​g\+f​g′\(fg\)^\{\\prime\}=f^\{\\prime\}g\+fg^\{\\prime\}on both sides\. The left\-hand side becomes the boundary term\[f​g\]ab\[fg\]\_\{a\}^\{b\}; rearranging yields integration by parts\.

#### C\.1\.2Higher\-dimensional case: divergence theorem version

###### Theorem C\.2\(Integration by Parts in Divergence Form / Green’s First Identity\)\.

Letφ\\varphibe a smooth function with compact support \(or vanishing on∂Ω\\partial\\Omega\), and let𝐅\\mathbf\{F\}be a smooth vector field\. Then:

∫ℝdφ​\(∇⋅𝐅\)​𝑑x=−∫ℝd∇φ⋅𝐅​d​x\\int\_\{\\mathbb\{R\}^\{d\}\}\\varphi\\,\(\\nabla\\cdot\\mathbf\{F\}\)\\,dx=\-\\int\_\{\\mathbb\{R\}^\{d\}\}\\nabla\\varphi\\cdot\\mathbf\{F\}\\,dx\(71\)

Derivation:Integrate the product rule∇⋅\(φ​𝐅\)=∇φ⋅𝐅\+φ​∇⋅𝐅\\nabla\\cdot\(\\varphi\\mathbf\{F\}\)=\\nabla\\varphi\\cdot\\mathbf\{F\}\+\\varphi\\,\\nabla\\cdot\\mathbf\{F\}on both sides\. The left\-hand side becomes the boundary integral∮φ​𝐅⋅𝐧​𝑑S=0\\oint\\varphi\\mathbf\{F\}\\cdot\\mathbf\{n\}\\,dS=0by the divergence theorem \(sinceφ\\varphihas compact support\); rearranging gives the result\.

Typical usage in the main text:

- •𝐅=ρ​∇V\\mathbf\{F\}=\\rho\\nabla V:∫∇φ⋅\(ρ​∇V\)​𝑑x=−∫φ​∇⋅\(ρ​∇V\)​𝑑x\\displaystyle\\int\\nabla\\varphi\\cdot\(\\rho\\nabla V\)\\,dx=\-\\int\\varphi\\,\\nabla\\cdot\(\\rho\\nabla V\)\\,dx
- •𝐅=ρ​η\\mathbf\{F\}=\\rho\\,\\eta\(η\\etais a perturbation vector field\):∫f​∇⋅\(ρ​η\)​𝑑x=−∫∇f⋅\(ρ​η\)​𝑑x\\displaystyle\\int f\\,\\nabla\\cdot\(\\rho\\eta\)\\,dx=\-\\int\\nabla f\\cdot\(\\rho\\eta\)\\,dx

#### C\.1\.3Transferring the Laplacian \(two integrations by parts\)

###### Theorem C\.3\(Green’s second identity\)\.

Letφ\\varphiandggbe smooth functions withφ\\varphicompactly supported\. Then:

∫ℝd\(Δ​φ\)​g​𝑑x=∫ℝdφ​\(Δ​g\)​𝑑x\\int\_\{\\mathbb\{R\}^\{d\}\}\(\\Delta\\varphi\)\\,g\\,dx=\\int\_\{\\mathbb\{R\}^\{d\}\}\\varphi\\,\(\\Delta g\)\\,dx\(72\)

Derivation:Apply \([71](https://arxiv.org/html/2606.24157#A3.E71)\) twice:

∫\(Δ​φ\)​g​𝑑x\\displaystyle\\int\(\\Delta\\varphi\)\\,g\\,dx=∫\(∇⋅∇φ\)​g​𝑑x=−∫∇φ⋅∇g​d​x\(first application, minus sign\)\\displaystyle=\\int\(\\nabla\\cdot\\nabla\\varphi\)\\,g\\,dx=\-\\int\\nabla\\varphi\\cdot\\nabla g\\,dx\\quad\\text\{\(first application, minus sign\)\}=\+∫φ​\(∇⋅∇g\)​𝑑x=∫φ​Δ​g​𝑑x\(second application,−⁣×⁣−⁣=⁣\+\)\\displaystyle=\+\\int\\varphi\\,\(\\nabla\\cdot\\nabla g\)\\,dx=\\int\\varphi\\,\\Delta g\\,dx\\quad\\text\{\(second application, $\-\\times\-=\+$\)\}
Intuition:The second\-order derivative can “jump” from one function to another without changing sign—because two integrations by parts produce two minus signs that cancel\.

#### C\.1\.4Integration by parts in time

The same procedure applies to integration in time\. Letφ​\(t,x\)\\varphi\(t,x\)vanish att=0t=0andt=Tt=T:

∫0T∫φ​∂tρ​d​x​d​t=−∫0T∫\(∂tφ\)​ρ​𝑑x​𝑑t\\int\_\{0\}^\{T\}\\\!\\int\\varphi\\,\\partial\_\{t\}\\rho\\,dx\\,dt=\-\\int\_\{0\}^\{T\}\\\!\\int\(\\partial\_\{t\}\\varphi\)\\,\\rho\\,dx\\,dtThis is used when deriving the weak form of the continuity equation\.

#### C\.1\.5Formula summary table

### C\.2Du Bois\-Reymond Lemma \(Fundamental Lemma of the Calculus of Variations\)

###### Theorem C\.4\(Du Bois\-Reymond Lemma / Fundamental Lemma of Calculus of Variations\)\.

Leth:ℝd→ℝh:\\mathbb\{R\}^\{d\}\\to\\mathbb\{R\}be a locally integrable function\. If for all compactly supported smooth functionsφ∈Cc∞​\(ℝd\)\\varphi\\in C\_\{c\}^\{\\infty\}\(\\mathbb\{R\}^\{d\}\)we have:

∫ℝdh​\(x\)​φ​\(x\)​𝑑x=0\\int\_\{\\mathbb\{R\}^\{d\}\}h\(x\)\\,\\varphi\(x\)\\,dx=0\(73\)thenh​\(x\)=0h\(x\)=0almost everywhere \(a\.e\.\)\. If furthermorehhis continuous, thenh​\(x\)=0h\(x\)=0holds everywhere\.

#### C\.2\.1Intuitive understanding

Finite\-dimensional analogy:If a vectorv∈ℝnv\\in\\mathbb\{R\}^\{n\}has zero inner product with every vector \(⟨v,w⟩=0​∀w\\langle v,w\\rangle=0\\;\\forall w\), thenv=0v=0\.

The Du Bois\-Reymond lemma is theinfinite\-dimensional versionof this fact:

- •“vector”→\\tofunctionh​\(x\)h\(x\)
- •“inner product”→\\tointegral∫h​φ​𝑑x\\int h\\varphi\\,dx
- •“zero inner product with all vectors”→\\to“zero integral against all test functionsφ\\varphi”
- •Conclusion:h=0h=0\(a\.e\.\)

#### C\.2\.2Proof sketch

Proof by contradiction: supposeh​\(x0\)\>0h\(x\_\{0\}\)\>0at some point\. By continuity,h\>0h\>0on some small ballBϵ​\(x0\)B\_\{\\epsilon\}\(x\_\{0\}\)aroundx0x\_\{0\}\. Choose a nonnegativeφ\\varphisupported insideBϵ​\(x0\)B\_\{\\epsilon\}\(x\_\{0\}\)\(e\.g\., a smooth bump function\)\. Then:

∫h​φ​𝑑x=∫Bϵ​\(x0\)h​φ​𝑑x\>0\\int h\\,\\varphi\\,dx=\\int\_\{B\_\{\\epsilon\}\(x\_\{0\}\)\}h\\,\\varphi\\,dx\>0This contradicts the hypothesis\. Hencehhcannot be positive at any point\. By the same argument it cannot be negative\. Thereforeh=0h=0\.

#### C\.2\.3Constrained version \(used in the main text\)

What appears in Step 1 of the JKO derivation in the main text is aconstrained version:

###### Theorem C\.5\(Version with mass conservation constraint\)\.

Lethhbe a continuous function\. If for all smooth perturbationsδ​ρ\\delta\\rhosatisfying∫δ​ρ​𝑑x=0\\int\\delta\\rho\\,dx=0we have:

∫h​\(x\)​δ​ρ​\(x\)​𝑑x=0\\int h\(x\)\\,\\delta\\rho\(x\)\\,dx=0thenh​\(x\)=consth\(x\)=\\mathrm\{const\}\(a constant\)\.

Why a constant rather than zero?

Because the test functionsδ​ρ\\delta\\rhoare constrained to satisfy∫δ​ρ=0\\int\\delta\\rho=0—they cannot be “arbitrary” but must have “zero total mass\.” This means we lose one “probing direction” \(the constant\-function direction\)\.

- •Unconstrained:∫h​φ=0​∀φ\\int h\\varphi=0\\;\\forall\\varphi⟹\\impliesh=0h=0
- •Constrained∫δ​ρ=0\\int\\delta\\rho=0:∫h​δ​ρ=0​∀δ​ρ​s\.t\.​∫δ​ρ=0\\int h\\,\\delta\\rho=0\\;\\forall\\,\\delta\\rho\\text\{ s\.t\. \}\\int\\delta\\rho=0⟹\\impliesh=consth=\\text\{const\}

Proof:For any two pointsx1,x2x\_\{1\},x\_\{2\}, takeδ​ρ=φ1−φ2\\delta\\rho=\\varphi\_\{1\}\-\\varphi\_\{2\}, whereφi\\varphi\_\{i\}is a nonnegative smooth bump function centered atxix\_\{i\}with∫φi=1\\int\\varphi\_\{i\}=1\. Then∫δ​ρ=0\\int\\delta\\rho=0, and by hypothesis:

∫h​φ1​𝑑x=∫h​φ2​𝑑x\\int h\\,\\varphi\_\{1\}\\,dx=\\int h\\,\\varphi\_\{2\}\\,dxAs the support of the bump functions shrinks toxix\_\{i\}\(φi→δxi\\varphi\_\{i\}\\to\\delta\_\{x\_\{i\}\}\), continuity ofhhgivesh​\(x1\)=h​\(x2\)h\(x\_\{1\}\)=h\(x\_\{2\}\)\. Sincex1,x2x\_\{1\},x\_\{2\}are arbitrary,hhis constant\.

#### C\.2\.4Where it is used in the main text

IntuitionOne\-sentence summary:The Du Bois\-Reymond lemma is the “starting tool” of the calculus of variations—it lets us pass from integral equalities \(weak/variational form\) to pointwise equalities \(strong form/PDE\)\.Without it, we can only say “the weighted average is zero”; with it, we can say “the function itself is zero \(or constant\)\.” This is the bridge from “integral equations” to “differential equations\.”

## Appendix DGibbs–Boltzmann Distribution

The formπ​\(x\)=1Z​e−β​V​\(x\)\\pi\(x\)=\\frac\{1\}\{Z\}e^\{\-\\beta V\(x\)\}appears repeatedly across all corners of science—whenever a system strikes a balance between “energy” and “randomness,” its stationary distribution is the Gibbs–Boltzmann distribution\.

### D\.1Core intuition

At thermal equilibrium, low\-energy states occur withexponentiallyhigher probability than high\-energy states\. The probability ratio between two states depends only on their energy difference:

π​\(x1\)π​\(x2\)=e−β​\(V​\(x1\)−V​\(x2\)\)\\frac\{\\pi\(x\_\{1\}\)\}\{\\pi\(x\_\{2\}\)\}=e^\{\-\\beta\(V\(x\_\{1\}\)\-V\(x\_\{2\}\)\)\}For everykB​Tk\_\{B\}Tunit of energy difference, the probabilities differ by a factor ofe≈2\.7e\\approx 2\.7\.

### D\.2The role of temperatureTT—the single control knob

- •T→0T\\to 0\(β→∞\\beta\\to\\infty\): the distribution degenerates to a Dirac delta; the system “freezes” at the lowest\-energy state
- •T→∞T\\to\\infty\(β→0\\beta\\to 0\): the distribution approaches uniform; the system completely ignores energy and wanders randomly
- •FiniteTT: a compromise between the two—low\-energy states are “preferred” but not exclusively

Intuition: temperature controls the ratio of “exploration” to “exploitation\.” Low temperature = pure exploitation \(stay at the optimum\); high temperature = pure exploration \(wander everywhere\)\.

### D\.3Applications across disciplines

### D\.4Core connection to the main text

The Fokker\-Planck equation∂tρ=Δ​ρ\+∇⋅\(ρ​∇V\)\\partial\_\{t\}\\rho=\\Delta\\rho\+\\nabla\\cdot\(\\rho\\nabla V\)describes precisely the process by which the system evolves from an arbitrary initial distributionρ0\\rho\_\{0\}toward the Gibbs–Boltzmann equilibriumπ∝e−V\\pi\\propto e^\{\-V\}\.

- •The free energyℱ​\(ρ\)=KL​\(ρ∥π\)\+const\\mathcal\{F\}\(\\rho\)=\\mathrm\{KL\}\(\\rho\\\|\\pi\)\+\\text\{const\}measures “how far from equilibrium”
- •The unique minimizer ofℱ\\mathcal\{F\}isπ\\pi\(KL​\(ρ∥π\)=0⇔ρ=π\\mathrm\{KL\}\(\\rho\\\|\\pi\)=0\\iff\\rho=\\pi\)
- •Each step of the JKO scheme minimizesℱ\+12​τ​W22\\mathcal\{F\}\+\\frac\{1\}\{2\\tau\}W\_\{2\}^\{2\}, discretely approximating this approach to equilibrium

The endpoint of the entire story is this elegante−Ve^\{\-V\}\.

### D\.5Why the exponential form?

Thee−β​Ve^\{\-\\beta V\}form of the Gibbs distribution is not an arbitrary choice but is dictated by themaximum entropy principle:

Subject to the constraint that the mean energy⟨V⟩=∫V​𝑑ρ=E\\langle V\\rangle=\\int V\\,d\\rho=Eis fixed, the distribution that maximizes the entropy−∫ρ​log⁡ρ\-\\int\\rho\\log\\rhoisρ∝e−β​V\\rho\\propto e^\{\-\\beta V\}, whereβ\\betais the Lagrange multiplier \(which happens to correspond to the inverse temperature\)\.

In other words:among all distributions with mean energyEE, the Gibbs distribution is the “most random” \(most uncertain\) one\.Nature selects it because there is no reason to prefer a more ordered state—this is the statistical interpretation of the second law of thermodynamics\.

## Appendix EFisher Information

The name “Fisher information” hastwo different meaningsin mathematics, originating from statistics and information theory/PDE respectively\. Below we introduce each in turn, explain how they are unified, and finally present the relative version used in the main text\.

### E\.1Fisher Information in Statistics \(Parametric Version\)

###### Definition E\.1\(Fisher Information — Statistical Version\)\.

Given a parametric family of probability distributions\{pθ​\(x\)\}θ∈ℝk\\\{p\_\{\\theta\}\(x\)\\\}\_\{\\theta\\in\\mathbb\{R\}^\{k\}\}, theFisher information matrix\(Fisher, 1925\) is defined as:

ℐ​\(θ\)i​j:=𝔼x∼pθ​\[∂log⁡pθ​\(x\)∂θi⋅∂log⁡pθ​\(x\)∂θj\]=−𝔼x∼pθ​\[∂2log⁡pθ​\(x\)∂θi​∂θj\]\\mathcal\{I\}\(\\theta\)\_\{ij\}:=\\mathbb\{E\}\_\{x\\sim p\_\{\\theta\}\}\\left\[\\frac\{\\partial\\log p\_\{\\theta\}\(x\)\}\{\\partial\\theta\_\{i\}\}\\cdot\\frac\{\\partial\\log p\_\{\\theta\}\(x\)\}\{\\partial\\theta\_\{j\}\}\\right\]=\-\\mathbb\{E\}\_\{x\\sim p\_\{\\theta\}\}\\left\[\\frac\{\\partial^\{2\}\\log p\_\{\\theta\}\(x\)\}\{\\partial\\theta\_\{i\}\\partial\\theta\_\{j\}\}\\right\]\(74\)In the one\-dimensional parameter case:ℐ​\(θ\)=𝔼x∼pθ​\[\(∂log⁡pθ∂θ\)2\]\\mathcal\{I\}\(\\theta\)=\\mathbb\{E\}\_\{x\\sim p\_\{\\theta\}\}\\left\[\\left\(\\frac\{\\partial\\log p\_\{\\theta\}\}\{\\partial\\theta\}\\right\)^\{2\}\\right\]\.

Intuition:ℐ​\(θ\)\\mathcal\{I\}\(\\theta\)measures “how much information the data carry about the parameterθ\\theta\.”

- •ℐ​\(θ\)\\mathcal\{I\}\(\\theta\)large→\\tochangingθ\\thetacauses a drastic change in the distribution→\\todata can easily distinguish differentθ\\theta→\\toeasy to estimateθ\\theta
- •ℐ​\(θ\)\\mathcal\{I\}\(\\theta\)small→\\tochangingθ\\thetabarely affects the distribution→\\todata are insensitive toθ\\theta→\\tohard to estimate
- •Cramér\-Rao lower bound: the variance of any unbiased estimatorθ^\\hat\{\\theta\}satisfiesVar​\(θ^\)≥ℐ​\(θ\)−1\\mathrm\{Var\}\(\\hat\{\\theta\}\)\\geq\\mathcal\{I\}\(\\theta\)^\{\-1\}

Example:pθ=𝒩​\(θ,σ2\)p\_\{\\theta\}=\\mathcal\{N\}\(\\theta,\\sigma^\{2\}\)\(mean as parameter\)\.log⁡pθ=−\(x−θ\)22​σ2\+C\\log p\_\{\\theta\}=\-\\frac\{\(x\-\\theta\)^\{2\}\}\{2\\sigma^\{2\}\}\+C,∂∂θ​log⁡pθ=x−θσ2\\frac\{\\partial\}\{\\partial\\theta\}\\log p\_\{\\theta\}=\\frac\{x\-\\theta\}\{\\sigma^\{2\}\},ℐ​\(θ\)=𝔼​\[\(x−θ\)2σ4\]=1σ2\\mathcal\{I\}\(\\theta\)=\\mathbb\{E\}\\left\[\\frac\{\(x\-\\theta\)^\{2\}\}\{\\sigma^\{4\}\}\\right\]=\\frac\{1\}\{\\sigma^\{2\}\}\. The smaller the variance, the larger the Fisher information—the more “concentrated” the data, the easier it is to estimate the mean\.

### E\.2Fisher Information in Information Theory/PDE \(Distribution Version\)

###### Definition E\.2\(Fisher Information — Distribution Version / de Bruijn Version\)\.

Given a probability densityρ\\rho\(not depending on any parameter\), itsFisher informationis defined as:

I​\(ρ\):=∫ℝdρ​\(x\)​\|∇xlog⁡ρ​\(x\)\|2​𝑑x=∫ℝd\|∇ρ​\(x\)\|2ρ​\(x\)​𝑑xI\(\\rho\):=\\int\_\{\\mathbb\{R\}^\{d\}\}\\rho\(x\)\\,\|\\nabla\_\{x\}\\log\\rho\(x\)\|^\{2\}\\,dx=\\int\_\{\\mathbb\{R\}^\{d\}\}\\frac\{\|\\nabla\\rho\(x\)\|^\{2\}\}\{\\rho\(x\)\}\\,dx\(75\)

Intuition:I​\(ρ\)I\(\\rho\)measures the “spatial irregularity” of the densityρ\\rho\.

- •The smoother/flatterρ\\rhois→\\tothe smaller∇log⁡ρ\\nabla\\log\\rho→\\tothe smallerI​\(ρ\)I\(\\rho\)
- •The sharper/more rapidly varyingρ\\rhois→\\tothe larger∇log⁡ρ\\nabla\\log\\rho→\\tothe largerI​\(ρ\)I\(\\rho\)
- •Uniform distribution:I=0I=0\(completely flat\)
- •Dirac delta:I=∞I=\\infty\(infinitely sharp\)

Relation to the heat equation\(de Bruijn identity\): ifρt\\rho\_\{t\}solves the heat equation∂tρ=Δ​ρ\\partial\_\{t\}\\rho=\\Delta\\rho, then:

dd​t​H​\(ρt\)=−I​\(ρt\)\(rate of entropy change=negative Fisher information\)\\frac\{d\}\{dt\}H\(\\rho\_\{t\}\)=\-I\(\\rho\_\{t\}\)\\qquad\\text\{\(rate of entropy change $=$ negative Fisher information\)\}whereH​\(ρ\)=−∫ρ​log⁡ρH\(\\rho\)=\-\\int\\rho\\log\\rhois the entropy\. Diffusion flattens the distribution \(entropy increases\), and the rate is controlled by the Fisher information\.

### E\.3Unification of the two versions

These two forms of Fisher information look different; they are unified as follows:

Key observation: take the parametric family to be the “translation family”pθ​\(x\)=ρ​\(x−θ\)p\_\{\\theta\}\(x\)=\\rho\(x\-\\theta\)\. Then:

∂∂θ​log⁡pθ​\(x\)\|θ=0=−∇xlog⁡ρ​\(x\)\\frac\{\\partial\}\{\\partial\\theta\}\\log p\_\{\\theta\}\(x\)\\bigg\|\_\{\\theta=0\}=\-\\nabla\_\{x\}\\log\\rho\(x\)Substituting into the statistical version of the Fisher information:

ℐ​\(0\)=𝔼x∼ρ​\[\|∇xlog⁡ρ​\(x\)\|2\]=I​\(ρ\)\\mathcal\{I\}\(0\)=\\mathbb\{E\}\_\{x\\sim\\rho\}\\left\[\|\\nabla\_\{x\}\\log\\rho\(x\)\|^\{2\}\\right\]=I\(\\rho\)The two coincide\!The distribution\-version Fisher information is precisely the statistical\-version Fisher information evaluated at the “translation parameter\.”

More generally:

- •The statistical version measures “how easily differentθ\\thetacan be distinguished in parameter space”
- •The distribution version measures “how rapidly the density varies across different locationsxxin physical space”
- •When “the parameter is the location,” the two are equivalent

Derivation details: why does∂∂θ​log⁡pθ​\(x\)\|θ=0=−∇xlog⁡ρ​\(x\)\\frac\{\\partial\}\{\\partial\\theta\}\\log p\_\{\\theta\}\(x\)\\big\|\_\{\\theta=0\}=\-\\nabla\_\{x\}\\log\\rho\(x\)?

Letpθ​\(x\)=ρ​\(x−θ\)p\_\{\\theta\}\(x\)=\\rho\(x\-\\theta\), solog⁡pθ​\(x\)=log⁡ρ​\(x−θ\)\\log p\_\{\\theta\}\(x\)=\\log\\rho\(x\-\\theta\)\.

∂∂θ​log⁡ρ​\(x−θ\)=−∇ylog⁡ρ​\(y\)\|y=x−θ⋅1=−∇log⁡ρ​\(x−θ\)\\frac\{\\partial\}\{\\partial\\theta\}\\log\\rho\(x\-\\theta\)=\-\\nabla\_\{y\}\\log\\rho\(y\)\\big\|\_\{y=x\-\\theta\}\\cdot 1=\-\\nabla\\log\\rho\(x\-\\theta\)Atθ=0\\theta=0:=−∇log⁡ρ​\(x\)=\-\\nabla\\log\\rho\(x\)\.

So the statistical version becomes:

ℐ​\(0\)=𝔼x∼ρ​\[\(−∇log⁡ρ​\(x\)\)2\]=∫ρ​\|∇log⁡ρ\|2​𝑑x=I​\(ρ\)✓\\mathcal\{I\}\(0\)=\\mathbb\{E\}\_\{x\\sim\\rho\}\\left\[\(\-\\nabla\\log\\rho\(x\)\)^\{2\}\\right\]=\\int\\rho\\,\|\\nabla\\log\\rho\|^\{2\}\\,dx=I\(\\rho\)\\quad\\checkmark

### E\.4Relative Fisher Information

###### Definition E\.3\(Relative Fisher Information\)\.

Given probability densitiesρ\\rhoand a target distributionπ\\pi, therelative Fisher informationis defined as:

I​\(ρ∥π\):=∫ρ​\(x\)​\|∇log⁡ρ​\(x\)π​\(x\)\|2​𝑑x=∫ρ​\(x\)​\|∇log⁡ρ​\(x\)−∇log⁡π​\(x\)\|2​𝑑xI\(\\rho\\\|\\pi\):=\\int\\rho\(x\)\\left\|\\nabla\\log\\frac\{\\rho\(x\)\}\{\\pi\(x\)\}\\right\|^\{2\}dx=\\int\\rho\(x\)\\left\|\\nabla\\log\\rho\(x\)\-\\nabla\\log\\pi\(x\)\\right\|^\{2\}dx\(76\)

Whenπ​\(x\)∝e−V​\(x\)\\pi\(x\)\\propto e^\{\-V\(x\)\}, we have∇log⁡π=−∇V\\nabla\\log\\pi=\-\\nabla V, so:

I​\(ρ∥π\)=∫ρ​\|∇log⁡ρ\+∇V\|2​𝑑x=∫ρ​\|sρ​\(x\)−sπ​\(x\)\|2​𝑑xI\(\\rho\\\|\\pi\)=\\int\\rho\\,\|\\nabla\\log\\rho\+\\nabla V\|^\{2\}\\,dx=\\int\\rho\\,\|s\_\{\\rho\}\(x\)\-s\_\{\\pi\}\(x\)\|^\{2\}\\,dxwheresρ=∇log⁡ρs\_\{\\rho\}=\\nabla\\log\\rhoandsπ=∇log⁡π=−∇Vs\_\{\\pi\}=\\nabla\\log\\pi=\-\\nabla V\.

Relation to ordinary Fisher information: whenπ\\piis the uniform distribution \(V=0V=0\):

I​\(ρ∥π\)=∫ρ​\|∇log⁡ρ\|2​𝑑x=I​\(ρ\)\(reduces to the ordinary version\)I\(\\rho\\\|\\pi\)=\\int\\rho\\,\|\\nabla\\log\\rho\|^\{2\}\\,dx=I\(\\rho\)\\qquad\\text\{\(reduces to the ordinary version\)\}
Intuition:I​\(ρ∥π\)I\(\\rho\\\|\\pi\)measures the discrepancy between the score ofρ\\rhoand the score ofπ\\pi\. Whenρ=π\\rho=\\pi, the two scores coincide andI=0I=0; the furtherρ\\rhodeviates fromπ\\pi, the largerIIbecomes\.

### E\.5Three identities

- •Measures “how far from equilibrium” \(dynamic sense\):I​\(ρ∥π\)I\(\\rho\\\|\\pi\)quantifies how fastρ\\rhois approachingπ\\pi\.I=0I=0means equilibrium has been reached; largeIImeans the system is still far from equilibrium\.
- •L2L^\{2\}discrepancy of scores:I​\(ρ∥π\)=∫ρ​\|sρ−sπ\|2​𝑑xI\(\\rho\\\|\\pi\)=\\int\\rho\\,\|s\_\{\\rho\}\-s\_\{\\pi\}\|^\{2\}\\,dx\. This is precisely the loss function used in score matching\! Training a score\-based diffusion model is essentially minimizing the Fisher information\.
- •“Steepness” of the gradient: by analogy with\|∇f\|2\|\\nabla f\|^\{2\}in Euclidean gradient flow,I​\(ρ∥π\)I\(\\rho\\\|\\pi\)measures the Wasserstein “slope” of the free energy at the currentρ\\rho\.

### E\.6Information inequalities

- •Log\-Sobolev inequality\(LSI\): ifπ\\pisatisfies an LSI with constantλ\>0\\lambda\>0, then: I​\(ρ∥π\)≥2​λ​KL​\(ρ∥π\)I\(\\rho\\\|\\pi\)\\geq 2\\lambda\\,\\mathrm\{KL\}\(\\rho\\\|\\pi\)Combined withdd​t​KL=−I\\frac\{d\}\{dt\}\\mathrm\{KL\}=\-I, this yields a Gronwall inequality: dd​t​KL​\(ρt∥π\)≤−2​λ​KL​\(ρt∥π\)⟹KL​\(ρt∥π\)≤e−2​λ​t​KL​\(ρ0∥π\)\\frac\{d\}\{dt\}\\mathrm\{KL\}\(\\rho\_\{t\}\\\|\\pi\)\\leq\-2\\lambda\\,\\mathrm\{KL\}\(\\rho\_\{t\}\\\|\\pi\)\\implies\\mathrm\{KL\}\(\\rho\_\{t\}\\\|\\pi\)\\leq e^\{\-2\\lambda t\}\\,\\mathrm\{KL\}\(\\rho\_\{0\}\\\|\\pi\)i\.e\.,exponential convergenceto equilibrium\.
- •Talagrand inequality:W22​\(ρ,π\)≤2λ​KL​\(ρ∥π\)W\_\{2\}^\{2\}\(\\rho,\\pi\)\\leq\\frac\{2\}\{\\lambda\}\\mathrm\{KL\}\(\\rho\\\|\\pi\)
- •Information hierarchy\(from strongest to weakest\): I​\(ρ∥π\)→LSIKL​\(ρ∥π\)→TalagrandW22​\(ρ,π\)I\(\\rho\\\|\\pi\)\\;\\xrightarrow\{\\text\{LSI\}\}\\;\\mathrm\{KL\}\(\\rho\\\|\\pi\)\\;\\xrightarrow\{\\text\{Talagrand\}\}\\;W\_\{2\}^\{2\}\(\\rho,\\pi\)Fisher information is the “strongest” measure, and the Wasserstein distance is the “weakest\.”

### E\.7Example

Letρ=𝒩​\(m,σ2\)\\rho=\\mathcal\{N\}\(m,\\sigma^\{2\}\)andπ=𝒩​\(0,1\)\\pi=\\mathcal\{N\}\(0,1\)\(one\-dimensional\)\.

- •∇log⁡ρ=−x−mσ2\\nabla\\log\\rho=\-\\frac\{x\-m\}\{\\sigma^\{2\}\},∇log⁡π=−x\\nabla\\log\\pi=\-x
- •∇log⁡ρπ=−x−mσ2\+x=x​\(1−1σ2\)\+mσ2\\nabla\\log\\frac\{\\rho\}\{\\pi\}=\-\\frac\{x\-m\}\{\\sigma^\{2\}\}\+x=x\(1\-\\frac\{1\}\{\\sigma^\{2\}\}\)\+\\frac\{m\}\{\\sigma^\{2\}\}
- •I​\(ρ∥π\)=𝔼x∼ρ​\[\(x​\(1−1σ2\)\+mσ2\)2\]I\(\\rho\\\|\\pi\)=\\mathbb\{E\}\_\{x\\sim\\rho\}\\left\[\\left\(x\(1\-\\frac\{1\}\{\\sigma^\{2\}\}\)\+\\frac\{m\}\{\\sigma^\{2\}\}\\right\)^\{2\}\\right\]

Expanding \(withx∼𝒩​\(m,σ2\)x\\sim\\mathcal\{N\}\(m,\\sigma^\{2\}\), lettinga=1−1σ2a=1\-\\frac\{1\}\{\\sigma^\{2\}\}andb=mσ2b=\\frac\{m\}\{\\sigma^\{2\}\}\):

I\\displaystyle I=𝔼​\[\(a​x\+b\)2\]=a2​𝔼​\[x2\]\+2​a​b​𝔼​\[x\]\+b2\\displaystyle=\\mathbb\{E\}\[\(ax\+b\)^\{2\}\]=a^\{2\}\\mathbb\{E\}\[x^\{2\}\]\+2ab\\mathbb\{E\}\[x\]\+b^\{2\}=a2​\(m2\+σ2\)\+2​a​b​m\+b2\\displaystyle=a^\{2\}\(m^\{2\}\+\\sigma^\{2\}\)\+2abm\+b^\{2\}=\(1−1σ2\)2​\(m2\+σ2\)\+2​\(1−1σ2\)​m2σ2\+m2σ4\\displaystyle=\\left\(1\-\\frac\{1\}\{\\sigma^\{2\}\}\\right\)^\{2\}\(m^\{2\}\+\\sigma^\{2\}\)\+2\\left\(1\-\\frac\{1\}\{\\sigma^\{2\}\}\\right\)\\frac\{m^\{2\}\}\{\\sigma^\{2\}\}\+\\frac\{m^\{2\}\}\{\\sigma^\{4\}\}
Checking special cases:

- •m=0,σ=1m=0,\\sigma=1\(ρ=π\\rho=\\pi\):a=0,b=0a=0,b=0,I=0I=0✓
- •m=0,σ≠1m=0,\\sigma\\neq 1:I=\(1−1σ2\)2​σ2=\(σ−1σ\)2I=\(1\-\\frac\{1\}\{\\sigma^\{2\}\}\)^\{2\}\\sigma^\{2\}=\(\\sigma\-\\frac\{1\}\{\\sigma\}\)^\{2\}, which is zero only whenσ=1\\sigma=1
- •σ=1,m≠0\\sigma=1,m\\neq 0:a=0,b=ma=0,b=m,I=m2I=m^\{2\}; the further the deviation, the largerII

## Appendix FGenerative Models: Score, Velocity, and Diffusion

This appendix elaborates on the score/velocity relationship mentioned in the main text and the mathematical details of major generative models\.

### F\.1The Complete Relationship among Score, Velocity, and Drift

Starting from the Fokker\-Planck velocity fieldvt=−∇V−∇log⁡ρtv\_\{t\}=\-\\nabla V\-\\nabla\\log\\rho\_\{t\}:

Key InsightDefinitions of the three quantities:•Score\(score function\):st​\(x\):=∇xlog⁡ρt​\(x\)s\_\{t\}\(x\):=\\nabla\_\{x\}\\log\\rho\_\{t\}\(x\), the gradient of the log\-density\. This is what diffusion models learn\.•Velocity\(velocity field\):vt​\(x\)v\_\{t\}\(x\), satisfying the continuity equation∂tρt\+∇⋅\(ρt​vt\)=0\\partial\_\{t\}\\rho\_\{t\}\+\\nabla\\cdot\(\\rho\_\{t\}v\_\{t\}\)=0\. This is what flow matching learns\.•Drift:f​\(x,t\)f\(x,t\), the deterministic part of the SDEd​Xt=f​d​t\+g​d​BtdX\_\{t\}=f\\,dt\+g\\,dB\_\{t\}\. This is given \(determined by the forward process\)\.For the SDEd​Xt=f​\(Xt,t\)​d​t\+g​\(t\)​d​BtdX\_\{t\}=f\(X\_\{t\},t\)\\,dt\+g\(t\)\\,dB\_\{t\}, the velocity field of the correspondingprobability flow ODEis:vt​\(x\)=f​\(x,t\)−12​g​\(t\)2​∇log⁡ρt​\(x\)v\_\{t\}\(x\)=f\(x,t\)\-\\frac\{1\}\{2\}g\(t\)^\{2\}\\,\\nabla\\log\\rho\_\{t\}\(x\)\(77\)In one sentence:velocityt=driftt−12​gt2×scoret\\boxed\{\\text\{velocity\}\_\{t\}=\\text\{drift\}\_\{t\}\-\\tfrac\{1\}\{2\}g\_\{t\}^\{2\}\\times\\text\{score\}\_\{t\}\}

Their interconversion:

- •Knowing score \+ drift⇒\\Rightarrowvelocity:vt=f−12​g2​stv\_\{t\}=f\-\\frac\{1\}\{2\}g^\{2\}s\_\{t\}
- •Knowing velocity \+ drift⇒\\Rightarrowscore:st=2​\(f−vt\)g2s\_\{t\}=\\frac\{2\(f\-v\_\{t\}\)\}\{g^\{2\}\}
- •Both can generate samples: score via Langevin dynamics or reverse SDE; velocity via ODE integration

### F\.2Geometric meaning of the score: the Wasserstein gradient of entropy

The score functions​\(x\)=∇xlog⁡ρ​\(x\)s\(x\)=\\nabla\_\{x\}\\log\\rho\(x\)has an elegant geometric interpretation—it is precisely the Wasserstein gradient of the negative entropy\.

Step 1: Define the negative entropy functional\.

Negentropy:

ℋ​\(ρ\)=∫ℝdρ​\(x\)​log⁡ρ​\(x\)​𝑑x\\mathcal\{H\}\(\\rho\)=\\int\_\{\\mathbb\{R\}^\{d\}\}\\rho\(x\)\\log\\rho\(x\)\\,dxA largerℋ\\mathcal\{H\}indicates a more concentrated distribution \(low entropy\); a smallerℋ\\mathcal\{H\}indicates a more spread\-out distribution \(high entropy\)\.

Step 2: Compute the functional derivativeδ​ℋδ​ρ\\frac\{\\delta\\mathcal\{H\}\}\{\\delta\\rho\}\.

dd​ϵ\|ϵ=0​ℋ​\(ρ\+ϵ​δ​ρ\)\\displaystyle\\frac\{d\}\{d\\epsilon\}\\bigg\|\_\{\\epsilon=0\}\\mathcal\{H\}\(\\rho\+\\epsilon\\,\\delta\\rho\)=∫\[δ​ρ⋅log⁡ρ\+δ​ρ\]​𝑑x=∫\(log⁡ρ\+1\)​δ​ρ​𝑑x\\displaystyle=\\int\\left\[\\delta\\rho\\cdot\\log\\rho\+\\delta\\rho\\right\]dx=\\int\(\\log\\rho\+1\)\\,\\delta\\rho\\,dxThereforeδ​ℋδ​ρ​\(x\)=log⁡ρ​\(x\)\+1\\frac\{\\delta\\mathcal\{H\}\}\{\\delta\\rho\}\(x\)=\\log\\rho\(x\)\+1\.

Step 3: From functional derivative to Wasserstein gradient\.

The core formula of Otto calculus: the velocity field corresponding to the Wasserstein gradient isv=∇δ​Gδ​ρv=\\nabla\\frac\{\\delta G\}\{\\delta\\rho\}\. Forℋ\\mathcal\{H\}:

gradW​ℋ=∇δ​ℋδ​ρ=∇\(log⁡ρ\+1\)=∇log⁡ρ\\mathrm\{grad\}\_\{W\}\\mathcal\{H\}=\\nabla\\frac\{\\delta\\mathcal\{H\}\}\{\\delta\\rho\}=\\nabla\(\\log\\rho\+1\)=\\nabla\\log\\rho
Step 4: Recognize the score\.

s​\(x\)=∇log⁡ρ=gradW​ℋ​\(ρ\)\\boxed\{s\(x\)=\\nabla\\log\\rho=\\mathrm\{grad\}\_\{W\}\\mathcal\{H\}\(\\rho\)\}That is:Score==Wasserstein gradient of the negative entropy\.

Key InsightLearning the score∇log⁡ρt\\nabla\\log\\rho\_\{t\}≡\\equivlearning the “gradient direction” in probability space\. A diffusion model learning the score is essentially learning “which direction to move in probability space so that the distribution converges to the data distribution as quickly as possible\.”

### F\.3A unified framework for mainstream models

All diffusion/flow\-based generative models follow the same framework:

1. 1\.Forward process: Design an SDE that gradually transforms the data distributionρdata\\rho\_\{\\text\{data\}\}into a simple distribution \(typically Gaussian𝒩​\(0,I\)\\mathcal\{N\}\(0,I\)\)
2. 2\.Reverse process: Learn the reverse dynamics to generate data from noise

The unified forward SDE takes the form:

d​Xt=f​\(Xt,t\)​d​t\+g​\(t\)​d​Bt,X0∼ρdatadX\_\{t\}=f\(X\_\{t\},t\)\\,dt\+g\(t\)\\,dB\_\{t\},\\qquad X\_\{0\}\\sim\\rho\_\{\\text\{data\}\}
Reverse SDE \(Anderson, 1982\):

d​Xt=\[f​\(Xt,t\)−g​\(t\)2​∇log⁡ρt​\(Xt\)\]​d​t\+g​\(t\)​d​B¯tdX\_\{t\}=\\left\[f\(X\_\{t\},t\)\-g\(t\)^\{2\}\\,\\nabla\\log\\rho\_\{t\}\(X\_\{t\}\)\\right\]dt\+g\(t\)\\,d\\bar\{B\}\_\{t\}
The corresponding probability flow ODE \(noise\-free version, with the same marginal distributions\):

d​Xtd​t=f​\(Xt,t\)−12​g​\(t\)2​∇log⁡ρt​\(Xt\)\\frac\{dX\_\{t\}\}\{dt\}=f\(X\_\{t\},t\)\-\\frac\{1\}\{2\}g\(t\)^\{2\}\\,\\nabla\\log\\rho\_\{t\}\(X\_\{t\}\)
#### F\.3\.1DDPM \(Denoising Diffusion Probabilistic Models\)

Ho et al\. \(2020\) proposed this discrete\-time framework\.

Forward process:xt=αt​xt−1\+1−αt​ϵtx\_\{t\}=\\sqrt\{\\alpha\_\{t\}\}\\,x\_\{t\-1\}\+\\sqrt\{1\-\\alpha\_\{t\}\}\\,\\epsilon\_\{t\}, with closed formxt=α¯t​x0\+1−α¯t​ϵx\_\{t\}=\\sqrt\{\\bar\{\\alpha\}\_\{t\}\}\\,x\_\{0\}\+\\sqrt\{1\-\\bar\{\\alpha\}\_\{t\}\}\\,\\epsilon\.

Training objective\(predicting noise\):

ℒDDPM=𝔼t,x0,ϵ​\[‖ϵ−ϵθ​\(α¯t​x0\+1−α¯t​ϵ,t\)‖2\]\\mathcal\{L\}\_\{\\text\{DDPM\}\}=\\mathbb\{E\}\_\{t,x\_\{0\},\\epsilon\}\\left\[\\\|\\epsilon\-\\epsilon\_\{\\theta\}\(\\sqrt\{\\bar\{\\alpha\}\_\{t\}\}\\,x\_\{0\}\+\\sqrt\{1\-\\bar\{\\alpha\}\_\{t\}\}\\,\\epsilon,\\;t\)\\\|^\{2\}\\right\]
Relationship to the score:∇log⁡ρt​\(xt\)=−ϵθ​\(xt,t\)1−α¯t\\nabla\\log\\rho\_\{t\}\(x\_\{t\}\)=\-\\frac\{\\epsilon\_\{\\theta\}\(x\_\{t\},t\)\}\{\\sqrt\{1\-\\bar\{\\alpha\}\_\{t\}\}\}\. DDPM is essentially learning the score\.

#### F\.3\.2NCSN / SMLD

Song & Ermon \(2019, 2020\): Directly trainsθ​\(x,σ\)≈∇log⁡ρσ​\(x\)s\_\{\\theta\}\(x,\\sigma\)\\approx\\nabla\\log\\rho\_\{\\sigma\}\(x\)and generate via annealed Langevin dynamics\.

ℒNCSN=𝔼σ,x,x~​\[‖sθ​\(x~,σ\)\+x~−xσ2‖2\]\\mathcal\{L\}\_\{\\text\{NCSN\}\}=\\mathbb\{E\}\_\{\\sigma,x,\\tilde\{x\}\}\\left\[\\\|s\_\{\\theta\}\(\\tilde\{x\},\\sigma\)\+\\frac\{\\tilde\{x\}\-x\}\{\\sigma^\{2\}\}\\\|^\{2\}\\right\]

#### F\.3\.3VE\-SDE \(Variance Exploding\)

Song et al\. \(2021\):d​Xt=σ˙2​\(t\)​d​BtdX\_\{t\}=\\sqrt\{\\dot\{\\sigma\}^\{2\}\(t\)\}\\,dB\_\{t\}\(f=0f=0, pure diffusion\)\. The variance “explodes” toσ2​\(T\)≫1\\sigma^\{2\}\(T\)\\gg 1\. The corresponding FP equation reduces to the heat equation∂tρ=12​σ˙2​Δ​ρ\\partial\_\{t\}\\rho=\\frac\{1\}\{2\}\\dot\{\\sigma\}^\{2\}\\,\\Delta\\rho\.

#### F\.3\.4VP\-SDE \(Variance Preserving\)

Song et al\. \(2021\):d​Xt=−β​\(t\)2​Xt​d​t\+β​\(t\)​d​BtdX\_\{t\}=\-\\frac\{\\beta\(t\)\}\{2\}X\_\{t\}\\,dt\+\\sqrt\{\\beta\(t\)\}\\,dB\_\{t\}\(OU process\)\. The variance remains bounded→1\\to 1\. The corresponding FP equation:∂tρ=β2​∇⋅\(ρ​x\)\+β2​Δ​ρ\\partial\_\{t\}\\rho=\\frac\{\\beta\}\{2\}\\nabla\\cdot\(\\rho\\,x\)\+\\frac\{\\beta\}\{2\}\\Delta\\rho\.

Probability flow ODE:d​xd​t=−β​\(t\)2​\[x\+∇log⁡ρt​\(x\)\]\\frac\{dx\}\{dt\}=\-\\frac\{\\beta\(t\)\}\{2\}\[x\+\\nabla\\log\\rho\_\{t\}\(x\)\]\.

#### F\.3\.5Flow Matching

Lipman et al\. \(2023\), Liu et al\. \(2023\): Instead of starting from an SDE, directly learn a deterministic velocity field that satisfies the continuity equation\.

Conditional Flow Matching training objective:

ℒFM=𝔼t,x1,x​\[‖vθ​\(x,t\)−x1−x1−t‖2\]\\mathcal\{L\}\_\{\\text\{FM\}\}=\\mathbb\{E\}\_\{t,x\_\{1\},x\}\\left\[\\\|v\_\{\\theta\}\(x,t\)\-\\frac\{x\_\{1\}\-x\}\{1\-t\}\\\|^\{2\}\\right\]wherex∼𝒩​\(t​x1,\(1−t\)2​I\)x\\sim\\mathcal\{N\}\(t\\,x\_\{1\},\(1\-t\)^\{2\}I\)\. Generation: solve the ODEx˙=vθ​\(x,t\)\\dot\{x\}=v\_\{\\theta\}\(x,t\)\.

Advantages: pure ODE \(fast sampling\), straight\-line paths \(easier to learn\), and freely designable interpolation schemes\.

### F\.4Unified comparison table

### F\.5A unified understanding from the Fokker\-Planck/JKO perspective

- •Score\-based models: The FP equation corresponding to the reverse SDE describes the density descending along the free energyℱ\\mathcal\{F\}\. Each denoising step is essentially one JKO step\.
- •Flow Matching: The velocity fieldvtv\_\{t\}satisfies the continuity equation\. For optimal transport paths,vtv\_\{t\}is precisely the velocity along theW2W\_\{2\}geodesic—a direct application of the Benamou\-Brenier formula\.
- •Common destination: Regardless of the method, the generation process “traces a path” in probability space from𝒩​\(0,I\)\\mathcal\{N\}\(0,I\)toρdata\\rho\_\{\\text\{data\}\}\. The only difference is whether the path is tortuous \(diffusion\) or a smooth geodesic \(flow matching\)\.

## Appendix GConvex Analysis Quick Review

This appendix provides a quick review of basic concepts from convex analysis, which appear repeatedly in the uniqueness and existence arguments for the JKO scheme\.

### G\.1Convex functions

###### Definition G\.1\(Convex function\)\.

A functionf:ℝd→ℝ∪\{\+∞\}f:\\mathbb\{R\}^\{d\}\\to\\mathbb\{R\}\\cup\\\{\+\\infty\\\}is aconvex functionif for allx,y∈ℝdx,y\\in\\mathbb\{R\}^\{d\}andt∈\[0,1\]t\\in\[0,1\]:

f​\(t​x\+\(1−t\)​y\)≤t​f​\(x\)\+\(1−t\)​f​\(y\)f\\bigl\(tx\+\(1\-t\)y\\bigr\)\\leq t\\,f\(x\)\+\(1\-t\)\\,f\(y\)

Geometric Intuition:The line segment between any two points on the graph lies above \(or coincides with\) the graph\. In other words, the function “curves upward\.”

Typical Examples:

- •f​\(x\)=x2f\(x\)=x^\{2\}\(parabola\),f​\(x\)=\|x\|f\(x\)=\|x\|\(V\-shape\),f​\(x\)=exf\(x\)=e^\{x\}
- •Inℝd\\mathbb\{R\}^\{d\}:f​\(x\)=‖x‖2f\(x\)=\\\|x\\\|^\{2\},f​\(x\)=maxi⁡xif\(x\)=\\max\_\{i\}x\_\{i\}
- •Non\-convex examples:f​\(x\)=sin⁡xf\(x\)=\\sin x,f​\(x\)=−x2f\(x\)=\-x^\{2\}

The key property of convex functions:every local minimum is a global minimum\. This is the fundamental reason why convex optimization is much easier than general optimization\.

### G\.2The Hessian matrix

###### Definition G\.2\(Hessian\)\.

TheHessian matrixof a scalar functionf:ℝd→ℝf:\\mathbb\{R\}^\{d\}\\to\\mathbb\{R\}is thed×dd\\times dmatrix of its second\-order partial derivatives:

\[∇2f​\(x\)\]i​j=∂2f∂xi​∂xj​\(x\)\\bigl\[\\nabla^\{2\}f\(x\)\\bigr\]\_\{ij\}=\\frac\{\\partial^\{2\}f\}\{\\partial x\_\{i\}\\,\\partial x\_\{j\}\}\(x\)

The Hessian describes the “degree and direction of curvature” of a function at a given point; it is the multidimensional generalization of the second derivativef′′​\(x\)f^\{\\prime\\prime\}\(x\)in one dimension\.

- •One dimension:f′′​\(x\)\>0f^\{\\prime\\prime\}\(x\)\>0means the function curves upward atxx\(convex\)
- •Multiple dimensions: the eigenvalues of∇2f​\(x\)\\nabla^\{2\}f\(x\)describe the curvature in each direction

### G\.3Positive definiteness and positive semidefiniteness

###### Definition G\.3\(Positive definite / semidefinite\)\.

For a symmetric matrixA∈ℝd×dA\\in\\mathbb\{R\}^\{d\\times d\}:

- •Positive definite\(denotedA≻0A\\succ 0\):v⊤​A​v\>0v^\{\\top\}Av\>0for allv≠0v\\neq 0\. Equivalently, all eigenvalues ofAAare\>0\>0\.
- •Positive semidefinite\(denotedA⪰0A\\succeq 0\):v⊤​A​v≥0v^\{\\top\}Av\\geq 0for allvv\. Equivalently, all eigenvalues are≥0\\geq 0\.

Connection to convexity:For a twice\-differentiable functionff:

f​is convex⇔∇2f​\(x\)⪰0​for all​xf\\text\{ is convex\}\\iff\\nabla^\{2\}f\(x\)\\succeq 0\\text\{ for all \}xA positive definite Hessian \(∇2f≻0\\nabla^\{2\}f\\succ 0\) means the function curves strictly upward in every direction\.

Intuition:Think ofA≻0A\\succ 0as “uphill in every direction\.” For the quadraticf​\(x\)=12​x⊤​A​xf\(x\)=\\frac\{1\}\{2\}x^\{\\top\}Ax,A≻0A\\succ 0means exactly thatffis “bowl\-shaped” \(with a unique minimum\), and the eigenvalues ofAAdetermine the steepness in each direction\.

### G\.4Strong convexity andλ\\lambda\-convexity

###### Definition G\.4\(Strongly convex\)\.

A functionffisλ\\lambda\-strongly convex\(λ\>0\\lambda\>0\) iff​\(x\)−λ2​‖x‖2f\(x\)\-\\frac\{\\lambda\}\{2\}\\\|x\\\|^\{2\}is still convex\. Equivalently:

f​\(t​x\+\(1−t\)​y\)≤t​f​\(x\)\+\(1−t\)​f​\(y\)−λ2​t​\(1−t\)​‖x−y‖2f\\bigl\(tx\+\(1\-t\)y\\bigr\)\\leq t\\,f\(x\)\+\(1\-t\)\\,f\(y\)\-\\frac\{\\lambda\}\{2\}\\,t\(1\-t\)\\\|x\-y\\\|^\{2\}For twice\-differentiableff, this is equivalent to∇2f​\(x\)⪰λ​I\\nabla^\{2\}f\(x\)\\succeq\\lambda Ifor allxx\.

Intuition:Strong convexity means the function is not only convex but has a “lower bound on its curvature”—it curves at least as much asλ2​‖x‖2\\frac\{\\lambda\}\{2\}\\\|x\\\|^\{2\}\.

- •The largerλ\\lambdais, the more the function “bends,” and the faster convergence is near the minimum
- •Convexity = the degenerate caseλ=0\\lambda=0
- •Strong convexity guarantees that the minimum isuniqueand that gradient descent converges at alinear rate

Appearance in this article:The condition∇2V⪰λ​I\\nabla^\{2\}V\\succeq\\lambda I\(λ\\lambda\-convexity\) on the potentialVVin the JKO scheme guarantees the uniqueness of the solution at each JKO step, as well as convergence of the discrete solutions to the continuous gradient flow\.

## Appendix HSupplementary Concepts

This appendix explains several terms that appear in the main text, providing a quick reference for readers who may be unfamiliar with these concepts\.

### H\.1Langevin dynamics

Langevin dynamicsis a stochastic differential equation originally proposed by Paul Langevin in 1908 \(Langevin, 1908\) to describe the Brownian motion of particles in a fluid\. In modern probability, statistical mechanics, and machine learning, the most common form is theoverdamped Langevin dynamics:

d​Xt=−∇V​\(Xt\)​d​t\+2​ε​d​BtdX\_\{t\}=\-\\nabla V\(X\_\{t\}\)\\,dt\+\\sqrt\{2\\varepsilon\}\\,dB\_\{t\}whereVVis the potential energy function,ε\>0\\varepsilon\>0is the temperature/noise strength, andBtB\_\{t\}is standard Brownian motion\. In the main text we often setε=1\\varepsilon=1for simplicity\.

Physical intuition:A particle moves in the potential energy landscapeVV, subject to two effects:

- •−∇V\-\\nabla V: deterministic drift, pushing the particle toward regions of low potential energy\.
- •2​ε​d​Bt\\sqrt\{2\\varepsilon\}\\,dB\_\{t\}: random thermal noise, preventing all particles from collapsing to a single minimizer ofVV\.

Stationary distribution\.The probability densityρt\\rho\_\{t\}ofXtX\_\{t\}satisfies the Fokker–Planck equation

∂tρt=∇⋅\(ρt​∇V\)\+ε​Δ​ρt\.\\partial\_\{t\}\\rho\_\{t\}=\\nabla\\cdot\(\\rho\_\{t\}\\nabla V\)\+\\varepsilon\\Delta\\rho\_\{t\}\.At stationarity, the deterministic force and the entropic force balance:

−∇V−ε​∇log⁡ρ∞=0\.\-\\nabla V\-\\varepsilon\\nabla\\log\\rho\_\{\\infty\}=0\.Therefore

∇log⁡ρ∞=−1ε​∇V,ρ∞​\(x\)∝e−V​\(x\)/ε\.\\nabla\\log\\rho\_\{\\infty\}=\-\\frac\{1\}\{\\varepsilon\}\\nabla V,\\qquad\\rho\_\{\\infty\}\(x\)\\propto e^\{\-V\(x\)/\\varepsilon\}\.Thus Langevin dynamics samples from the Gibbs/Boltzmann distribution with temperatureε\\varepsilon\.

Free\-energy interpretation\.The same equation is the Wasserstein gradient flow of

ℱε​\(ρ\)=∫V​ρ​𝑑x\+ε​∫ρ​log⁡ρ​d​x\.\\mathcal\{F\}\_\{\\varepsilon\}\(\\rho\)=\\int V\\rho\\,dx\+\\varepsilon\\int\\rho\\log\\rho\\,dx\.The first term pulls mass toward low potential energy\. The second term is negative entropy; it creates the diffusion force−ε​∇log⁡ρ\-\\varepsilon\\nabla\\log\\rho\. This is the precise sense in which Langevin dynamics minimizes free energy rather than potential energy alone\.

Key InsightGradient descent onVValone,X˙t=−∇V​\(Xt\),\\dot\{X\}\_\{t\}=\-\\nabla V\(X\_\{t\}\),finds low\-energy points\. Langevin dynamics,d​Xt=−∇V​\(Xt\)​d​t\+2​ε​d​Bt,dX\_\{t\}=\-\\nabla V\(X\_\{t\}\)\\,dt\+\\sqrt\{2\\varepsilon\}\\,dB\_\{t\},samples from the full distributionρ∞∝e−V/ε\\rho\_\{\\infty\}\\propto e^\{\-V/\\varepsilon\}\. The noise is not a nuisance: it is what realizes the entropy term in the free energy\.

Discretization: the unadjusted Langevin algorithm\.Euler–Maruyama discretization gives

Xk\+1=Xk−η​∇V​\(Xk\)\+2​ε​η​Zk,Zk∼𝒩​\(0,I\)\.X\_\{k\+1\}=X\_\{k\}\-\\eta\\nabla V\(X\_\{k\}\)\+\\sqrt\{2\\varepsilon\\eta\}\\,Z\_\{k\},\\qquad Z\_\{k\}\\sim\\mathcal\{N\}\(0,I\)\.This is called theunadjusted Langevin algorithm\(ULA\)\. If one adds a Metropolis–Hastings accept/reject correction, the method becomes MALA \(Metropolis\-adjusted Langevin algorithm\), which removes discretization bias under suitable assumptions\.

In machine learning\.

- •Energy\-based models: If a model definespθ​\(x\)∝e−Vθ​\(x\)p\_\{\\theta\}\(x\)\\propto e^\{\-V\_\{\\theta\}\(x\)\}, Langevin dynamics can sample from it using only∇Vθ\\nabla V\_\{\\theta\}\.
- •SGLD: Stochastic Gradient Langevin Dynamics replaces∇V\\nabla Vby a minibatch stochastic gradient\. Adding calibrated Gaussian noise turns optimization\-like SGD into approximate posterior sampling\.
- •Score\-based models: Annealed Langevin dynamics uses learned scores∇log⁡pσ​\(x\)\\nabla\\log p\_\{\\sigma\}\(x\)at multiple noise levels\. Since∇log⁡p​\(x\)=−∇V​\(x\)\\nabla\\log p\(x\)=\-\\nabla V\(x\)whenp∝e−Vp\\propto e^\{\-V\}, score\-based sampling can be viewed as Langevin dynamics written in score language\.
- •Diffusion models: Reverse SDE sampling is a time\-inhomogeneous, score\-driven Langevin\-type dynamics\. Probability flow ODE sampling removes the explicit Brownian noise but keeps the same density evolution by absorbing the diffusion effect into a deterministic score velocity\.

Relation to this article\.Langevin dynamics is the particle\-level stochastic process; the Fokker–Planck equation is its density\-level PDE; the free energy is its Lyapunov function; and the JKO scheme is the implicit Euler discretization of this free\-energy gradient flow in Wasserstein space\.

### H\.2Kolmogorov forward equation

TheKolmogorov forward equation\(Kolmogorov, 1931\) is another name for the Fokker\-Planck equation\. The two are identical\. For the SDEd​Xt=f​\(Xt,t\)​d​t\+g​\(t\)​d​BtdX\_\{t\}=f\(X\_\{t\},t\)\\,dt\+g\(t\)\\,dB\_\{t\}\(where the diffusion coefficientggdepends only on time, not onXtX\_\{t\}—all models considered in this article are of this type\):

∂tρ=−∇⋅\(f​ρ\)\+g22​Δ​ρ\\partial\_\{t\}\\rho=\-\\nabla\\cdot\(f\\rho\)\+\\frac\{g^\{2\}\}\{2\}\\Delta\\rhoThis describes the time evolution of the probability densityρt​\(x\)\\rho\_\{t\}\(x\)ofXtX\_\{t\}\. \(For the more general case of state\-dependent diffusionG​\(Xt\)G\(X\_\{t\}\), the equation becomes∂tρ=−∇⋅\(f​ρ\)\+12​∑i​j∂i∂j\[\(G​G⊤\)i​j​ρ\]\\partial\_\{t\}\\rho=\-\\nabla\\cdot\(f\\rho\)\+\\frac\{1\}\{2\}\\sum\_\{ij\}\\partial\_\{i\}\\partial\_\{j\}\[\(GG^\{\\top\}\)\_\{ij\}\\rho\]\.\)

Why two names?Historical reasons: Kolmogorov derived this equation in 1931 from a purely mathematical perspective, while Fokker and Planck independently discovered it earlier \(1913–1917\) from a physics perspective\. Physicists call it the Fokker\-Planck equation; probabilists call it the Kolmogorov forward equation\.

There is also a correspondingKolmogorov backward equation, which describes the probability evolution “looking back from the endpoint to the starting point,” corresponding to the PDE satisfied by expected values of the stochastic process\.

### H\.3Ornstein\-Uhlenbeck \(OU\) process

TheOU process\(Uhlenbeck & Ornstein, 1930\) is one of the most important linear SDEs:

d​Xt=−θ​Xt​d​t\+σ​d​Bt,θ\>0dX\_\{t\}=\-\\theta\\,X\_\{t\}\\,dt\+\\sigma\\,dB\_\{t\},\\qquad\\theta\>0
Intuition:Spring \+ noise\. The particle is pulled back to the origin by a spring \(−θ​Xt\-\\theta X\_\{t\}is the restoring force\), while simultaneously being subject to random perturbations\.

Key properties:

- •The stationary distribution is Gaussian:X∞∼𝒩​\(0,σ22​θ\)X\_\{\\infty\}\\sim\\mathcal\{N\}\\bigl\(0,\\frac\{\\sigma^\{2\}\}\{2\\theta\}\\bigr\)
- •From any initial condition,XtX\_\{t\}converges to the stationary distribution at an exponential rate
- •It is a “compromise” between Brownian motion \(θ=0\\theta=0\) and pure deterministic decay \(σ=0\\sigma=0\)

In this article:The forward process of VP\-SDE,d​Xt=−β​\(t\)2​Xt​d​t\+β​\(t\)​d​BtdX\_\{t\}=\-\\frac\{\\beta\(t\)\}\{2\}X\_\{t\}\\,dt\+\\sqrt\{\\beta\(t\)\}\\,dB\_\{t\}, is an OU process with time\-varying coefficients\. It gradually “pulls” the data distribution toward the standard Gaussian𝒩​\(0,I\)\\mathcal\{N\}\(0,I\)\.

### H\.4Otto Calculus

Otto calculusis a formal computational framework introduced by Felix Otto in 2001, whose core idea is:

Treat the space of probability measures𝒫2​\(ℝd\)\\mathcal\{P\}\_\{2\}\(\\mathbb\{R\}^\{d\}\)as an infinite\-dimensional Riemannian manifold and do calculus on it\.

Specifically, Otto calculus consists of the following “conventions”:

1. 1\.Tangent space: The “tangent vectors” at a distributionρ\\rhoare gradient fields∇ψ\\nabla\\psi\(velocity fields satisfying the continuity equationρ˙=−∇⋅\(ρ​∇ψ\)\\dot\{\\rho\}=\-\\nabla\\cdot\(\\rho\\nabla\\psi\)\)
2. 2\.Riemannian metric: The inner product of two tangent vectors∇ψ1,∇ψ2\\nabla\\psi\_\{1\},\\nabla\\psi\_\{2\}is defined as⟨∇ψ1,∇ψ2⟩ρ=∫∇ψ1⋅∇ψ2​ρ​d​x\\langle\\nabla\\psi\_\{1\},\\nabla\\psi\_\{2\}\\rangle\_\{\\rho\}=\\int\\nabla\\psi\_\{1\}\\cdot\\nabla\\psi\_\{2\}\\,\\rho\\,dx
3. 3\.Gradient: The Wasserstein gradient of a functionalG​\(ρ\)G\(\\rho\)is the tangent vector \(velocity field\)gradW​G=∇δ​Gδ​ρ\\mathrm\{grad\}\_\{W\}G=\\nabla\\frac\{\\delta G\}\{\\delta\\rho\}\. The induced direction of density evolution is∂tρ=−∇⋅\(ρ​gradW​G\)=−∇⋅\(ρ​∇δ​Gδ​ρ\)\\partial\_\{t\}\\rho=\-\\nabla\\cdot\\bigl\(\\rho\\,\\mathrm\{grad\}\_\{W\}G\\bigr\)=\-\\nabla\\cdot\\bigl\(\\rho\\nabla\\frac\{\\delta G\}\{\\delta\\rho\}\\bigr\)

Why is it called “calculus”?Because with the conventions above, many complex infinite\-dimensional computations reduce to “plugging into formulas”—just as ordinary calculus frees you from going back to the limit definition every time\.

Rigorous vs\. formal:Otto calculus is*formal*—𝒫2\\mathcal\{P\}\_\{2\}is not strictly a smooth manifold \(it is infinite\-dimensional and has a boundary\)\. However, its computational results can be verified by rigorous theory \(the Ambrosio\-Gigli\-Savaré framework\)\. In practice, computing directly with Otto’s formalism is usually more efficient\.

### H\.5Geodesics

###### Definition H\.1\(Geodesic — informal\)\.

On a metric space or Riemannian manifold, ageodesicconnecting two pointsxxandyyis the “shortest path”—a curve whose length equals the distance betweenxxandyy\.

Examples in various spaces:

- •Euclidean spaceℝd\\mathbb\{R\}^\{d\}: geodesics are straight line segments
- •SphereS2S^\{2\}: geodesics are great circle arcs \(the shortest flight routes\)
- •Wasserstein space𝒫2​\(ℝd\)\\mathcal\{P\}\_\{2\}\(\\mathbb\{R\}^\{d\}\): geodesics are McCann displacement interpolationsρt=\(\(1−t\)​id\+t​T\)\#​μ\\rho\_\{t\}=\\bigl\(\(1\-t\)\\mathrm\{id\}\+t\\,T\\bigr\)\_\{\\\#\}\\mu, whereTTis the optimal transport map fromμ\\mutoν\\nu

Connection to Benamou\-Brenier:W2​\(μ,ν\)W\_\{2\}\(\\mu,\\nu\)is precisely the length of the geodesic connectingμ\\muandν\\nu, and the Benamou\-Brenier formulaW22=inf∫01∫\|vt\|2​ρt​𝑑x​𝑑tW\_\{2\}^\{2\}=\\inf\\int\_\{0\}^\{1\}\\int\|v\_\{t\}\|^\{2\}\\rho\_\{t\}\\,dx\\,dtseeks the shortest path among all paths satisfying the continuity equation\.

### H\.6Riesz representation theorem

TheRiesz representation theoremhas several versions; the one used in this article is the Hilbert space version:

###### Theorem H\.2\(Riesz representation — Hilbert space version\)\.

LetHHbe a Hilbert space andφ:H→ℝ\\varphi:H\\to\\mathbb\{R\}a continuous linear functional onHH\. Then there exists a uniquef∈Hf\\in Hsuch that:

φ​\(h\)=⟨f,h⟩Hfor all​h∈H\\varphi\(h\)=\\langle f,h\\rangle\_\{H\}\\quad\\text\{for all \}h\\in H

Intuition:In a space equipped with an inner product, every “linear measurement device” \(continuous linear functional\) can be realized as “taking the inner product with some fixed vector\.”

Role in this article:The Wasserstein gradientgradW​ℱ\\mathrm\{grad\}\_\{W\}\\mathcal\{F\}is precisely the “Riesz representative” of the functional derivativeδ​ℱδ​ρ\\frac\{\\delta\\mathcal\{F\}\}\{\\delta\\rho\}—it converts the functional derivative \(a cotangent vector\) under theL2L^\{2\}inner product into a gradient \(a tangent vector\) under the Wasserstein metric\. This is exactly the infinite\-dimensional manifestation of the principle from Appendix[B](https://arxiv.org/html/2606.24157#A2)that going from “differential→\\togradient” requires a metric\.

### H\.7Lyapunov functions

###### Definition H\.3\(Lyapunov function — informal\)\.

For a dynamical systemx˙=F​\(x\)\\dot\{x\}=F\(x\), aLyapunov functionV​\(x\)V\(x\)satisfies:

1. 1\.V​\(x\)≥0V\(x\)\\geq 0, andV​\(x∗\)=0V\(x^\{\*\}\)=0only at the equilibrium pointx∗x^\{\*\}
2. 2\.It is monotonically decreasing along trajectories of the system:dd​t​V​\(x​\(t\)\)≤0\\frac\{d\}\{dt\}V\(x\(t\)\)\\leq 0

Intuition:A Lyapunov function is like “energy”—if you can find a quantity that only decreases during the system’s evolution, then the system must evolve toward the point where this quantity is zero \(the equilibrium\)\. It is the standard tool for proving stability\.

In this article:The free energyℱ​\(ρ\)=KL​\(ρ∥π\)\\mathcal\{F\}\(\\rho\)=\\mathrm\{KL\}\(\\rho\\\|\\pi\)is the Lyapunov function of the Fokker\-Planck dynamical system—it decreases monotonically along solutions of the FP equation \(dd​t​ℱ​\(ρt\)≤0\\frac\{d\}\{dt\}\\mathcal\{F\}\(\\rho\_\{t\}\)\\leq 0\) and attains its minimum value of zero at the stationary distributionρ=π\\rho=\\pi\.

### H\.8KKT conditions

TheKKT conditions\(Karush\-Kuhn\-Tucker conditions\) are necessary optimality conditions for constrained optimization problems, generalizing the “gradient equals zero” condition from unconstrained optimization\.

Consider the problem:min⁡f​\(x\)\\min f\(x\)subject tog​\(x\)=0g\(x\)=0\. The Lagrangian isL​\(x,λ\)=f​\(x\)\+λ​g​\(x\)L\(x,\\lambda\)=f\(x\)\+\\lambda\\,g\(x\)\.

The KKT conditions require:

∇xL=∇f​\(x\)\+λ​∇g​\(x\)=0,g​\(x\)=0\\nabla\_\{x\}L=\\nabla f\(x\)\+\\lambda\\,\\nabla g\(x\)=0,\\qquad g\(x\)=0
Intuition:On the constraint surface, the gradient of the objective function must be parallel to the normal direction of the constraint surface \(otherwise one could continue descending along the constraint surface\)\. The multiplierλ\\lambdameasures “how tight” the constraint is\.

In this article:The derivation of the optimality conditions for the JKO scheme \(Section 6\) is essentially the infinite\-dimensional version of the KKT conditions—the functional derivative equals zero, together with the Lagrange multiplier corresponding to the mass conservation constraint∫ρ=1\\int\\rho=1\.