On Nonconvex Optimization for Machine Learning: Gradients, Stochasticity, and Saddle Points

Introduction

One of the principal discoveries in machine learning in recent years is an empirical one—that simple algorithms often suffice to solve difficult real-world learning problems. Machine learning algorithms generally arise via formulations as optimization problems, and, despite a massive classical toolbox of sophisticated optimization algorithms and a major modern effort to further develop that toolbox, the simplest algorithms—gradient descent, which dates to the 1840s (Cauchy, 1847) and stochastic gradient descent, which dates to the 1950s (Robbins and Monro, 1951)—reign supreme in machine learning.

This empirical discovery is appealing in many ways. First, at the scale of modern machine learning applications—often involving many millions of data points and millions of parameters—complex algorithms are generally infeasible, so that the only hope is that simple algorithms might be not only feasible but successful. Second, simple algorithms are easier to implement, debug, and maintain. Third, as the field of machine learning transforms into a real-world engineering discipline, it will be necessary to develop solid theoretical foundations for entire systems that employ machine learning algorithms at their core, and such an effort seems less daunting if the basic ingredients are simple.

These developments were presaged and supported by optimization researchers such as Nemirovskii, Nesterov, and Polyak who, from the 1960s until the present day, have pursued an in-depth study of first-order, gradient-based algorithms, developing novel algorithms and accompanying theory (Nemirovskii and Yudin, 1983; Nesterov, 1998; Polyak, 1963). Their results have included lower bounds and algorithms that achieve those lower bounds. This line of work has made clear that even simple algorithms require delicate theoretical treatment when they are studied in large-scale settings. Thus, much of the focus has been on the setting of convex optimization where many of the complexities have been stripped away. This has allowed the development of an elegant theory, and has provided a solid jumping-off point for further analysis that has brought additional computational constraints into play—including distributed platforms, fault tolerance, communication bottlenecks, and asynchronous computation (Recht et al., 2011; Zhang et al., 2012; Smith et al., 2018).

The most notable machine-learning success stories, however, have generally involved nonconvex optimization formulations, and a gap has arisen between theory and practice. Attempts to fill this gap include Choromanska et al. (2014) in the setting of learning multi-layer neural networks, Bandeira et al. (2016); Mei et al. (2017) for synchronization and MaxCut, Boumal et al. (2016) for smooth semidefinite programs, Bhojanapalli et al. (2016) for matrix sensing, Ge et al. (2016) for matrix completion, and Ge et al. (2017) for robust principal component analysis. But there remains a need to develop general theory that relates the convergence of machine learning algorithms to geometry and dynamics.

Most of the theory in the optimization literature has focused on the relationship between the number of iterations of the algorithm and a suitable notion of accuracy. Dimension is often neglected in such analyses, in part because in the convex setting even the simplest algorithms, including gradient descent, are provably independent of dimension. In developing algorithmic theory for nonconvex optimization formulations of machine learning problems, however, it is critically important to study iteration complexity as a function of of dimension, which can be in the millions. Moreover, we cannot resort to asymptotics—we are interested in problems at all scales.

In the current paper we have three goals. The first is to show that significant progress has been made in recent years in the theoretical analysis of algorithms for nonconvex machine learning. The second is to extend that line of analysis to handle both stochastic and non-stochastic algorithms in a single framework. In both cases we upper bound the iteration complexity as a function of both accuracy and dimension. The third is to exhibit a simple proof that exposes the core of the phenomenon that determines the dimension dependence.

Nonconvex optimization problems are intractable in general. Progress has been in machine learning by noting that in many problems the principal difficulty is not local minima, either because there are no spurious local minima (we review a list of such problems in Section 2) or because empirical work has shown that the local minima that are found by local gradient-based algorithms tend to be effective in terms of the ultimate goal of machine learning, which is performance on a test set. The problem then becomes one of avoiding saddle points, which are ubiquitous in machine learning architectures. Saddle points slow down gradient-based algorithms and in millions of dimensions they are potentially a major bottleneck for such algorithms. The theoretical problem becomes that of characterizing the iteration complexity of avoiding saddle points, as a function of target accuracy and dimension.

The current paper provides a positive answer to this question. We shows that suitably-perturbed versions of gradient descent and stochastic gradient descent escape saddle points in a number of iterations that is only polylogarithmic in dimension. More technically, defining a notion of $\epsilon$-second-order stationarity (see Section 2), which rules out saddle points, to be contrasted with classical $\epsilon$-first-order stationarity, which simply means near vanishing of the gradient, and which therefore does not rule out saddle points, we show that:

In this section we discuss related work on convergence guarantees for finding second-order stationary points. Some key comparisons are summarized in Table 1, and an augmented table is provided in Appendix A.

Classical approaches to finding second-order stationary points assume access to second-order information, in particular the Hessian matrix of second derivatives. Examples of such approaches include the cubic regularization method (Nesterov and Polyak, 2006) and trust-region methods (Curtis et al., 2014), both of which require $\mathcal{O}(\epsilon^{-1.5})$ queries of gradients and Hessians. This favorable convergence rate is, however, obtained at a high cost per iteration, owing to the fact that Hessian matrices scale quadratically with respect to dimension. In practice researchers have turned to first-order methods, which only utilize gradients and are therefore are substantially cheaper per iteration.

Given the preference among practitioners for simple, single-loop algorithms, and the striking empirical successes obtained with such algorithms, it is important to pin down the theoretical properties of such algorithms. In such analyses, the results for second-order and Hessian-vector algorithms serve as baselines. Of key interest is the convergence rate not merely as a function of the accuracy $\epsilon$ but also as a function of the dimension $d$. Indeed, while second-order algorithms can use the structure of the Hessian to readily avoid unhelpful directions even in high-dimensional spaces. Without the Hessian there is a concern that algorithms may scale poorly as a function of dimension.

Stochastic setting with Lipschitz gradient.

General stochastic setting.

There is significantly less work in the general setting in which the stochastic gradients are no longer guaranteed to be Lipschitz. In fact, only the results of Ge et al. (2015) and Daneshmand et al. (2018) apply here, and both of them require at least $\Omega(d^{4})$ gradient queries to find second-order stationary points. The current paper brings this dependence down to linear dimension dependence. See the last three lines in Table 1 for a summary of the results in the stochastic case.

Other settings.

Finally, there are also several recent results in the setting in which objective functions can be written as a finite sum of individual functions. We refer readers to Reddi et al. (2017); Allen-Zhu and Li (2017) and Lei et al. (2018) and the references therein for further reading.

2 Organization

In Section 2, we review some algorithmic and mathematical preliminaries. Section 3 presents several examples of nonconvex problems in machine learning, demonstrating how second-order stationarity can ensure approximate global optimality. In Section 4, we present the algorithms that we analyze and present our main theoretical results for perturbed GD and SGD. In Section 5, we present the proof for the non-stochastic case (perturbed GD), which illustrates some of our key ideas. The proof for the stochastic setting is presented in the Appendix. We conclude in Section 6.

Background

In this section, we introduce our notation and present definitions and assumptions. We also overview existing results in nonconvex optimization, in both the deterministic and stochastic settings.

2 Nonconvex optimization and gradient descent

In this paper, we are interested in solving general unconstrained optimization problems of the form:

where $f$ is a smooth function which can be nonconvex. In particular we assume that $f$ has Lipschitz gradients and Lipschitz Hessians, which ensures that the gradient and Hessian can not change too rapidly.

A twice-differentiable function $f$ is $\rho$-Hessian Lipschitz if:

Our point of departure is the classical Gradient Descent (GD) algorithm, whose update takes following form:

where $\eta>0$ is a step size or learning rate. Since the problem of finding a global optimum for general nonconvex functions is NP-hard, the classical literature in optimization has resorted to a local surrogate—first-order stationarity.

For a differentiable function $f$, $\mathbf{x}$ is a first-order stationary point if $\nabla f(\mathbf{x})=\mathbf{0}$.

For a differentiable function $f$, $\mathbf{x}$ is an $\epsilon$-first-order stationary point if $\|{\nabla f(\mathbf{x})}\|\leq\epsilon$.

It is of major importance that gradient descent converges to a first-order stationary point in a number of iterations that is independent of dimension. This fact, referred to as “dimension-free convergence” in the optimization literature, is captured in the following classical theorem.

Note that in this formulation, the last iterate is not guaranteed to be a stationary point. However, it is not hard to figure out which iterate is the stationary point by calculating the norm of the gradient at every iteration.

A first-order stationary point can be a local minimum, a local maximum or even a saddle point:

For a differentiable function $f$, a stationary point $\mathbf{x}$ is a

local minimum, if there exists $\delta>0$ such that $f(\mathbf{x})\leq f(\mathbf{y})$ for any $\mathbf{y}$ with $\|{\mathbf{y}-\mathbf{x}}\|\leq\delta$.

local maximum, if there exists $\delta>0$ such that $f(\mathbf{x})\geq f(\mathbf{y})$ for any $\mathbf{y}$ with $\|{\mathbf{y}-\mathbf{x}}\|\leq\delta$.

For minimization problems, both saddle points and local maxima are clearly undesirable. Our focus will be “saddle points,” although our results also apply directly to local maxima as well. Unfortunately, distinguishing saddle points from local minima for smooth functions is still NP-hard in general (Nesterov, 2000). To avoid these hardness results, we focus on a subclass of saddle points.

For a twice-differentiable function $f$, $\mathbf{x}$ is a strict saddle point if $\mathbf{x}$ is a stationary point and $\lambda_{\min}(\nabla^{2}f(\mathbf{x}))<0$.

A generic saddle point must satisfy that $\lambda_{\min}(\nabla^{2}f(\mathbf{x}))\leq 0$. Being “strict” simply rules out the case where $\lambda_{\min}(\nabla^{2}f(\mathbf{x}))=0$. We reformulate our goal as that of finding stationary points that are not strict saddle points.

For twice-differentiable function $f(\cdot)$, $\mathbf{x}$ is a second-order stationary point if

For a $\rho$-Hessian Lipschitz function $f(\cdot)$, $\mathbf{x}$ is an $\epsilon$-second-order stationary point if:

Our definition again makes use of an $\epsilon$-ball around the stationary point so that we can discuss rates, and the condition on the Hessian in Definition 4 uses the Hessian Lipschitz parameter $\rho$ to retain a single accuracy parameter and to match the units of the gradient and Hessian, following the convention of Nesterov and Polyak (2006).

Although second-order stationarity is only a necessary condition for being a local minimum, a line of recent work in the machine learning literature shows that for many popular models in machine learning, all $\epsilon$-second-order stationary points are approximate global minima. Thus for these models finding second-order stationary points is sufficient for solving those problems. See Section 3 for references and discussion of these results.

3 Stochastic approximation

Other than being an unbiased estimator of true gradient, another standard assumption on the stochastic gradients is that their variance is bounded by some number $\sigma^{2}$:

When we are interested in high-probability bounds, we make the following stronger assumption on the tail of the distribution.

Prior work shows that stochastic gradient descent converges to first-order stationary points in a number of iterations that is independent of dimension.

On the Sufficiency of Second-Order Stationarity

In this section we show that for a wide class of nonconvex problems in machine learning and signal processing, all second-order stationary points are global minima. Thus, for this class of problems, finding second-order stationary points efficiently is equivalent to solving the problem. We focus on the underlying global geometry that these problems have in common.

Problems for which second-order stationary points are global minima include tensor decomposition (Ge et al., 2015), dictionary learning (Sun et al., 2016a), phase retrieval (Sun et al., 2016b), synchronization and MaxCut (Bandeira et al., 2016; Mei et al., 2017), smooth semidefinite programs (Boumal et al., 2016), and many problems related to low-rank matrix factorization, such as matrix sensing (Bhojanapalli et al., 2016), matrix completion (Ge et al., 2016) and robust principale component analysis (Ge et al., 2017). In particular, these papers show that by adding appropriate regularization terms, and under mild conditions, there are two key geometric properties satisfied by the corresponding objective functions: (a) All local minima are global minima. There might be multiple local minima due to permutation, but they are all equally good; (b) All saddle points have at least one direction with strictly negative curvature, thus are strict saddle points. we summarize the consequences of these properties in the following proposition, for which we omit the proof as it follows essentially by definition.

If a function $f$ satisfies (a) all local minima are global minima; (b) all saddle points (including local maxima) are strict saddle points, then all second-order stationary points are global minima.

This implies that the core problem for these nonconvex machine-learning applications is to find second-order stationary points efficiently. If we can prove that some simple variants of GD and SGD converges to second-order stationary points efficiently, then we immediately establish global convergence results for these algorithms for all of the above applications (i.e. convergence from arbitrary initialization).

Denote the eigenvalues and eigenvectors of $\mathbf{M}$ as $(\lambda_{i},\mathbf{v}_{i})$ for $i=1,\ldots,d$, and assume there is a gap between the first and second eigenvalues: $\lambda_{1}>\lambda_{2}\geq\lambda_{3}\geq\ldots\geq\lambda_{d}\geq 0$. In this case, the global optimal solutions are $\mathbf{x}=\pm\sqrt{\lambda_{1}}\mathbf{v}_{1}$ giving the top eigenvector direction.

The objective function (3) is nonconvex as a function of $\mathbf{x}$. In order to optimize this objective via gradient-descent methods, we need to analyze the global landscape of the objective function. Its gradient and Hessian are of the form:

Therefore, all stationary points satisfy the equation $\mathbf{M}\mathbf{x}=\|{\mathbf{x}}\|^{2}\mathbf{x}$. Thus they are $\mathbf{0}$ and $\pm\sqrt{\lambda_{i}}\mathbf{v}_{i}$ for $i=1,\ldots,d$. We already know that $\pm\sqrt{\lambda_{1}}\mathbf{v}_{1}$ are global minima, thus they are also local minima and are equivalent up to a sign difference. For the remaining stationary points $\mathbf{x}^{\dagger}$, we note that their Hessian always has strict negative curvature along the $\mathbf{v}_{1}$ direction: $\mathbf{v}_{1}^{\top}\nabla^{2}f(\mathbf{x}^{\dagger})\mathbf{v}_{1}\leq\lambda_{2}-\lambda_{1}<0$. Thus these points are strict saddle points. Having established the preconditions for Proposition 11, we are able to conclude:

Assume that $\mathbf{M}$ is a positive semidefinite matrix whose top two eigenvalues are $\lambda_{1}>\lambda_{2}\geq 0$. For the problem of minimizing the objective Eq. (3), all second-order stationary points are global optima.

Further analysis can be carried out to establish the $\epsilon$-robust version of Corollary 12. Informally, it can be shown that under technical conditions, for polynomially small $\epsilon$, all $\epsilon$-second-order stationary point are close to global optima. We refer the readers to Ge et al. (2017) for the formal statement.

Main Results

In this section, we present our main results on the efficiency of GD and SGD in the nonconvex setting. We first study the case where the exact gradients are accessible, where we focus on an algorithm that we refer to as Perturbed Gradient Descent (PGD). We then turn to the stochastic setting, and present the results for Perturbed SGD and its mini-batch version.

We begin by considering the case in which exact gradients are available, such that GD can be implemented. For convex problems, GD is efficient, but, as can be seen in Eq.(1), GD makes a non-zero step only when the gradient is non-zero, and thus in the nonconvex setting it will be stuck at saddle points if initialized there. We thus consider a simple variant of GD which adds randomness to the iterates at each step (Algorithm 1). The question is whether such a simple procedure can be efficient, particularly in terms of its dimension dependence.

If we wish to output an $\epsilon$-second-order stationary point, it suffices to run PGD for double the number of iterations in Theorem 13. A simple change to the proof shows that half of the iterates will be $\epsilon$-second-order stationary points in this case, so that if we output an iterate uniformly at random, with at least a constant probability it will be an $\epsilon$-second-order stationary point.

We have chosen the distribution of the perturbations to be Gaussian in Algorithm 1 for simplicity. This choice is not necessary. The key properties needed for the perturbation distributions are (a) that the tail of the distribution is sufficiently light such that an appropriate concentration inequality holds, and (b) the variance in every direction is bounded below.

Comparing Theorem 13 to the classical result in Theorem 5, our result shows that PGD finds second-order stationary points in almost the same time as GD finds first-order stationary points, up to only logarithmic factors. Therefore, strict saddle points are computationally benign for first-order gradient methods.

Comparing to Theorem 5, we see that Theorem 13 makes an additional smoothness assumption. This assumption is essential in separating strict saddle points from second-order stationary points.

2 Stochastic setting

In machine learning, the stochastic gradient $\mathbf{g}$ is often obtained as an exact gradient of a smooth function: $\mathbf{g}(\cdot;\theta)=\nabla f(\cdot;\theta)$. We formalize this assumption.

We are now ready to present a guarantee on the efficiency of PSGD (Algorithm 2) for finding a second-order stationary point. We make the following choices of parameters for Algorithm 2:

Let the function $f$ satisfy Assumption A and assume that the stochastic gradient $\mathbf{g}$ satisfies Assumption B (or Assumption C optionally). For any $\epsilon,\delta>0$, the PSGD algorithm (Algorithm 2), with parameter $(\eta,r)$ chosen as in Eq. (4), will visit an $\epsilon-$second-order stationary point at least once in the following number of iterations, with probability at least $1-\delta$:

We note that Remark 14 and Remark 15 apply directly to Theorem 16.

Finally, Theorem 16 can be easily extended to the mini-batch setting, with parameters chosen as:

Let the function $f$ satisfy Assumption A and assume that the stochastic gradient $\mathbf{g}$ satisfies Assumption B (or C optionally). Then, for any $\epsilon,\delta,m>0$, the mini-batch PSGD algorithm (Algorithm 3), with parameters $(\eta,r)$ chosen as in Eq. (5), will visit an $\epsilon-$second-order stationary point at least once in the following number of iterations, with probability at least $1-\delta$:

Theorem 17 says that if the mini-batch size $m$ is not too large—$m\leq\mathfrak{N}$, where $\mathfrak{N}$ is defined in Eq. (4)—then mini-batch PSGD will reduce the number of iterations linearly, while not increasing the total number of stochastic gradient queries.

Proofs

We provide full proofs of our main results, Theorem 16 and Theorem 17, in Appendix B. (Theorem 13 follows directly from the proof of Theorem 16 by setting $\sigma=0$). These proofs require novel concentration inequalities and other tools from stochastic analysis to handle the relatively complex way in which stochasticity interacts with geometry in the neighborhood of saddle points. In the current section we circumvent some of these complexities by presenting a conceptually straightforward proof for an algorithm that is a variant of PGD. This algorithm, summarized in Algorithm 4, removes some of the stochasticity of PGD by restricting the way in which perturbation noise is added. The algorithm is more complex than PGD, but the proof is streamlined, allowing the core concepts underlying the full proof to be conveyed more simply.

The following theorem is the specialization of Theorem 13 to the setting of Algorithm 4.

We proceed to the proof of the theorem. We first specify the choice of hyperparameters $\eta$, $r$, and $\mathscr{T}$, and two quantities $\mathscr{F}$ and $\mathscr{S}$ which are frequently used:

Our high-level proof strategy is a proof by contradiction: when the current iterate is not an $\epsilon$-second order stationary point, it must either have a large gradient or have a strictly negative Hessian, and we prove that in either case, PGD must yield a significant decrease in function value in a controlled number of iterations. Also, since the function value can not decrease more than $f(\mathbf{x}_{0})-f^{\star}$, we know that the total number of iterates that are not $\epsilon$-second order stationary points can not be very large.

First, we bound the rate of decrease when the gradient is large.

We now show that if the starting point has a strictly negative eigenvalue of the Hessian, then adding a perturbation and following by gradient descent will yield a significant decrease in function value in $\mathscr{T}$ iterations.

where $\mathbf{x}_{\mathscr{T}}$ is the $\mathscr{T}^{\textrm{th}}$ gradient descent iterate starting from $\mathbf{x}_{0}$.

In order to prove this, we need to prove two lemmas, and the major simplification over Jin et al. (2017a) comes from the following lemma which says that if function value does not decrease too much over $t$ iterations, then all iterates $\{\mathbf{x}_{\tau}\}_{\tau=0}^{t}$ will remain in a small neighborhood of $\mathbf{x}_{0}$.

Under the setting of Lemma 19, for any $t\geq\tau>0$:

Given the gradient update, $\mathbf{x}_{t+1}=\mathbf{x}_{t}-\eta\nabla f(\mathbf{x}_{t})$, we have that for any $\tau\leq t$:

where step (1) uses Cauchy-Schwarz inequality, and step (2) is due to Lemma 19. ∎

Second, we show that the region in which GD will get remain in a small local neighborhood for at least $\mathscr{T}$ iterations if initialized there (which we refer to as the “stuck region”) is thin. We show this by tracking any pair of points that differ only in an escaping direction and are at least $\omega$ far apart. We show that at least one sequence of the two GD sequences initialized at these points is guaranteed to escape the saddle point with high probability, so that the width of the stuck region along an escaping direction is at most $\omega$.

The claim is true for the base case $t=0$ as $\|{\mathbf{q}(0)}\|=0\leq\|{\hat{\mathbf{x}}_{0}}\|/2=\|{\mathbf{p}(0)}\|/2$. Now suppose the induction claim is true up to $t$. Denote $\lambda_{\min}(\nabla^{2}f(\mathbf{x}_{0}))=-\gamma$. Note that $\hat{\mathbf{x}}_{0}$ lies in the direction of the minimum eigenvector of $\nabla^{2}f(\mathbf{x}_{0})$. Thus for any $\tau\leq t$, we have:

where the second-to-last inequality uses $t+1\leq\mathscr{T}$. By our choice of the hyperparameter in Eq. (6), we have $2\eta\rho\mathscr{S}\mathscr{T}\leq 1/2$, which finishes the inductive proof.

where step (1) uses the fact $(1+x)^{1/x}\geq 2$ for any $x\in(0,1]$. This contradicts the localization property of Eq. (7), which finishes the proof. ∎

Equipped with Lemma 21 and Lemma 22, we are ready to prove Lemma 20.

On the event $\{\mathbf{x}_{0}\not\in\mathcal{X}_{\text{stuck}}\}$, due to our parameter choice in Eq. (6), we have:

With Lemma 19 and Lemma 20 in hand, it is not hard to establish Theorem 18.

First, we set the total number of iterations $T$ to be:

We then argue that, with probability $1-\delta$, Algorithm 4 will add a perturbation at most $T/(4\mathscr{T})$ times. This is because if otherwise, we can appeal to Lemma 20 every time we add a perturbation, and conclude:

which can not happen. Finally, excluding those iterations that are within $\mathscr{T}$ steps after adding perturbations, we still have $3T/4$ steps left. They are either large gradient steps, $\|{\nabla f(\mathbf{x}_{t})}\|\geq\epsilon$, or $\epsilon$-second order stationary points. Within them, we know that the number of large gradient steps cannot be more than $T/4$. This is true because if otherwise, by Lemma 19:

which again cannot happen. Therefore, we conclude that at least $T/2$ of the iterates must be $\epsilon-$second order stationary points. ∎

Conclusions

We have shown that simple perturbed versions of GD and SGD escape saddle points and find second-order stationary points in essentially the same time that classical GD and SGD take to find first-order stationary points. The overheads are only logarithmic factors in dimensionality in both the non-stochastic setting and the stochastic setting with Lipschitz stochastic gradient. In the general stochastic setting, the overhead is a linear factor in dimension.

Combined with previous literature that shows that all second-order stationary points are global optima for a broad class of nonconvex optimization problems in machine learning and signal processing, our results directly provide efficient guarantees for solving those nonconvex problem via simple local search approaches. We now discuss several possible future directions, and further connections to other fields.

Carmon et al. (2017a) have presented lower bounds that imply that GD achieves the optimal rate for finding stationary points for gradient Lipschitz functions. In our setting, we additionally assume that the Hessian is Lipschitz. This implies that GD is no longer necessarily an optimal algorithm. While our results show that variants of GD are efficient in this setting, one would additionally like to know whether they are close to optimality.

Escaping high-order saddle points.

The current paper focuses on escaping strict saddle points and finding second-order stationary points. More generally, one can define $n$th-order stationary points as points that satisfies the KKT necessary conditions for being local minima up to $n$th-order derivatives. It becomes more challenging to find $n$th-order stationary points as $n$ increases, since it is necessary to escape higher-order saddle points. In terms of worst-case efficiency, Nesterov (2000) rules out the possibility of efficient algorithms for finding $n$th-order stationary points for all $n\geq 4$, showing that the problem is NP-hard. Anandkumar and Ge (2016) present a third-order algorithm that finds third-order stationary points in polynomial time. It remains open whether simple variants of GD can also find third-order stationary points efficiently. It is unlikely that the overhead will still be small or only logarithmic in this case, but it is not clear what to expect for the overhead. A related question is to identify applications where third-order stationarity is needed, in addition to second-order stationarity, to achieve global optimality.

Connection to gradient Langevin dynamics.

The Bayesian counterpart of SGD is the Langevin Monte Carlo (LMC) algorithm (Roberts et al., 1996), which performs the following iterative update:

Here $\beta$ is known as the inverse temperature. When the step size $\eta$ goes to zero, the distribution of the LMC iterates is known to converge to the stationary distribution $\mu(\mathbf{x})\propto e^{-\beta f(\mathbf{x})}$ (Roberts et al., 1996).

While the LMC algorithm is superficially very close to stochastic gradient descent, the goals for the two algorithms are quite different.

Convergence: While the focus of the optimization literature is to find stationary points, the goal of the LMC algorithm is to converge to a stationary distribution (i.e., to mix rapidly).

Scaling of Noise: The scaling of the stochasticity in SGD—and in particular the size of the perturbation that we consider in the current paper—is much smaller than the scaling considered in the LMC literature. Running our algorithm is equivalent to running LMC with temperature $\beta^{-1}\propto d^{-1}$. In this low-temperature or small-noise regime, the algorithm can no longer mix efficiently for smooth nonconvex functions, as it requires $\Omega(e^{d})$ steps in the worst case (Bovier et al., 2004). However, with this small amount of noise, the algorithm can still perform local search efficiently, and can find a second-order stationary point in a small number of iterations, as shown in Theorem 13.

Recent work of Zhang et al. (2017) studied the time that LMC takes to hit a second-order stationary point as a criterion for convergence, instead of the traditional mixing time to a stationary distribution. In this analysis, the runtime is no longer exponential, but it is still polynomially dependent on dimension $d$ with large degree.

On the necessity of adding perturbations.

We have shown that adding perturbations to the iterations of GD or SGD allows these algorithms to escape saddle points efficiently. As an alternative, one can also simply run GD with random initialization, and try to escape saddle points using only the randomness due to the initialization. Although this alternative algorithm exhibits asymptotic convergence (Lee et al., 2016), it does not yield efficient convergence in general. Du et al. (2017) shows that even with fairly natural random initialization schemes and non-pathological functions, randomly initialized GD can be significantly slowed by saddle points, taking exponential time to escape them.

Acknowledgements

We thank Tongyang Li and Quanquan Gu for valuable discussions. This work was supported in part by the Mathematical Data Science program of the Office of Naval Research under grant number N00014-18-1-2764. Rong Ge also acknowledges the funding support from NSF CCF-1704656, NSF CCF-1845171 (CAREER), Sloan Fellowship and Google Faculty Research Award.

References

Appendix A Tables of Related Work

In Table 2 and Table 3, we present a detailed comparison of our results with other related work in both non-stochastic and stochastic settings. See Section 1.1 for the full text descriptions. We note that our algorithms are simple variants of standard GD and SGD, which are the simplest among all the algorithms listed in the table.

Appendix B Proofs for Stochastic Setting

In this section, we provide proofs for our main results—Theorem 16 and Theorem 17. Theorem 13 can be proved as a special case of Theorem 16 by taking $\sigma=0$.

where $\mathfrak{N}$ and the log factor $\iota$ are defined as:

Here, $\mu$ is a sufficiently large absolute constant to be determined later. Note also that throughout this section $c$ denotes an absolute constant that does not depend on the choice of $\mu$. Its value may change from line to line.

B.2 Descent lemma

We first prove that the change in the function value can be always decomposed into the decrease due to the magnitudes of gradients and the possible increase due to randomness in both the stochastic gradients and the perturbations.

Since Algorithm 2 is Markovian, the operations in each iteration do not depend on time step $t$. Thus, it suffices to prove Lemma 23 for special case $t_{0}=0$. Recall the update equation:

Summing over this inequality, we have the following:

For the second term on the right-hand side, applying Lemma 39, there exists an absolute constant $c$, such that, with probability $1-2e^{-\iota}$:

For the third term on the right-hand side of Eq. (9), applying Lemma 38, with probability $1-2e^{-\iota}$:

The descent lemma allows us to show the following “Improve or Localize” phenomenon for perturbed SGD. That is, with high probability over a small number of iterations, either the function value decreases significantly, or the iterates stay within a small local region.

Under the same setting as in Lemma 23, with probability at least $1-8dt\cdot e^{-\iota}$, the sequence PSGD$(\eta,r)$ (Algorithm 2) satisfies:

By a similar argument as in the proof of Lemma 23, it suffices to prove Lemma 24 in the special case $t_{0}=0$. According to Lemma 23, with probability $1-4e^{-\iota}$, for some absolute constant $c$:

Therefore, for any fixed $\tau\leq t$, with probability $1-8d\cdot e^{-\iota}$,

where in step (1) we use the Cauchy-Schwarz inequality and Lemma 36. Finally, applying a union bound for all $\tau\leq t$, we finish the proof. ∎

B.3 Escaping saddle points

Lemma 23 shows that large gradients contribute to the fast decrease of the function value. In this subsection, we will show that starting in the vicinity of strict saddle points will also enable PSGD to decrease the function value rapidly. Concretely, this entire subsection will be devoted to proving the following lemma:

Since Algorithm 2 is Markovian, the operations in each iterations do not depend on time step $t$. Thus, it suffices to prove Lemma 25 for special case $t_{0}=0$. To prove this lemma, we first need to introduce the concept of a coupling sequence.

Throughout this subsection, we let $\mathcal{H}\mathrel{\mathop{:}}=\nabla^{2}f(\mathbf{x}_{0})$, $\mathbf{e}_{1}$ be the minimum eigendirection of $\mathcal{H}$, and $\gamma\mathrel{\mathop{:}}=\lambda_{\min}(\mathcal{H})$. We also let $\mathcal{P}_{-1}$ be the projection onto the subspace complement of $\mathbf{e}_{1}$.

Consider sequences $\{\mathbf{x}_{i}\}$ and $\{\mathbf{x}^{\prime}_{i}\}$ that are obtained as separate runs of the PSGD (algorithm 2), both starting from $\mathbf{x}_{0}$. They are coupled if both sequences share the same randomness $\mathcal{P}_{-1}\xi_{\tau}$ and $\theta_{\tau}$, while in $\mathbf{e}_{1}$ direction we have $\mathbf{e}_{1}^{\top}\xi_{\tau}=-\mathbf{e}_{1}^{\top}\xi^{\prime}_{\tau}$.

The first thing we can show is that if the function values of both sequences do not exhibit a sufficient decrease, then both sequences are localized in a small ball around $\mathbf{x}_{0}$ within $\mathscr{T}$ iterations.

This lemma follows from applying Lemma 24 on both sequences and using a union bound. ∎

The overall proof strategy for Lemma 25 is to show that localization happens with a very small probability, thus at least one of the sequence must have sufficient descent. In order to prove this, we study the dynamics of the difference of the coupling sequence.

Consider coupling sequences $\{\mathbf{x}_{i}\}$ and $\{\mathbf{x}^{\prime}_{i}\}$ as in Definition 26 and let $\hat{\mathbf{x}}_{t}\mathrel{\mathop{:}}=\mathbf{x}_{i}-\mathbf{x}^{\prime}_{i}$. Then $\hat{\mathbf{x}}_{t}=-\mathbf{q}_{h}(t)-\mathbf{q}_{sg}(t)-\mathbf{q}_{p}(t)$, where:

Recall $\zeta_{i}=\mathbf{g}(\mathbf{x}_{i};\theta_{i})-\nabla f(\mathbf{x}_{i})$, thus, we have the update formula:

Taking the difference between $\{\mathbf{x}_{i}\}$ and $\{\mathbf{x}^{\prime}_{i}\}$:

At a high level, we will show that with constant probability, $\mathbf{q}_{p}(t)$ is the dominating term which controls the behavior of the dynamics, and $\mathbf{q}_{h}(t)$ and $\mathbf{q}_{sg}(t)$ will stay small compared to $\mathbf{q}_{p}(t)$. To show this, we prove the following three lemmas.

Under the notation of Lemma 28 and Lemma 29, letting $-\gamma\mathrel{\mathop{:}}=\lambda_{\min}(\mathcal{H})$, we have $\forall t>0$:

There exists an absolute constant $c_{\max}$ such that, for any $\iota\geq c_{\max}$, under the notation of Lemma 28 and Lemma 29, and letting $-\gamma\mathrel{\mathop{:}}=\lambda_{\min}(\mathcal{H})$, we have:

For simplicity we denote $\mathfrak{E}$ as the event $\{\forall\tau\leq t:\max\{\|{\mathbf{x}_{\tau}-\mathbf{x}_{0}}\|^{2},\|{\mathbf{x}^{\prime}_{\tau}-\mathbf{x}_{0}}\|^{2}\}\leq\mathscr{S}^{2}\}$. We use induction to prove following claim for any $t\in[0,\mathscr{T}]$:

Then Lemma 31 follows directly from combining Lemma 27 and the induction claim.

Clearly for the base case $t=0$, the claim holds trivially, as $\mathbf{q}_{sg}(0)=\mathbf{q}_{h}(0)=\mathbf{0}$. Suppose the claim holds for $t$, then by Lemma 30, with probability at least $1-2\mathscr{T}e^{-\iota}$, we have for any $\tau\leq t$:

where the last step is due to $\eta\rho\mathscr{S}\mathscr{T}=1/\iota$ by Eq. (8). By picking $\iota$ larger than the absolute constant $40c$, we have $c\eta\rho\mathscr{S}\mathscr{T}\leq 1/40$.

Recall also that $\hat{\zeta}_{\tau}|\mathcal{F}_{\tau-1}$ is the summation of a $\text{nSG}(\sigma)$ random vector and a $\text{nSG}(c\cdot r)$ random vector. By Lemma 36, we know that with probability at least $1-4de^{-\iota}$:

Finally, combining both cases, and by our choice of step size $\eta,r$ as in Eq. (8) with $\iota$ large enough:

and the induction follows by the triangle inequality and a union bound. ∎

We are ready to prove Lemma 25, which is the focus of this subsection.

B.4 Proof of Theorem 16

Lemma 23 and Lemma 25 describe the speed of decrease in the function values when either large gradients or strictly negative curvatures are present. Combining them gives the proof for our main theorem.

First, we set the total number of iterations $T$ to be:

We will show that the following two claims hold simultaneously with probability $1-\delta$:

At most $T/4$ iterates have large gradient; i.e., $\|{\nabla f(\mathbf{x}_{t})}\|\geq\epsilon$;

At most $T/4$ iterates are close to saddle points; i.e., $\|{\nabla f(\mathbf{x}_{t})}\|\leq\epsilon$ and $\lambda_{\min}(\nabla^{2}f(\mathbf{x}_{t}))\leq-\sqrt{\rho\epsilon}$.

Therefore, at least $T/2$ iterates are $\epsilon$-second order stationary point. We prove the two claims separately.

Suppose that within $T$ steps, we have more than $T/4$ iterates for which gradient is large (i.e., $\|{\nabla f(\mathbf{x}_{t})}\|\geq\epsilon$). Recall that by Lemma 23 we have with probability $1-4e^{-\iota}$:

Claim 2.

We first define the stopping times that allow us to invoke Lemma 25:

Clearly, $z_{i}$ is a stopping time, and it is the $i$th time in the sequence along which we can apply Lemma 25. We also let $M$ be the random variable $M=\max\{i|z_{i}+\mathscr{T}\leq T\}$. We can decompose the decrease $f(\mathbf{x}_{T})-f(\mathbf{x}_{0})$ as follows:

For the first term $T_{1}$, by Lemma 25 and a supermartingale concentration inequality, for each fixed $m\leq T$:

Since the random variable $M\leq T/\mathscr{T}\leq T$, by a union bound, we know that with probability $1-5d\mathscr{T}^{2}T^{2}\cdot\log(\mathscr{S}\sqrt{d}/(\eta r))e^{-\iota}$:

For the second term, by a union bound and Lemma 23 for all $0\leq t_{1},t_{2}\leq T$, with probability $1-4T^{2}e^{-\iota}$:

Therefore, if within $T$ steps we have more than $T/4$ saddle points, then $M\geq T/4\mathscr{T}$, and with probaility $1-10d\mathscr{T}^{2}T^{2}\cdot\log(\mathscr{S}\sqrt{d}/(\eta r))e^{-\iota}$:

This will gives $f(\mathbf{x}_{T})\leq f(x_{0})-0.4T\mathscr{F}/\mathscr{T}<f^{\star}$ which is not possible.

B.5 Proof of Theorem 17

Our proofs for PSGD easily generalize to the mini-batch setting.

Appendix C Concentration Inequalities

In this section, we present the concentration inequalities required for this paper. Please refer to the technical note [Jin et al., 2019] for the proofs of Lemmas 33, 34, 36 and 37.

Recall the definition of a norm-subGaussian random vector.

Note that a bounded random vector and a subGaussian random vector are two special cases of a norm-subGaussian random vector.

There exists an absolute constant $c$ so that following random vectors are $\text{nSG}(c\cdot\sigma)$.

Second, we have that if $\mathbf{X}$ is norm-subGaussian, then its norm square is subExponential, and its component along a single direction is subGaussian.

For concentration, we are interested in the properties of norm-subGaussian martingale difference sequences. Concretely, they are sequences satisfying the following conditions.

Similar to subGaussian random variables, we can also prove a Hoeffding-type inequality for norm-subGaussian random vectors which is tight up to a $\log(d)$ factor.

When $\{\sigma_{i}\}$ is also random, we have the following.

Finally, we can also provide concentration inequalities for the sum of norm squares of norm-subGaussian random vectors, and for the sum of inner products of norm-subGaussian random vectors with another set of random vectors.

For any $i\in[n]$ and fixed $\lambda>0$, since $\mathbf{u}_{i}\in\mathcal{F}_{i-1}$, according to Lemma 34 there exists a constant $c$ such that $\langle\mathbf{u}_{i},\mathbf{X}_{i}\rangle|\mathcal{F}_{i-1}$ is $c\cdot\|{\mathbf{u}_{i}}\|\sigma_{i}$-subGaussian. Thus:

Therefore, consider the following quantity: