JKO gradient flow (Wasserstein)

Wiki card — concept node for exp-families-stability.

(i) Formal statement

Let \(\mathcal{P}_2(\mathbb{R}^d)\) denote probability measures with finite second moment, equipped with the quadratic Wasserstein distance \(W_2\). Define the free energy

\[ \mathcal{F}(\rho) \;:=\; \int_{\mathbb{R}^d} \rho\,\log\rho\, dx \;+\; \int V\, \rho\, dx, \]

for a potential \(V : \mathbb{R}^d \to \mathbb{R}\) (possibly \(V \equiv 0\)).

Theorem (Jordan–Kinderlehrer–Otto 1998). The Fokker–Planck equation \[ \partial_t \rho \;=\; \tfrac12 \Delta \rho \;+\; \nabla\cdot(\rho \nabla V) \] is the \(W_2\)-gradient flow of \(\mathcal{F}\): \[ \boxed{\;\rho^{n+1} \;=\; \arg\min_{\rho \in \mathcal{P}_2}\Bigl\{\tfrac{1}{2\tau} W_2^2(\rho, \rho^n) \;+\; \mathcal{F}(\rho)\Bigr\}\;} \tag{JKO} \] converges, as \(\tau \to 0\), to the Fokker–Planck solution.

In particular, for \(V \equiv 0\), the heat equation \(\partial_t \rho = \tfrac12 \Delta\rho\) is the \(W_2\)-gradient flow of (negative) differential entropy: \(\mathcal{F}(\rho) = \int \rho \log\rho\).

(ii) Role in Q1 / Q2 / Q3

• Q1 (geometric recast, Wheeler). Stability of \(\mathcal{E}_\varphi\) under heat flow = \(\mathcal{E}_\varphi\) is an invariant submanifold of the \(W_2\)- gradient flow of entropy. The tangent space to \(\mathcal{E}_\varphi\) at \(p_\theta\) is \(\{(\delta\theta)^\top \nabla\varphi \cdot p_\theta : \delta\theta \in \mathbb{R}^r\}\) (up to mass-conservation projection). Invariance = the gradient vector \(-\nabla\mathcal{F}|_{p_\theta}\) lies in this tangent space.
• Q2 (Otto calculus on Gaussians). Restricting \(W_2\)-gradient flow to the Gaussian submanifold gives the Bures–Wasserstein metric on covariance matrices; the gradient flow of entropy becomes a Lyapunov ODE for \(\Sigma\) which, in precision coordinates \(K = \Sigma^{-1}\), is the Riccati \(\dot K = -K^2\) of Q2.
• Q3. Hamilton–Jacobi (6) is the Hamiltonian formulation of the same dynamics: the pair (position = \(\eta\) = \(m\)-coordinates, momentum = \(\theta\) = \(e\)-coordinates) governed by the Lagrangian of \(W_2\)-gradient flow. Otto 2001 makes this precise for porous-medium equations; for the heat equation, Q3 (6) is the classical shadow.

(iii) References

• Jordan, Kinderlehrer, Otto (1998) — The Variational Formulation of the Fokker–Planck Equation. doi:10.1137/S0036141096303359.
• Otto (2001) — The Geometry of Dissipative Evolution Equations: The Porous Medium Equation. doi:10.1081/PDE-100002243.
• Ambrosio, Gigli, Savaré (2008) — Gradient Flows in Metric Spaces and in the Space of Probability Measures. doi:10.1007/978-3-7643-8722-8.

(iv) Worked miniature — JKO on 1D Gaussians

Restrict to \(\rho = \mathcal{N}(\mu, \sigma^2)\) on \(\mathbb{R}\). Then:

• \(W_2^2(\mathcal{N}(\mu_0, \sigma_0^2), \mathcal{N}(\mu_1, \sigma_1^2)) = (\mu_0 - \mu_1)^2 + (\sigma_0 - \sigma_1)^2\) (Bures–Wasserstein on 1D Gaussians reduces to a Euclidean distance on \((\mu, \sigma)\)).
• \(\int \rho \log\rho = -\tfrac12 \log(2\pi e \sigma^2)\).

JKO with \(V = 0\): minimise \(\tfrac{1}{2\tau}\bigl((\mu - \mu_n)^2 + (\sigma - \sigma_n)^2\bigr) - \tfrac12 \log(2\pi e \sigma^2)\) over \((\mu, \sigma)\).

Critical-point conditions:

• \(\partial_\mu\): \((\mu - \mu_n)/\tau = 0\), hence \(\mu_{n+1} = \mu_n\).
• \(\partial_\sigma\): \((\sigma - \sigma_n)/\tau = 1/\sigma\), i.e. \(\sigma^2 - \sigma_n\sigma - \tau = 0\), so \(\sigma_{n+1} = \tfrac{\sigma_n + \sqrt{\sigma_n^2 + 4\tau}}{2}\).

Continuous limit \(\tau \to 0\): \(\dot\mu = 0\), \(\dot\sigma = 1/\sigma\), i.e. \(\sigma^2(t) = \sigma^2(0) + 2 \cdot \tfrac{t}{2} = \sigma^2(0) + t\) (scaling by factor 2 absorbed into the \(\tfrac12\Delta\) convention; with the standard convention \(\partial_t \rho = \tfrac12 \partial_{xx}\rho\), \(\sigma^2(t) = \sigma^2(0) + t\)). ✓

Information-geometric takeaway. On the Gaussian submanifold, JKO reduces to a Euclidean gradient flow on \((\mu, \sigma)\), producing Q2's Riccati flow exactly. The submanifold is invariant, confirming Wheeler's reformulation: \(\mathcal{E}_{\mathrm{Gauss}}\) is a totally geodesic locus for the Wasserstein gradient flow of entropy.