The Randomized Midpoint Method for Log-Concave Sampling

TL;DR

This paper introduces a faster MCMC algorithm based on a new discretization framework for Langevin diffusion, significantly improving sampling efficiency from log-concave distributions in high dimensions.

Contribution

The paper proposes a novel discretization framework for stochastic differential equations and an accelerated MCMC algorithm for log-concave sampling, outperforming previous methods in convergence speed.

Findings

Achieves $ ilde{O}(rac{ ext{condition number}^{7/6}}{ ext{error}^{1/3}} + rac{ ext{condition number}}{ ext{error}^{2/3}})$ steps

Performs faster than the previous best algorithm requiring $ ilde{O}( ext{condition number}^{1.5}/ ext{error})$ steps

Can be parallelized to require only $O( ext{condition number} imes ext{log}(1/ ext{error}))$ parallel steps

Abstract

Sampling from log-concave distributions is a well researched problem that has many applications in statistics and machine learning. We study the distributions of the form $p^{*}\propto\exp(-f(x))$, where $f:\mathbb{R}^{d}\rightarrow\mathbb{R}$ has an $L$-Lipschitz gradient and is $m$-strongly convex. In our paper, we propose a Markov chain Monte Carlo (MCMC) algorithm based on the underdamped Langevin diffusion (ULD). It can achieve $\epsilon\cdot D$ error (in 2-Wasserstein distance) in $\tilde{O}\left(\kappa^{7/6}/\epsilon^{1/3}+\kappa/\epsilon^{2/3}\right)$ steps, where $D\overset{\mathrm{def}}{=}\sqrt{\frac{d}{m}}$ is the effective diameter of the problem and $\kappa\overset{\mathrm{def}}{=}\frac{L}{m}$ is the condition number. Our algorithm performs significantly faster than the previously best known algorithm for solving this problem, which requires $\tilde{O}\left(\kappa^{1.5}/\epsilon\right)$ steps. Moreover, our algorithm can be easily parallelized to require only $O(\kappa\log\frac{1}{\epsilon})$ parallel steps. To solve the sampling problem, we propose a new framework to discretize stochastic differential equations. We apply this framework to discretize and simulate ULD, which converges to the target distribution $p^{*}$. The framework can be used to solve not only the log-concave sampling problem, but any problem that involves simulating (stochastic) differential equations.