Is There an Analog of Nesterov Acceleration for MCMC?

Introduction

While optimization methodology has provided much of the underlying algorithmic machinery that has driven the theory and practice of machine learning in recent years, sampling-based methodology, in particular Markov chain Monte Carlo (MCMC), remains of critical importance, given its role in linking algorithms to statistical inference and, in particular, its ability to provide notions of confidence that are lacking in optimization-based methodology. However, the classical theory of MCMC is largely asymptotic and the theory has not developed as rapidly in recent years as the theory of optimization.

Recently, however, a literature has emerged that derives nonasymptotic rates for MCMC algorithms [see, e.g., 9, 12, 10, 8, 6, 14, 27, 28, 2, 5]. This work has explicitly aimed at making use of ideas from optimization; in particular, whereas the classical literature on MCMC focused on reversible Markov chains, the recent literature has focused on non-reversible stochastic processes that are built on gradients [see, e.g., 24, 26, 3, 1]. In particular, the gradient-based Langevin algorithm has been shown to be a form of gradient descent on the space of probabilities [see, e.g., 19, 44].

What has not yet emerged is an analog of acceleration. Recall that the notion of acceleration has played a key role in gradient-based optimization methods . In particular, Nesterov’s accelerated gradient descent (AGD) method, an instance of the general family of “momentum methods,” provably achieves a faster convergence rate than gradient descent (GD) in a variety of settings . Moreover, it achieves the optimal convergence rate under an oracle model of optimization complexity in the convex setting .

This motivates us to ask: Is there an analog of Nesterov acceleration for gradient-based MCMC algorithms? And does it provably accelerate the convergence rate of these algorithms?

This paper answers these questions in the affirmative by showing that an underdamped form of the Langevin algorithm performs accelerated gradient descent. Critically, our work is based on the use of Kullback-Leibler (KL) divergence as the metric. We build on previous work that has studied the underdamped Langevin algorithm and has used coupling methods to establish convergence of the algorithm in the Wasserstein distance [see, e.g., 8, 7, 11]. Our work establishes a direct linkage between the underdamped Langevin algorithm and Nesterov acceleration by working directly in the objective functional, the KL divergence. Combining ideas from optimization theory and diffusion processes, we construct a Lyapunov functional that couples the convergence in the momentum and the original variables. We then prove the overall convergence rate by leveraging the hypocoercivity structure of the underdamped Langevin algorithm . For target distributions satisfying a log-Sobolev inequality, we find that the underdamped Langevin algorithm accelerates the convergence rate of the classical Langevin algorithm from $d/\epsilon$ to $\sqrt{d/\epsilon}$ in terms of KL divergence (See Theorem 1 for formal statement).

Preliminaries

We start by laying out the problem setting, including our assumptions on the target distribution that we sample from, properties of the KL divergence with respect to other measure of differences between probability distributions, and the notion of gradient on the space of probabilities.

We use this KL divergence as an objective functional in an optimization-theoretic formulation of convergence to $\mathbf{p}^{*}(\theta)$.

We assume that $\mathbf{p}^{*}$ satisfies the following conditions.

As a concrete example, these assumptions are satisfied in the “locally nonconvex” case studied by , with nonconvex region of radius $R$ and strong convexity $m$; see also Assumption (a)–(c) in Appendix A. Note that instantiates both the log-Sobolev constant $\rho$ and the normalization constants $C_{N}$ in terms of the smoothness and conditioning of $U$, showing that $\rho\geq\frac{m}{2}e^{-16L_{G}R^{2}}$. Here we additionally establish (see Fact 1) that $C_{N}\leq\frac{1}{2}\ln\frac{4\boldsymbol{\pi}}{m}$, and $C_{M}\leq 32\frac{L_{G}^{2}}{m^{2}}L_{G}R^{2}$.

2 KL divergence and relation to other metrics

By Pinsker’s inequality, we can upper bound the total variation distance by the KL divergence:

Since $\mathbf{p}^{*}$ satisfies the log-Sobolev inequality (A1) with constant $\rho$ and has a Lipschitz smoothness property, by the Talagrand inequality (Theorem 1 of ), we can upper bound the Wasserstein-$2$ distance (defined in Eq. (2)) by the KL divergence:

3 Gradients on the space of probabilities

See [23, Definition 10.1.1] for more details. This strong subdifferential provides us the proper notion of “gradient.” In particular, for functionals with enough regularity, the strong subdifferential of $\mathcal{L}$ taken at $\mathbf{p}$ can be expressed as $\nabla_{\theta}\frac{\delta\mathcal{L}}{\delta\mathbf{p}}$, where $\frac{\delta}{\delta\mathbf{p}}$ is the functional derivative taken at $\mathbf{p}$ and $\nabla_{\theta}$ is the ordinary gradient operator in the space of $\theta$ [23, Lemma 10.4.1].

Underdamped Langevin Algorithm as Accelerated Gradient Descent

where $B_{t}$ is a standard Brownian motion. The evolution of the probability density function $\mathbf{p}_{t}$ of the random variable $\theta_{t}$ follows the transport of probability mass along a vector flow $v_{t}$ in the state space:

where the vector flow can be calculated as: $v_{t}(\theta)=b(\theta)-\nabla\ln\mathbf{p}_{t}(\theta)$. This can be compared with the following Liouville equation:

On the other hand, we formulate the “gradient” of the KL divergence corresponding to the vector flow point of view. For the objective functional $\mathcal{F}[\mathbf{p}_{t}]$, its time change when $\theta_{t}$ follows Eq. (3) is:

or, equivalently, $b^{GD}(\theta)=-\nabla U(\theta)$ in Eq. (3).

Along this gradient descent flow, $v_{t}^{GD}$, the time evolution of the KL divergence is

If $\mathbf{p}^{*}(\theta)$ satisfies Assumption A1 then taking $g=\frac{\mathbf{p}_{t}}{\mathbf{p}^{*}}$ in the log-Sobolev inequality yields:

Note the resemblance of this bound to the Polyak-Łojasiewicz condition used in optimization theory for studying the convergence of gradient methods—in both cases the difference in objective value from the current iterate to the optimum is upper bounded by the squared norm of the gradient of the objective. With the log-Sobolev inequality, we obtain that

2 Accelerated gradient descent in KL divergence: A continuous perspective

The corresponding continuity equation defined by this vector field is

This implies that the vector field can be implemented via the following stochastic differential equation

which is the underdamped Langevin dynamics .

This only demonstrates the contractive property in the $r$ coordinates (note that the gradient is only in $r$ in Line (18)) and does not directly provide a linear convergence rate over time. To quantify the convergence rate for this accelerated gradient descent dynamics with respect to the KL divergence objective, we need to couple the convergence in $\theta$ coordinates to that in $r$. To this end, we follow recent work in the optimization literature and design a Lyapunov functional which makes use of a quadratic form of the gradient of the distance $\mathcal{D}$ between the current iteration $\mathbf{p}_{t}$ and the stationary solution $\mathbf{p}^{*}$:

Interestingly, similar forms appear in the analyses of both accelerated gradient descent dynamics and hypocoercive diffusion operators .

We then make use of this Lyapunov functional to obtain a linear convergence rate for the accelerated gradient descent dynamics with respect to the KL divergence.

Under Assumptions A1–A3, the time evolution of the Lyapunov functional $\mathcal{L}$ with respect to the continuous time vector flow $v_{t}^{AGD}$ in Eq. (13) with $\gamma=2$ and $\xi=2L_{G}$ is upper bounded as:

This establishes linear convergence of the continuous process with a rate of $\frac{\rho}{10}$.

2.2 Accelerated gradient descent dynamics for optimization

It is worth noting that the derivation in the previous subsection has a close correspondence to recent analyses of the accelerated gradient descent dynamics in convex optimization . Indeed, when optimizing a strongly convex function $U(\theta)$ on a Euclidean space with the accelerated gradient descent dynamics, the continuous limit of the algorithm is expressed as an ordinary differential equation :

We also extend the original objective function $U(\theta)$ to $H(x)=U(\theta)+\frac{\xi}{2}\|r\|_{2}^{2}$ to capture the overall dynamical behavior in the space of $x$. With the definition of this extended objective function $H$, we can simplify the expression of the dynamics:

To quantify convergence for the strongly convex objective $U$, considers a Lyapunov function of the form $l(x)=H(\theta)+\left<\nabla_{x}^{T}D_{h}(x),S\nabla_{x}D_{h}(x)\right>$, where $D_{h}(x)=\frac{1}{2}\left\|\theta-\theta^{*}\right\|^{2}+\frac{1}{2}\left\|r\right\|^{2}$ is the squared distance from $(\theta,r)$ to the optimum of $H$, $(\theta^{*},0)$.

Comparing the dynamics of Eq. (23) versus Eq. (10) and the convergence analyses for them, we observe that the underdamped Langevin diffusion defined in Eq. (14) is precisely accelerated gradient descent with respect to the KL divergence.

3 Underdamped Langevin via second-order discretization

While the continuous-time perspective yields insight into the convergence rates achievable by acceleration, for these insights to apply to discrete-time algorithms it is necessary to understand the effects of discretization. In optimization, an emerging literature has begun to show how to design discretization procedures that retain accelerated rates from continuous time . The literature in MCMC has not yet formalized lower bounds on convergence rates that allow characterizations of acceleration, in either continuous time or discrete time, but there are results that exhibit the importance of discretization for convergence. In particular, higher order (and more accurate) discretization schemes are found to accelerate convergence .

In this section we show how to design a discretization for the an underdamped Langevin algorithm that yields accelerated rates. Following , we discretize the time dimension underlying Eq. (14) into intervals of equal length $h$ (at the end of the $k$-th iteration, we have $t=kh$). Then in the $(k+1)$-th step, we define a continuous dynamics in the interval of $\tau\in[kh,(k+1)h]$ by conditioning on the initial value of $x_{kh}$:

In Appendix B we derive explicit formulas for $x_{\tau}$ given $x_{kh}$. These are used to generate the $(k+1)$-th iterate. In particular, define the hyperparameters $\gamma=2$, $\xi=2L_{G}$, and set the step size as follows:

where $\widetilde{C_{N}}=C_{N}+\frac{1}{2}\ln\frac{L_{G}}{2\pi}$. The discretized vector field is

This leads to a high-order discretization scheme that is defined explicitly in Appendix B and summarized in Algorithm 1.

By way of comparison, the Euler-Maruyama discretization scheme corresponds to:

After integration, we obtain that for $\tau\in[kh,(k+1)h]$:

There are other higher-order discretization schemes that can be considered in addition to our scheme in Eq. (30). In particular, note that $v_{t}^{AGD}$ decomposes into two parts:

Convergence of the Underdamped Langevin Algorithm

From Fig. 1, we see that the underdamped Langevin algorithm, Eq. (65), seems to have a similar profile to accelerated gradient descent; it uses oscillatory behavior to increase the convergence rate. In this section, we rigorously establish acceleration, by proving that the convergence of the underdamped Langevin algorithm is of order $\mathcal{O}\left(\sqrt{{d/\epsilon}}\right)$ in terms of KL divergence.

Let the KL divergence from $\mathbf{p}_{t}(\theta)$ to $\mathbf{p}^{*}(\theta)$ be the target functional to minimize:

If we further assume that the function $U$ is locally nonconvex with radius $R$ and has global strong convexity $m$ (Assumption (a)–(c)), we obtain an explicit dependence of the convergence time $K$ on other constants:

where $\rho=\min\left\{\frac{m}{2}e^{-16L_{G}R^{2}},1\right\}$.

We devote the remainder of Section 4 to the proof of Theorem 1. As advertised, the proof decomposes into a continuous-time analysis and a discretization analysis. We first establish the convergence rate of the continuous underdamped Langevin dynamics in Proposition 1 to quantify the instantaneous contraction provided by the dynamics. We then study the discretization error of the underdamped Langevin algorithm in each step. Combining these two results and integrating over the time steps leads us to the final conclusion.

We begin by formulating the instantaneous change of the probability density $\mathbf{p}(x_{\tau})$ within each step of the underdamped Langevin algorithm. The time evolution of $\mathbf{p}(x_{\tau}|x_{kh})$ following the discretized vector flow $\hat{v}_{\tau}^{AGD}$ for $\tau\in[{kh},(k+1)h]$ is as follows:

We have thus separated the time evolution of $\mathbf{p}(x_{\tau})$ into two parts: the continuous component and the discretization error component.

We now analyze term (43a) and term (43b) separately, returning later to combine the analyses and obtain the overall convergence rate.

We use Lemma 7 in the Appendix to expand term (43a) and quantify the convergence of $\mathcal{L}$ with respect to the continuous vector flow $v_{\tau}^{AGD}$:

where $M_{C}$ is defined in Eq. (82). The two terms on the right-hand side of Eq. (44) are both less than or equal to zero. We will use the first term to cancel similar terms in the discretization error and use the second term to drive the convergence of the process (by way of the log-Sobolev inequality).

For term (43b) capturing the discretization error, we provide an upper bound in the following proposition.

Under Assumption A2, when $\tau-kh\leq\frac{1}{8L_{G}}$, $\gamma=2$, and $\xi=2L_{G}$, term (43b) is upper bounded as:

Roughly speaking, Proposition 2 upper bounds the instantaneous contribution of the discretization error by the terms appearing in Eq. (44) (the contraction of the continuous process), the variance of $\theta_{\tau}-\theta_{kh}$ (the progress of $\theta$ within one step), and constant terms that depend on the step size. After combining Proposition 2 with Proposition 1, the only nonnegative terms that remain are the variance of $\theta_{\tau}-\theta_{kh}$ and other constant terms.

We devote the rest of this subsection to the proof of Proposition 2. We first expand term (43b) using the definitions of the functional $\mathcal{L}$ as well as the discrete and continuous vector flows $\hat{v}_{\tau}^{AGD}$ and $v_{\tau}^{AGD}$.

For $\tau-kh\leq\frac{1}{8L_{G}}$, the time evolution of the Lyapunov functional $\mathcal{L}$ with respect to the discretization error $\hat{v}_{\tau}^{AGD}-v_{\tau}^{AGD}$ is:

It can be observed that of the three terms (45a)–(45c) in Lemma 3, there are two types of term: Terms (45a) and (45b) only involve first-order derivatives, $\nabla_{\#}\ln\frac{\mathbf{p}_{\tau}(x_{\tau})}{\mathbf{p}^{*}(x_{\tau})}$ (for $\#$ labeling $\theta$ or $r$); while term (45c) involves a second-order derivative, $\nabla_{x}\nabla_{r}\ln\frac{\mathbf{p}_{\tau}(x_{\tau})}{\mathbf{p}^{*}(x_{\tau})}$.

For terms (45a) and (45b), we make use of Young’s inequality to obtain upper bounds:

The main difficulty is in bounding term (45c), which is the object of the following lemma.

Under Assumption A2, we provide an explicit bound for term (45c). When $\tau-kh\leq\frac{1}{8L_{G}}$, $\gamma=2$, and $\xi=2L_{G}$,

Applying Eq. (46a)–(46b) and Lemma 4 to Eq. (45a)–(45c), we bound the overall discretization error and finish the proof of Proposition 2 as follows:

2 Convergence of the underdamped Langevin algorithm

Combining Propositions 1 and 2, which establish the convergence rates of the continuous underdamped Langevin dynamics and the discretization error, we find that the overall time evolution of the Lyapunov functional $\mathcal{L}$ within each step of the underdamped Langevin algorithm can be upper bounded as follows:

In this section, we will further analyze terms (47a)–(47c) to obtain the overall convergence rate of the underdamped Langevin algorithm. We will need to quantify the convergence contributed by term (47a) and upper bound the extra discretization error in terms (47b)–(47c) as the algorithm progresses. After these two steps, choosing a suitable step size will finish the proof of Theorem 1.

We begin by using the log-Sobolev inequality to relate term (47a) to the Lyapunov functional $\mathcal{L}(\mathbf{p}_{t})$. A key step is lower bounding matrix $M$ which is done in the following Lemma 5 (the proof of which is deferred to Appendix E).

We can thus upper bound term (47a) using this lower bound on $M$ in conjunction with the log-Sobolev inequality, Eq. (5):

Consequently, Eq. (47a)–(47c) simplify to:

This implies that without the extra discretization error of terms (52b)–(52c), the Markov process converges exponentially (similarly as for the continuous dynamics) with a rate of $\rho/30$, proportional to the log-Sobolev constant.

Assume that function $U$ satisfies Assumption A1–A3, where $\rho$ denotes the minimum of the log-Sobolev constant and $1$. Assume that we take $\gamma=2$, $\xi=2L_{G}$, and

To establish this uniform upper bound, we use an inductive argument—we prove that if the above bound holds for $t\leq kh$, then, given the effect of contraction and the discretization error in $[kh,\tau]$, the bound will still hold for any $\tau\in[kh,(k+1)h]$. We defer the complete proof of Lemma 6 to Appendix E.

Applying Grönwall’s lemma, we arrive at a bound for the Lyapunov functional at every step:

We now use the definition of the step size $h$ and the upper bound on the initial value $\mathcal{L}[\mathbf{p}_{0}]$ from Lemma 12 to obtain the number of iterations for Algorithm 1 to converge to within $\epsilon$ of the target distribution $\mathbf{p}^{*}$:

If the function $U$ further satisfies assumptions A1—A3 (that $U$ is nonconvex inside a region of radius $R$ and $m$-strongly convex outside of it), we can instantiate the constants $\rho\geq\frac{m}{2}e^{-16L_{G}R^{2}}$, $\widetilde{C_{N}}=C_{N}+\frac{1}{2}\ln\frac{L_{G}}{2\pi}\leq\frac{1}{2}\ln\frac{2L_{G}}{m}$, and $C_{M}\leq 32\frac{L_{G}^{2}}{m^{2}}L_{G}R^{2}$, and study the computational complexity in more detail. The number of iterations required becomes:

Emphasizing the dimension dependency, we have:

Discussion

We have shown that there is an analog of Nesterov accelerated gradient method for MCMC—it is the underdamped Langevin algorithm. We demonstrated this by adopting a view of sampling algorithms as optimizing over the space of probability measures, with KL divergence as the objective functional. By constructing an appropriate Lyapunov functional, we were able to prove that the underdamped Langevin algorithm has an accelerated convergence rate compared to the classical overdamped Langevin algorithm.

A line of recent results leverage richer stochastic dynamics to obtain better pre-conditioning and employ higher-order discretization schemes . They observe that in practice such dynamics increase stability and in turn results in faster convergence of the algorithm.

Our particular approach involves multiplying the strong sub-differential of the KL divergence by a symplectic matrix and a positive semidefinite matrix. An interesting direction for future research would be to consider other, more general choices. Indeed, a general construction of underdamped stochastic processes would involve taking a vector field $v_{t}$ to have the following form:

where $D(x)$ is a positive semidefinite diffusion matrix, and $Q(x)$ is a skew-symmetric curl matrix. This has the form of a generic dynamics for smooth optimization. It can be checked that when $\mathbf{p}_{t}(x)=\mathbf{p}^{*}(x)$, $v_{t}=0$. Therefore, $\mathbf{p}^{*}$ is a stationary distribution when $\mathbf{p}_{t}$ follows the vector flow $v_{t}$:

It has been previously proved that any continuous Markov process with the stationary distribution $\mathbf{p}^{*}$ which satisfies an integrability condition can be represented in the form of Eq. (55).

To simulate the dynamics of $v_{t}$ on the state space of $x$, we can realize it as a stochastic process with an Itô diffusion:

where $\Gamma_{i}(x)=\sum_{j}\frac{\partial}{\partial x_{j}}\left[D(x)+Q(x)\right]_{i,j}$. Eq. (56) corresponds to the probability density of $x_{t}$ following a stochastic differential equation:

Using notation from statistical mechanics, we can represent Eq. (58) in a more compact form using a (Ginzburg-Landau) dissipative bracket and a generalized Poisson bracket to generate the stochastic process $\frac{d}{dt}\mathbf{p}_{t}(x)$ with $\nabla\frac{\delta\mathcal{F}}{\delta\mathbf{p}_{t}}$. Define the dissipative bracket $\{\cdot,\cdot\}$ as

and the generalized Poisson bracket $[\cdot,\cdot]$ as

By taking $\mathcal{G}=\mathcal{F}$ as the KL-divergence, we can calculate its time derivative as:

Some attempts have been made in this direction in the stochastic optimization literature for a class of constant $D$ and $Q$ matrices . For the generic case, explores an approach based on Stein factors; this seems like a particularly promising avenue to explore further.

Acknowledgements

We would like to thank Jianfeng Lu, Chi Jin, and Nilesh Tripuraneni for many helpful discussions and insights. This work was partially supported by Army Research Office grant W911NF-17-1-0304, and National Science Foundation Grant NSF-IIS-1740855, NSF-IIS-1909365, and NSF-IIS-1619362.

References

Appendix A Local Nonconvexity Assumption

$U(\theta)$ is $m$-strongly convex for $\left\|\theta\right\|>R$.

$U(\theta)$ is $L_{G}$-Lipschitz smooth and Hessian $L_{H}$-Lipschitz.

For convenience, let $\nabla U(0)=0$ (i.e., zero is a local extremum).

From , we know that $\rho\geq\dfrac{m}{2}e^{-16L_{G}R^{2}}$. We prove that the constants in Assumption A3 are also upper bounded by functions of $m$, $L_{G}$, and $R$.

In other words, constants in Assumption A3 are bounded as: $C_{N}\leq\dfrac{1}{2}\ln\dfrac{4\pi}{m}$, and $C_{M}\leq 32\dfrac{L_{G}^{2}}{m^{2}}L_{G}R^{2}$.

Appendix B Explicit Iteration Rule for Algorithm 1

We provide an explicit iteration formula for $x_{\tau}$ given $x_{kh}$ in Eq. (24). Given $x_{kh}$ at the previous iteration, $x_{\tau}$ can be calculated as:

Therefore, the update rule in Algorithm 1 can be expressed as:

In Algorithm 1, the hyperparameters are set to be: $\gamma=2$, $\xi=2L_{G}$, and

where $\widetilde{C_{N}}=C_{N}+\dfrac{1}{2}\ln\dfrac{L_{G}}{2\pi}$.

Appendix C Convergence of the Continuous Process

To simplify the notations in the proofs, we let $a=\dfrac{1}{L_{G}}$, $b=\dfrac{1}{4L_{G}}$, and $c=\dfrac{2}{L_{G}}$, so that

We first compute the time evolution of the Lyapunov function $\mathcal{L}$ with respect to the continuous time vector flow $v_{t}^{AGD}$ in Eq. (13).

The time derivative of the Lyapunov functional $\mathcal{L}$ with respect to the continuous time vector flow $v_{t}^{AGD}$ in Eq. (13) with $\gamma=2$ and $\xi=2L_{G}$ is:

We then upper bound the time derivative of $\mathcal{L}$ by a negative factor times itself to obtain linear convergence rate.

For $L_{G}$-Lipschitz smooth $U$, matrix $M_{C}$ defined in Eq. (81) satisfy:

Since the matrix $S$ is positive definite, we can directly bound the evolution of the Lyapunov functional $\mathcal{L}$ as

Using the log-Sobolev inequality in Assumption A1, we directly obtain:

which implies the linear convergence of the continuous process with a rate of $\dfrac{\rho}{10}$. $\blacksquare$

Denote $h(\mathbf{p}_{t})=\sqrt{\dfrac{\mathbf{p}_{t}}{\mathbf{p}^{*}}}$. Then

The variational derivative of $\mathcal{L}[\mathbf{p}_{t}]$ can be thus calculated as:

the adjoint operator can be expressed as:

The vector flow $v_{t}$ can also be expressed in terms of $h(\mathbf{p}_{t})$ as:

for $a=\dfrac{1}{L_{G}}$, $b=\dfrac{1}{4L_{G}}$, $c=\dfrac{2}{L_{G}}$, $\gamma=2$, $\xi=2L_{G}$, and $\lambda=\dfrac{\rho}{10}$. That is equivalent to having:

To guarantee that $l\geq 0$, we need that $\forall\Lambda_{j}\in[-L_{G},L_{G}]$,

Since the linear function $\dfrac{a}{2}\Lambda_{j}-\alpha-\sigma$ of $\Lambda_{j}$ is increasing; the quadratic function $\dfrac{b^{2}}{4}\Lambda_{j}^{2}+\left(\dfrac{a}{2}\alpha-b\beta\right)\Lambda_{j}+\beta^{2}-\alpha\sigma$ of $\Lambda_{j}$ is convex, we simply need the inequality to be satisfied at the end points:

We verify these inequalities by plugging in the setting of $a=\dfrac{1}{L_{G}}$, $b=\dfrac{1}{4L_{G}}$, $c=\dfrac{2}{L_{G}}$, $\gamma=2$, $\xi=2L_{G}$, and $\lambda=\dfrac{\rho}{10}$ in the definition of $\alpha$, $\beta$, and $\sigma$. We obtain:

We then deal with the three terms one by one.

Here, $\nabla_{\theta}$ commutes with $\nabla_{r}$ and $(\nabla_{r})^{*}$.

where we have used $\left<\cdot,\cdot\right>_{F}$ to also denote Frobenius inner product between matrices.

Line (136) can be simplified by using the representation of the vector flow in Eq. (13):

Since $[\nabla_{r},B][h]=\xi\nabla_{\theta}h$ and $[\nabla_{\theta},B][h]=-\nabla^{2}U(\theta)\nabla_{r}h$, Eq. (162) becomes

Appendix D Discretization Error

As in the continuous case, define ${h}=\sqrt{\dfrac{\mathbf{p}_{\tau}(x_{\tau})}{\mathbf{p}^{*}(x_{\tau})}}$, and denote $a=\dfrac{1}{L_{G}}$, $b=\dfrac{1}{4L_{G}}$, $c=\dfrac{2}{L_{G}}$. First note that

Similar to the continuous case, the term in Line (183) separates into four terms:

We first simplify Lines (184) and (185) and then deal with Lines (186) and (187).

Therefore, Lines (184)–(187) combines to be:

It can be seen that the expectation in Line (188) can be rewritten as $x_{kh}$ conditioning on $x_{\tau}$:

Then for $\nu\leq\dfrac{1}{8L_{G}}$ (and $\gamma=2$, and $\xi=2L_{G}$),

Taking Lemma 10 as given, we can separate Term (45c) into two:

We then make use of the properties of Frobenius inner product to upper bound Terms (192c) and (192g) by the Frobenius norms:

As a result, we obtain that for Term (192c),

To obtain the final bound, we simplify Eq. (204) by demonstrating the following fact.

For $0\leq\nu\leq\min\left\{\dfrac{1}{\gamma\xi},\dfrac{1}{\sqrt{2eL_{G}\xi}}\right\}$,

Since $\nu\leq\dfrac{1}{8L_{G}}\leq\min\left\{\dfrac{1}{\gamma\xi},\dfrac{1}{\sqrt{2eL_{G}\xi}}\right\}$, and $\left\|\nabla^{2}U(\theta_{\tau})-\nabla^{2}U(\theta_{kh})\right\|_{F}\leq L_{H}\left\|\theta_{\tau}-\theta_{kh}\right\|$, we plug the above inequalities into Terms (192c) and (192g) and arrive at our conclusion:

where $\Gamma\left(\mathbf{p}(x_{kh}|x_{\tau}),\mathbf{p}(\hat{x}_{n}|x_{\tau}+hv)\right)$ is any joint distribution of $x_{kh}$ and $\hat{x}_{n}$ with marginal distributions being $\mathbf{p}(x_{kh}|x_{\tau})$ and $\mathbf{p}(\hat{x}_{n}|x_{\tau}+hv)$ – any coupling between the two random variables.

Recall from (65) that the relation between $x_{\tau}$ and $x_{kh}$ is:

where the Gaussian random variable $W_{x}$ takes the same value as that in Eq. (207). Then we get that for any pair of $\left(x_{kh},\hat{x}_{n}\right)$ following this joint law,

$\bar{\theta}$ a convex combination of $\theta_{kh}$ and $\hat{\theta}_{n}$, and

Appendix E Overall Convergence of the Underdamped Langevin Algorithm

for $a=\dfrac{1}{L_{G}}$, $b=\dfrac{1}{4L_{G}}$, $c=\dfrac{2}{L_{G}}$, $\gamma=2$, $\xi=2L_{G}$, and $\lambda=\dfrac{\rho}{30}$. That is equivalent to having:

To guarantee that $l\geq 0$, we need that $\forall\Lambda_{j}\in[-L_{G},L_{G}]$,

Since the linear function $\dfrac{a}{2}\Lambda_{j}-\alpha-\sigma$ of $\Lambda_{j}$ is increasing; the quadratic function $\dfrac{b^{2}}{4}\Lambda_{j}^{2}+\left(\dfrac{a}{2}\alpha-b\beta\right)\Lambda_{j}+\beta^{2}-\alpha\sigma$ of $\Lambda_{j}$ is convex, we simply need the inequality to satisfy at the end points:

We verify these inequalities by plugging in the setting of $a=\dfrac{1}{L_{G}}$, $b=\dfrac{1}{4L_{G}}$, $c=\dfrac{2}{L_{G}}$, $\gamma=2$, $\xi=2L_{G}$, and $\lambda=\dfrac{\rho}{30}$, in the definition of $\alpha$, $\beta$, and $\sigma$. Then for $L_{G}\geq 2\rho$, we obtain that

For the expectation of ${\left\|\theta_{\tau}-\theta_{kh}\right\|^{2}}$ taken over the joint distribution of $(x_{\tau},x_{kh})$, we use the definition of $x_{\tau}$ in our Equation (24) to expand it (by way of Jensen’s inequality):

Assume that function $U$ satisfies Assumption A1–A3, where $\rho$ denotes the minimum of the log-Sobolev constant and $1$. If we take $\gamma=2$, $\xi=2L_{G}$, and

where $\epsilon\leq d\dfrac{L_{G}}{\rho}$. Then for $r_{s}$ following Equation (24), $\forall s\geq 0$,

We defer the proof of Lemma 11 to Sec. E.1.

Let $\mathbf{p}_{0}(x)=\mathbf{p}_{0}(\theta)\mathbf{p}_{0}(r)$, where

For $\mathbf{p}^{*}(x)\propto\left(-U(\theta)-\dfrac{\xi}{2}\left\|r\right\|^{2}\right)$, if $U(\theta)$ follows Assumptions A1–A3, then we can define $\widetilde{C_{N}}=C_{N}+\dfrac{1}{2}\ln\dfrac{L_{G}}{2\pi}$ and obtain that

With the setting of $\xi=2L_{G}$, we can also obtain that

We further expand this inequality by using the extended Talagrand inequality, Eq. (1), which applies to the joint density function $\mathbf{p}^{*}(\theta,r)\propto\exp\left(-U(\theta)-\dfrac{\xi}{2}\left\|r\right\|^{2}\right)$ with log-Sobolev constant greater than or equal to $\rho$ and Lipschitz smoothness of $U+\dfrac{\xi}{2}\left\|r\right\|^{2}$ less than or equal to $4L_{G}$:

It can be verified that for $\epsilon\leq 2d$ and $\rho\leq 1$, $h$ is indeed smaller than $\dfrac{1}{8L_{G}}$. Thus Lemma 13, in conjunction with the induction hypothesis, gives us a rough bound that $\forall s\in[kh,(k+1)h]$,

Applying the extended Talagrand inequality, Eq. (1), we obtain that

Let $x_{s}$ follow the underdamped Langevin algorithm 1 with parameters $\xi=2L_{G}$, $\gamma=2$, and the step size $h=(k+1)h-kh$ given in Eq. (75). Also let $p_{s}$ be the probability distribution of $x_{s}$. Assume that Eq. (243) (given by the induction hypothesis in conjunction with Lemma 13) holds for any $s\in[kh,(k+1)h]$. Then for $\epsilon\leq 2d$ and $\rho\leq 1$, $\forall s\in[kh,(k+1)h]$,

Applying Grönwall’s Lemma in Eq. (245), we obtain that the objective functional $\mathcal{L}$ will not increase by more than $\epsilon/2$ throughout the progress of the algorithm:

From Lemma 12, we know that $\mathcal{L}[\mathbf{p}_{0}]\leq\left(\widetilde{C_{N}}+1\right)d+C_{M}$. Therefore, for $\epsilon\leq 2d$,

Plugging Eq. (246) into Eq. (244), we obtain our final result that

We begin from the discretized dynamics of underdamped Langevin diffusion Eq. (30) to calculate that $\forall s\in[kh,(k+1)h]$,

where the last step follows from plugging in the setting of $\gamma=2$ and $\xi=2L_{G}$ and using Young’s inequality. Multiplying $e^{-2L_{G}s}>0$ on both ends of Eq. (252), we obtain that $\forall s$,

Applying the fundamental theorem of calculus and multiplying $e^{2L_{G}\tau}>0$ on both sides, we have that

It can then be checked that when $\tau-kh\leq h\leq\dfrac{1}{8L_{G}}$, the factor $\left(e^{2L_{G}(\tau-kh)}-1\right)\leq\dfrac{1}{2}$, and that

to Eq. (52a)–(52c), we obtain that for $\xi=2L_{G}$, $\gamma=2$, and $\forall\tau\in[kh,(k+1)h]$,

Using the definition of $h=\dfrac{1}{56}\dfrac{1}{\sqrt{L_{G}}}\min\left\{\dfrac{1}{24}\dfrac{\rho}{L_{G}},\dfrac{\sqrt{L_{G}}\rho}{L_{H}}\right\}\cdot\min\left\{\left(\widetilde{C_{N}}+2\right)^{-1/2}\sqrt{\dfrac{\epsilon}{d}},\sqrt{\dfrac{\epsilon}{C_{M}}}\right\}$ in Eq. (75), we know that

Plugging this setting into the last term of Eq. (254), we obtain that for $\epsilon\leq 2d$ and $\rho\leq 1$,

Consequently, the time derivative of the Lyapunov functional $\mathcal{L}$ is bounded as:

Appendix F Proofs for Auxiliary Facts

The latter case follows directly from Assumptions (b) and (c). For the former case where $\|\theta\|\geq\dfrac{8L_{G}}{m}R$, define $\vartheta=\dfrac{R}{\|\theta\|}\theta$. Since $\|\vartheta\|=R$,

since $\|\theta\|\geq\dfrac{8L_{G}}{m}R$. Again, using Assumptions (b) and (c), $U(\vartheta)\geq-\dfrac{L_{G}}{2}R^{2}$, which leads to the result that $U(\theta)\geq\dfrac{m}{4}\|\theta\|^{2}$.

Therefore, $U(\theta)\geq\dfrac{m}{4}\|\theta\|^{2}-32\dfrac{L_{G}^{2}}{m^{2}}L_{G}R^{2}$ and

Hence $C_{N}\leq\dfrac{1}{2}\ln\dfrac{4\pi}{m}$ and $C_{M}\leq 32\dfrac{L_{G}^{2}}{m^{2}}L_{G}R^{2}$. $\blacksquare$

and provide bound for it when $0\leq(\tau-kh)\leq\min\left\{\dfrac{1}{\gamma\xi},\dfrac{1}{\sqrt{2eL_{G}\xi}}\right\}$.

First note that for $0\leq(\tau-kh)\leq\dfrac{1}{\gamma\xi}$,

We then prove Fact 2 by separating the following term: