JKO Scheme: Gradient Flows in Wasserstein Space

• The JKO scheme is a time discretization method that uses an implicit Euler step in Wasserstein space to approximate gradient flows via energy minimization.
• It applies to a wide range of PDEs—such as Fokker–Planck and aggregation–diffusion models—ensuring stability and convergence through variational principles.
• Recent developments include entropic regularization and neural network implementations that enhance computational efficiency in high-dimensional applications.

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 (Halmos et al., 18 Nov 2025, Marino et al., 2019, Xu et al., 2022, Aksenov et al., 19 Nov 2024).

The functional $\mathcal{F}$ typically encompasses physical (or model-based) entropy, potential energy, interaction energies, or more complex terms (e.g., total variation (Carlier et al., 2017), 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 (Rankin et al., 27 Feb 2024), or with alternative metrics as in the Kantorovich–Fisher–Rao splitting (Gallouët et al., 2016).

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 (Halmos et al., 18 Nov 2025, Marino et al., 2019).

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 (Halmos et al., 18 Nov 2025, Moretti et al., 2016). 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 (Marino et al., 2019, Coudreuse, 10 Oct 2025).

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 (Halmos et al., 18 Nov 2025).

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$ (Marino et al., 29 May 2025). Exact minimization at each step is not required: under summable error conditions (either Wasserstein distance or energy gap), convergence and rates persist (Marino et al., 29 May 2025).

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) (Carrillo et al., 2017, Marino et al., 2019, Elbar, 19 Oct 2024, Coudreuse, 10 Oct 2025). Discrete Li–Yau–Hamilton inequalities and maximum/minimum principles are established, yielding Lipschitz bounds and quantitative Harnack-type inequalities (Coudreuse, 10 Oct 2025, Carlier et al., 2017).

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 (Rankin et al., 27 Feb 2024). 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 (Baradat et al., 18 Feb 2025). 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 (Gallouët et al., 2016). For systems with source terms (e.g., birth-death in chemotaxis), a mass-shift is performed before the transport-proximal step (Valencia-Guevara, 2020).

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 (Hraivoronska et al., 18 Apr 2025). Entropic regularization enables scalable Eulerian solvers via fixed-point and Anderson-accelerated methods, leveraging low-rank tensor decompositions for high-dimensional Bayesian inference (Aksenov et al., 19 Nov 2024).

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 (Xu et al., 2022, Choi et al., 8 Feb 2024, Persiianov et al., 2 Jun 2025).

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 (Wu et al., 11 Jan 2025).

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 (Halmos et al., 18 Nov 2025, Xu et al., 2022)
Aggregation–Diffusion	$$\int U(\rho) + \int V\rho + \iint W(x-y)\rho(x)\rho(y)$$	Nonlinearities, singular kernels (Marino et al., 2019, Coudreuse, 10 Oct 2025)
Keller–Segel	Coupling with Poisson potential	Blowup and subcritical mass regimes (Carrillo et al., 2017, Elbar, 19 Oct 2024)
Landau Equation	Landau metric $d_L$, Boltzmann entropy	Particle schemes with neural parameterization (Huang et al., 18 Sep 2024)
TV–JKO	TV(ρ) as energy, optional lower bound	Fourth-order PDE limit, BV and maximum principle (Carlier et al., 2017)
KFR Splitting	Wasserstein and Fisher–Rao steps	Mass variation, inf-convolution structure (Gallouët et al., 2016)
General Metric	Smooth cost c(x, y), manifold setting	Riemannian Fokker–Planck, Bregman divergence (Rankin et al., 27 Feb 2024)

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 (Carrillo et al., 2017, Marino et al., 2019).
• 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 (Elbar, 19 Oct 2024, Coudreuse, 10 Oct 2025).
• Modulus of continuity and Fisher information decrease monotonically under the scheme for nonlinear diffusions, and concave moduli propagate through all steps (Caillet et al., 4 Jul 2024).
• 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 (Coudreuse, 10 Oct 2025).

6. Numerical Techniques and Practical Implementation

Efficient algorithms for JKO steps employ:

• Benamou–Brenier dynamic reformulation and Sinkhorn accelerations for entropic regularization (Baradat et al., 18 Feb 2025, Aksenov et al., 19 Nov 2024).
• Eulerian grid-based solvers with low-rank tensor-train compression, enabling high-dimensional Bayesian inverse problems and posterior sampling (Aksenov et al., 19 Nov 2024).
• Neural parametrizations of transport maps and energy functionals, trained block-wise or end-to-end for generative modeling and inverse problems (Xu et al., 2022, Choi et al., 8 Feb 2024, Persiianov et al., 2 Jun 2025).
• Implementation of inexact proximal steps—energy or distance up to summable error—without loss of convergence guarantees (Marino et al., 29 May 2025).

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 (Halmos et al., 18 Nov 2025).

---

Key references: (Halmos et al., 18 Nov 2025, Marino et al., 2019, Xu et al., 2022, Marino et al., 29 May 2025, Aksenov et al., 19 Nov 2024, Wu et al., 11 Jan 2025, Elbar, 19 Oct 2024, Coudreuse, 10 Oct 2025, Rankin et al., 27 Feb 2024, Baradat et al., 18 Feb 2025, Persiianov et al., 2 Jun 2025, Carlier et al., 2017, Carrillo et al., 2017, Huang et al., 18 Sep 2024, Gallouët et al., 2016, Hraivoronska et al., 18 Apr 2025, Valencia-Guevara, 2020, Caillet et al., 4 Jul 2024, Choi et al., 8 Feb 2024).