---
title: 'JKO Scheme: Gradient Flows in Wasserstein Space'
url: https://www.emergentmind.com/topics/jko-scheme
type: topic
---

# JKO Scheme: Gradient Flows in Wasserstein Space

The Jordan–Kinderlehrer–Otto (JKO) Scheme is a canonical time discretization of gradient flows in the Wasserstein space of probability measures. Formulated originally for the Fokker–Planck equation, it generalizes to a wide variety of dissipative PDEs with variational structure, spanning diffusion, aggregation–diffusion, reaction–advection–diffusion, granular media, plasma dynamics, and statistical learning. The JKO approach enables both rigorous analysis of solution existence and novel computational algorithms grounded in optimal transport theory, convex analysis, and variational optimization.

## 1. Core Variational Principle and Formulation

The JKO scheme defines an implicit Euler step in Wasserstein geometry for any lower-semicontinuous, geodesically convex energy \( \mathcal F \) on \( \mathcal{P}_2(\mathbb{R}^d) \):
\[
\rho^{n+1} = \arg\min_{\rho \in \mathcal{P}_2(\mathbb{R}^d)} \left\{ \frac{1}{2\tau} W_2^2(\rho, \rho^n) + \mathcal{F}(\rho) \right\},
\]
where \( W_2 \) is the 2-Wasserstein distance, \( \tau > 0 \) is the time step, and \( \{\rho^n\} \) is the discrete solution sequence [2511.14827], [1911.10999], [2212.14424], [2411.12430].

The functional \( \mathcal{F} \) typically encompasses physical (or model-based) entropy, potential energy, interaction energies, or more complex terms (e.g., total variation [1703.00243], or Landau entropy for plasma). For systems with additional structures, \( \mathcal{F} \) and the transport cost may be generalized, e.g., to non-Euclidean manifolds or more general costs [2402.17681], or with alternative metrics as in the Kantorovich–Fisher–Rao splitting [1602.04457].

Gradient flow in \( W_2 \) can be formally recovered from the continuous-time limit \( \tau \to 0 \), yielding
\[
\partial_t \rho = \nabla \cdot \left[ \rho \nabla \frac{\delta \mathcal{F}}{\delta \rho} \right],
\]
in the sense of distributions [2511.14827], [1911.10999].

## 2. Structure, Properties, and Interpretation

**Implicit Euler and Stability**  
The JKO scheme is a variational implicit Euler (proximal-point) method in metric spaces. Unlike explicit (forward) schemes, it enjoys unconditional stability for all \( \tau > 0 \) when \( \mathcal{F} \) is displacement-convex [2511.14827], [1611.09029]. The discrete energy-dissipation inequality at each step,
\[
\mathcal{F}(\rho^{n+1}) + \frac{1}{2\tau}W_2^2(\rho^{n+1}, \rho^n) \leq \mathcal{F}(\rho^n),
\]
reflects non-increase of energy and ensures regularity and compactness vital for passage to the continuum limit [1911.10999], [2510.09231].

**First- and Second-order Expansion: Implicit Bias**  
The JKO update approximates, to first order in \( \tau \), Wasserstein gradient flow for \( \mathcal{F} \). At second order, it is the gradient flow for the *modified* energy
\[
J^\tau(\rho) = \mathcal{F}(\rho) - \frac{\tau}{4} \int \left\| \nabla \frac{\delta \mathcal{F}}{\delta \rho} \right\|^2 \rho(dx),
\]
i.e., it induces a canonical deceleration determined by the *squared metric slope* of \( \mathcal{F} \), with explicit forms as the Fisher information or Fisher–Hyvärinen divergence for entropy or KL functionals [2511.14827].

**Convergence and Optimality**  
For \( \lambda \)-displacement convex functionals, the scheme converges (narrowly in \( \mathcal{P}_2 \)), admitting rates \( O(1/\sigma_n) \) in energy for total time \( \sigma_n = \sum_{i=1}^{n-1}\tau_i \) [2505.23517]. Exact minimization at each step is not required: under summable error conditions (either Wasserstein distance or energy gap), convergence and rates persist [2505.23517].

**Strong Regularity and Quantitative Bounds**  
JKO iterates propagate \( L^\infty \) and Sobolev regularity, with explicit bounds for key PDE models (Fokker–Planck, Keller–Segel, aggregation–diffusion) [1709.04201], [1911.10999], [2410.15095], [2510.09231]. Discrete Li–Yau–Hamilton inequalities and maximum/minimum principles are established, yielding Lipschitz bounds and quantitative Harnack-type inequalities [2510.09231], [1703.00243].

## 3. Extensions and Algorithmic Innovations

**General Cost Functions and Geometric Setups**  
The transport cost in the JKO step can be replaced by general smooth costs \( c(x, y) \) satisfying a mixed Hessian condition, producing gradient flows on arbitrary Riemannian manifolds, including those with Bregman or other costs [2402.17681]. The induced metric in the continuity equation is directly determined by the cost's mixed Hessian.

**Entropic and Schrödinger Regularization**  
Replacing \( W_2 \) with entropic-regularized costs (Schrödinger problems, solved via Sinkhorn) smooths each step and provides computational tractability in high dimensions [2502.12666]. A scaling limit \( \epsilon/\tau \to \alpha \) yields extra linear diffusion in the limit PDE, and the classical gradient flow is restored as \( \epsilon/\tau \to 0 \).

**Splitting Schemes for Reaction, Unbalanced Mass, or Source Terms**  
The JKO scheme is adapted for dynamics such as the Kantorovich–Fisher–Rao (KFR) metric, which entails a two-step procedure: a conservative (mass-preserving) Wasserstein step, then a reaction (mass-changing) Fisher–Rao step [1602.04457]. For systems with source terms (e.g., birth-death in chemotaxis), a mass-shift is performed before the transport-proximal step [2004.12368].

**Full Discretization and Algorithms**  
Natural space discretizations using grid-based atomic measures with suitable scaling (\( h/\tau \to 0 \)) are shown to converge to the continuous gradient flow solution [2504.13513]. Entropic regularization enables scalable Eulerian solvers via fixed-point and Anderson-accelerated methods, leveraging low-rank tensor decompositions for high-dimensional Bayesian inference [2411.12430].

**Deep Learning and Neural Approaches**  
The JKO scheme underpins neural generative models (JKO-iFlow, S-JKO, iJKOnet), connecting block-wise normalizing flows and Wasserstein gradient flow. Sequential residual neural ODEs and adversarial minimax optimization estimate transport maps and energy functionals, achieving state-of-the-art results in high-dimensional synthetic and real generative tasks with improved scalability [2212.14424], [2402.05443], [2506.01502].

**Computational–Statistical Analysis and Parameter Learning**  
Statistical extensions combine parameter estimation (offline and online) with the JKO scheme. Joint asymptotics yield stochastic PDE (SPDE) central limit theory for the error between statistical and population JKO flows, allowing quantification of discretization and parameter-estimation-induced fluctuations [2501.06408].

## 4. Model-Specific JKO Schemes

| Model / Setting          | JKO Functional and Constraints                           | Special Features/References                          |
|--------------------------|---------------------------------------------------------|------------------------------------------------------|
| Fokker–Planck            | \( \mathrm{KL}(\rho\|\rho_*) \), quadratic cost         | Classic, geodesic convexity [2511.14827], [2212.14424] |
| Aggregation–Diffusion    | \( \int U(\rho) + \int V\rho + \iint W(x-y)\rho(x)\rho(y) \) | Nonlinearities, singular kernels [1911.10999], [2510.09231] |
| Keller–Segel             | Coupling with Poisson potential                         | Blowup and subcritical mass regimes [1709.04201], [2410.15095] |
| Landau Equation          | Landau metric \( d_L \), Boltzmann entropy              | Particle schemes with neural parameterization [2409.12296] |
| TV–JKO                   | TV(ρ) as energy, optional lower bound                   | Fourth-order PDE limit, BV and maximum principle [1703.00243] |
| KFR Splitting            | Wasserstein and Fisher–Rao steps                        | Mass variation, inf-convolution structure [1602.04457] |
| General Metric           | Smooth cost c(x, y), manifold setting                   | Riemannian Fokker–Planck, Bregman divergence [2402.17681] |

## 5. Regularity, Convergence, and Compactness

The JKO scheme propagates quantitative regularity under mild assumptions:
- ***L^p and L^\infty bounds*** are established for key drift-diffusion, aggregation, and chemotaxis equations, often matching PDE-level bounds [1709.04201], [1911.10999].
- ***Propagated Sobolev regularity*** (e.g., \( L^\infty_t W^{1,p}_x \), \( L^2_t H^2_x \) for Fokker–Planck and Keller–Segel) is recovered from discrete-level functional inequalities [2410.15095], [2510.09231].
- ***Modulus of continuity*** and Fisher information decrease monotonically under the scheme for nonlinear diffusions, and concave moduli propagate through all steps [2407.03758].
- ***Strong convergence*** in suitable topology (e.g., \( L^2_{\text{loc}}((0,T); H^2) \) for Fokker–Planck) is achieved via compactness, discrete Gronwall, and functional inequalities [2510.09231].

## 6. Numerical Techniques and Practical Implementation

Efficient algorithms for JKO steps employ:
- *Benamou–Brenier* dynamic reformulation and Sinkhorn accelerations for entropic regularization [2502.12666], [2411.12430].
- *Eulerian grid-based solvers* with low-rank tensor-train compression, enabling high-dimensional Bayesian inverse problems and posterior sampling [2411.12430].
- *Neural parametrizations* of transport maps and energy functionals, trained block-wise or end-to-end for generative modeling and inverse problems [2212.14424], [2402.05443], [2506.01502].
- Implementation of *inexact proximal steps*—energy or distance up to summable error—without loss of convergence guarantees [2505.23517].

## 7. Significance and Impact Across Domains

The JKO scheme is foundational for the modern theory of metric-measure gradient flows and has reshaped understanding of PDEs, probability, and data science. Its theoretical robustness (energy dissipation, unconditional stability, preservation of maximum/minimum principles) and flexibility (handling singular energies, mass systems, geometric generalization, and statistical uncertainty) underpin both rigorous PDE analysis and scalable computational methodologies. It continues to drive new developments at the intersection of optimal transport, mathematical physics, machine learning, and statistical inference [2511.14827].

---
**Key references:** [2511.14827], [1911.10999], [2212.14424], [2505.23517], [2411.12430], [2501.06408], [2410.15095], [2510.09231], [2402.17681], [2502.12666], [2506.01502], [1703.00243], [1709.04201], [2409.12296], [1602.04457], [2504.13513], [2004.12368], [2407.03758], [2402.05443].

Source: https://www.emergentmind.com/topics/jko-scheme