Tuning-Free Sampling via Optimization on the Space of Probability Measures

TL;DR

This paper develops tuning-free gradient-based sampling algorithms derived from Wasserstein gradient flows, providing theoretical guarantees and competitive performance without the need for step size tuning.

Contribution

It introduces a family of tuning-free sampling algorithms for probability measures, with theoretical convergence guarantees and practical benchmarking results.

Findings

Achieves convergence rates comparable to optimally tuned algorithms.

Provides tuning-free variants of popular sampling methods like ULA and SGLD.

Performs competitively in various tasks without step size tuning.

Abstract

We introduce adaptive, tuning-free step size schedules for gradient-based sampling algorithms obtained as time-discretizations of Wasserstein gradient flows. The result is a suite of tuning-free sampling algorithms, including tuning-free variants of the unadjusted Langevin algorithm (ULA), stochastic gradient Langevin dynamics (SGLD), mean-field Langevin dynamics (MFLD), Stein variational gradient descent (SVGD), and variational gradient descent (VGD). More widely, our approach yields tuning-free algorithms for solving a broad class of stochastic optimization problems over the space of probability measures. Under mild assumptions (e.g., geodesic convexity and locally bounded stochastic gradients), we establish strong theoretical guarantees for our approach. In particular, we recover the convergence rate of optimally tuned versions of these algorithms up to logarithmic factors, in both nonsmooth and smooth settings. We then benchmark the performance of our methods against comparable existing approaches. Across a variety of tasks, our algorithms achieve similar performance to the optimal performance of existing algorithms, with no need to tune a step size parameter.