JKO Solution Operator

• JKO solution operator is a variational framework that defines discrete approximations of Wasserstein gradient flows by minimizing a convex energy functional with a Wasserstein proximity term.
• It ensures energy dissipation and unconditional stability through an implicit Euler scheme, making it effective for nonlinear diffusion, aggregation, and reaction problems.
• Extensions and numerical methods, including neural operator frameworks and operator-splitting, enable its practical application in complex, high-dimensional problems.

The Jordan–Kinderlehrer–Otto (JKO) solution operator defines a central step in the discrete-time approximation of Wasserstein gradient flows, furnishing a variational time-marching framework for a wide variety of nonlinear diffusion, aggregation, reaction, and geometry-driven evolution problems. At each time step, the JKO scheme maps a given measure or density to the unique—or, in certain degenerate cases, minimal-energy—solvent of a convex functional augmented by a Wasserstein proximity term. This operator realizes an implicit Euler scheme in the infinite-dimensional manifold of probability measures equipped with the $2$-Wasserstein metric, yielding unconditionally stable, energy-dissipative, and convergent approximations under broad assumptions.

1. Variational Definition and General Properties

Let $(\mathcal{P}_2(\Omega), W_2)$ denote the space of Borel probability measures on a domain $\Omega \subset \mathbb{R}^d$ (or a Riemannian manifold), with finite second moment and quadratic Wasserstein distance $W_2$. Given an energy functional $$F: \mathcal{P}_2(\Omega) \to \mathbb{R} \cup \{+\infty\}$$—typically lower semicontinuous and geodesically convex—and a previous iterate $\mu \in \mathcal{P}_2(\Omega)$, the JKO solution operator $J_\tau$ is defined as: $$J_\tau(\mu) = \arg\min_{\nu \in \mathcal{P}_2(\Omega)} \left\{ F(\nu) + \frac{1}{2\tau} W_2^2(\nu, \mu) \right\}$$ for discrete time-step $\tau > 0$ (Halmos et al., 18 Nov 2025, Marino et al., 29 May 2025, Benamou et al., 2014). The minimization is over all admissible measures or classes where $F$ is properly defined. Under $\lambda$-geodesic convexity, existence and uniqueness of the minimizer are classical; the sequence $\{\mu^k\}$ with $\mu^{k+1} = J_\tau(\mu^k)$ converges (as $\tau \to 0$) to the Wasserstein gradient flow of $F$ (Marino et al., 29 May 2025).

Key properties of the operator include:

• Energy dissipation: $$F(J_\tau(\mu)) + \frac{1}{2\tau} W_2^2(J_\tau(\mu), \mu) \leq F(\mu)$$ for each step (Halmos et al., 18 Nov 2025).
• Unconditional stability: Iterates converge to minimizers of $F$ for any positive step-size, provided geodesic convexity (Halmos et al., 18 Nov 2025).
• $L^\infty$, $L^p$, Sobolev, and BV bounds: Discrete analogues of Bakry–Émery, McCann contraction, and gradient-entropy inequalities propagate through the scheme (Marino et al., 2019, Carrillo et al., 2017, Elbar, 2024).
• Discrete chain rules and monotonicity: Quantities such as Fisher information and moduli of continuity are nonincreasing across steps for entropic flows (Caillet et al., 2024).

2. Euler–Lagrange Conditions and Structural Formulae

For absolutely continuous minimizers, the first-order optimality condition is typically encoded by the existence of a Kantorovich potential $\varphi$ (related to the optimal transport map) and the first variation $\frac{\delta F}{\delta \nu}$: $$\frac{\varphi}{\tau} + \frac{\delta F}{\delta \nu}(\nu) = \mathrm{const}, \quad \text{on } \operatorname{supp}(\nu).$$ Equivalently, if $\nu$ has density $\rho^{k+1}$, and the previous iterate $\mu$ has density $\rho^k$, then (in many classical cases),

$$\nabla \varphi(x) = -\tau \nabla \frac{\delta F}{\delta \rho} (\rho^{k+1})(x)$$

or, for Wasserstein-2 flows of entropy and interaction energies, the transport map $T(x) = x - \nabla\varphi(x)$ pushes $\rho^{k+1}$ forward to $\rho^k$ (Caillet et al., 2024, Marino et al., 2019, Carlier et al., 2017, Elbar, 2024). For the TV-JKO scheme, the Euler–Lagrange system additionally involves dual variables for total-variation and nonnegativity constraints (Carlier et al., 2017).

Such conditions enable derivation of maximum/minimum principles, as well as refined regularity and entropy-dissipation estimates at the discrete level.

3. Discretization and Numerical Realization

The practical implementation of the JKO operator can range from classical finite-dimensional optimization formulations to high-dimensional machine learning approaches and operator-splitting schemes:

• Discrete Monge–Ampère-type variational problems: One step is reformulated as a convex minimization over finite sets using power diagrams or discrete subdifferential polytopes, with convergence guaranteed by $\Gamma$-convergence (Benamou et al., 2014).
• Operator splitting: For functionals decomposed as $\mathcal{F} = \mathcal{F}_1 + \mathcal{F}_2$, iterative application of resolvent steps for each term (e.g., the proximal ULA algorithm alternating between entropy and potential steps) approximates a full JKO time-step efficiently (Bernton, 2018, Gallouët et al., 2016).
• Entropic and kinetic regularization: Dynamic (Benamou–Brenier) and Eulerian PDE reformulations, sometimes with controlled diffusion or entropy-penalization, mitigate high-dimensional complexity. Low-rank tensor-train decompositions and Anderson acceleration have been used for feasibility in Bayesian inverse problems (Aksenov et al., 2024).

Recent neural operator frameworks parameterize the JKO map directly (as a neural displacement field or potential), trained with self-supervised or adversarial objectives from synthetic or empirical trajectories (Lee et al., 2023, Feng et al., 9 Jan 2026, Persiianov et al., 2 Jun 2025).

4. Extensions: General Metrics, Inexactness, and Geometry

Generalizations of the JKO operator encompass:

• General transport costs: The proximity term $W_2^2$ is replaced by a cost $\mathcal{T}_c(\mu,\nu)$, where $c(x, y)$ is a smooth cost satisfying metric-compatibility conditions. If $c$ induces a Riemannian or Hessian-metric, the discrete solution operator converges to the corresponding geometric Fokker–Planck evolution (Rankin et al., 2024). Bregman divergences yield explicit representations in Hessian manifolds.
• Kantorovich–Fisher–Rao (KFR) frameworks: JKO splitting strategies handle flows with reaction and dissipation on positive Radon measures—separately evolving transport (Wasserstein) and reaction (Fisher–Rao) substeps (Gallouët et al., 2016).
• Inexact and approximate solution steps: Proximal-point steps computed only approximately (with controlled error in Wasserstein distance or functional value) satisfy convergence to the continuous minimizer, provided the error sequence is summable in a suitable sense (Marino et al., 29 May 2025). This robustness is crucial for scalable computational approaches relying on interior-point, variational, or neural solvers.

5. Analytical and Regularity Results

The JKO operator admits discrete maximum/minimum principles, $L^\infty$ bounds, and compactness properties mirroring those of the continuous flow. Specifically:

• For the heat equation (entropy gradient flow), the Fisher information is nonincreasing per JKO step, and the modulus of continuity of $\log \rho$ is preserved across iterations (Caillet et al., 2024).
• For chemotaxis (Keller–Segel) flows, the operator yields $L^\infty$ and Sobolev bounds up to blow-up time, with explicit time-of-existence estimates matching the continuous aggregation regime (Carrillo et al., 2017).
• For the total variation Wasserstein flow, the TV-JKO operator satisfies maximum principles in all dimensions, minimum principles in low-dimensional or symmetric settings, and global uniform bounds under initial positivity. The limiting PDE is a fourth-order nonlinear evolution driven by the divergence of the TV subdifferential field (Carlier et al., 2017).
• For the granular-medium and Fokker–Planck equations, discrete Li–Yau–Hamilton matrix inequalities grant Lipschitz/Harnack control, while $L^2_{t}H^2_{x}$ strong convergence is established in the limit $\tau \to 0$ (Coudreuse, 10 Oct 2025, Elbar, 2024).

6. Neural and Inverse Learning of JKO Operators

Neural methods realize approximate JKO operators ("learned displacement fields" or "operator-nets") via optimization against typical JKO loss functionals or min-max adversarial objectives. In "Learn-to-Evolve", iterative self-supervised alternation between trajectory generation and operator update allows recovery of accurate and robust discrete solution operators with strong out-of-distribution generalization in aggregation, porous medium, and Fokker–Planck settings (Feng et al., 9 Jan 2026). Adversarial inverse optimization approaches (as in iJKOnet) ensure recovery of energy landscapes governing observed (possibly unpaired) evolving distributions, with explicit theoretical error bounds for convex potential energies (Persiianov et al., 2 Jun 2025).

7. Implicit Bias and Higher-Order Structure

Beyond the first-order approximation to the Wasserstein gradient flow, the JKO operator exhibits a canonical second-order "implicit bias." Backward error analysis reveals that the iterates match the gradient flow of a modified energy $$J^\eta(\rho) = J(\rho) - \frac{\eta}{4} \int \|\nabla \frac{\delta J}{\delta \rho}\|^2 \rho$$, which penalizes high metric curvature in the driving energy (Halmos et al., 18 Nov 2025). For entropy, the bias corresponds to Fisher information, and for KL divergence, to the Fisher–Hyvärinen divergence.

This structural correction yields additional regularization and improved stability, not present in naive time-discretization schemes, and is reflected in the superior performance and stability guarantees of the JKO solution operator.

---

The JKO solution operator thus constitutes a rigorous, unifying, and flexible tool for the discrete-time approximation of a broad range of dissipative, constrained, and nonlinear PDEs and serves as an organizing principle for classical, numerical, and learning-based methods in Wasserstein gradient flows. The persistence of quantitative contractivity, stability, and structural bias explain both its theoretical appeal and its practical efficacy in modern computational and data-driven contexts.