On the Convergence of Wasserstein Gradient Descent for Sampling

TL;DR

This paper introduces a Wasserstein gradient descent framework for sampling from complex distributions, providing theoretical convergence guarantees and demonstrating practical effectiveness in high-dimensional Bayesian inference.

Contribution

It develops a new theoretical foundation for Wasserstein gradient descent in measure spaces and proposes a particle-based algorithm with score matching for practical sampling.

Findings

WGD converges for specific subclasses of Wasserstein space.

The particle-based WGD effectively approximates complex target distributions.

WGD outperforms standard MCMC and variational methods in challenging scenarios.

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.