Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

On the Convergence of Wasserstein Gradient Descent for Sampling

Van Chien Ta, Thi Mai Hong Chu, Minh-Ngoc Tran

TL;DR

The work addresses sampling by reframing it as KL minimization $F(\mu)=\mathrm{KL}(\mu\|\pi)$ on the 2-Wasserstein space $\mathbb{W}_2(\mathbb{R}^d)$ and solving it with Wasserstein gradient descent (WGD). It identifies two regularity regimes, $(\alpha,\beta)$-regular and $(c_1,c_2)$-regular measures, under which WGD converges (with exact or estimated gradients) and provides corresponding theoretical guarantees with carefully chosen step sizes. A practical particle-based algorithm using score matching to estimate $\nabla\log\mu_t$ is developed, including an annealing strategy and stopping criteria, and validated on challenging targets such as Bayesian logistic regression, banana-shaped, and eggbox distributions. The results suggest WGD as a flexible, scalable alternative for Bayesian inference in high-dimensional or multimodal settings, while highlighting improvements that could come from faster score estimators and adaptive step-size schemes.

Abstract

This paper studies the optimization of the KL functional on the Wasserstein space of probability measures, and develops a sampling framework based on Wasserstein gradient descent (WGD). We identify two important subclasses of the Wasserstein space for which the WGD scheme is guaranteed to converge, thereby providing new theoretical foundations for optimization-based sampling methods on measure spaces. For practical implementation, we construct a particle-based WGD algorithm in which the score function is estimated via score matching. Through a series of numerical experiments, we demonstrate that WGD can provide good approximation to a variety of complex target distributions, including those that pose substantial challenges for standard MCMC and parametric variational Bayes methods. These results suggest that WGD offers a promising and flexible alternative for scalable Bayesian inference in high-dimensional or multimodal settings.

On the Convergence of Wasserstein Gradient Descent for Sampling

TL;DR

The work addresses sampling by reframing it as KL minimization on the 2-Wasserstein space and solving it with Wasserstein gradient descent (WGD). It identifies two regularity regimes, -regular and -regular measures, under which WGD converges (with exact or estimated gradients) and provides corresponding theoretical guarantees with carefully chosen step sizes. A practical particle-based algorithm using score matching to estimate is developed, including an annealing strategy and stopping criteria, and validated on challenging targets such as Bayesian logistic regression, banana-shaped, and eggbox distributions. The results suggest WGD as a flexible, scalable alternative for Bayesian inference in high-dimensional or multimodal settings, while highlighting improvements that could come from faster score estimators and adaptive step-size schemes.

Abstract

This paper studies the optimization of the KL functional on the Wasserstein space of probability measures, and develops a sampling framework based on Wasserstein gradient descent (WGD). We identify two important subclasses of the Wasserstein space for which the WGD scheme is guaranteed to converge, thereby providing new theoretical foundations for optimization-based sampling methods on measure spaces. For practical implementation, we construct a particle-based WGD algorithm in which the score function is estimated via score matching. Through a series of numerical experiments, we demonstrate that WGD can provide good approximation to a variety of complex target distributions, including those that pose substantial challenges for standard MCMC and parametric variational Bayes methods. These results suggest that WGD offers a promising and flexible alternative for scalable Bayesian inference in high-dimensional or multimodal settings.
Paper Structure (11 sections, 8 theorems, 104 equations, 3 figures, 3 algorithms)

This paper contains 11 sections, 8 theorems, 104 equations, 3 figures, 3 algorithms.

Table of Contents

  1. Introduction
  2. $(\alpha,\beta)$-regular measures
  3. WGD with Estimated Wasserstein Gradient
  4. $(c_1,c_2)$-regular measures
  5. WGD with Estimated Wasserstein Gradient
  6. Numerical examples
  7. Implemenation strategy
  8. Bayesian logistic regression
  9. Banana-shaped distribution
  10. Eggbox distribution
  11. Conclusion

Key Result

Proposition 1

Assume that the target $\pi\in\mathcal{P}_{(\alpha,\beta)}^r(\mathbb{R}^d)$. For every $\mu\in\mathcal{P}_{(\alpha,\beta)}^r(\mathbb{R}^d)$, let $\nu=(\text{Id}-\epsilon \nabla_\mu F)_{\#}\mu$, with $\nabla_\mu F=\nabla \log\frac{\mu}{\pi}$ the Wasserstein gradient of $F$. Then, we have where $C$ is a finite constant depending only on $\alpha,\beta$ and dimension $d$.

Figures (3)

  • Figure 1: Approximate marginal posterior distributions obtained using WGD (dashed lines) and MCMC (solid lines).
  • Figure 2: Contour of the banana-shaped distribution (solid lines) together with the particles generated by the annealing WGD algorithm. The dashed lines are the contours of the Gaussian VB approximation
  • Figure 3: Eggbox targets and their WGD approximations

Theorems & Definitions (17)

  • Definition 1: $(\alpha,\beta)$-regular measures
  • Proposition 1
  • Theorem 1
  • Proposition 2
  • Theorem 2
  • Definition 2: $(c_1,c_2)$-regular measures
  • Theorem 3
  • proof
  • Proposition 3
  • Theorem 4
  • ...and 7 more