Accelerated Regularized Wasserstein Proximal Sampling Algorithms

TL;DR

This work presents Accelerated Regularized Wasserstein Proximal (ARWP) sampling, a momentum-enabled extension of regularized Wasserstein proximal flows for drawing samples from Gibbs distributions. By replacing the intractable score with a RWProx-based surrogate and incorporating a second-order dynamics, ARWP achieves faster asymptotic contraction and improved discrete-time mixing compared to kinetic Langevin and BRWP, while keeping computational cost similar. The authors provide a thorough analysis for quadratic (Gaussian) targets, including continuous-time and linearized discrete-time rates, and validate the approach on 1D and 2D examples, Rosenbrock, multimodal mixtures, and Bayesian neural networks. The results demonstrate better tail exploration, structured particle dynamics, and improved generalization in some non-log-concave Bayesian tasks, signaling a robust path toward scalable, accelerated score-based sampling. The framework builds a bridge between mean-field control, Wasserstein proximal theory, and accelerated information flows in a practical, kernel-free setting.

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.

Figures (11)

• Figure 1: Contour plots of the covariance error of discrete-time ARWP \ref{['eqs:symplecticARWP']}, for target distribution $\Lambda = {\mathcal{N}}(0, I_1)$, with initialization ${\mathcal{N}}(0,10^{-3})$ (top) and ${\mathcal{N}}(0,4)$ (bottom). Error is plotted against the damping parameter $a\in [10^0, 10^2]$ and step-size $\eta \in [10^{-4}, 10^{-0.5}]$, and fixed $T=0.2$. The gray line indicates the optimal damping parameter in continuous time, given by \ref{['eq:optimalDamping']}. The black region in the top right corner indicates (empirical) divergence, occurring when $a\eta > 2$.
• Figure 2: Convergence in KL divergence for the 2D Gaussian, run with 100 particles over 100 iterations. We observe that the deterministic methods ARWP and BRWP enjoy particle-wise convergence, indicated by the smaller oscillations between iterations. The accelerated Langevin methods ILA and KLMC continue to evolve due to the Brownian motion in the velocity. We observe that while BRWP has a faster initial convergence rate, both ARWP-Nesterov and ARWP-Heavy-ball reach their steady states faster. This is consistent with classical optimization results.
• Figure 3: Particle positions after 100 iterations for the 2D Gaussian with condition number $\kappa=50$, run with 100 particles. We observe that both the accelerated and non-accelerated Langevin algorithms look more randomly sampled, as the particles do not interact. Moreover, the proposed ARWP method has a similarly structured but slightly messier terminal position compared to BRWP. Both ARWP and BRWP particle positions converge and do not move.
• Figure 4: Particle evolution for Rosenbrock distribution at iterations 50, 200, 500. Top to bottom: (1) ARWP-Nesterov ($T = \eta = 0.02$), (2) "underdamped ILA" ($\eta = 0.05$, damping parameter $=2$), (3) KLMC ($\eta = 0.01$, damping parameter $=5$), and (4) ULA ($\eta = 0.01$). We observe that ARWP and ILA are better at exploring the tails than KLMC. However, ILA and ULA both have particles straying away from the main parabola due to time-discretization bias.
• Figure 5: KL divergence between the particles and the underlying distribution for the Gaussian mixture. We observe that ARWP method converges faster than BRWP, and the particle cloud stabilizes with lower KL divergence than the corresponding Langevin methods. KLMC experiences mode collapse, which persists through hyperparameter changes.

Theorems & Definitions (24)

• Definition 1: santambrogio2015optimalambrosio2005gradient
• Definition 2
• Remark 1
• Remark 2
• Lemma 1
• Proposition 1
• proof
• Remark 3
• Proposition 2
• proof : Sketch