mlsanityData to 7 Oct 2026

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

arXiv 2501.06408 ↗First posted 11 Jan 2025stat.ML cs.LGcited by 1

Not on the map. Why?

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.

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

Builds on

20 of 55 references.

References · 20

  1. 2016Variational Inference: A Review for Statisticianscited by 1,466
  2. 2016Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithmcited by 498
  3. 2019Rapid Convergence of the Unadjusted Langevin Algorithm: Isoperimetry Sufficescited by 224
  4. 2018Analysis of Langevin Monte Carlo via convex optimizationcited by 126
  5. 2018Sampling as optimization in the space of measures: The Langevin dynamics as a composite optimization problemcited by 123
  6. 2015On the properties of variational approximations of Gibbs posteriorscited by 121
  7. 2021Proximal Optimal Transport Modeling of Population Dynamicscited by 91
  8. 2020Structured Logconcave Sampling with a Restricted Gaussian Oraclecited by 85
  9. 2022Variational inference via Wasserstein gradient flowscited by 84
  10. 2019On the geometry of Stein variational gradient descentcited by 82
  11. 2019The Randomized Midpoint Method for Log-Concave Samplingcited by 77
  12. 2021Large-Scale Wasserstein Gradient Flowscited by 76
8 more
  1. 2018On sampling from a log-concave density using kinetic Langevin diffusionscited by 61
  2. 2021Optimizing Functionals on the Space of Probabilities with Input Convex Neural Networkscited by 53
  3. 2021Variational Wasserstein gradient flowcited by 53
  4. 2019Accelerating Langevin Sampling with Birth-deathcited by 45
  5. 2018Policy Optimization as Wasserstein Gradient Flowscited by 33
  6. 2023Learning Gaussian Mixtures Using the Wasserstein-Fisher-Rao Gradient Flowcited by 25
  7. 2021Efficient Gradient Flows in Sliced-Wasserstein Spacecited by 21
  8. 2023The limited-memory recursive variational Gaussian approximation (L-RVGA)cited by 6

Built on by

Cited by year

Citing papers per year: 2025: 1.

1 paper.