Underdamped Langevin MCMC with third order convergence

TL;DR

This paper introduces a novel third order convergent numerical method for underdamped Langevin diffusion, achieving faster sampling error reduction under certain smoothness conditions, with practical validation on Bayesian logistic regression tasks.

Contribution

It presents the first gradient-only ULD method with third order convergence, improving sampling efficiency under specific smoothness assumptions.

Findings

Achieves $ ext{W}_2$ error $oldsymbol{ ext{O}(rac{ ext{d}^{1/2}}{ ext{ε}^{1/3}})}$ steps with third derivative Lipschitz continuity.

Performs competitively on real-world datasets compared to existing Langevin MCMC and NUTS.

Provides non-asymptotic error bounds for the proposed algorithm.

Abstract

In this paper, we propose a new numerical method for the underdamped Langevin diffusion (ULD) and present a non-asymptotic analysis of its sampling error in the 2-Wasserstein distance when the $d$-dimensional target distribution $p(x)\propto e^{-f(x)}$ is strongly log-concave and has varying degrees of smoothness. Precisely, under the assumptions that the gradient and Hessian of $f$ are Lipschitz continuous, our algorithm achieves a 2-Wasserstein error of $\varepsilon$ in $\mathcal{O}(\sqrt{d}/\varepsilon)$ and $\mathcal{O}(\sqrt{d}/\sqrt{\varepsilon})$ steps respectively. Therefore, our algorithm has a similar complexity as other popular Langevin MCMC algorithms under matching assumptions. However, if we additionally assume that the third derivative of $f$ is Lipschitz continuous, then our algorithm achieves a 2-Wasserstein error of $\varepsilon$ in $\mathcal{O}(\sqrt{d}/\varepsilon^{\frac{1}{3}})$ steps. To the best of our knowledge, this is the first gradient-only method for ULD with third order convergence. To support our theory, we perform Bayesian logistic regression across a range of real-world datasets, where our algorithm achieves competitive performance compared to an existing underdamped Langevin MCMC algorithm and the popular No U-Turn Sampler (NUTS).