Accelerated Regularized Wasserstein Proximal Sampling Algorithms

TL;DR

This paper introduces ARWP, an accelerated sampling algorithm for Gibbs distributions using a second-order score-based ODE and Wasserstein proximal methods, demonstrating faster convergence and better exploration in experiments.

Contribution

The paper proposes ARWP, a novel accelerated sampling algorithm combining second-order ODEs with Wasserstein proximal methods, improving mixing rates and exploration over existing methods.

Findings

ARWP achieves higher contraction rates than kinetic Langevin algorithms.

Numerical experiments show ARWP's faster mixing and tail exploration.

ARWP particles demonstrate better generalization in Bayesian neural network tasks.

Abstract

We consider sampling from a Gibbs distribution by evolving a finite number of particles using a particular score estimator rather than Brownian motion. To accelerate the particles, we consider a second-order score-based ODE, similar to Nesterov acceleration. In contrast to traditional kernel density score estimation, we use the recently proposed regularized Wasserstein proximal method, yielding the Accelerated Regularized Wasserstein Proximal method (ARWP). We provide a detailed analysis of continuous- and discrete-time non-asymptotic and asymptotic mixing rates for Gaussian initial and target distributions, using techniques from Euclidean acceleration and accelerated information gradients. Compared with the kinetic Langevin sampling algorithm, the proposed algorithm exhibits a higher contraction rate in the asymptotic time regime. Numerical experiments are conducted across various low-dimensional experiments, including multi-modal Gaussian mixtures and ill-conditioned Rosenbrock distributions. ARWP exhibits structured and convergent particles, accelerated discrete-time mixing, and faster tail exploration than the non-accelerated regularized Wasserstein proximal method and kinetic Langevin methods. Additionally, ARWP particles exhibit better generalization properties for some non-log-concave Bayesian neural network tasks.