A convergence rate for the entropic JKO scheme

1. Problem Statement and Motivation

Context:
The Jordan-Kinderlehrer-Otto (JKO) scheme is a variational method used to construct discrete-time approximations of gradient flows in the space of probability measures equipped with the Wasserstein metric ($W_2$). It is widely used for solving Partial Differential Equations (PDEs) arising in physics, biology, and machine learning.

The Challenge:
While the classical JKO scheme is theoretically sound, numerically computing the Wasserstein distance $W_2$ is computationally expensive (often requiring solving a linear programming problem or a complex optimization).

The Solution (Entropic Regularization):
To address the computational cost, the "Entropic JKO scheme" replaces the Wasserstein distance with its entropic regularization, known as the Schrödinger cost ($Sch_{\alpha\tau}$). This allows the use of the highly efficient Sinkhorn algorithm.

• The Trade-off: Previous work (e.g., [4]) established that if the regularization parameter $\varepsilon = \alpha\tau$ scales linearly with the time step $\tau$, the entropic scheme converges to a modified PDE:
  $$\partial_t \rho - \text{div}\left(\rho \nabla \frac{\delta F}{\delta \rho}(\rho)\right) = \frac{\alpha}{2} \Delta \rho$$
  This introduces an artificial diffusion term $\frac{\alpha}{2} \Delta \rho$.
• The Gap: While convergence was known, there were no explicit convergence rates quantifying the error between the entropic scheme and the true solution of the original PDE (where $\alpha=0$), nor was there a direct bound between the classical and entropic schemes.

Objective:
The authors aim to provide an explicit convergence rate for the Entropic JKO scheme towards the solution of the original PDE as both the time step $\tau \to 0$ and the regularization parameter $\alpha \to 0$.

2. Methodology and Mathematical Framework

The paper relies on the theory of gradient flows in the Wasserstein space, specifically utilizing the Evolution Variational Inequality (EVI) and properties of $\lambda$-convexity along generalized geodesics.

Key Assumptions (Hypothesis 1.1):
The functional $F$ (the energy driving the flow) must satisfy:

1. Lower Semicontinuity: With respect to narrow convergence in $P_2(\mathbb{R}^d)$.
2. $\lambda$-Convexity: Along generalized geodesics (a strong form of convexity ensuring stability).
3. Controlled Growth: A bound on the variation of $F$ under the heat flow (related to the Laplacian of the potential).

Core Strategy:
The authors employ a stability analysis comparing the classical JKO step ($J^0_\tau$) and the entropic JKO step ($J^\alpha_\tau$).

1. Decomposition: They decompose the distance between the $k+1$-th iterates of the two schemes using the triangle inequality:
   $$W_2(J^0_{k+1}, J^\alpha_{k+1}) \leq W_2(J^0(J^0_k), J^0(J^\alpha_k)) + W_2(J^0(J^\alpha_k), J^\alpha(J^\alpha_k))$$

   • Term (I): The contraction of the classical scheme (already known from [3]).
   • Term (II): The stability of the scheme under perturbation (the novel contribution).

2. Discrete EVI: They utilize the discrete EVI for $\lambda$-convex functionals to bound Term (II). This requires showing that the minimizer of the entropic scheme is a "good competitor" for the classical scheme's minimization problem.

3. Competitor Construction: To bound the difference in costs, they construct a specific competitor for the entropic problem by taking the classical minimizer and convolving it with a heat kernel ($\sigma_{\alpha\tau}$). This exploits the heuristic that the flow of $F + \frac{\alpha}{2}H$ (where $H$ is entropy) can be approximated by flowing $F$ then flowing $H$.

4. Recursive Bounds: They derive a recursive inequality for the squared Wasserstein distance and solve it using a discrete Grönwall lemma adapted for squared distances.

3. Key Contributions and Results

A. Main Theorem: Convergence Rate (Theorem 1.3)

The paper establishes an explicit bound on the Wasserstein distance between the $n$-th iterate of the classical JKO scheme ($J^0_{n,\tau}$) and the entropic JKO scheme ($J^\alpha_{n,\tau}$) after time $t = n\tau$.

• Case $\lambda = 0$ (Flat/Convex):
  $$W_2(J^0_{n,\tau}, J^\alpha_{n,\tau}) \leq \sqrt{2\tau(F(\mu_0) - F(J^0_{n,\tau}))} + \sqrt{n\tau\alpha \left( H(\mu_0) - H(J^\alpha_{n,\tau}) + Kn\tau \right)}$$
• Case $\lambda \neq 0$ (General $\lambda$-convex):
  The bound involves exponential factors depending on $\lambda$ and $\tau$, but retains the same structure: a term depending on the energy drop of the classical scheme and a term depending on the entropy drop of the entropic scheme, scaled by $\sqrt{\alpha}$.

B. Corollary: Convergence to the True PDE Solution (Corollary 1.3.1)

By combining Theorem 1.3 with known convergence rates of the classical JKO scheme to the PDE solution $\rho_0(t)$, the authors derive the total error bound:
$$W_2(J^\alpha_{n,t/n}(\mu_0), \rho_0(t)) \leq C \left( \sqrt{\alpha} + \frac{1}{\sqrt{n}} \right)$$
This proves that the entropic scheme converges to the true solution with a rate of $O(\sqrt{\alpha} + \sqrt{\tau})$.

C. Optimality Analysis (Section 5)

The authors investigate whether the $\sqrt{\alpha}$ rate is optimal.

• Continuous Level: They compare their bound to a known bound between the continuous PDE solutions (Theorem 1.4). They show their discrete bound is within a factor of 2 of the continuous bound, suggesting the order of magnitude is correct.
• Discrete Level: Using a specific "toy model" where $F(\rho) = \int \frac{\lambda}{2}|x|^2 d\rho$ (quadratic potential) and initial data is Gaussian, they compute exact solutions.
  • They prove that for this model, the derived inequalities become equalities (up to terms vanishing as $\tau \to 0$).
  • This confirms that the $\sqrt{\alpha}$ dependence is optimal and not an artifact of the proof technique (e.g., the splitting argument or the Grönwall lemma).

4. Significance and Impact

1. Theoretical Justification: The paper provides the first rigorous, explicit convergence rate for the Entropic JKO scheme towards the unregularized PDE solution. It clarifies that the artificial diffusion introduced by entropic regularization vanishes at a rate of $\sqrt{\alpha}$.
2. Practical Guidance: For practitioners using Sinkhorn-based JKO schemes (common in optimal transport and generative modeling), the result quantifies the trade-off between computational speed (controlled by $\alpha$) and accuracy. It suggests that to achieve an error of $\epsilon$, one must choose $\alpha \sim \epsilon^2$ and $\tau \sim \epsilon^2$.
3. Methodological Innovation: The proof technique, which treats the entropic scheme as a perturbation of the classical scheme and uses a "heat kernel competitor," offers a new tool for analyzing regularized gradient flows.
4. Optimality: By proving the bound is sharp in the Gaussian case, the authors rule out the possibility that the $\sqrt{\alpha}$ rate is merely a loose estimate, setting a benchmark for future improvements.

5. Conclusion

Baradat and Cherf successfully bridge the gap between the theoretical convergence of the Entropic JKO scheme and practical error bounds. They demonstrate that while entropic regularization introduces a diffusion error, this error is controllable and converges to zero at a rate of $\sqrt{\alpha}$ as the regularization parameter vanishes, provided the time step also vanishes appropriately. This work solidifies the theoretical foundation for using entropic regularization in gradient flow simulations.