High-Order Langevin Diffusion Yields an Accelerated MCMC Algorithm

Introduction

Recent years have seen substantial progress in the theoretical analysis of algorithms for large-scale statistical inference. For both the optimization algorithms used to compute frequentist point estimates and the sampling algorithms that underpin Bayesian inference, nonasymptotic rates of convergence have been obtained and, increasingly, those rates include dimension dependence [see, e.g., 10, 13, 11, 9, 7, 14, 20, 5]. In particular, for the gradient-based algorithms that have become the state-of-the-art in many large-scale applications, the dimension dependence is generally linear or sublinear, providing strong theoretical support for the deployment of these algorithms in large-scale problems.

Although progress has been made in both optimization and sampling, the latter has lagged the former, arguably because of the inherent stochasticity of the sampling paradigm. Indeed, much of the recent progress in both paradigms has involved taking a continuous-time point of view, whereby algorithms are obtained as discretizations of underlying continuous dynamical systems, and this line of attack is more challenging for sampling methods. For optimization algorithms the continuous dynamics can be represented as ordinary differential equations (ODEs) , whereas the underlying dynamics are characterized as stochastic differential equations (SDEs) in the case of sampling algorithms . The non-smooth nature of the Brownian motion underlying these SDEs raises fundamental challenges in carrying out the discretization that is needed to transfer the continuous-time results to discrete time.

We focus on densities that can written in the form $p^{*}(\theta)\propto\exp(-U(\theta))$, where the potential function $U:{}^{d}\rightarrow\real$ is strongly convex and Lipschitz smooth. There is a substantial body of past work on sampling problems of this type; among other results, it has been shown that a discretization of the second-order Langevin diffusion has mixing time that scales as $O(\sqrt{d})$, which compares favorably to the best known $O(d)$ scaling of the first-order Langevin diffusion . The results were further improved by Shen and Lee , who used a uniform random time to construct the low-bias estimator for an integral, leading to an improved discretization scheme with $O(d^{1/3}/\varepsilon^{2/3})$ mixing time. Furthermore, Cao et al. show for second-order Langevin diffusions, this rate cannot be further improved to achieve $\varepsilon$ discretization error.

However, if additional and relatively strong assumptions are imposed on the density, then even faster rates of convergence can be obtained . Here we ask whether it is possible to accelerate convergence of sampling algorithms beyond the $O(d^{1/3})$ barrier in the general setting without imposing additional assumptions beyond strong convexity and Lipschitz smoothness. The main contribution of our paper is an affirmative answer to this question.

Let us provide some context for our line of attack and our contributions. As is well known from past work , the continuous-time Langevin dynamics converge to the target distribution at an exponential rate, with no dependence on dimension. However, to be implemented with digital computation, the continuous-time dynamics must be approximated with a discrete-time scheme, leading to numerical error that does scale with the dimension and the conditioning of the problem. Direct application of higher-order schemes to the Langevin diffusion is hindered by the non-smoothness of the Brownian motion. One way to circumvent this problem is to augment the traditional Langevin diffusion by moving to higher-order continuous dynamics. In particular, recent work has studied the the second-order (or underdamped) Langevin algorithm, which lifts the original $d$-dimensional space to a $2d$-dimensional space consisting of vectors of the form $x=(\theta,r)\in{}^{d}\times{}^{d}$, and considers a $2d$-dimensional collection of SDEs in these variables .

There is a natural hierarchy of such lifted schemes, and this paper is based on proposing and analyzing a carefully designed third-order lifting, to be described in Section 2.4. We provide a careful analysis of a particular discretization of this third-order scheme, establishing non-asymptotic bounds on mixing time for particular classes of potential functions.

Our presentation begins with an analysis of potential functions that have the following ridge-separable form:

where $\{u_{i}\}_{i=1}^{n}$ are a collection of univariate functions, and $\{a_{i}\}_{i=1}^{n}$ are a given collection of vectors in d. Many log-concave sampling problems that arise in statistics and machine learning involve potential functions of this form. In particular, posterior sampling in Bayesian generalized linear models, including Bayesian logistic regression and one-layer neural networks, can be written in the form (1). It is worth noticing that we do not impose any additional assumptions on the vectors $\{a_{i}\}_{i=1}^{n}$. The ridge-separable form is needed only to make sure a one-dimensional integral to have close-form solution.

Given a distribution of the form (1), with $U$ strongly convex and smooth, we prove that $O(d^{\frac{1}{4}}/\varepsilon^{\frac{1}{2}})$ steps suffice to make the Wasserstein distance between the sample and target distributions less than $\varepsilon$. This is the first time that the $O(d^{1/3})$ barrier for the log-concave sampling problem has been overcome without additional structural assumptions on the data, even for the special case of Bayesian logistic regression. The dependency on $\varepsilon$ is also improved relative to the current state of the art. It is worth noticing that our analysis allows for arbitrary vectors $\{a_{i}\}_{i=1}^{n}$ and functions $\{u_{i}\}_{i=1}^{n}$, as long as the smoothness and strong convexity of $U$ are guaranteed. This is in sharp contrast to previous work that requires incoherence conditions on the data vectors $\{a_{i}\}_{i=1}^{n}$ and/or high-order smoothness conditions on the component functions $\{u_{i}\}_{i=1}^{n}$.

We then tackle the more general setting in which the function $U$ need not be ridge-separable as in equation (1). Assuming only that we are given access to gradients from a black-box gradient oracle, we show that the dimension dependency of our algorithm is $O(\sqrt{d})$, but the dependency on final accuracy $\varepsilon$ associated with this term can be adaptive to the smoothness assumptions satisfied by $U$. In particular, we establish an upper bound on the mixing time of $O\left(d^{\frac{1}{4}}/\varepsilon^{\frac{1}{2}}+d^{\frac{1}{2}}/\varepsilon^{\frac{1}{\alpha-1}}\right)$, under $\alpha$-th order smoothness of $U$. Thus, when the potential function $U$ satisfies high-order smoothness conditions and a high-accuracy solution is needed, this bound is favorable compared to existing $O(\sqrt{d}/\varepsilon)$ rates.

The remainder of the paper is organized as follows. Section 2 is devoted to background on the Langevin diffusion, and various higher-order variants. In Section 3, we describe the third-order Langevin scheme analyzed in this paper, and state our two main results: Theorem 1 for the special case of ridge-separable functions, and Theorem 2 for general functions under black-box gradient access with additional smoothness. Section 4 is devoted to the proofs of our main results, with more technical aspects of the arguments deferred to the appendices. We conclude with a discussion in Section 5.

Background and problem formulation

In this section, we first introduce the class of sampling problems that are our focus, before turning to specific sampling algorithms that are based on discretizations of diffusion processes. We begin with the classical first-order discretization of Langevin diffusion, and then introduce the higher-order discretization that is the principal object of study of our work.

We consider the problem of drawing samples from a distribution with density written in the form $p^{*}(\theta)\propto\exp(-U(\theta))$. The potential function $U:{}^{d}\rightarrow\real$ is assumed to be strongly convex and smooth in the following sense:

The function $U$ is differentiable, and $m$-strongly convex and $L$-smooth:

where $m$ and $L$ are positive constants.

We say that the potential is $(m,L)$-convex-smooth when this sandwich relation holds. The condition number of the problem is given by the ratio $\kappa:=\frac{L}{m}\in[1,\infty)$.

Given an iterative algorithm that generates a random vector $\theta^{(k)}$ at each step $k$, we use $\pi^{(k)}$ to denote the law of $\theta^{(k)}$. We are interested in the convergence $\pi^{(k)}$ to the measure $\pi^{*}$ defined by the target density $p^{*}$. In order to quantify closeness of the measures $\pi^{(k)}$ and $\pi^{*}$, we use the Wasserstein-$2$ distance (see the book for background). Given a pair of distributions $p$ and $q$ on d, a coupling $\gamma$ is a joint distribution over the product space ${}^{d}\times{}^{d}$ that has $p$ and $q$ as its marginal distributions. We let $\Gamma(p,q)$ denote the space of all possible couplings of $p$ and $q$. With this notation, the Wasserstein-$2$ distance is given by

Given this definition, we obtain the notion of the $\varepsilon$-mixing time—it is the minimum number of steps the algorithm takes to converge to within $\varepsilon$-close of the target measure $\pi^{*}$ in $\mathcal{W}_{2}^{2}$ distance:

2 First-order Langevin algorithm

We refer to the stochastic process represented by the following stochastic differential equation as continuous-time Langevin dynamics:

Here the reader should recall that $L$ is the smoothness parameter from Assumption 1.

It is well known that continuous-time Langevin dynamics converges to the target distribution $p^{*}$ exponentially quickly; moreover, under Assumption 1, it is known that the convergence rate is independent of the dimension $d$. However, after discretization using the Euler scheme, the resulting algorithm—a discrete-time stochastic process—has a mixing rate that scales as $\mathcal{O}(d)$ . As this result makes clear, the principal difficulty in high-dimensional sampling problems based on Langevin diffusion is the numerical error that arises from the integration of the continuous-time dynamics. A classical response to this problem is to introduce higher-order discretizations, but in the setting of Langevin diffusion a major challenge arises—the non-smoothness of the Brownian motion makes it difficult to control high-order numerical errors.

3 Underdamped (second-order) Langevin dynamics

One way to proceed is to augment the dynamics to yield smoother trajectories that are more readily discretized. For example, the second-order Langevin algorithm, also known as the underdamped Langevin algorithm, lifts the original $d$-dimensional space to a $2d$-dimensional space consisting of vectors of the form $x=(\theta,r)\in{}^{d}\times{}^{d}$, and defines the following $2d$-dimensional collection of SDEs:

where $\xi>0$ is an algorithmic parameter.

In the second-order Langevin dynamics determined by the system (6), the trajectory $\theta_{t}$ has one additional order of smoothness compared to the Brownian motion $B_{t}^{r}$. As a result, it is possible to introduce higher-order discretizations for the augmented dynamics. Examples of such discretizations include Hamiltonian Monte Carlo and underdamped Langevin algorithms , both of which are derived from equation (6). These methods can provably accelerate convergence; in particular, the underdamped Langevin algorithm provides a convergence rate of $\mathcal{O}(\sqrt{d})$ when the objective function $U$ is strongly convex and Lipschitz smooth.

4 A third-order scheme

It is natural to ask whether one can further accelerate the convergence of Langevin algorithms via higher-order dynamics, where we expand the ambient space and drive the variable of interest via higher-order integration of an SDE. In order to pursue this question, let us recall a generic recipe for constructing Markov dynamics with a desired stationary distribution. Consider the family of SDEs of the form

where ${\rm D}$ is a constant positive semidefinite matrix, and ${\rm Q}$ is a constant skew-symmetric matrix. It can be shown that for any choice of the matrices $({\rm D},{\rm Q})$ respecting these constraints, the SDE in equation (7) has $p^{*}(x)\propto\exp(-H(x))$ as its invariant distribution.

Within this general framework, note that the second-order Langevin dynamics (6) are obtained by setting $x_{t}=(\theta_{t},r_{t})$, and choosing

Note that the positive semidefinite matrix ${\rm D}$ has a zero top-left block matrix (corresponding to the $\theta$ coordinates), which means that $\theta_{t}$ is not directly coupled with the Brownian motion.

This observation motivates us to choose an even more singular ${\rm D}$ matrix. Beginning from the general equation (7), let $x_{t}=(\theta_{t},p_{t},r_{t})$, and define the function $H(x_{t})=U(\theta_{t})+\frac{L}{2}\left\|p_{t}\right\|^{2}+\frac{L}{2}\left\|r_{t}\right\|^{2}$, along with the matrices

Given these definitions, we set up a third-order form of Langevin dynamics as follows:

The trajectory of $\theta_{t}$ under these third-order dynamics is substantially smoother than the corresponding trajectory under the underdamped Langevin dynamics; this higher degree of smoothness provides more control over discretization errors. In particular, in our numerical analysis of equation (8), we exploit the fact that the Brownian motion and $\nabla U$ are passed into the time derivative of two different variables, $r_{t}$ and $p_{t}$, respectively. This allows us to adopt a splitting scheme that takes advantage of the structure of $U$ and thereby provides an improvement in convergence rate relative to past work. Indeed, in Section 3 we prove that faster convergence is achieved with a proper choice of integration scheme.

Main results

In this section, we describe our higher-order Langevin algorithm, and state two theorems that characterize its convergence rate.

We propose an algorithm, akin to the Langevin or underdamped Langevin algorithm, that at every iteration generates a normal random variable centered according to the previous iterate (see Algorithm 1). The algorithm is constructed via a three-stage discretization scheme of the continuous-time Markov dynamics (8). See Section 3.4 for a detailed discussion of the discretization scheme.

Recalling the $(m,L)$-strong-convexity-smoothness condition given in Assumption 1, we see that the potential function $U$ has a unique global minimizer $\theta^{*}\in{}^{d}$ such that $\nabla U(\theta^{*})=0$. We initialize our algorithm at a point $\theta_{0}$ satisfying $\left\|\theta^{*}-\theta_{0}\right\|_{2}\leq\frac{1}{L}$. Such a point can be found in $O(\sqrt{\kappa}\log(d\kappa))$ gradient evaluations using accelerated gradient methods . Our algorithm generates a sequence of vector triples $x^{(k)}=(\theta^{(k)},p^{(k)},r^{(k)})$ for $k=1,2,\ldots$ in a recursive manner. Any instance of the algorithm is specified by a stepsize parameter $\eta>0$, two positive auxiliary parameters $\gamma$ and $\xi$, and a function $\Delta U:{}^{d}\times{}^{d}\rightarrow\real$. We provide specific choices of these parameters and the function $\Delta U$ in the theory to follow.

Given the iterate $x^{(k)}$ at step $k$, the next iterate $x^{(k+1)}$ is obtained by drawing from a multivariate Gaussian distribution with mean $\bm{\mu}(x^{(k)})$, where

The constants $\mu_{12}$–$\mu_{33}$, as well as $\sigma_{11}$–$\sigma_{33}$ above are entirely determined by the triple $(\gamma,\xi,\eta)$; see Appendix C for their explicit definitions.

We make a few remarks about the algorithm:

The vector $\Delta U(\theta,p)$ is chosen to be either an exact or approximate value of the integral $\int_{0}^{\eta}\nabla U(\theta+tp)dt$. As we discuss in the two versions of the main theorem, different choices of $\Delta U(\theta,p)$ are available depending on the starting assumptions, and each such choice leads to a different mixing time bound.

In each iteration, we only need to compute $\Delta U(\theta,p)$ once. Below we provide choices of the function $\Delta U$ for which this step has equivalent computational cost with a gradient evaluation.

When the stepsize $\eta$ is small, the leading terms in $(\mu(x^{(k)}),\Sigma)$ are the same as the dynamics in equation (8). However, the high-order correction terms are essential for achieving accelerated rates. This high-order scheme allows us to separate $\nabla U$, the only nonlinear part of the equation, and carry out a direct integration on a deterministic path.

While our description allows for different choices of the parameters $\gamma$ and $\xi$, in our analysis, we adopt the choices $\gamma=\kappa$ and $\xi=2\kappa$.

2 Guarantees for ridge-separable potentials

We begin by describing and analyzing a version of our algorithm applicable when the potential function $U$ is of ridge-separable form (1). In this case, the integral $\int_{0}^{\eta}\nabla U(\theta+tp)dt$ can be computed exactly using the Newton-Leibniz formula. This fact allows us to run Algorithm 1 with the choice

We claim that an $O(d^{\frac{1}{3}})$ mixing time can be achieved in this way. More precisely, we have:

Let $U$ be an $(m,L)$-convex-smooth potential of the form (1). Given a desired Wasserstein accuracy $\varepsilon\in(0,1)$, suppose that we run Algorithm 1 with stepsize $\eta=c\kappa^{-\frac{11}{4}}d^{-\frac{1}{4}}L^{\frac{1}{4}}\varepsilon^{\frac{1}{2}}$, using the function $\Delta U(\theta,p)$ defined in equation (10). Then there is a universal constant $C$ such that the mixing time is bounded as

Note that the result holds true for any potential function of the form equation (1), regardless of the distribution of the data points. In particular, we do not assume any form of incoherence assumptions as in ; nor do we assume any conditions on the norm of vectors $a_{i}$. Furthermore, only the strong convexity and smoothness assumptions are used, without requiring high-order smoothness assumptions. Many log-concave sampling problems of practical interest in statistical applications arise from a Gibbs measure defined by generalized linear potential functions. Under this setup, our result significantly improves the previous best known $O(d^{1/3}/\varepsilon^{2/3})$ rate in the dependency on both $\varepsilon$ and $d$. Finally, it is also worth noticing that the ridge-separable form is needed only to ensure the close-form expression (10). For a function that does not satisfy Eq (1) but allows the close-form integration of $\Delta U(\theta,p)$, Theorem 1 also applies.

As a caveat, we note that the stepsize used in running Algorithm 1 depends on knowledge of $\kappa$ and $L$, which might not be unavailable in practice. An important direction for future work is to provide an automated mechanism for stepsize selection with similar guarantees.

3 Guarantees under black-box gradient access

We now turn to the more general setting, in which the potential function $U$ is no longer ridge separable (1). Suppose moreover that we have access to $U$ only via a black-box gradient oracle, meaning that we can compute the gradient $\nabla U(\theta)$ at any point $\theta$ of our choice. Under these assumptions, the closed form integrator described in Algorithm 1 is no longer available. However, by using Lagrange interpolation as an approximation, we can still derive a practical high-order algorithm that yields a faster mixing time. In particular, while the mixing time scales as $O(\sqrt{d})$, as with lower-order methods, we show that the $\varepsilon$-dependency term can be adaptive to the degree of smoothness of the function $U$.

When the objective $U$ does not take the form of a generalized linear function, we use Lagrange interpolating polynomials with Chebyshev nodes to approximate the function:

The Chebyshev polynomial interpolation operator $\Phi$ takes as inputs a scalar $t$, and a function $z:[0,\eta]\rightarrow{}^{d}$, and returns the scalar $\Phi(t;z):=\sum_{i=1}^{\alpha}z(s_{i})\prod_{j\neq i}\frac{t-s_{i}}{s_{j}-s_{i}}$. Note that the integral of this function over $t$ can be computed in closed form.

For each pair $(\theta,p)$ define the mapping $\psi_{\theta,p}(s)=\nabla U(\theta+sp)$ from to d. Applying the interpolation polynomial to this mapping, we define

Note that $t\mapsto\Phi(t;\psi_{\theta,p})$ is a polynomial function, and hence the integral over $t$ can be computed exactly. Computing this integral requires $\alpha$ gradient evaluations in total, along with $O(d)$ additional computational cost. Thus, when the smoothness $\alpha$ is viewed as a constant, the computational complexity is order-equivalent to a gradient evaluation.

For some $\alpha\geq 3$, the potential function $U$ is $\alpha$-th order differentiable, and the associated tensor of derivatives satisfies the bound

Note that in the special case $\alpha=3$, Assumption 2 corresponds to a Lipschitz condition on the Hessian function, as has been used in past analysis of sampling algorithms.

In general, under a smoothness assumption of order $\alpha\geq 2$, we have the following guarantee:

Consider a potential $U$ satisfying Assumptions 1 and and 2 for some $\alpha\geq 2$. Given a desired Wasserstein accuracy $\varepsilon\in(0,1)$, suppose that we run Algorithm 1 with stepsize

We observe that the dimension dependence becomes $d^{1/2}$, but the corresponding $\epsilon$ dependence is reduced to $\epsilon^{-1/(\alpha-1)}$, where higher-order smoothness leads to better accuracy dependence.

4 Derivation of the discretization

In this section we provide a detailed derivation of Algorithm 1 as a discretization scheme for the continuous-time process (8). Our overall approach involves a combination of a splitting method and a high-order integration scheme. More precisely, we reduce the problem of one-step simulation of equation (8) to (approximately) computing the integral of $\nabla U$ along a straight line, using a three-stage discretization scheme. Specifically, letting $\hat{g}_{s}$ be an approximation for $\nabla U\left(\theta^{(k)}+\frac{s-k\eta}{\eta}p^{(k)}\right)$ and letting $\Delta U(\theta^{(k)},p^{(k)}):=\int_{k\eta}^{(k+1)\eta}\hat{g}_{s}ds$, the discretization error bound depends on the accuracy with which $\hat{g}_{s}$ approximates $\nabla U$. Depending on the assumptions imposed on the target distribution, various choices of $\hat{g}$ can be used, the exact integration of which leads to different choices of $\Delta U$. Theorems 1 and 2 correspond to two instances of this general approach.

We begin by constructing estimators $\hat{\theta}_{(k+1)\eta}$ and $\hat{r}_{(k+1)\eta}$ following the Ornstein-Uhlenbeck process:

The values $(\hat{\theta}_{(k+1)\eta},\hat{r}_{(k+1)\eta})$ are then used to calculate a high-accuracy result by adding a correction term. We use the function $\hat{g}_{s}$ as an approximation of the gradient

It is worth noting that $\hat{g}_{t}$ approximates a function along a fixed curve determined by $x_{k\eta}$, and has no interaction with other variables nor the noise. This makes it possible to obtain high-accuracy solutions to the equation by integration of a deterministic and known function.

In the second stage, we solve the system of differential equations

Note that the Brownian motion $(B_{t}^{r}:t\geq 0)$ used in process (15) must be the same as that used in process (13), so that the two processes must be solved jointly. As shown in Appendix C, we can carry out the integrations in closed form, so as to obtain the explicit quantities required to implement Algorithm 1.

Proofs

In this section, we provide the proofs of our main results. We begin in Section 4.1 by stating and proving a result (Proposition 1) on the exponential convergence of the third-order dynamics in continuous time. Section 4.2 is devoted to our proofs of Theorems 1 and Theorem 2 on the behavior of the discrete-time algorithm. In all cases, we defer the proofs of more technical results to the appendices.

We begin by studying the process $\{x_{t}\}_{t\geq 0}$ defined by the continuous-time third-order dynamics in equation (8), with the particular goal of showing convergence in the Wasserstein-$2$ distance. In all of our analysis—in this section as well as others—we make the choices $\gamma=\kappa$ and $\xi=2\kappa$ in defining the dynamics.

It is known that the limiting stationary distribution of the process $x_{t}=(\theta_{t},p_{t})$ has a product form:

Our goal is to show that the distribution of $x_{t}=(\theta_{t},p_{t})$ converges at an exponential rate in the Wasserstein-$2$ distance, as previously defined in equation (3), to this expanded target distribution.

In order to do so, we consider two processes following the third-order dynamics (8), where the process $\{x_{t}\}_{t\geq 0}$ and $\{x^{*}_{t}\}_{t\geq 0}$ are started, respectively, from the initial distributions $p_{0}$ and $p^{*}$. We then couple these two processes in a synchronous coupling. In order to establish a convergence rate, we make use of the following Lyapunov function:

With this setup, our main result on the continuous-time dynamics is the following:

Let $x_{0}$ and $x^{*}_{0}$ follow the laws of $p_{0}$ and $p^{*}$, respectively. Then the process $\{x_{t}\}_{t\geq 0}$ defined by the dynamics (8) satisfies the bound

As shown in Lemma 3, to be stated in the next section, the eigenvalues of $S$ lie in the interval $[1/(5\kappa),\;\kappa^{2}+10]$. Thus, Proposition 1 implies convergence in the Wasserstein-$2$ distance at an exponential rate.

The remainder of this section is devoted to the proof of Proposition 1. The first step in the proof involves establishing a differential inequality for the Lyapunov function $\mathcal{L}_{t}$ via the coupling technique.

The proof of Lemma 1, given in Section 4.1.1, is based on the synchronous coupling technique, in which two processes are coupled based on the same underlying Brownian motion.

Taking Lemma 1 as given, we can now complete the proof of Proposition 1. Applying Grönwall’s lemma to equation (17) yields

which establishes the bound in Proposition 1.

We prove Lemma 1 by choosing a synchronous coupling $\bar{\zeta}\in\Gamma\left(p_{t}(x_{t}|x_{0}),p^{*}_{t}(x^{*}_{t}|x^{*}_{0})\right)$ for the laws of $x_{t}$ and $x^{*}_{t}$. (A synchronous coupling simply means that we use the same Brownian motion $B_{t}^{r}$ in defining both $x_{t}$ and $x^{*}_{t}$.) We then obtain that for any pair $(x_{t},x^{*}_{t})\sim\bar{\zeta}$,

Consequently, the derivative of the function $t\mapsto(x_{t}-x^{*}_{t})^{\rm T}S(x_{t}-x^{*}_{t})$ is given by

In order to proceed, we need to relate the eigenvalues of the matrix $SM_{t}+M_{t}^{\rm T}S$ to those of $S$. The following lemmas allow us to carry out this conversion:

Using these two lemmas, it follows that for any pair of random variables $(x_{t},x^{*}_{t})\sim\bar{\zeta}$, we have

where inequality (i) follows from Lemma 2 and inequality (ii) follows from the upper bound on the eigenvalues of $S$ in Lemma 3. This completes the proof of Lemma 1.

Finally, we turn to the proofs of Lemmas 2 and 3.

1.2 Proof of Lemma 2

In order to simplify notation, we first define $\{l_{k}\ |\ k=1,\cdots d\}$ as the eigenvalues of $H_{t}/L$. Since $H_{t}$ is the Hessian of the potential function $U$ (which is $m$-strongly convex and $L$-Lipschitz smooth), its eigenvalues are always bounded above and below: $1/\kappa\leq l_{k}\leq 1$.

Next we calculate the eigenvalues $\{\lambda_{k}\}$ of $SM_{t}+M_{t}^{\rm T}S$: $\lambda_{0}=-1$ and

where we define $g_{1}(\kappa)$ in Appendix B.1 to simplify the notation.

We first upper bound $\lambda_{k}^{+}$. Since $(l_{k}-1/\kappa)\sqrt{g_{1}(\kappa)}\geq 0$, we find that

Turning to the bound on $\lambda_{k}^{-}$, we can rewrite it as:

Since $\left(2-\sqrt{g_{1}(\kappa)}\right)\leq 0$ (proven in Lemma 4) and $l_{k}\leq 1$, we can upper bound $\lambda_{k}^{-}$ with an expression independent of $l_{k}$:

The function $g_{1}$ given in equation (33a) has the following properties:

1.3 Proof of Lemma 3

where the coefficients $g_{2}(\kappa)$ and $g_{3}(\kappa)$ are defined in Appendix B.1. Since $S$ is a symmetric matrix, all the roots of equation (27) are real.

For any $\kappa\geq 1$, $f(x)$ defined in equation (27) satisfy that $\forall x\leq\frac{4}{5\kappa^{6}}$, $f(x)<0$, and that $\forall x\geq\frac{4}{\kappa^{3}}+\frac{40}{\kappa^{5}}$, $f(x)>0$.

We defer the proof of this lemma to Appendix B.3. Note that Lemma 5 implies that all real roots of equation (27) lie in the range of $\left[\frac{4}{5\kappa^{6}},\frac{4}{\kappa^{3}}+\frac{40}{\kappa^{5}}\right]$, which completes the proof.

2 Proofs of discrete-time results

We now turn to the proofs of our two main results—namely, Theorems 1 and Theorem 2—that establish the behavior of the discrete-time algorithm. We begin with a general roadmap for the proofs, along with a key auxiliary result (Proposition 2) common to both arguments.

In the last step we transform the $\left\|\cdot\right\|_{2}$ norm into the $\left\|\cdot\right\|_{S}$ norm, which is possible because the process $(x_{t}:t\geq 0)$ is contractive under the $\left\|\cdot\right\|_{S}$ norm.

To analyze the one-step discretization error, we have the following key lemma, which holds in general for any approximator of $\nabla U$. Note that Proposition 2 is used in the proof of both theorems.

Comparing the constructions in equations (13)–(15), we note that:

2.2 Proof of Theorem 1

Note that by the $L$-smoothness condition (Assumption 1), we have

The moments can further be controlled as:

Therefore, with $\eta<c\kappa^{-11/4}$, we have $2C\kappa^{6}\eta^{5}|\!|\!|S^{-1}|\!|\!|_{{op}}\leq\frac{\eta}{40\kappa^{2}+400}$, and consequently:

Choosing the parameters accordingly completes the proof.

2.3 Proof of Theorem 2

Now we describe the proof of Theorem 2, a general result for functions with high-order smoothness. The proof is also based on synchronous coupling. In addition to exploiting Proposition 2, we also use the following standard result for Chebyshev node interpolation :

For Lagrange interpolating polynomials, the time derivative is the finite difference between interpolation points. These differences can be further bounded by the time derivative of the original process $\nabla U(\theta+tp)$—viz.

Similar to the proof of Theorem 1, since the weights in Lagrangian interpolation at Chebyshev nodes are non-negative, using Proposition 2, we obtain the bound:

Discussion

In this paper, we focus on accelerating the convergence of gradient-based MCMC algorithms in high-dimensional spaces. We break the problem into two parts: a splitting scheme that reduces the problem of SDE discretization to that of integration along a fixed straight line; a third-order Langevin dynamics in continuous time which allows this fine-grained discretization analysis while satisfying exponentially fast convergence.

For the second problem, we construct a third-order Langevin dynamics so that the trajectories are smoother and the integration of $\nabla U$ is separated from the Brownian motion part. We then apply utilize this dynamics to design MCMC algorithms adaptive to underlying structures of the problem. A mixing time of order $\mathcal{O}\left(d^{1/4}/\varepsilon^{1/2}\right)$ is achieved for ridge-separable potentials, which cover a large class of machine learning models. Under black-box oracle model, a rate of order $O(d^{1/4}/\varepsilon^{1/2}+d^{1/2}/\varepsilon^{1/(\alpha-1)})$ is achieved for $k$-th order smooth objective functions $U$.

An important future direction is to further investigate higher-order splitting scheme with the use of higher-order dyanmics, to further reduce the dimension dependency for the mixing time of MCMC. We conjecture that the exponent on $d$ can be further reduced, with a trade-off between the dependency on the dimension and condition number.

This work was partially supported by Office of Naval Research Grant ONR-N00014-18-1-2640, Army Research Office grant W911NF-17-1-0304, and National Science Foundation Grants NSF-IIS-1740855, NSF-IIS-1909365, and NSF-IIS-1619362. We thank Santosh Vempala and Chris Junchi Li for helpful discussions.

References

Appendix A Proof of Proposition 2

In this appendix, we prove Proposition 2, as previously stated in Section 4.2.1. Recall that this result provides a bound on the discretization error with certain choice of $\hat{g}_{t}$ used in equation (14). Our proof of this bound is based on direct coupling estimates.

Comparing the two processes along the path, we obtain:

Introducing the function $J(u)=u^{T}Su$, for the one-step analysis with $t\in[k\eta,(k+1)\eta]$, we have:

The remainder of the proof is devoted to bounding the terms $\{I_{j}\}_{j=1}^{3}$.

A.2 Some auxiliary results

The straight-line approximation error of the curve is uniformly bounded as

Our second auxiliary lemma relates the squared Euclidean norm of the interpolation process (cf. equations (13)–(15)) with the squared norm at discrete time steps.

Our third auxiliary lemma upper bounds the higher order moments of the stochastic process generated by Algorithm 1. These bounds are useful for controlling certain higher order derivatives along the path.

We return to prove all of these claims in Section A.5.

A.3 Bounding the three terms

Taking the auxiliary lemmas as given for now, we now bound each of the terms $\{I_{j}\}_{j=1}^{3}$ in succession.

Therefore, applying Cauchy-Schwartz to the integral, we arrive at:

The second term is easy to control by the properties of OU processes:

The upper bound involve the expected supremum of squared gradient norm along the path of $\hat{\theta}_{s}$, which is already obtained in Section A.3.1:

Putting them together, for $t\in[k\eta,(k+1)\eta]$ we have:

A.4 Obtaining the final bound

Plugging the moment upper bounds in equation (32) to the estimates for $I_{1}$, $I_{2}$ and $I_{3}$, for $\eta<\min\left(1/\gamma,1/\xi\right)$, we obtain:

A.5 Proof of auxiliary lemmas

We now return to prove the auxiliary results that were stated and used in the previous sections—namely, Lemmas 7, 8 and 9.

A.5.2 Proof of Lemma 8

The two terms appearing in the above upper bound are both easy to control:

Putting them together, with $\eta<\min\left(1/\gamma,1/\xi\right)$, we have:

A.5.3 Proof of Lemma 9

Since $\hat{g}$ belongs to the convex hull of the curve $\nabla U(\theta+tp)$, as assumed in the statement of the lemma, we have:

and using the smoothness of $U$, we can easily see that

Collecting the main terms and bound the rest of terms directly using the norm of $x^{(k)}$, we obtain:

By Appendix C, it is easy to see that $\left\|\zeta_{k}\right\|_{2}\leq C\eta^{2}\kappa^{2}\left\|x^{(k)}\right\|_{2}$.

For $\eta<c\kappa^{-7}$ with some universal constant $c>0$, we have:

Now we turn to deal with the stochastic part. By our construction, it is easy to verify that $|\!|\!|\Sigma|\!|\!|_{{op}}\leq C\kappa^{2}\eta$ for some universal constant $C>0$.

Letting $\xi_{k}\sim\mathcal{N}(0,\frac{1}{L}I)$ be independent from $x^{(k)}$, we have:

Noting that $\left\|x^{(0)}\right\|_{2}\leq\frac{1}{L}$, by solving the recursion inequalities, we obtain:

Appendix B Auxiliary results for Lemmas 2 and 3

This appendix is dedicated to proofs of two auxiliary results—namely, Lemmas 4 and 5—that underlie the proofs of Lemmas 2 and 3.

The functions $g_{1}$, $g_{2}$ and $g_{3}$ are given by:

B.2 Proof of Lemma 4

By inspecting the definition of $g_{1}$ in equation (33a), we see that $\forall\kappa\geq 1$, $g_{1}(\kappa)$ always lies in the interval $[4,+\infty)$, which implies the inequality $2-\sqrt{g_{1}(\kappa)}\leq 0$.

As for the second inequality, we first group the $\sqrt{g_{1}(\kappa)}$ term and rewrite inequality (26b) into the following equivalent form:

Since both the left and right sides are non-negative, we can square both sides, thereby obtaining

Expanding the left hand side, we obtain the equivalent form of equation (26b):

Since $\kappa\geq 1$, we can divide by $\kappa$ on both sides and obtain an equivalent inequality with a polynomial function of $1/\kappa$:

We can rewrite $g_{6}(1/\kappa)$ as a polynomial of $\left(\frac{1}{\kappa}-\frac{1}{2}\right)$:

For all $\kappa$ such that $1/\kappa\in(0,1]$, we have

which is strictly positive. This completes the proof of Lemma 4.

B.3 Proof of Lemma 5

In order to prove the bounds in this lemma, we first examine the monotonicity properties of the cubic function

For any $\kappa\geq 1$, the function $f$ from equation (34) has the following properties:

It is monotonically increasing over the interval $(\frac{4}{\kappa^{3}}+\frac{40}{\kappa^{5}},\infty)$.

It is monotonically increasing over the interval $(-\infty,\frac{4}{5\kappa^{6}})$.

Therefore, $\max\limits_{x\leq\frac{4}{5\kappa^{6}}}f(x)=f\left(\frac{4}{5\kappa^{6}}\right)$ and $\min\limits_{x\geq\frac{4}{\kappa^{3}}+\frac{40}{\kappa^{5}}}f(x)=f\left(\frac{4}{\kappa^{3}}+\frac{40}{\kappa^{5}}\right)$. Then we simply need to prove that $f\left(\frac{4}{5\kappa^{6}}\right)<0$ and $f\left(\frac{4}{\kappa^{3}}+\frac{40}{\kappa^{5}}\right)>0$ to obtain the result.

using the fact that $1/\kappa\in(0,1]$. Moreover, we also have

which is negative. Therefore, we conclude that for any $\kappa>1$, the cubic function $f$ satisfies the inequalities

We divide our proof into separate parts, corresponding to claims (a) and (b) in the lemma statement. For both parts, we establish monotonicity of the function $f$ from equation (34) by studying its derivative

By inspection, for large enough $x$, the quadratic function $f^{\prime}$ is positive, and hence the function $f$ is monotonically increasing in this range. Concretely, we claim that $f^{\prime}$ remains positive for all $x>\frac{4}{\kappa^{3}}+\frac{40}{\kappa^{5}}$.

Our strategy is to compute the solutions $x^{*}_{\pm}$ to the quadratic equation $f^{\prime}(x)=0$, and prove that the larger one $x^{*}_{+}$ satisfies the lower bound

In detail, the two solutions to the quadratic equation $f^{\prime}(x)=0$ are given by the pair $x^{*}_{\pm}=\frac{1}{3}\left(g_{4}(\kappa)\pm\sqrt{g_{5}(\kappa)}\right)$, where we define

From the fact that $\tfrac{1}{\kappa}\in(0,1]$, it follows that

Combining equations (38a) and (38b) yields the bound (36), and hence completes the proof of part (a).

Recall the solutions to the quadratic equation $f^{\prime}(x)=0$ that we computed in the previous section. For this part, it suffices to show that the smaller solution $x^{*}_{-}=\frac{1}{3}\left(g_{4}(\kappa)-\sqrt{g_{5}(\kappa)}\right)$ satisfies the bound

Equation (38b) states that $\sqrt{g_{5}(\kappa)}\leq\tfrac{4}{\kappa^{3}}+\tfrac{12}{\kappa^{5}}$, and hence for any $1/\kappa\in(0,1]$,

Since $g_{4}(\kappa)+\sqrt{g_{5}(\kappa)}$ is positive for $\kappa\geq 1$, we can substitute the bound (41) into equation (40), thereby finding that

which completes the proof of the bound (39).

Appendix C Details of Algorithm 1

This section is devoted to explicit definition of all constants involved in Algorithm 1, as well as a derivation of how the algorithm’s updates follows from integrating equations (13)–(15). We begin by defining the constants precisely:

Given these definitions, we now demonstrate how to obtain Algorithm 1 via integrating equations (13) through to (15).

The first step involving the Ornstein-Uhlenbeck process can be explicitly solved as (let $B_{0}=0$):

If function $U$ takes the form of $U(\theta)=\sum_{i}u_{i}(y_{i}^{\rm T}\theta)$, then we directly take $\hat{g}_{s}=\nabla U(\hat{\theta}_{s})$ and calculate the integral:

Using the Newton-Leibniz formula, we obtain that

Therefore, we obtain $\hat{p}_{t}$ from explicit integration:

Putting together the pieces, we find that

Summing the terms together, we obtain that