# Computational and Statistical Asymptotic Analysis of the JKO Scheme for Iterative Algorithms to update distributions

arXiv: https://arxiv.org/abs/2501.06408 · HTML: https://mlsanity.com/paper/2501.06408/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 20 of its 55 references that are in the map.
- First posted on arXiv 11 Jan 2025.

## Builds on (most cited first, 20 of 20)

- [Variational Inference: A Review for Statisticians](https://mlsanity.com/paper/1601.00670/) (2016)
- [Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm](https://mlsanity.com/paper/1608.04471/) (2016)
- [Rapid Convergence of the Unadjusted Langevin Algorithm: Isoperimetry Suffices](https://mlsanity.com/paper/1903.08568/) (2019)
- [Analysis of Langevin Monte Carlo via convex optimization](https://mlsanity.com/paper/1802.09188/) (2018)
- [Sampling as optimization in the space of measures: The Langevin dynamics as a composite optimization problem](https://mlsanity.com/paper/1802.08089/) (2018)
- [On the properties of variational approximations of Gibbs posteriors](https://mlsanity.com/paper/1506.04091/) (2015)
- [Proximal Optimal Transport Modeling of Population Dynamics](https://mlsanity.com/paper/2106.06345/) (2021)
- [Structured Logconcave Sampling with a Restricted Gaussian Oracle](https://mlsanity.com/paper/2010.03106/) (2020)
- [Variational inference via Wasserstein gradient flows](https://mlsanity.com/paper/2205.15902/) (2022)
- [On the geometry of Stein variational gradient descent](https://mlsanity.com/paper/1912.00894/) (2019)
- [The Randomized Midpoint Method for Log-Concave Sampling](https://mlsanity.com/paper/1909.05503/) (2019)
- [Large-Scale Wasserstein Gradient Flows](https://mlsanity.com/paper/2106.00736/) (2021)
- [On sampling from a log-concave density using kinetic Langevin diffusions](https://mlsanity.com/paper/1807.09382/) (2018)
- [Optimizing Functionals on the Space of Probabilities with Input Convex Neural Networks](https://mlsanity.com/paper/2106.00774/) (2021)
- [Variational Wasserstein gradient flow](https://mlsanity.com/paper/2112.02424/) (2021)
- [Accelerating Langevin Sampling with Birth-death](https://mlsanity.com/paper/1905.09863/) (2019)
- [Policy Optimization as Wasserstein Gradient Flows](https://mlsanity.com/paper/1808.03030/) (2018)
- [Learning Gaussian Mixtures Using the Wasserstein-Fisher-Rao Gradient Flow](https://mlsanity.com/paper/2301.01766/) (2023)
- [Efficient Gradient Flows in Sliced-Wasserstein Space](https://mlsanity.com/paper/2110.10972/) (2021)
- [The limited-memory recursive variational Gaussian approximation (L-RVGA)](https://mlsanity.com/paper/2303.14195/) (2023)

## Built on by (most cited first, 1 of 1)

- [Learning of Population Dynamics: Inverse Optimization Meets JKO Scheme](https://mlsanity.com/paper/2506.01502/) (2025)

## Abstract (arXiv)

The seminal paper of Jordan, Kinderlehrer, and Otto introduced what is now widely known as the JKO scheme, an iterative algorithmic framework for computing distributions. This scheme can be interpreted as a Wasserstein gradient flow and has been successfully applied in machine learning contexts, such as deriving policy solutions in reinforcement learning. In this paper, we extend the JKO scheme to accommodate models with unknown parameters. Specifically, we develop statistical methods to estimate these parameters and adapt the JKO scheme to incorporate the estimated values. To analyze the adopted statistical JKO scheme, we establish an asymptotic theory via stochastic partial differential equations that describes its limiting dynamic behavior. Our framework allows both the sample size used in parameter estimation and the number of algorithmic iterations to go to infinity. This study offers a unified framework for joint computational and statistical asymptotic analysis of the statistical JKO scheme. On the computational side, we examine the scheme's dynamic behavior as the number of iterations increases, while on the statistical side, we investigate the large-sample behavior of the resulting distributions computed through the scheme. We conduct numerical simulations to evaluate the finite-sample performance of the proposed methods and validate the developed asymptotic theory.