Global Convergence of Stochastic Gradient Hamiltonian Monte Carlo for Non-Convex Stochastic Optimization: Non-Asymptotic Performance Bounds and Momentum-Based Acceleration

Xuefeng Gao, Mert Gürbüzbalaban, Lingjiong Zhu

Introduction

We consider the stochastic non-convex optimization problem

Because the population distribution D\mathcal{D} is unknown, a common popular approach is to consider the empirical risk minimization problem

based on the dataset z:=(z1,z2,…,zn)∈Zn\mathbf{z}:=(z_{1},z_{2},\dots,z_{n})\in\mathcal{Z}^{n} as a proxy to the problem (1.1) and minimize the empirical risk

instead, where the expectation is taken with respect to any randomness encountered during the algorithm to generate xx.We note that in our notation Z\mathbf{Z} is a random vector, whereas z\mathbf{z} is deterministic vector associated to a dataset that corresponds to a realization of the random vector Z\mathbf{Z}. Many algorithms have been proposed to solve the problem (1.1) and its finite-sum version (1.2). Among these, gradient descent, stochastic gradient and their variance-reduced or momentum-based variants come with guarantees for finding a local minimizer or a stationary point for non-convex problems. In some applications, convergence to a local minimum can be satisfactory ([GLM17, DLT+18]). However, in general, methods with global convergence guarantees are also desirable and preferable in many settings ([HLSS16, ŞimşekliYN+18]).

It has been well known that sampling from a distribution which concentrates around a global minimizer of FF is a similar goal to computing an approximate global minimizer of FF. For example such connections arise in the study of simulated annealing algorithms in optimization which admit several asymptotic convergence guarantees (see e.g. [Gid85, Haj85, GM91, KGV83, BT93, BLNR15, BM99]). Recent studies made such connections between the fields of statistics and optimization stronger, justifying and popularizing the use of Langevin Monte Carlo-based methods in stochastic non-convex optimization and large-scale data analysis further (see e.g. [CCS+17, Dal17, RRT17, CCG+16, ŞimşekliBCR16, ŞimşekliYN+18, WT11, Wib18]).

Stochastic gradient algorithms based on Langevin Monte Carlo are popular variants of stochastic gradient which admit asymptotic global convergence guarantees where a properly scaled Gaussian noise is added to the gradient estimate. Two popular Langevin-based algorithms that have demonstrated empirical success are stochastic gradient Langevin dynamics (SGLD) ([WT11, CDC15]) and stochastic gradient Hamiltonian Monte Carlo (SGHMC) ([CFG14, CDC15, Nea10, DKPR87]) and their variants to improve their efficiency and accuracy ([AKW12, MCF15, PT13, DFB+14, Wib18]). In particular, SGLD can be viewed as the analogue of stochastic gradient in the Markov Chain Monte Carlo (MCMC) literature whereas SGHMC is the analogue of stochastic gradient with momentum (see e.g. [CFG14]). SGLD iterations consist of

On the other hand, the SGHMC algorithm is based on the underdamped (a.k.a. second-order or kinetic) Langevin diffusion

(see e.g. [HN04, Pav14]) where Γz\Gamma_{\mathbf{z}} is the normalizing constant:

Hence, the xx-marginal distribution of stationary distribution πz(dx,dv)\pi_{\mathbf{z}}(dx,dv) is exactly the invariant distribution of the overdamped Langevin diffusion.With slight abuse of notation, we use πz(dx)\pi_{\mathbf{z}}(dx) to denote the xx-marginal of the equilibrium distribution πz(dx,dv)\pi_{\mathbf{z}}(dx,dv). SGHMC dynamics correspond to the discretization of the underdamped Langevin SDE where the gradients are replaced with their unbiased estimates. Although various discretizations of the underdamped Langevin SDE has also been considered and studied ([CDC15, LMS15]), the following first-order Euler scheme is the simplest approach that is easy to implement, and a common scheme among the practitioners ([TTV16, CCG+16, CDC15]):

In this paper, we focus on the unadjusted dynamics (without Metropolis-Hastings type of correction) that works well in many applications ([CFG14, CDC15]), as Metropolis-Hastings correction is typically computationally expensive for applications in machine learning and large-scale optimization when the size of the dataset nn is large and low to medium accuracy is enough in practice (see e.g. [WT11, CFG14]).

There is also an alternative discretization to (1.8)-(1.9), recently proposed by [CCBJ18] which leads to state-of-the-art estimates in the special case that improves upon the Euler discretization when the objective is strongly convex ([CCBJ18]). To introduce this alternative discretization by [CCBJ18], we first define a sequence of functions ψk\psi_{k} by ψ0(t)=e−γt\psi_{0}(t)=e^{-\gamma t} and ψk+1(t)=∫0tψk(s)ds\psi_{k+1}(t)=\int_{0}^{t}\psi_{k}(s)ds, k≥0k\geq 0. The iterates (X^k,V^k)(\hat{X}_{k},\hat{V}_{k}) are then defined by the following recursion:

where (ξk+1,ξk+1′)(\xi_{k+1},\xi^{\prime}_{k+1}) is a 2d2d-dimensional centered Gaussian vector so that (ξj,ξj′)(\xi_{j},\xi^{\prime}_{j})’s are independent and identically distributed (i.i.d.) and independent of the initial condition, and for any fixed jj, the random vectors ((ξj)1,(ξj′)1)((\xi_{j})_{1},(\xi^{\prime}_{j})_{1}), ((ξj)2,(ξj′)2)((\xi_{j})_{2},(\xi^{\prime}_{j})_{2}), …\ldots ((ξj)d,(ξj′)d)((\xi_{j})_{d},(\xi^{\prime}_{j})_{d}) are i.i.d. with the covariance matrix:

In the rest of the paper, we refer to Euler discretization (1.8)-(1.9) as SGHMC1 whereas the alternative discretization (1.10)-(1.11) as SGHMC2.

Recently, [EGZ19] show that the underdamped SDE converges to its stationary distribution faster than that of the best known convergence rate of overdamped SDE in the 2-Wasserstein metric under some assumptions, where FzF_{\mathbf{z}} can be non-convex. Their result is for the continuous-time underdamped dynamics. This raises the natural question whether the discretized underdamped dynamics (SGHMC), can lead to better guarantees than the SGLD method for solving stochastic non-convex optimization problems. Indeed, experimental results show that SGHMC can outperform SGLD dynamics in many applications (see e.g. [EGZ19, CDC15, CFG14]). Although asymptotic convergence guarantees for SGHMC exist (see e.g. [CFG14] [MSH02, Section 3], [LMS15]), there is a lack of finite-time explicit performance bounds for solving non-convex stochastic optimization problems with SGHMC in the literature including risk minimization problems.

Our main contributions can be summarized as follows:

We provide for the first time the non-asymptotic provable guarantees for SGHMC to find approximate minimizers of both empirical and population risks with explicit constants. We establish the results under some regularity and growth assumptions for the component functions f(x,z)f(x,z) and the noise in the gradients, but we do not assume ff is strongly convex in any region.

We show that for a class of non-convex problems, SGHMC2 can improve upon the (vanilla) SGLD algorithm in terms of the gradient complexity, i.e. the total number of stochastic gradients required to achieve a global minimum. Here, “improvement” means the best available bounds for SGHMC2, which we prove in our paper, are better than the best available bounds for SGLD for some class of problems; see Section 5 for details. As a consequence, our analysis gives further theoretical justification to the success of momentum-based methods for solving non-convex machine learning problems, empirically observed in practice (see e.g. [SMDH13]).

We illustrate the applications of our theoretical results using two examples including binary linear classification and robust ridge regression.

On the technical side, we adapt the proof techniques of [RRT17] developed for the overdamped dynamics to the underdamped dynamics and combine it with the analysis of [EGZ19] which quantifies the convergence rate of the underdamped Langevin SDE to its equilibrium. The main new technical results we derive in this paper, relative to these studies, include controlling the discretization errors between SGHMC and the continuous-time underdamped Langevin SDE, and bounding the moments of underdamped dynamics.

2 Related Work and Comparison to Existing Literature

In a recent work, [ŞimşekliYN+18] obtained a finite-time performance bound for the ergodic average of the SGHMC iterates in the presence of delays in gradient computations. Their analysis highlights the dependency of the optimization error on the delay in the gradient computations and the stepsize explicitly, however it hides some implicit constants which can be exponential both in β\beta and dd in the worst case. A comparison with the SGLD algorithm is also not given. On the contrary, in our paper, we make all the constants explicit. This allows us to make gradient complexity comparisons with respect to overdamped MCMC approaches such as SGLD.

A related paper [XCZG18] applies variance reduction techniques to overdamped MCMC to improve performance when the empirical risk can be non-convex satisfying the same dissipativity assumption considered in our paper. However, these results do not give guarantees for the risk minimization problem (1.1). Furthermore, such variance reduction techniques require objectives in the form of a finite sum and do not apply to the streaming data setting when each data point is used only once. In this work, we obtain guarantees for both the risk minimization problem and the empirical risk minimization and our results apply to the streaming data setting. Also, the convergence guarantees provided in [XCZG18] depends on a spectral gap-related parameter that is not provided explicitly; whereas all our results are explicit and this allows us to have explicit performance comparisons between the upper bounds of SGLD and SGHMC algorithms.

We also note that underdamped Langevin MCMC (also known as Hamiltonian MCMC) and its practical applications have also been analyzed further in a number of recent works (see e.g. [LV18, BBLG17, Bet17, BBG14, MPS18]). In particular, [MPS18] provide a characterization of the conductance of Hamiltonian Monte Carlo (HMC) in continuous time using Liouville’s theorem and invoking the Cheeger’s inequality, they obtain upper and lower bounds on the spectral gap of HMC in continuous-time. Although the formula provided in [MPS18] for the conductance of HMC is elegant, it is not an explicit formula. In our analysis, our focus is to obtain performance bounds with explicit constants and therefore we build on the coupling techniques of [EGZ19] which leads to explicit constants for the class of problems we consider.

We also note that [MPS18] consider sampling from the target distribution 12N(−1,σ2)+12N(1,σ2)\frac{1}{2}\mathcal{N}(-1,\sigma^{2})+\frac{1}{2}\mathcal{N}(1,\sigma^{2}) in dimension one and estimate the spectral gap of HMC in the regime as σ→0\sigma\to 0 . This is a mixture of two Gaussians with the same variance σ2\sigma^{2} centered at −1-1 and 11 respectively where they argue that for this specific example HMC does not lead to much improvement over the Random Walk approach for sampling. In our paper, our results apply to more general targets that are not necessarily mixture of Gaussians. However, if we consider sampling from the distribution 12N(−a,σ2)+12N(a,σ2)\frac{1}{2}\mathcal{N}(-a,\sigma^{2})+\frac{1}{2}\mathcal{N}(a,\sigma^{2}) as a→∞a\to\infty for σ2\sigma^{2} fixed, Proposition 11 is applicable and it implies that HMC will be more efficient than overdamped Langevin dynamics in terms of dependency to aa (which measures the distance between the modes) in the sense that the mixing time will be O(a)\mathcal{O}(a) in HMC whereas it will be O(a2)\mathcal{O}(a^{2}) in Random Walk. This does not contradict results of [MPS18] because we consider different scaling regimes: We fix σ>0\sigma>0 and let a→∞a\to\infty whereas [MPS18] fix a=1a=1 and let σ→0\sigma\to 0.

There are also some connections of our work to existing momentum-based optimization algorithms. More specifically, if the term with dB(t)dB(t) involving the Brownian noise is removed in the underdamped SDE (1.5)–(1.6), this results in a second-order ODE in X(t)X(t). Momentum-based algorithms for strongly convex objectives such as Polyak’s heavy ball method as well as Nesterov’s accelerated gradient method can be both viewed as (alternative) discretizations of this ODE (see e.g. [Pol87, SBC14, SDJS18, WRJ16]). It is known ([SBC14, SDJS18, WRJ16]) that Nesterov’s accelerated gradient method tracks this second-order ODE (also referred to as the Nesterov’s ODE in the literature), whereas the first-order non-accelerated methods such as the classical gradient descent are known to track a first-order ODE in X(t)X(t) called the gradient flow dynamics. Furthermore, existing analysis shows that Nesterov’s ODE converges to its equilibrium faster (in time) than the first-order gradient flow ODE in terms of upper bounds and this speed-up is also inherited by the discretized dynamics. Roughly speaking, our results can be interpreted as the analogue of these results in the non-convex optimization setting where we deal with SDEs instead of ODEs building on the theory of Markov processes and show that SGHMC tracks the second-order (underdamped) Langevin SDE closely and inherits its favorable convergence guarantees (in terms of upper bounds on the expected suboptimality) compared to that of overdamped Langevin SDE.

Acceleration of first-order gradient or stochastic gradient methods and their variance-reduced versions for finding a local stationary point (a point with a gradient less than ε\varepsilon in norm) are also studied in the literature (see e.g. [CDHS18, Nes83, GL16, JT19, AZH16]). It has also been shown that under some assumptions momentum-based accelerated methods can escape saddle points faster (see e.g. [OW19, LCZZ18]). In contrast, in this work, our focus is obtaining performance guarantees for convergence to global minimizers instead.

Preliminaries and Assumptions

The function ff is continuously differentiable, takes non-negative real values, and there exist constants A0,B≥0A_{0},B\geq 0 so that

For each z∈Zz\in\mathcal{Z}, the function f(⋅,z)f(\cdot,z) is MM-smooth:

For each z∈Zz\in\mathcal{Z}, the function f(⋅,z)f(\cdot,z) is (m,b)(m,b)-dissipative:

There exists a constant δ∈[0,1)\delta\in[0,1) such that for every z{\mathbf{z}}:

The probability law μ0\mu_{0} of the initial state (X0,V0)(X_{0},V_{0}) satisfies:

where V\mathcal{V} is a Lyapunov function to be used repeatedly for the rest of the paper:

and γ\gamma is the friction coefficient as in (1.5), λ\lambda is a positive constant less than min⁡(1/4,m/(M+γ2/2))\min(1/4,m/(M+\gamma^{2}/2)), and α=λ(1−2λ)/12\alpha=\lambda(1-2\lambda)/12.

We note that the Lyapunov function V\mathcal{V} is used in [EGZ19] to study the rate of convergence to equilibrium for underdamped Langevin diffusion, which itself is motivated by e.g. [MSH02]. It follows from the above assumptions (applying Lemma 25) that there exists a constant A∈(0,∞)A\in(0,\infty) so that

This drift condition, which will be used later, guarantees the stability and the existence of Lyapunov function V\mathcal{V} for the underdamped Langevin diffusion in (1.5)–(1.6), see [EGZ19].

Main Results for SGHMC1 Algorithm

Our first result shows SGHMC1 iterates (Xk,Vk)(X_{k},V_{k}) in (1.8)–(1.9) track the underdamped Langevin SDE in the sense that the expectation of the empirical risk FzF_{\mathbf{z}} with respect to the probability law of (Xk,Vk)(X_{k},V_{k}) conditional on the sample z\mathbf{z}, denoted by μk,z\mu_{k,\mathbf{z}}, and the stationary distribution πz\pi_{\mathbf{z}} of the underdamped SDE is small when kk is large enough. The difference in expectations decomposes as a sum of two terms J0(z,ε)\mathcal{J}_{0}(\mathbf{z},\varepsilon) and J1(ε)\mathcal{J}_{1}(\varepsilon) while the former term quantifies the dependency on the initialization and the dataset z\mathbf{z} whereas the latter term is controlled by the discretization error and the amount of noise in the gradients which depends on the parameter δ\delta. We also note that the parameter μ∗\mu_{*} (see Table 1) in our bounds governs the speed of convergence to the equilibrium of the underdamped Langevin diffusion.

Consider the SGHMC1 iterates (Xk,Vk)(X_{k},V_{k}) defined by the recursion (1.8)–(1.9) from the initial state (X0,V0)(X_{0},V_{0}) which has the law μ0\mu_{0}. If Assumption 1 is satisfied, then for β,ε>0\beta,\varepsilon>0, we have

with σ\sigma defined by (A.20) provided that

Here Hρ\mathcal{H}_{\rho} is a semi-metric for probability distributions defined by (A.12). All the constants are made explicit and are summarized in Table 1.

The proof of Theorem 2 will be presented in details in Section A in the Appendix. In the following subsections, we discuss that this theorem combined with some basic properties of the equilibrium distribution πz\pi_{\mathbf{z}} leads to a number of results which provide performance guarantees for both the empirical risk and population risk minimization.

In order to obtain guarantees for the empirical risk given in (1.3), in light of Theorem 2, one has to control the quantity

which is a measure of how much the x−x-marginal of the equilibrium distribution πz\pi_{\mathbf{z}} concentrates around a global minimizer of the empirical risk. As β\beta goes to infinity, it can be verified that this quantity goes to zero. For finite β\beta, [RRT17] (see Proposition 11) derives an explicit bound of the form

(which is also provided in the Appendix for the sake of completeness, see Lemma 28). This combined with Theorem 2 immediately leads to the following performance bound for the empirical risk minimization. The proof is omitted.

Under the setting of Theorem 2, the empirical risk minimization problem admits the performance bounds:

provided that conditions (3.3) and (3.4) hold where the terms J0(z,ε)\mathcal{J}_{0}(\mathbf{z},\varepsilon), J1(ε)\mathcal{J}_{1}(\varepsilon) and J2\mathcal{J}_{2} are defined by \eqrefeq:J0z\eqref{eq:J0z}, \eqrefeq:J1\eqref{eq:J1} and \eqrefdef−J2\eqref{def-J2} respectively.

2 Performance bound for the population risk minimization

By exploiting the fact that the x−x-marginal of the invariant distribution for the underdamped dynamics is the same as it is in the overdamped case, it can be shown that the generalization error F(Xk)−FZ(Xk)F(X_{k})-F_{\mathbf{Z}}(X_{k}) is no worse than that of the available bounds for SGLD given in [RRT17], and therefore, we have the following corollary. A more detailed proof will be given in Section A in the Appendix.

Under the setting of Theorem 2, the expected population risk of XkX_{k} (the iterates in (1.9)) is bounded by

where σ\sigma is defined by (A.20), H‾ρ(μ0)\overline{\mathcal{H}}_{\rho}(\mu_{0}) is defined by (A.18), J1(ε)\mathcal{J}_{1}(\varepsilon) and J2\mathcal{J}_{2} are defined by (3.2) and (3.5) respectively and cLSc_{LS} is a constant satisfying

and λ∗\lambda_{\ast} is the uniform spectral gap for overdamped Langevin dynamics In [RRT17], their formula for λ∗\lambda_{\ast} missed β−1\beta^{-1} factor.:

3 Generalization error of SGHMC1 in the one pass regime

where XπX^{\pi} is the Gibbs output, i.e. its distribution conditional on Z=z\mathbf{Z}=\mathbf{z} is given by πz\pi_{\mathbf{z}}. If every sample is used once, i.e. if only one pass is made over the dataset, then the second term in (3.9) disappears. As a consequence, the generalization error is controlled by the bound

The following result provides a bound on this quantity. The proof is similar to the proof of Theorem 2 and its corollaries, and hence omitted.

provided that (3.3) and (3.4) hold where XπX^{\pi} is the output of the underdamped Langevin dynamics, i.e. its distribution conditional on Z=z\mathbf{Z}=\mathbf{z} is given by πz\pi_{\mathbf{z}} and J‾0(ε)\overline{\mathcal{J}}_{0}(\varepsilon) is defined by (3.6). Then, it follows from (3.10) that if each data point is used once, the expected generalization error satisfies

Main Results for SGHMC2 Algorithm

Recall the SGHMC2 algorithm (X^k,V^k)(\hat{X}_{k},\hat{V}_{k}) defined in (1.10)-(1.11), and denote the probability law of (X^k,V^k)(\hat{X}_{k},\hat{V}_{k}) conditional on the sample z\mathbf{z} by μ^k,z(dx,dv)\hat{\mu}_{k,\mathbf{z}}(dx,dv). Similar to our analysis for SGHMC1, we can derive similar performance guarantees for SGHMC2 in terms of empirical risk, population risk and the generalization error. The main difference is that the term J1(ε)\mathcal{J}_{1}(\varepsilon) is controlled by the accuracy of the discretization and has to be replaced by another term J^1(ε){\hat{\mathcal{J}}}_{1}(\varepsilon), as SGHMC2 algorithm is based on an alternative discretization. In particular, the performance bounds we get for SGHMC2 are tighter than SGHMC1, as will be elaborated further in the Section 5.

Consider the SGHMC2 iterates (X^k,V^k)(\hat{X}_{k},\hat{V}_{k}) defined by the recursion (1.10)–(1.11) from the initial state (X0,V0)(X_{0},V_{0}) which has the law μ0\mu_{0}. If Assumption 1 is satisfied, then for β,ε>0\beta,\varepsilon>0, we have

where J0(z,ε)\mathcal{J}_{0}(\mathbf{z},\varepsilon) is defined in (3.1) and

with σ\sigma defined by (A.20) provided that

Here Hρ\mathcal{H}_{\rho} is a semi-metric for probability distributions defined by (A.12). All the constants are made explicit and are summarized in Table 1 and Table 2.

The proof of Theorem 6 is given in Section B in the Appendix. Relying on Theorem 6, one can readily derive the following result on the performance bound for the empirical risk minimization with the SGHMC2 algorithm. The proof follows a similar argument as discussed in Section 3.1, and is omitted.

Under the setting of Theorem 6, the empirical risk minimization problem admits the performance bounds:

provided that conditions (4.2) and (4.3) hold where the terms J0(z,ε)\mathcal{J}_{0}(\mathbf{z},\varepsilon), J^1(ε)\hat{\mathcal{J}}_{1}(\varepsilon) and J2\mathcal{J}_{2} are defined by \eqrefeq:J0z\eqref{eq:J0z}, \eqrefeq:J1:Jordan\eqref{eq:J1:Jordan} and \eqrefdef−J2\eqref{def-J2} respectively.

Next, we present the performance bound for the population risk minimization with the SGHMC2 algorithm. Similar as in Section 3.2, to control the population risk during SGHMC2 iterations, one needs to control the difference between the finite sample size problem (1.2) and the original problem (1.1) in addition to the empirical risk. This leads to the following result. The details of the proof are given in Section B in the Appendix.

Under the setting of Theorem 6, the expected population risk of X^k\hat{X}_{k} (the iterates in (1.11)) is bounded by

where J‾0(ε)\overline{\mathcal{J}}_{0}(\varepsilon), J^1(ε)\hat{\mathcal{J}}_{1}(\varepsilon), J2\mathcal{J}_{2}, J3(n)\mathcal{J}_{3}(n) are defined in (3.6), (4.1), (3.5) and (3.7).

Finally, we present a result on the generalization error of the SGHMC2 algorithm in the one pass regime. The proof follows from Theorem 6 and the discussion for SGHMC1 algorithm in Section 3.3, and hence is omitted.

provided that (4.2) and (4.3) hold where XπX^{\pi} is the output of the underdamped Langevin dynamics, i.e. its distribution conditional on Z=z\mathbf{Z}=\mathbf{z} is given by πz\pi_{\mathbf{z}} and J‾0(ε)\overline{\mathcal{J}}_{0}(\varepsilon) is defined by (3.6). Then, it follows from (3.10) that if each data point is used once, the expected generalization error satisfies

Performance comparison with respect to SGLD algorithm

The constants λ∗\lambda_{*} (see (3.8)) and μ∗\mu_{*} (see Table 1) are exponentially small in both β\beta and dd in the worst case, but under some extra assumptions the dependency on dd can be polynomial (see e.g. [CCBJ18]) although the exponential dependence to β\beta is unavoidable in the presence of multiple minima in general (see [BGK05]). One can readily see that KSGHMC2K_{SGHMC2} has better dependency on ϵ\epsilon than KSGHMC1K_{SGHMC1}, and infer from (5.1)–(5.2) that the performance of SGHMC2 is better than SGHMC1. Hence, in the rest of the section, we will only focus on the comparison between SGHMC2 and SGLD.

We see that the generalization error for SGHMC2 (5.2) is bounded by

as μ∗\mu_{*} is small, and if we ignore the log⁡log⁡(1/ε)\sqrt{\log\log(1/\varepsilon)} factor We emphasize that the effect of the last term log⁡log⁡(1/ε)\sqrt{\log\log(1/\varepsilon)} appearing in (5.4) is typically negligible compared to other parameters. For instance even if ε=2−216\varepsilon=2^{-2^{16}} is double-exponentially small, we have log⁡log⁡(1/ε)≤4\sqrt{\log\log(1/\varepsilon)}\leq 4., then, we get

iterations of the SGHMC2 algorithm whereas the corresponding bound for SGLD from [RRT17, Theorem 1] is

iterations of the SGLD algorithm. Note that KSGHMC2K_{SGHMC2} and KSGLDK_{SGLD} do not have the same dependency to ε\varepsilon up to log⁡\log factors (the former scales with ε\varepsilon as log⁡2(1/ε)ε−2\log^{2}(1/\varepsilon)\varepsilon^{-2} and the latter log⁡5(1/ε)ε−4\log^{5}(1/\varepsilon)\varepsilon^{-4}), and this improvement on ε\varepsilon dependency is due to better diffusion approximation of SGHMC2 (see Lemma 22) compared to SGLD and the exponential integrability estimate we have in Lemma 17 which improves the estimate in [RRT17] and using the same argument, one can improve the log⁡5(1/ε)/ε4\log^{5}(1/\varepsilon)/\varepsilon^{4} term in (5.6) to log⁡3(1/ε)/ε4\log^{3}(1/\varepsilon)/\varepsilon^{4}.

To make the comparison to SGLD simpler, we notice that in both expressions (5.5) and (5.6), we see a term scaling with δ1/4\delta^{1/4} due to the gradient noise level δ\delta (δ\delta is fixed in the one-pass setting), and we fix the error in (5.5) and (5.6) without the δ\delta term to be the same order, and then compare the number of iterations KSGHMC2K_{SGHMC2} and KSGLDK_{SGLD}. More precisely, given ε^>0\hat{\varepsilon}>0 and we choose ε>0\varepsilon>0 such that (d+β)3/2βμ∗ε=ε^\frac{(d+\beta)^{3/2}}{\beta\mu_{*}}\varepsilon=\hat{\varepsilon} in (5.5) so that the generalization error for SGHMC2 is

Similarly, the generalization error for SGLD is

When λ∗\lambda_{*} and μ∗\mu_{*} are on the same order or μ∗\mu_{*} is larger, since typically β≥1\beta\geq 1, the term involving δ\delta in the generalization error for SGHMC2 above is (smaller) better than the counterpart for SGLD, and this is guaranteed to be achieved in a less number of iterations ignoring the log factors and universal constants for KSGHMC2K_{SGHMC2} in (5.7) and KSGLDK_{SGLD} in (5.8).

Empirical risk minimization.

The empirical risk minimization bound given in Corollary 7 has an additional term J2\mathcal{J}_{2} compared to the J‾0(ε)\overline{\mathcal{J}}_{0}(\varepsilon) and J^1(ε)\hat{\mathcal{J}}_{1}(\varepsilon) terms appearing in the one-pass generalization bounds. Note also that J0(z,ε)≤J‾0(ε)\mathcal{J}_{0}(\mathbf{z},\varepsilon)\leq\overline{\mathcal{J}}_{0}(\varepsilon). As a consequence, SGHMC2 algorithm has expected empirical risk

We next briefly discuss the comparisons of SGHMC2 and SGLD based on the total number of stochastic gradient evaluations (gradient complexity), and we compare with a recent work [XCZG18] which established a faster convergence result and improved the gradient complexity for SGLD in the mini-batch setting compared with [RRT17]. Here, the total number of stochastic gradient evaluations of an algorithm is defined as the number of stochastic gradients calculated per iteration (which is equal to the batch size in the mini-batch setting) times the total number of iterations. [XCZG18] showed that it suffices to take

stochastic gradient evaluations, ignoring the log⁡\log factors in the parameters ε^,μ∗,d\hat{\varepsilon},\mu_{*},d and hiding factors in β\beta that can be made explicit. To see (5.12), we infer from (5.9) that for fixed precision ε^>0\hat{\varepsilon}>0 and dimension dd, by ignoring the log factors and β\beta, we can choose ε\varepsilon so that d3/2ε/μ∗=ε^d^{3/2}\varepsilon/\mu_{*}=\hat{\varepsilon} and choose the gradient noise level δ\delta so that d3/2δ1/4/μ∗=ε^d^{3/2}\delta^{1/4}/\sqrt{\mu_{*}}=\hat{\varepsilon}. So the number of SGHMC2 iterations is

On the other hand, the mini-batch size to achieve gradient noise level δ\delta is given by 1/δ1/\delta (see [RRT17]), which is equal to d6/(μ∗2ε^4)d^{6}/(\mu_{*}^{2}\hat{\varepsilon}^{4}). Hence, we obtain (5.12) which is the product of the mini-batch size and number of iterations.

It is hard to compare λ^\hat{\lambda} in (5.11) and μ∗\mu_{*} in (5.12) in general since λ^\hat{\lambda} is the spectral gap of the discrete overdamped Langevin dynamics (i.e. SGLD with zero gradient noise) without a simple closed-form formula. However, when the stepsize is small enough, we expect λ^\hat{\lambda} will be similar to λ∗\lambda_{*}, which is the spectral gap of the continuous-time overdamped Langevin diffusion. As a consequence, when the stepsize η\eta is small enough (which is the case for instance, when target accuracy ε^\hat{\varepsilon} is small enough), we will have λ^≈λ∗\hat{\lambda}\approx\lambda_{*} and 1μ∗=O(1λ∗)=O(1λ^)\frac{1}{\mu_{*}}=\mathcal{O}\left(\sqrt{\frac{1}{\lambda_{*}}}\right)=\mathcal{O}\left(\sqrt{\frac{1}{\hat{\lambda}}}\right) for the class of non-convex functions we discuss in Proposition 11 and Example 10. For this class of problems, comparing (5.11) and (5.12), we see that we obtain an improvement in the spectral gap parameter (μ∗4\mu_{*}^{4} vs. λ^5\hat{\lambda}^{5}), however ε^\hat{\varepsilon} and dd dependency of the bound (5.11) is better than (5.12).

Population risk minimization.

If samples are recycled and multiple passes over the dataset is made, then one can see from Corollary 4 that there is an extra term J3\mathcal{J}_{3} that needs to be added to the bounds given in (5.9) and (5.10). This term satisfies

If this term is dominant compared to other terms J‾0,J1\overline{\mathcal{J}}_{0},\mathcal{J}_{1} and J2\mathcal{J}_{2}, for instance this may happen if the number of samples nn is not large enough, then the performance guarantees for population risk minimization via SGLD and SGHMC2 will be similar. Otherwise, if nn is large and β\beta is chosen in a way to keep the J2\mathcal{J}_{2} term on the order J0‾\overline{\mathcal{J}_{0}}, then similar improvement can be achieved.

The parameters λ∗\lambda_{*} (see (3.8)) and μ∗\mu_{*} (see Table 1) govern the convergence rate to the equilibrium of the overdamped and underdamped Langevin SDE, they can be both exponentially small in dimension dd and in β\beta. They appear naturally in the complexity estimates of SGHMC2 and SGLD method as these algorithms can be viewed as discretizations of Langevin SDEs (when the discretization step is small and the gradient noise δ=0\delta=0, the discrete dynamics will behave similarly as the continuous dynamics). Next, to get further intuition, first we discuss some toy examples of non-convex functions below where 1μ∗=O(1λ∗)\frac{1}{\mu_{*}}=\mathcal{O}\left(\sqrt{\frac{1}{\lambda_{*}}}\right). For these examples if the other parameters (β,d,δ)(\beta,d,\delta) are fixed, then SGHMC2 can lead to an improvement upon the SGLD performance. We will then show in Proposition 11 that these examples generalize to a more general class of non-convex functions.

where a>0a>0 is a scaling parameter which is illustrated in the left panel of Figure 1. For this example, there are two minima that are apart at a distance R=O(a)\mathcal{R}=\mathcal{O}(a). For simplicity, we assume there is only one sample, i.e. z=(z1)\mathbf{z}=(z_{1}) and Fz(x)=f(x,z1)=fa(x)F_{\mathbf{z}}(x)=f(x,z_{1})=f_{a}(x). We consider the non-convex optimization problem (1.2) with both the SGHMC2 algorithm and the SGLD algorithm. [EGZ19] showed that μ∗≥Θ(1a)\mu_{*}\geq\Theta(\frac{1}{a}) for this example whereas λ∗≤Θ(1a2)\lambda_{*}\leq\Theta(\frac{1}{a^{2}}) making the constants hidden by the Θ\Theta explicit. This shows that the contraction rate of the underdamped diffusion μ∗\mu_{*} is (faster) larger than that of the overdamped diffusion λ∗\lambda_{*} by a square root factor when aa is large where all the constants can be made explicit. Such results extend to a more general class of non-convex functions with multiple-wells and higher dimensions as long as the gradient of the objective satisfies a growth condition (see Example 1.1, Example 1.13 in [EGZ19] for a further discussion).

is the asymmetric double well potential in dimension one. It follows from Theorem 19 (see also [EGZ19]) that the contraction rate satisfies μ∗=Θ(a−1) ,\mu_{*}=\Theta\left(a^{-1}\right)\,, whereas it follows from Theorem 1.2 in [BGK05] that λ∗=Θ(1/a2)\lambda_{*}=\Theta(1/a^{2}). This shows that when the separation between minima, or alternatively the scaling factor aa is large enough, μ∗\mu_{*} is larger than λ∗\lambda_{*} by a square root factor up to constants.

The behavior in these toy examples can be generalized to more general non-convex objectives with a finite-sum structure satisfying Assumption 1. Proposition 11 below gives a class of functions where μ∗\mu_{*} is on the order of the square root of λ∗\lambda_{*}. The proof will be presented in details in Section F.

Suppose that the functions fa(x,z)f_{a}(x,z) indexed by aa satisfies Assumption 1 (i)-(iii) with m=m1a−2m=m_{1}a^{-2}, M=M1a−2M=M_{1}a^{-2} and B=B1a−1B=B_{1}a^{-1} for some fixed constants m1m_{1}, M1M_{1}, and B1B_{1}. Then, we have as a→∞,a\rightarrow\infty,

This result is more general than the previous example. In particular, if f(x,z)f(x,z) satisfies Assumption 1 (i)-(iii) with m,M,Bm,M,B replaced by m1,M1,B1m_{1},M_{1},B_{1}, then fa(x,z):=f(x/a,z)f_{a}(x,z):=f(x/a,z) satisfies Assumption 1 (i)-(iii) with m=m1a−2m=m_{1}a^{-2}, M=M1a−2M=M_{1}a^{-2} and B=B1a−1B=B_{1}a^{-1}. Proposition 11 essentially says that if we consider the normalized empirical risk objective Fz(x/a)=1n∑i=1nf(x/a,zi)F_{\mathbf{z}}(x/a)=\frac{1}{n}\sum_{i=1}^{n}f(x/a,z_{i}) where aa is a (normalization) scaling parameter and f(x,z)f(x,z) satisfies Assumption 1, then for large enough values of aa, the empirical risk surface will be relatively flat and the convergence rate of momentum variant SGHMC2 to an ε\varepsilon-neighborhood of the global minimum will be governed by the parameter μ∗\mu_{*} which will be larger than that of the parameter λ∗\lambda_{*} of SGLD when aa is sufficiently large. This will lead to improved performance bounds for SGHMC2 compared to known performance bounds for SGLD.

Applications

We note that several non-convex stochastic optimization problems of interest can satisfy Assumption 1 under appropriate noise assumptions for the underlying dataset. For example, Lasso problems with non-convex regularizers (see e.g. [HLM+17]), non-convex formulations of the phase retrieval problem around global minimum (see e.g. [ZZLC17]) or non-convex stochastic optimization problems defined on a compact set including but not limited to dictionary learning over the sphere (see e.g. [SQW16]), training deep learning models subject to norm constraints in the model parameters (see e.g. [ALG19]). In this section, we discuss some applications of our results where we provide two specific examples.

where λr>0\lambda_{r}>0 is a regularization parameter that may depend on the number of samples nn. By Lagrangian duality, this problem is equivalent to the constrained optimization problem

for some RR, which has also been considered in the literature (see e.g. [MBM18, FSS18, WCX19]). For non-convex σc(⋅)\sigma_{c}(\cdot), this problem is also non-convex in general. We consider minimizing the objective (6.1) in the mini-batch setting where the gradients in SGHMC iterations are estimated from nbn_{b} data points sampled with replacement, i.e. the gradient is estimated as

where zjz_{j} are i.i.d. with a uniform distribution over the set of indices {1,2,…,n}\{1,2,\dots,n\}. We also consider the following assumption for the threshold function σc\sigma_{c} which are satisfied by many choices of σc\sigma_{c} in practice. A prominent example is the logistic (or sigmoid) function in which case σc(z)=1/(1+e−z)\sigma_{c}(z)=1/(1+e^{-z}) which is also used in deep learning. Another possible choice is the probit function which corresponds to σc(t)=Φ(t)\sigma_{c}(t)=\Phi(t) where Φ\Phi is the cumulative distribution function of the standard normal distribution.

We show in the next lemma that if Assumption 12 holds, then Assumption 1 holds with explicit constants A0,B,M,m,bA_{0},B,M,m,b and σc\sigma_{c} that we can precise. The proof can be found in the Appendix.

In the setting of binary linear classification, consider the SGHMC method applied to the objective (6.1) where gradients are estimated according to (6.2) where the probability law μ0\mu_{0} of the initial state has compact support. If Assumption 12 holds; then Assumption 1 hold for any δ∈[14nb,1)\delta\in[\frac{1}{4n_{b}},1) with the following constants:

stochastic gradient evaluations to converge to an ε^\hat{\varepsilon} neighborhood of an almost ERM ignoring the log⁡\log factors in the parameters ε^,μ∗,d\hat{\varepsilon},\mu_{*},d and hiding other constants that can be made explicit based on Lemma 13.We also note that under further assumptions on the statistical nature of the input and if the number of data points is large enough, it can be shown that the objective (6.1) admits a unique minimizer and the objective is strongly convex in some regions [MBM18]. However, our assumptions here are weaker, therefore such arguments are not directly applicable.

2 Robust Ridge Regression

(see e.g. [MBM18]) and exponential squared loss [WJHZ13]: ρexp(t)=1−e−∥t∥2/t0\rho_{exp}(t)=1-e^{-\|t\|^{2}/t_{0}}, where t0>0t_{0}>0 is a tuning parameter. In the following, similar to [WCX19], we assume that the data AinA_{in} is bounded and the threshold function and its derivatives up to order two are bounded, similar to [MBM18]. This assumption for ρ\rho is satisfied in several cases, including Tukey’s bisquare loss and exponential squares loss mentioned above.

The following lemma shows that under Assumption 14, our assumptions (Assumption 1) for analyzing SGHMC methods hold with proper initialization.

In the setting of robust regression, consider the objective (6.1) where gradients are estimated according to (6.2) where the probability law μ0\mu_{0} of the initial state has compact support. If Assumption 12 holds; then Assumption 1 hold for both SGHMC1 and SGHMC2 methods for any choice of δ∈[14nb,1)\delta\in[\frac{1}{4n_{b}},1) with the following constants:

Similarly, we conclude from Lemma 15 that our main results for SGHMC1 and SGHMC2 algorithms described in Sections 3–5 apply to the problem of robust regression under Assumption 14.

Outline of the Proof

To obtain the main results in this paper, we adapt the proof techniques of [RRT17] developed for the overdamped dynamics to the underdamped dynamics and combine it with the analysis of [EGZ19] which quantifies the convergence rate of the underdamped Langevin SDE to its equilibrium. In an analogy to the fact that momentum-based first-order optimization methods require a different Lyapunov function and a quite different set of analysis tools (compared to their non-accelerated variants) to achieve fast rates (see e.g. [LFM18, SBC14, Nes83]), our analysis of the momentum-based SGHMC1 and SGHMC2 algorithms requires studying a different Lyapunov function V\mathcal{V} defined in (2.1) that also depends on the objective ff as opposed to the classic Lyapunov function H(x)=∥x∥2\mathcal{H}(x)=\|x\|^{2} arising in the study of the SGLD algorithm (see e.g. [MSH02, RRT17]). This fact introduces some challenges for the adaptation of the existing analysis techniques for SGLD to SGHMC. For this purpose, we take the following steps:

First, we show that SGHMC1 and SGHMC2 iterates track the underdamped Langevin diffusion closely in the 2-Wasserstein metric. As this metric requires finiteness of second moments, we first establish uniform (in time) L2L^{2} bounds for both the underdamped Langevin SDE and SGHMC1 and SGHMC2 iterates (see Lemma 16 and Lemma 21 in Appendix), exploiting the structure of the Lyapunov function V\mathcal{V}. Second, we obtain a bound for the Kullback-Leibler divergence between the discrete and continuous underdamped dynamics making use of the Girsanov theorem, which is then converted to bounds in the 2-Wasserstein metric by an application of an optimal transportation inequality of [BV05]. This step requires proving a certain exponential integrability property of the underdamped Langevin diffusion (Lemma 17 in Appendix). We show in Lemma 17 that the exponential moments grow at most linearly in time, which strictly improves the exponential growth in time in Lemma 4 in [RRT17]. The method that is used in the proof of Lemma 17 in Appendix can indeed be adapted to improve the exponential integrability and hence the overall estimates in [RRT17] for SGLD as well. As a result, the method improves upon the ε\varepsilon dependence of the number of iterates (see equations (5.5) and (5.6)).

Second, we apply the seminal result of [EGZ19] which showed that the continuous-time underdamped Langevin SDE is geometrically ergodic with an explicit rate μ∗\mu_{*} in the 2-Wasserstein metric. In order to get explicit performance guarantees, we derive new bounds that make the dependence of the constants to the initialization in [EGZ19] explicit (see Lemma 20 in Appendix).

As the xx-marginal of the equilibrium distribution πz(dx,dv)\pi_{\mathbf{z}}(dx,dv) of the underdamped Langevin SDE concentrates around the global minimizers of FzF_{\mathbf{z}} for β\beta appropriately chosen, and we can control the error between the discrete-time SGHMC1 and SGHMC2 dynamics and the underdamped SDE by choosing the step size accordingly; this leads to performance bounds for the empirical risk minimizations for SGHMC1 and SGHMC2 algorithms in Corollary 3 and Corollary 7. For controlling the population risk during SGHMC iterations, in addition to the empirical risk, one has to control the generalization error F(Xk)−FZ(Xk)F(X_{k})-F_{\mathbf{Z}}(X_{k}) that accounts for the differences between the finite sample size problem (1.2) and the original problem (1.1). By exploiting the fact that the x−x-marginal of the invariant distribution for the underdamped dynamics is the same as it is in the overdamped case, we control the generalization error in Corollary 4 and Corollary 8 which is no worse than that of the available bounds for SGLD given in [RRT17].

Conclusion

SGHMC is a momentum-based popular variant of stochastic gradient where a controlled amount of isotropic Gaussian noise is added to the gradient estimates for optimizing a non-convex function. We obtained first-time finite-time guarantees for the convergence of SGHMC1 and SGHMC2 algorithms to the ε\varepsilon-global minimizers under some regularity assumption on the non-convex objective ff. We also show that on a class of non-convex problems, SGHMC2 can be faster than overdamped Langevin MCMC approaches such as SGLD in the sense that the best available bounds for SGHMC2, which we prove in our paper, are better than the best available bounds for SGLD. This effect is due to the momentum term in the underdamped SDE. Furthermore, our results show that momentum-based acceleration is possible on a class of non-convex problems under some conditions if we compare known upper bounds between SGLD and SGHMC. Finally, we mention a few limitations in our work that may lead to some future research directions. In our paper, the performance dependence on dimension is exponential in general. In the future, we will investigate for what class of (non-convex) target functions ff we can obtain performance bound independent of dimension dd or has polynomial dependence on dd. In addition, our results suggest that momentum-based SGHMC methods will work particularly well when the (non-convex) target functions have relatively flat landscapes. In the future, we will investigate whether we can obtain theoretical results for SGHMC on a wider class of non-convex problems.

Acknowledgements

We thank Agostino Capponi, Xiuli Chao, Wenbin Chen, Jim Dai, Murat A. Erdogdu, Fuqing Gao, Jianqiang Hu, Jin Ma, Sanjoy Mitter, Asuman Ozdaglar, Pablo Parrilo, Umut Şimşekli, and S. R. S. Varadhan for helpful discussions. Xuefeng Gao acknowledges support from Hong Kong RGC Grants 24207015 and 14201117. Mert Gürbüzbalaban’s research is supported in part by the grants NSF DMS-1723085 and NSF CCF-1814888. Lingjiong Zhu is grateful to the support from the grant NSF DMS-1613164.

References

Appendix A Proof of Theorem 2 and Corollary 4

We first present several technical lemmas that will be used in our analysis and review existing results for the underdamped Langevin SDE. The proof of these lemmas will be deferred to Section C.

Our analysis for analyzing the convergence speed of the SGHMC1 algorithm and its comparison to the underdamped Langevin SDE is based on the 2-Wasserstein distance and this requires the L2L^{2} norm of the iterates to be finite. In the next lemma, we show that L2L^{2} norm of the both discrete and continuous dynamics are uniformly bounded over time with explicit constants. The main idea is to make use of the properties of the Lyapunov function V\mathcal{V} which is designed originally for the continuous-time process and show that the discrete dynamics can also be controlled by it.

For 0<η≤min⁡{γK2(d/β+A/β),γλ2K1,2γλ}0<\eta\leq\min\left\{\frac{\gamma}{K_{2}}(d/\beta+A/\beta),\frac{\gamma\lambda}{2K_{1}},\frac{2}{\gamma\lambda}\right\}, where

Since SGHMC1 is a discretization of the underdamped SDE (except that noise is also added to the gradients), we expect SGHMC1 to follow the underdamped SDE dynamics. It is natural to seek for bounds between the probability law μz,k\mu_{\mathbf{z},k} of the SGHMC1 algorithm at step kk with time step η\eta and that of the underdamped SDE at time t=kηt=k\eta which we denote by νz,kη\nu_{{\mathbf{z}},k\eta}. In our analysis, we first control the Kullback-Leibler (KL) divergence between these two, and then convert these bounds into bounds in terms of the 2-Wasserstein metric, applying an optimal transportation inequality by [BV05]. Note that Bolley and Villani theorem has also been successfully applied to analyzing the SGLD dynamics in [RRT17]. However, the analysis in [RRT17] does not directly apply to our setting as underdamped dynamics require a different Lyapunov function. This step requires an exponential integrability property of the underdamped SDE process which we establish next, before stating our result in Lemma 18 about the diffusion approximation of the SGHMC1 iterates.

We showed in the above Lemma 17 that the exponential moments grow at most linearly in time tt, which is a strict improvement from the exponential growth in time tt in [RRT17]. As a result, in the following Lemma 18 for the diffusion approximation, our upper bound is of the order (kη)3/2log⁡(kη)(δ1/4+η1/4)+kηη(k\eta)^{3/2}\sqrt{\log(k\eta)}(\delta^{1/4}+\eta^{1/4})+k\eta\sqrt{\eta} compared to kη(δ1/4+η1/4)k\eta(\delta^{1/4}+\eta^{1/4}) in [RRT17]. The method that is used in the proof of Lemma 17 for the underdamped dynamics can indeed be adapted to the case of the overdamped dynamics to improve the results in [RRT17].

where C0C_{0}, C1C_{1} and C2C_{2} are given by:

We consider the underdamped SDE and bound the 2-Wasserstein distance W2(νz,t,πz)\mathcal{W}_{2}(\nu_{z,t},\pi_{\mathbf{z}}) to the equilibrium for a fix arbitrary time t≥0t\geq 0. Crucial to the analysis is [EGZ19], which quantifies the convergence to equilibrium for underdamped Langevin diffusions. We first review the results from [EGZ19]. Let us recall from (2.1) the definition of the Lyapunov function V(x,v)\mathcal{V}(x,v):

Note that Hρ\mathcal{H}_{\rho} is a semi-metric, but not necessarily a metric. A simplified version of the main result from [EGZ19] which will be used in our setting is given below.

There exist constants α1,ϵ1∈(0,∞)\alpha_{1},\epsilon_{1}\in(0,\infty) and a continuous non-decreasing function h:[0,∞)→[0,∞)h:[0,\infty)\rightarrow[0,\infty) with h(0)=0h(0)=0 such that we have

We remark that the definitions of Λ,α1\Lambda,\alpha_{1} in (A.15) are coupled and there exists α1∈(0,∞)\alpha_{1}\in(0,\infty) so that Λ,α1\Lambda,\alpha_{1} in (A.15) are well defined; see Theorem 2.3. in [EGZ19]. In order to get explicit performance bounds, we also derive an upper bound for Hρ(μ0,πz)\mathcal{H}_{\rho}(\mu_{0},\pi_{\mathbf{z}}) in the next lemma. It is based on the (integrability properties) structure of the stationary distribution πz\pi_{\mathbf{z}} and the Lyapunov function V\mathcal{V} that controls the L2L^{2} norm of the initial distribution μ0\mu_{0}.

If parts (i)(i), (ii)(ii), (iii)(iii) and (iv)(iv) of Assumption 1 hold, then we have

A.2 Proof of Theorem 2

As the function FzF_{\mathbf{z}} satisfies the conditions in Lemma 26 in Section E with c1=Mc_{1}=M and c2=Bc_{2}=B (Lemma 25 in Section E), and the probability measures μk,z,πz\mu_{k,{\mathbf{z}}},\pi_{\mathbf{z}} have finite second moments (Lemma 16), we can apply Lemma 26 and deduce that

Here, one can obtain from Lemma 16 and Theorem 19 (convergence in 2-Wasserstein distance implies convergence of second moments) that

Then for any η\eta satisfying the condition in Lemma 16 and η≤(ε(log⁡(1/ε))3/2)4\eta\leq\left(\frac{\varepsilon}{(\log(1/\varepsilon))^{3/2}}\right)^{4}, we have

A.3 Proof of Corollary 4

With a slight abuse of notations, consider the random elements (X^,V^)(\hat{X},\hat{V}) and (X^∗,V^∗)(\hat{X}^{*},\hat{V}^{*}) with Law((X^,V^)∣Z=z)=μz,k\text{Law}((\hat{X},\hat{V})|\mathbf{Z}={\mathbf{z}})=\mu_{{\mathbf{z}},k} and Law((X^∗,V^∗)∣Z=z)=πz\text{Law}((\hat{X}^{*},\hat{V}^{*})|\mathbf{Z}={\mathbf{z}})=\pi_{\mathbf{z}}. Then we can decompose the expected population risk of X^\hat{X} (which has the same distribution as XkX_{k}) as follows:

The first term in (A.21) can be written as:

where PnP^{n} is the product measure of independent random variables Z1,…,ZnZ_{1},\ldots,Z_{n}. Then it follows from Theorem 2 and Lemma 20 that

Next, we bound the second and third terms in (A.21). Note that

Specifically, the second term in (A.21) can be bounded as

by applying Lemma 27, and the last term in (A.21) can be bounded as

where x∗x^{*} is any minimizer of F(x)F(x), i.e., F(x∗)=F∗F(x^{*})=F^{*}, and the last step is due to Lemma 28. The proof is complete.

Appendix B Proof of Theorem 6 and Corollary 8

The proof of Theorem 6 (Corollary 8) is similar to the proof of Theorem 2 (Corollary 4). There are two key new results that we need to establish: a uniform (in time) L2L^{2} bound for the SGHMC2 iterates (X^k,V^k)(\hat{X}_{k},\hat{V}_{k}), and the diffusion approximation that characterizes the 2-Wasserstein distance between the SGHMC2 iterates and the continuous-time underdampled Langevin diffusion. We summarize these two results in the following two lemmas and defer their proofs to Section D. With these two lemmas, Theorem 6 and Corollary 8 readily follow and we omit the proof details.

For 0<η≤min⁡{1,γK^2(d/β+A/β),γλ2K^1,2γλ}0<\eta\leq\min\left\{1,\frac{\gamma}{\hat{K}_{2}}(d/\beta+A/\beta),\frac{\gamma\lambda}{2\hat{K}_{1}},\frac{2}{\gamma\lambda}\right\}, where

where K1K_{1}, K2K_{2} are defined in (A.3) and (A.4), and

where CxdC_{x}^{d} and CvdC_{v}^{d} are defined in (A.5) and (A.6).

Next, let us provide a diffusion approximation between the SGHMC2 algorithm (X^k,V^k)(\hat{X}_{k},\hat{V}_{k}) and the continuous time underdamped diffusion process (X(kη),V(kη))(X(k\eta),V(k\eta)), and we use μ^z,k\hat{\mu}_{\mathbf{z},k} to denote the law of (X^k,V^k)(\hat{X}_{k},\hat{V}_{k}) and νz,k\nu_{\mathbf{z},k} to denote the law of (X(kη),V(kη))(X(k\eta),V(k\eta)).

where C0C_{0} is defined in (A.8) and C^1\hat{C}_{1} is given by:

where γ^\hat{\gamma} is defined in (A.11).

Appendix C Proofs of Lemmas in Section A

(i) We first prove the continuous–time case. The main idea is to use the following Lyapunov function (see (2.1)) introduced in [EGZ19] for the underdamped Langevin diffusion:

Lemma 1.3 in [EGZ19] showed that if the drift condition in (2.2) holds, then

where L\mathcal{L} is the infinitesimal generator of the underdamped Langevin diffusion (X,V)(X,V) defined in (1.5)–(1.6):

To show part (i), we first note that for λ≤14,\lambda\leq\frac{1}{4},

and we will provide an upper bound for L(t)L(t).

and hence ∫0teγλs(βV(s)+βγ2X(s))⋅2γβ−1B(s)\int_{0}^{t}e^{\gamma\lambda s}\left(\beta V(s)+\frac{\beta\gamma}{2}X(s)\right)\cdot\sqrt{2\gamma\beta^{-1}}B(s) is a martingale. Then we can infer from (C.1) and (C.5) that for any t≥0t\geq 0,

In combination with (C.1), we obtain that (X,V)(X,V) are uniformly (in time) L2L^{2} bounded. Indeed, we have

(ii) Next, we prove the uniform (in time) L2L^{2} bounds for (Xk,Vk)(X_{k},V_{k}). Let us recall the dynamics:

We show below that one can find explicit constants K1,K2>0K_{1},K_{2}>0, such that

We proceed in several steps in upper bounding L2(k+1)L_{2}(k+1).

First, by using the independence of Vk−η[γVk+gk(Xk,Uz,k)]V_{k}-\eta[\gamma V_{k}+g_{k}(X_{k},U_{z,k})] and ξk\xi_{k}, we can obtain from (C.8) that

where we have used part (iv) of Assumption 1 and Lemma 25 in Section E in the Appendix. By using ∣x∣≤∣x∣2+12|x|\leq\frac{|x|^{2}+1}{2}, we immediately get

where the last inequality is due to the M−M-smoothness of FzF_{\mathbf{z}}. This implies

where we have used part (iv) of Assumption 1 in the inequality above.

Combining the equations (C.11), (C.12), (C.13) and (C.14), we get

where we used the drift condition (2.2) in the last inequality, and

We can upper bound Ek\mathcal{E}_{k} as follows:

Since λ≤14\lambda\leq\frac{1}{4}, we obtain from (C.1) and (C.10) that

where we recall from (A.3) and (A.4) that

Moreover, since λ≤14\lambda\leq\frac{1}{4}, we infer from the definition of L2(k)L_{2}(k) in (C.10) that

Together with (C.15) and (C.17), we deduce that

For 0<η≤min⁡{γK2(d/β+A/β),γλ2K1}0<\eta\leq\min\left\{\frac{\gamma}{K_{2}}(d/\beta+A/\beta),\frac{\gamma\lambda}{2K_{1}}\right\}, we get

and we have ρ∈[0,1)\rho\in[0,1), where we used the assumption that η≤2γλ\eta\leq\frac{2}{\gamma\lambda}. It follows that

The result then follows from the inequality above and (C.16).

C.2 Proof of Lemma 17

From (C.1)–(C.3), we can directly obtain that

Since LeαV=[LαV+γβ−1∥∇vαV∥2]eαV\mathcal{L}e^{\alpha\mathcal{V}}=\left[\mathcal{L}\alpha\mathcal{V}+\gamma\beta^{-1}\|\nabla_{v}\alpha\mathcal{V}\|^{2}\right]e^{\alpha\mathcal{V}}, we have showed that

Applying an exponential integrability result, e.g. Corollary 2.4. in [CHJ13], we get

Next, applying Itô’s formula to e14αV(X(t),V(t))e^{\frac{1}{4}\alpha\mathcal{V}(X(t),V(t))}, we obtain

where we used (C.1) and (C.22). Thus, ∫0t12(βV(s)+βγ2X(s))e14αV(X(s),V(s))⋅dB(s)\int_{0}^{t}\frac{1}{2}\left(\beta V(s)+\frac{\beta\gamma}{2}X(s)\right)e^{\frac{1}{4}\alpha\mathcal{V}(X(s),V(s))}\cdot dB(s) is a martingale. By taking expectations on both hand sides of (C.23), we get

From (C.18), (C.19) and (C.20), we can infer that

where in the last inequality we used the facts that V≥0\mathcal{V}\geq 0 and 14αγ(d+A)−316αγλV≥0\frac{1}{4}\alpha\gamma(d+A)-\frac{3}{16}\alpha\gamma\lambda\mathcal{V}\geq 0 if and only if V≤4(d+A)3λ\mathcal{V}\leq\frac{4(d+A)}{3\lambda}. Therefore, it follows from (C.24) that

C.3 Proof of Lemma 18

The proof is inspired by the proof of Lemma 7 in [RRT17] although more delicate in our setting. Note that the main technical difficulty here is that the underdamped Langevin diffusion is a hypoelliptic diffusion, i.e. the diffusion matrix of the stochastic differential equation defining the multidimensional diffusion process is not of full rank, but its solutions admit a smooth density, see [DS19]. In our case, there is no Brownian noise in dX(t)dX(t) term in (1.6) and the underdamped Langevin diffusion (1.5)-(1.6) is hypoelliptic. Consider the following continuous-time interpolation of (Xk,Vk)(X_{k},V_{k}):

where we used part (ii) of Assumption 1 Cauchy-Schwarz inequality.

We can also bound the second term in (C.29):

where the first inequality follows from part (iv) of Assumption 1.

Finally, let us bound the third term in (C.29) as follows:

Hence, together with Lemma 16, we conclude that that

From the exponential integrability of the measure νz,kη\nu_{\mathbf{z},k\eta} in Lemma 17, we have

Note that η≤1\eta\leq 1 so that 2γ2ηCvd+(4+2δ)η(M2Cxd+B2)+2γβ−1≤(C2)22\gamma^{2}\eta C_{v}^{d}+(4+2\delta)\eta\left(M^{2}C_{x}^{d}+B^{2}\right)+2\gamma\beta^{-1}\leq(C_{2})^{2}, where C2C_{2} is defined in (A.10). Then, we have

By using (x+y)2≤2(x2+y2)(x+y)^{2}\leq 2(x^{2}+y^{2}), we get

where C0C_{0} and C1C_{1} are defined in (A.8) and (A.9). The result then follows from the fact that x+y≤x+y\sqrt{x+y}\leq\sqrt{x}+\sqrt{y} for non-negative real numbers xx and yy.

where we used the assumption η≤1\eta\leq 1 so that 2γ2ηCvd+(4+2δ)η(M2Cxd+B2)+2γβ−1≤(C2)22\gamma^{2}\eta C_{v}^{d}+(4+2\delta)\eta\left(M^{2}C_{x}^{d}+B^{2}\right)+2\gamma\beta^{-1}\leq(C_{2})^{2} in the last inequality above, where C2C_{2} is defined in (A.10). Therefore,

C.4 Proof of Lemma 20

Next, let us notice that by the concavity of the function hh, we have (see [EGZ19])

Moreover, by the definition of V\mathcal{V} in (2.1) and Lemma 25, we deduce that

It has been shown in [RRT17, Section 3.5] that

In addition, from the explicit expression of πz(dx,dv)\pi_{\mathbf{z}}(dx,dv) in (1.7), we have

Hence, the conclusion follows from (C.4).

Appendix D Proofs of Lemmas in Section B

Before we proceed to the proof of Lemma 21, let us state two technical lemmas, which will be used in the proof of Lemma 21. Recall ψ0(t)=e−γt\psi_{0}(t)=e^{-\gamma t} and ψk+1(t)=∫0tψk(s)ds\psi_{k+1}(t)=\int_{0}^{t}\psi_{k}(s)ds, and (ξk+1,ξk+1′)(\xi_{k+1},\xi^{\prime}_{k+1}) is a 2d2d-dimensional centered Gaussian vector from the SGHMC2 iterates (X^k,V^k)(\hat{X}_{k},\hat{V}_{k}) given in (1.10)–(1.11). Using the definitions, it is straightforward to establish these two lemmas, so we omit the details of their proofs.

Now, we are ready to prove Lemma 21, i.e. the uniform (in time) L2L^{2} bounds for (X^k,V^k)(\hat{X}_{k},\hat{V}_{k}) defined in (1.10)–(1.11). We can rewrite the dynamics of the SGHMC2 iterates as follows:

By following the proofs of the L2L_{2} uniform bound for SGHMC1 iterates, we get

where K1K_{1} and K2K_{2} are given in (A.3) and (A.4).

where we used the fact that c11=dc_{11}=d. Moreover,

By applying the assumption η≤1\eta\leq 1, we have

where the constants Q1,Q2,Q3Q_{1},Q_{2},Q_{3} are given in (B.3)–(B.5). Let us recall that for λ≤14\lambda\leq\frac{1}{4},

where K^1:=K1+4Q11−2λ+8Q2(1−2λ)γ2,\hat{K}_{1}:=K_{1}+\frac{4Q_{1}}{1-2\lambda}+\frac{8Q_{2}}{(1-2\lambda)\gamma^{2}}, and K^2:=K2+Q3,\hat{K}_{2}:=K_{2}+Q_{3}, we get

This implies L^2(k+1)≤ρL^2(k)+K\hat{L}_{2}(k+1)\leq\rho\hat{L}_{2}(k)+K, where ρ:=1−ηγλ/2∈[0,1)\rho:=1-\eta\gamma\lambda/2\in[0,1), where we used the assumption η≤2γλ\eta\leq\frac{2}{\gamma\lambda}, and K:=2ηγ(d/β+A/β).K:=2\eta\gamma(d/\beta+A/\beta). It follows that

The uniform L2L^{2} bounds then readily follow.

D.2 Proof of Lemma 22

We follow similar steps as in the proof of Lemma 7 in [RRT17]. Recall that with the same initialization, the SGHMC2 iterates (X^k,V^k)(\hat{X}_{k},\hat{V}_{k}) has the same distribution as (X^(kη),V^(kη))(\hat{X}(k\eta),\hat{V}(k\eta)) where (X^(⋅),V^(⋅))(\hat{X}(\cdot),\hat{V}(\cdot)) is a continuous-time process satisfying

We first bound the first term in (D.18). Before we proceed, let us notice that for any kη≤s<(k+1)ηk\eta\leq s<(k+1)\eta,

where we used (D.14), the assumption η≤1\eta\leq 1 and Lemma 21.

We can also bound the second term in (D.18):

where the first inequality follows from part (iv) of Assumption 1, and we also used Lemma 21. Hence, we conclude that

To complete the proof, we can follow similar steps as in the proof of Lemma 18. By using the estimate in (D.20), the result from [BV05], and the exponential integrability of the measure νz,kη\nu_{\mathbf{z},k\eta} in Lemma 17, we can infer that

where C0C_{0} and C^1\hat{C}_{1} are defined in (A.8) and (B.8). The result then follows from the fact that x+y≤x+y\sqrt{x+y}\leq\sqrt{x}+\sqrt{y} for non-negative real numbers xx and yy.

Appendix E Supporting Lemmas

In this section, we present several supporting lemmas from the existing literature. These lemmas are used in our proofs, so we include them here for the sake of completeness. The first lemma shows that ff admits lower and upper bounds that are quadratic functions.

The next lemma shows a 2-Wasserstein continuity result for functions of quadratic growth. This lemma was also used in [RRT17] to study the SGLD dynamics.

for some constants c1>0c_{1}>0 and c2≥0c_{2}\geq 0. Then,

The next lemma shows a uniform stability of πz\pi_{\mathbf{z}}. Note that the x−x-marginal of πz(dx,dv)\pi_{\mathbf{z}}(dx,dv) for the underdamped diffusion is the same as the stationary distribution for the overdamped diffusion studied in [RRT17]. For two n−n-tuples z=(z1,…,zn),z‾=(z‾1,…,z‾n)∈Zn\mathbf{z}=(z_{1},\ldots,z_{n}),\overline{\mathbf{z}}=(\overline{z}_{1},\ldots,\overline{z}_{n})\in\mathcal{Z}^{n}, we say z\mathbf{z} and z‾\overline{\mathbf{z}} differ only in a single coordinate if card∣{i:zi≠z‾i}∣=1|\{i:z_{i}\neq\overline{z}_{i}\}|=1.

For any two z,z‾∈Zn{\mathbf{z}},\overline{\mathbf{z}}\in\mathcal{Z}^{n} that differ only in a single coordinate,

where λ∗\lambda_{\ast} is the uniform spectral gap for overdamped Langevin dynamics:

The next lemma show that for large values of β\beta, the x−x-marginal of the stationary distribution πz(dx,dv)\pi_{\mathbf{z}}(dx,dv) is concentrated at the minimizer of FzF_{\mathbf{z}}. Note in Proposition 11 of [RRT17], they have the assumption β≥2/m\beta\geq 2/m, which seems to be only used to derive their Lemma 4, but not used in deriving their Proposition 11.

Appendix F Proof of Proposition 11

Let us first prove that λ∗=O(a−2)\lambda_{\ast}=\mathcal{O}(a^{-2}). We first recall that λ∗\lambda_{\ast} is the uniform spectral gap for overdamped Langevin dynamics:

with m=m1a−2m=m_{1}a^{-2}, M=M1a−2M=M_{1}a^{-2}, and B=B1a−1B=B_{1}a^{-1}.

Next, let us take the test function g1(x):=∥x∥2g_{1}(x):=\|x\|^{2}. And we further define

Next, by the definition of c1c_{1} in (F.2) and the bounds in (F.1), we get

where we used m=m1a−2m=m_{1}a^{-2}, M=M1a−2M=M_{1}a^{-2}, B=B1a−1B=B_{1}a^{-1} and g1(ax)=a2g1(x)=a2∥x∥2g_{1}(ax)=a^{2}g_{1}(x)=a^{2}\|x\|^{2}. Hence, we conclude that λ∗=O(a−2)\lambda_{\ast}=\mathcal{O}(a^{-2}).

Next, let us prove that μ∗=Θ(a−1)\mu_{\ast}=\Theta(a^{-1}). We recall that μ∗\mu_{*} the convergence rate for underdamped Langevin dynamics is given by:

where λ,A\lambda,A come from the drift condition (2.2), and from [GGZ20], we can take

Note that μ∗\mu_{\ast} depends on the objective function FzF_{\mathbf{z}} only via the parameters from its properties, which is independent of z\mathbf{z}. Recall that m=m1a−2m=m_{1}a^{-2}, M=M1a−2M=M_{1}a^{-2}, B=B1a−1B=B_{1}a^{-1}. We define γ=:γ1a−1\gamma=:\gamma_{1}a^{-1} so that γ1\gamma_{1} is independent of aa and

where we can check that λ\lambda, Λ\Lambda are independent of aa. Then, we can see from (F.4) that μ∗\mu_{\ast} is linear in a−1a^{-1} so that we have μ∗=Θ(a−1)\mu_{\ast}=\Theta(a^{-1}). The proof is complete.

Appendix G Explicit dependence of constants on key parameters

We recall the constants from Table 1. It is easy to see that

In addition, in view of (G.1), it follows that

The structure of the initial distribution μ0(dx,dv)\mu_{0}(dx,dv) would affect the overall dependence on β,d\beta,d. Since we assumed in Section 5 that μ0(dx,dv)\mu_{0}(dx,dv) is supported on a Euclidean ball with radius being a universal constant, then the Lyapunov function V\mathcal{V} in (2.1) is linear in β\beta. We can then obtain

Moreover, by the definition of C^1\hat{C}_{1} in (B.8), we get

Appendix H Proof of Lemma 13 and Lemma 15

Since the distribution of AinA_{in} has compact support, we have ∥ai∥≤D\|a_{i}\|\leq D for some D>0D>0. Let si:=⟨ai,x⟩s_{i}:=\langle a_{i},x\rangle. By taking the gradient of f(x,zi)f(x,z_{i}) with respect to xx, we obtain

where we used the triangle inequality and the Cauchy-Schwartz inequality. Then, it is straightforward to check that we obtain ⟨∇f(x,zi),x⟩≥m∥x∥2−b\langle\nabla f(x,z_{i}),x\rangle\geq m\|x\|^{2}-b for

and therefore part (iii) of Assumption 1 holds. Also for any z=(a,y)z=(a,y), ∣f(0,z)∣=∣(y−σ(0))2∣≤A0|f(0,z)|=|(y-\sigma(0))^{2}|\leq A_{0} for A0=(1+∥σ(0)∥)2A_{0}=(1+\|\sigma(0)\|)^{2}. Similarly, ∥∇f(0,z)∥=∥−2(y−σ(0))σ′(0)a∥≤B1\|\nabla f(0,z)\|=\|-2(y-\sigma(0))\sigma^{\prime}(0)a\|\leq B_{1} for

Therefore, part (i) of Assumption 1 holds for any B≥B1B\geq B_{1}. We also have the Hessian matrix

where IdI_{d} is the d×dd\times d identity matrix. Hence, ∥∇2f(x,zi)∥≤M1\|\nabla^{2}f(x,z_{i})\|\leq M_{1} where

where we used Cauchy-Schwarz inequality. This implies

for any δ∈[14nb,1)\delta\in[\frac{1}{4n_{b}},1), M≥M2:=4λrM\geq M_{2}:=4\lambda_{r}, B≥B2:=4M2B\geq B_{2}:=4M_{2} where we used (H.5) and the fact that uju_{j} are i.i.d. with mean zero. If we choose, for instance, M=M1+M2M=M_{1}+M_{2}, B=max⁡(B1,B2)=B2B=\max(B_{1},B_{2})=B_{2}; we observe that part (i) and (iv) of Assumptions 1 hold.

H.2 Proof of Lemma 15

We set ri=yi−⟨ai,x⟩r_{i}=y_{i}-\langle a_{i},x\rangle and follow a similar approach to the proof of Lemma 13.

Therefore, part (iii) of Assumption 1 holds. We have also

for any ziz_{i}. Therefore, part (i) of Assumption 1 holds with A0=∥ρ∥∞A_{0}=\|\rho\|_{\infty} and B=∥ρ′∥∞DB=\|\rho^{\prime}\|_{\infty}D. Since

where IdI_{d} is the d×dd\times d identity matrix, we also have

Therefore, part (ii) of Assumption 1 holds for any M≥∥ρ′′∥∞D2+λrM\geq\|\rho^{\prime\prime}\|_{\infty}D^{2}+\lambda_{r}. We have also

where we used Cauchy-Schwarz inequality. This implies

for any δ∈[14nb,1)\delta\in[\frac{1}{4n_{b}},1) and M2≥4λrM_{2}\geq 4\lambda_{r} and B≥4∥ρ′∥∞DB\geq 4\|\rho^{\prime}\|_{\infty}D where we used (H.8) and the fact that vjv_{j} are i.i.d. with mean zero. We conclude that Assumption 1 work for M=∥ρ′′∥∞D2+5λrM=\|\rho^{\prime\prime}\|_{\infty}D^{2}+5\lambda_{r} and B=4∥ρ′∥∞DB=4\|\rho^{\prime}\|_{\infty}D.