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 is unknown, a common popular approach is to consider the empirical risk minimization problem
based on the dataset 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 .We note that in our notation is a random vector, whereas is deterministic vector associated to a dataset that corresponds to a realization of the random vector . 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 is a similar goal to computing an approximate global minimizer of . 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 is the normalizing constant:
Hence, the -marginal distribution of stationary distribution is exactly the invariant distribution of the overdamped Langevin diffusion.With slight abuse of notation, we use to denote the -marginal of the equilibrium distribution . 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 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 by and , . The iterates are then defined by the following recursion:
where is a -dimensional centered Gaussian vector so that ’s are independent and identically distributed (i.i.d.) and independent of the initial condition, and for any fixed , the random vectors , , 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 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 and the noise in the gradients, but we do not assume 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 and 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 in dimension one and estimate the spectral gap of HMC in the regime as . This is a mixture of two Gaussians with the same variance centered at and 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 as for fixed, Proposition 11 is applicable and it implies that HMC will be more efficient than overdamped Langevin dynamics in terms of dependency to (which measures the distance between the modes) in the sense that the mixing time will be in HMC whereas it will be in Random Walk. This does not contradict results of [MPS18] because we consider different scaling regimes: We fix and let whereas [MPS18] fix and let .
There are also some connections of our work to existing momentum-based optimization algorithms. More specifically, if the term with involving the Brownian noise is removed in the underdamped SDE (1.5)–(1.6), this results in a second-order ODE in . 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 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 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 is continuously differentiable, takes non-negative real values, and there exist constants so that
For each , the function is -smooth:
For each , the function is -dissipative:
There exists a constant such that for every :
The probability law of the initial state satisfies:
where is a Lyapunov function to be used repeatedly for the rest of the paper:
and is the friction coefficient as in (1.5), is a positive constant less than , and .
We note that the Lyapunov function 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 so that
This drift condition, which will be used later, guarantees the stability and the existence of Lyapunov function for the underdamped Langevin diffusion in (1.5)–(1.6), see [EGZ19].
Main Results for SGHMC1 Algorithm
Our first result shows SGHMC1 iterates in (1.8)–(1.9) track the underdamped Langevin SDE in the sense that the expectation of the empirical risk with respect to the probability law of conditional on the sample , denoted by , and the stationary distribution of the underdamped SDE is small when is large enough. The difference in expectations decomposes as a sum of two terms and while the former term quantifies the dependency on the initialization and the dataset whereas the latter term is controlled by the discretization error and the amount of noise in the gradients which depends on the parameter . We also note that the parameter (see Table 1) in our bounds governs the speed of convergence to the equilibrium of the underdamped Langevin diffusion.
Consider the SGHMC1 iterates defined by the recursion (1.8)–(1.9) from the initial state which has the law . If Assumption 1 is satisfied, then for , we have
with defined by (A.20) provided that
Here 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 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 marginal of the equilibrium distribution concentrates around a global minimizer of the empirical risk. As goes to infinity, it can be verified that this quantity goes to zero. For finite , [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 , and are defined by , and respectively.
2 Performance bound for the population risk minimization
By exploiting the fact that the 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 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 (the iterates in (1.9)) is bounded by
where is defined by (A.20), is defined by (A.18), and are defined by (3.2) and (3.5) respectively and is a constant satisfying
and is the uniform spectral gap for overdamped Langevin dynamics In [RRT17], their formula for missed factor.:
3 Generalization error of SGHMC1 in the one pass regime
where is the Gibbs output, i.e. its distribution conditional on is given by . 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 is the output of the underdamped Langevin dynamics, i.e. its distribution conditional on is given by and 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 defined in (1.10)-(1.11), and denote the probability law of conditional on the sample by . 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 is controlled by the accuracy of the discretization and has to be replaced by another term , 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 defined by the recursion (1.10)–(1.11) from the initial state which has the law . If Assumption 1 is satisfied, then for , we have
where is defined in (3.1) and
with defined by (A.20) provided that
Here 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 , and are defined by , and 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 (the iterates in (1.11)) is bounded by
where , , , 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 is the output of the underdamped Langevin dynamics, i.e. its distribution conditional on is given by and 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 (see (3.8)) and (see Table 1) are exponentially small in both and in the worst case, but under some extra assumptions the dependency on can be polynomial (see e.g. [CCBJ18]) although the exponential dependence to is unavoidable in the presence of multiple minima in general (see [BGK05]). One can readily see that has better dependency on than , 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 is small, and if we ignore the factor We emphasize that the effect of the last term appearing in (5.4) is typically negligible compared to other parameters. For instance even if is double-exponentially small, we have ., 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 and do not have the same dependency to up to factors (the former scales with as and the latter ), and this improvement on 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 term in (5.6) to .
To make the comparison to SGLD simpler, we notice that in both expressions (5.5) and (5.6), we see a term scaling with due to the gradient noise level ( is fixed in the one-pass setting), and we fix the error in (5.5) and (5.6) without the term to be the same order, and then compare the number of iterations and . More precisely, given and we choose such that in (5.5) so that the generalization error for SGHMC2 is
Similarly, the generalization error for SGLD is
When and are on the same order or is larger, since typically , the term involving 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 in (5.7) and in (5.8).
Empirical risk minimization.
The empirical risk minimization bound given in Corollary 7 has an additional term compared to the and terms appearing in the one-pass generalization bounds. Note also that . 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 factors in the parameters and hiding factors in that can be made explicit. To see (5.12), we infer from (5.9) that for fixed precision and dimension , by ignoring the log factors and , we can choose so that and choose the gradient noise level so that . So the number of SGHMC2 iterations is
On the other hand, the mini-batch size to achieve gradient noise level is given by (see [RRT17]), which is equal to . Hence, we obtain (5.12) which is the product of the mini-batch size and number of iterations.
It is hard to compare in (5.11) and in (5.12) in general since 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 will be similar to , which is the spectral gap of the continuous-time overdamped Langevin diffusion. As a consequence, when the stepsize is small enough (which is the case for instance, when target accuracy is small enough), we will have and 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 ( vs. ), however and 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 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 and , for instance this may happen if the number of samples is not large enough, then the performance guarantees for population risk minimization via SGLD and SGHMC2 will be similar. Otherwise, if is large and is chosen in a way to keep the term on the order , then similar improvement can be achieved.
The parameters (see (3.8)) and (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 and in . 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 , 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 . For these examples if the other parameters 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 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 . For simplicity, we assume there is only one sample, i.e. and . We consider the non-convex optimization problem (1.2) with both the SGHMC2 algorithm and the SGLD algorithm. [EGZ19] showed that for this example whereas making the constants hidden by the explicit. This shows that the contraction rate of the underdamped diffusion is (faster) larger than that of the overdamped diffusion by a square root factor when 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 whereas it follows from Theorem 1.2 in [BGK05] that . This shows that when the separation between minima, or alternatively the scaling factor is large enough, is larger than 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 is on the order of the square root of . The proof will be presented in details in Section F.
Suppose that the functions indexed by satisfies Assumption 1 (i)-(iii) with , and for some fixed constants , , and . Then, we have as
This result is more general than the previous example. In particular, if satisfies Assumption 1 (i)-(iii) with replaced by , then satisfies Assumption 1 (i)-(iii) with , and . Proposition 11 essentially says that if we consider the normalized empirical risk objective where is a (normalization) scaling parameter and satisfies Assumption 1, then for large enough values of , the empirical risk surface will be relatively flat and the convergence rate of momentum variant SGHMC2 to an -neighborhood of the global minimum will be governed by the parameter which will be larger than that of the parameter of SGLD when 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 is a regularization parameter that may depend on the number of samples . By Lagrangian duality, this problem is equivalent to the constrained optimization problem
for some , which has also been considered in the literature (see e.g. [MBM18, FSS18, WCX19]). For non-convex , 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 data points sampled with replacement, i.e. the gradient is estimated as
where are i.i.d. with a uniform distribution over the set of indices . We also consider the following assumption for the threshold function which are satisfied by many choices of in practice. A prominent example is the logistic (or sigmoid) function in which case which is also used in deep learning. Another possible choice is the probit function which corresponds to where 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 and 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 of the initial state has compact support. If Assumption 12 holds; then Assumption 1 hold for any with the following constants:
stochastic gradient evaluations to converge to an neighborhood of an almost ERM ignoring the factors in the parameters 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]: , where is a tuning parameter. In the following, similar to [WCX19], we assume that the data is bounded and the threshold function and its derivatives up to order two are bounded, similar to [MBM18]. This assumption for 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 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 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 defined in (2.1) that also depends on the objective as opposed to the classic Lyapunov function 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) 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 . 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 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 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 -marginal of the equilibrium distribution of the underdamped Langevin SDE concentrates around the global minimizers of for 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 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 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 -global minimizers under some regularity assumption on the non-convex objective . 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 we can obtain performance bound independent of dimension or has polynomial dependence on . 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 norm of the iterates to be finite. In the next lemma, we show that 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 which is designed originally for the continuous-time process and show that the discrete dynamics can also be controlled by it.
For , 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 of the SGHMC1 algorithm at step with time step and that of the underdamped SDE at time which we denote by . 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 , which is a strict improvement from the exponential growth in time in [RRT17]. As a result, in the following Lemma 18 for the diffusion approximation, our upper bound is of the order compared to 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 , and are given by:
We consider the underdamped SDE and bound the 2-Wasserstein distance to the equilibrium for a fix arbitrary time . 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 :
Note that 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 and a continuous non-decreasing function with such that we have
We remark that the definitions of in (A.15) are coupled and there exists so that 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 in the next lemma. It is based on the (integrability properties) structure of the stationary distribution and the Lyapunov function that controls the norm of the initial distribution .
If parts , , and of Assumption 1 hold, then we have
A.2 Proof of Theorem 2
As the function satisfies the conditions in Lemma 26 in Section E with and (Lemma 25 in Section E), and the probability measures 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 satisfying the condition in Lemma 16 and , we have
A.3 Proof of Corollary 4
With a slight abuse of notations, consider the random elements and with and . Then we can decompose the expected population risk of (which has the same distribution as ) as follows:
The first term in (A.21) can be written as:
where is the product measure of independent random variables . 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 is any minimizer of , i.e., , 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) bound for the SGHMC2 iterates , 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 , where
where , are defined in (A.3) and (A.4), and
where and are defined in (A.5) and (A.6).
Next, let us provide a diffusion approximation between the SGHMC2 algorithm and the continuous time underdamped diffusion process , and we use to denote the law of and to denote the law of .
where is defined in (A.8) and is given by:
where 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 is the infinitesimal generator of the underdamped Langevin diffusion defined in (1.5)–(1.6):
To show part (i), we first note that for
and we will provide an upper bound for .
and hence is a martingale. Then we can infer from (C.1) and (C.5) that for any ,
In combination with (C.1), we obtain that are uniformly (in time) bounded. Indeed, we have
(ii) Next, we prove the uniform (in time) bounds for . Let us recall the dynamics:
We show below that one can find explicit constants , such that
We proceed in several steps in upper bounding .
First, by using the independence of and , 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 , we immediately get
where the last inequality is due to the smoothness of . 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 as follows:
Since , we obtain from (C.1) and (C.10) that
where we recall from (A.3) and (A.4) that
Moreover, since , we infer from the definition of in (C.10) that
Together with (C.15) and (C.17), we deduce that
For , we get
and we have , where we used the assumption that . 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 , we have showed that
Applying an exponential integrability result, e.g. Corollary 2.4. in [CHJ13], we get
Next, applying Itô’s formula to , we obtain
where we used (C.1) and (C.22). Thus, 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 and if and only if . 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 term in (1.6) and the underdamped Langevin diffusion (1.5)-(1.6) is hypoelliptic. Consider the following continuous-time interpolation of :
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 in Lemma 17, we have
Note that so that , where is defined in (A.10). Then, we have
By using , we get
where and are defined in (A.8) and (A.9). The result then follows from the fact that for non-negative real numbers and .
where we used the assumption so that in the last inequality above, where is defined in (A.10). Therefore,
C.4 Proof of Lemma 20
Next, let us notice that by the concavity of the function , we have (see [EGZ19])
Moreover, by the definition of 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 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 and , and is a -dimensional centered Gaussian vector from the SGHMC2 iterates 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) bounds for defined in (1.10)–(1.11). We can rewrite the dynamics of the SGHMC2 iterates as follows:
By following the proofs of the uniform bound for SGHMC1 iterates, we get
where and are given in (A.3) and (A.4).
where we used the fact that . Moreover,
By applying the assumption , we have
where the constants are given in (B.3)–(B.5). Let us recall that for ,
where and we get
This implies , where , where we used the assumption , and It follows that
The uniform 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 has the same distribution as where is a continuous-time process satisfying
We first bound the first term in (D.18). Before we proceed, let us notice that for any ,
where we used (D.14), the assumption 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 in Lemma 17, we can infer that
where and are defined in (A.8) and (B.8). The result then follows from the fact that for non-negative real numbers and .
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 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 and . Then,
The next lemma shows a uniform stability of . Note that the marginal of for the underdamped diffusion is the same as the stationary distribution for the overdamped diffusion studied in [RRT17]. For two tuples , we say and differ only in a single coordinate if card.
For any two that differ only in a single coordinate,
where is the uniform spectral gap for overdamped Langevin dynamics:
The next lemma show that for large values of , the marginal of the stationary distribution is concentrated at the minimizer of . Note in Proposition 11 of [RRT17], they have the assumption , 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 . We first recall that is the uniform spectral gap for overdamped Langevin dynamics:
with , , and .
Next, let us take the test function . And we further define
Next, by the definition of in (F.2) and the bounds in (F.1), we get
where we used , , and . Hence, we conclude that .
Next, let us prove that . We recall that the convergence rate for underdamped Langevin dynamics is given by:
where come from the drift condition (2.2), and from [GGZ20], we can take
Note that depends on the objective function only via the parameters from its properties, which is independent of . Recall that , , . We define so that is independent of and
where we can check that , are independent of . Then, we can see from (F.4) that is linear in so that we have . 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 would affect the overall dependence on . Since we assumed in Section 5 that is supported on a Euclidean ball with radius being a universal constant, then the Lyapunov function in (2.1) is linear in . We can then obtain
Moreover, by the definition of in (B.8), we get
Appendix H Proof of Lemma 13 and Lemma 15
Since the distribution of has compact support, we have for some . Let . By taking the gradient of with respect to , we obtain
where we used the triangle inequality and the Cauchy-Schwartz inequality. Then, it is straightforward to check that we obtain for
and therefore part (iii) of Assumption 1 holds. Also for any , for . Similarly, for
Therefore, part (i) of Assumption 1 holds for any . We also have the Hessian matrix
where is the identity matrix. Hence, where
where we used Cauchy-Schwarz inequality. This implies
for any , , where we used (H.5) and the fact that are i.i.d. with mean zero. If we choose, for instance, , ; we observe that part (i) and (iv) of Assumptions 1 hold.
H.2 Proof of Lemma 15
We set and follow a similar approach to the proof of Lemma 13.
Therefore, part (iii) of Assumption 1 holds. We have also
for any . Therefore, part (i) of Assumption 1 holds with and . Since
where is the identity matrix, we also have
Therefore, part (ii) of Assumption 1 holds for any . We have also
where we used Cauchy-Schwarz inequality. This implies
for any and and where we used (H.8) and the fact that are i.i.d. with mean zero. We conclude that Assumption 1 work for and .