Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Shang Wu, Yazhen Wang

TL;DR

This work extends the JKO scheme to models with unknown parameters by developing a statistical JKO framework that combines offline and online parameter estimation with Wasserstein-gradient-flow dynamics. The authors derive joint computational-statistical asymptotics, showing that the scaled errors converge to stochastic PDE-driven limits, with clear rates and Brownian forcing in the online setting. A Bures-Wasserstein specialization reveals tractable mean- and covariance-based dynamics for Gaussian distributions, and the paper provides numerical validation in a simple one-dimensional setting. The results offer a unified theory for analyzing iterative distribution updates under parameter uncertainty with practical implications for learning-based sampling and inference.

Abstract

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.

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

TL;DR

This work extends the JKO scheme to models with unknown parameters by developing a statistical JKO framework that combines offline and online parameter estimation with Wasserstein-gradient-flow dynamics. The authors derive joint computational-statistical asymptotics, showing that the scaled errors converge to stochastic PDE-driven limits, with clear rates and Brownian forcing in the online setting. A Bures-Wasserstein specialization reveals tractable mean- and covariance-based dynamics for Gaussian distributions, and the paper provides numerical validation in a simple one-dimensional setting. The results offer a unified theory for analyzing iterative distribution updates under parameter uncertainty with practical implications for learning-based sampling and inference.

Abstract

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.
Paper Structure (28 sections, 16 theorems, 312 equations, 4 figures)

This paper contains 28 sections, 16 theorems, 312 equations, 4 figures.

Table of Contents

  1. Introduction
  2. Reviews on Langevin diffusions and gradient flows in the Wasserstein space
  3. The JKO scheme
  4. The plain scheme
  5. The statistical JKO scheme
  6. The statistical JKO scheme in the offline case
  7. The statistical JKO scheme in the online case
  8. Asymptotic theory of the statistical JKO scheme in the offline case
  9. Asymptotic theory of the stochastic JKO scheme in the online case
  10. Application on Bures-Wasserstein gradient flow
  11. Numerical studies
  12. Conclusion
  13. Sketch of the Proof in Section 4
  14. Proof of Theorem 4.1
  15. Sketch of the Proof in Section 5
  16. ...and 13 more sections

Key Result

Theorem 4.1

Assume the conditions (A1),(A2),(A3),(A4). Then $\hat{V}_\delta^n (t,x)$ converges to $V(t,x)$ in the sense that for any $\xi\in{\cal C}_0^\infty(\mathbb{R}_+\times\mathbb{R}^d)$, as $n\rightarrow\infty$, $\delta\rightarrow 0$, $\delta\sqrt{n}\rightarrow 0$, where $V(t, x)$ satisfies the following PDE, $\tau(x)=\nabla_\theta'\Psi_\theta(x)\gamma_\theta$, $\gamma_\theta^2$ is the asymptotic varia

Figures (4)

  • Figure 1: Function $\hat{\rho}_{\delta,1}^m(T,x)$ and $\rho(T,x)$
  • Figure 2: Function $\hat{V}_{\delta,1}^m(T,x)$ and $V_1(T,x)$
  • Figure 3: Contour plot of $\hat{V}_{\delta,1}^m(t,x)$ (left) and $V_1(t,x)$ (right)
  • Figure 4: Variance of $t/\delta(\hat{\theta}_\delta(t)-\hat{\theta}(t))$ for $\delta_m=0.0001m$, $m=1,2,\dots,100$

Theorems & Definitions (33)

  • Theorem 4.1
  • Remark 4.1
  • Theorem 5.1
  • Remark 5.1
  • Remark 5.2
  • Theorem 5.2
  • Remark 5.3
  • Proposition 5.1
  • Remark 5.4
  • Proposition 5.2
  • ...and 23 more