E SCAPING S ADDLE P OINTS F ASTER WITH S TOCHASTIC M OMENTUM
A BSTRACT
Stochastic gradient descent (SGD) with stochastic momentum is popular in nonconvex stochastic optimization and particularly for the training of deep neural networks.
In standard SGD, parameters are updated by improving along the path of the gradient at the current iterate on a batch of examples, where the addition of a “momentum” term biases the update in the direction of the previous change in parameters.
In non-stochastic convex optimization one can show that a momentum adjustment provably reduces convergence time in many settings, yet such results have been elusive in the stochastic and non-convex settings.
At the same time, a widely-observed empirical phenomenon is that in training deep networks stochastic momentum appears to signiﬁcantly improve convergence time, variants of it have ﬂourished in the development of other popular update methods, e.g. ADAM (Kingma & Ba (2015)), AMSGrad (Reddi et al.
(2018b)), etc.
Yet theoretical justiﬁcation for the use of stochastic momentum has remained a signiﬁcant open question.
In this paper we propose an answer: stochastic momentum improves deep network training because it modiﬁes SGD to escape saddle points faster and, consequently, to more quickly ﬁnd a second order stationary point.
Our theoretical results also shed light on the related question of how to choose the ideal momentum parameter–our analysis suggests that β ∈ [0 , 1) should be large (close to 1), which comports with empirical ﬁndings.
We also provide experimental ﬁndings that further validate these conclusions.
I NTRODUCTION
SGD with stochastic momentum has been a de facto algorithm in nonconvex optimization and deep learning.
It has been widely adopted for training machine learning models in various applications.
Modern techniques in computer vision (e.g.Krizhevsky et al.
; He et al.
; Cubuk et al.
; Gastaldi (2017)), speech recognition (e.g. Amodei et al.
(2016)), natural language processing (e.g. Vaswani et al.
(2017)), and reinforcement learning (e.g. Silver et al.
(2017)) use SGD with stochastic momentum to train models.
The advantage of SGD with stochastic momentum has been widely observed (Hoffer et al.
; Loshchilov & Hutter (2019); Wilson et al.
Sutskever et al.
(2013) demonstrate that training deep neural nets by SGD with stochastic momentum helps achieving in faster convergence compared with the standard SGD (i.e. without momentum).
The success of momentum makes it a necessary tool for designing new optimization algorithms in optimization and deep learning.
For example, all the popular variants of adaptive stochastic gradient methods like Adam (Kingma & Ba (2015)) or AMSGrad (Reddi et al.
(2018b)) include the use of momentum.
Despite the wide use of stochastic momentum (Algorithm 1) in practice, 1 justiﬁcation for the clear empirical improvements has remained elusive, as has any mathematical guidelines for actually setting the momentum parameter—it has been observed that large values (e.g. β = 0 .
9 ) work well in practice.
It should be noted that Algorithm 1 is the default momentum-method in popular software 1: Required: Step size parameter η and {+momentum parameter β . 
 Init: w 0 ∈ R d and m − 1 = 0 ∈ R d . 3: for t = 0 to T do 4:+}
Given current iterate w t , obtain stochastic gradient g t := ∇ f ( w t ; ξ t ) . Update stochastic momentum m t := βm t − 1 + g t . 6: Update iterate w t +1 :
= w t − ηm t .
7: end for packages such as PyTorch and Tensorﬂow.
In this paper we provide a theoretical analysis for SGD with momentum.
We identify some mild conditions that guarantees SGD with stochastic momentum will provably escape saddle points faster than the standard SGD, which provides clear evidence for the beneﬁt of using stochastic momentum.
For stochastic heavy ball momentum, a weighted average of stochastic gradients at the visited points is maintained.
The new update is computed as the current update minus a step in the direction of the momentum.
Our analysis shows that these updates can amplify a component in an escape direction of the saddle points.
In this paper, we focus on ﬁnding a second-order stationary point for smooth non-convex optimization by SGD with stochastic heavy ball momentum.
Speciﬁcally, we consider the stochastic nonconvex optimization problem, min w ∈ R d f ( w ) :=
E ξ ∼D [ f ( w ; ξ )] , where we overload the notation so that f ( w ; ξ ) represents a stochastic function induced by the randomness ξ while f ( w ) is the expectation of the stochastic functions.
An ( (cid:15), (cid:15) ) -second-order stationary point w satisﬁes (cid:107)∇ f ( w ) (cid:107) ≤ (cid:15) and ∇ 2 f ( w ) (cid:23) − (cid:15)I.
Obtaining a second order guarantee has emerged as a desired goal in the nonconvex optimization community.
Since ﬁnding a global minimum or even a local minimum in general nonconvex optimization can be NP hard (Anandkumar & Ge (2016); Nie (2015); Murty & Kabadi (1987);
Nesterov (2000)), most of the papers in nonconvex optimization target at reaching an approximate second-order stationary point with additional assumptions like Lipschitzness in the gradients and the
Hessian (e.g. Allen-Zhu & Li (2018); Carmon & Duchi (2018); Curtis et al.
; Daneshmand et al.
Du et al.
Fang et al.
Ge et al.
; Jin et al.
Kohler & Lucchi (2017);
Lei et al.
; Lee et al.
; Levy (2016); Mokhtari et al.
Nesterov & Polyak (2006); Reddi et al.
(2018a);
Staib et al.
Tripuraneni et al.
Xu et al.
2 We follow these related works for the goal and aim at showing the beneﬁt of the use of the momentum in reaching an ( (cid:15), (cid:15) ) -second-order stationary point.
We introduce a required condition, akin to a model assumption made in (Daneshmand et al.
, that ensures the dynamic procedure in Algorithm 2 produces updates with suitable correlation with the negative curvature directions of the function f .
Deﬁnition 1.
Assume, at some time t , that the Hessian H t = ∇ 2 f ( w t ) has some eigenvalue smaller than − (cid:15) and (cid:107)∇ f ( w t ) (cid:107) ≤ (cid:15) .
Let v t be the eigenvector corresponding to the smallest eigenvalue of ∇ 2 f ( w t ) .
The stochastic momentum m t satisﬁes Correlated Negative Curvature (CNC) at t with parameter γ >
if
[Equation]
As we will show, the recursive dynamics of SGD with heavy ball momentum helps in amplifying the escape signal γ , which allows it to escape saddle points faster.
Contribution: We show that, under CNC assumption and some minor constraints
that upper-bound parameter β , if SGD with momentum has properties called Almost Positively Aligned with Gradient (APAG), Almost Positively Correlated with Gradient (APCG), and Gradient Alignment or Curvature
Exploitation (GrACE), deﬁned in the later section, then it takes T
= O
((1 − β ) log(1 / (1 − β ) (cid:15) )
(cid:15) − 10 ) iterations to return an ( (cid:15), (cid:15) )
second order stationary point.
Alternatively, one can obtain an ( (cid:15), √ (cid:15) )
second order stationary point in T
=
O
((1 − β ) log(1 / (1 − β ) (cid:15) )
(cid:15) − 5 ) iterations.
Our theoretical result demonstrates that a larger momentum parameter β can help in escaping saddle points faster.
As saddle points are pervasive in the loss landscape of optimization and deep learning (Dauphin et al.
(2014); Choromanska et al. (2015))
, the result sheds light on explaining why SGD with momentum enables training faster in optimization and deep learning.
Notation: In this paper we use E t [ · ] to represent conditional expectation E [ ·| w 1 , w 2 , . .
. , w t ] , which is about ﬁxing the randomness upto but not including t and notice that w t was determined at t − 1 .
B ACKGROUND 2. A THOUGHT EXPERIMENT
Let us provide some high-level intuition about the beneﬁt of stochastic momentum with respect to avoiding saddle points.
In an iterative update scheme, at some time t 0 the parameters w t 0 can enter a saddle point region , that is a place where Hessian ∇ 2 f ( w t 0 ) has a non-trivial negative eigenvalue, say λ min ( ∇ 2 f ( w t 0 )) ≤ − (cid:15) , and the gradient ∇ f ( w t 0 ) is small in norm, say (cid:107)∇ f ( w t 0 )
(cid:107) ≤ (cid:15) .
The challenge here is that gradient updates may drift only very slowly away from the saddle point, and may not escape this region; see (Du et al.
; Lee et al.
for additional details.
On the other hand, if the iterates were to move in one particular direction, namely along v t 0 the direction of the smallest eigenvector of ∇ 2 f ( w t 0 ) , then a fast escape is guaranteed under certain constraints on the step size η ; see e.g. (Carmon et al.
While the negative eigenvector could be computed directly, this 2nd-order method is prohibitively expensive and hence we typically aim to rely on gradient methods.
With this in mind, Daneshmand et al.
(2018), who study non-momentum SGD, make an assumption akin to our CNC property described above that each stochastic gradient g t 0 is strongly non-orthogonal to v t 0 the direction of large negative curvature.
This sufﬁces to drive the updates out of the saddle point region.
In the present paper we study stochastic momentum, and our CNC property requires that the update direction m t 0 is strongly non-orthogonal to v t 0 ; more precisely, E t
[ (cid:104) m t 0 , v t 0 (cid:105) 2 ] ≥ γ > 0 .
We are able to take advantage of the analysis of (Daneshmand et al.
(2018)) to establish that updates begin to escape a saddle point region for similar reasons.
Further, this effect is ampliﬁed in successive iterations through the momentum update when β is close to 1.
Assume that at some w t 0 we have m twhich possesses signiﬁcant correlation with the negative curvature direction v t 0 , then on successive rounds m t
+1 is quite close to βm t 0 , m t
+2 is quite close to β 2 m t 0 , and so forth; see Figure 1 for an example.
This provides an intuitive perspective on how momentum might help accelerate the escape process.
Yet one might ask does this procedure provably contribute to the escape process and, if so, what is the aggregate performance improvement of the momentum?
We answer the ﬁrst question in the afﬁrmative, and we answer the second question essentially by showing that momentum can help speed up saddle-point escape by a multiplicative factor of 1 − β .
On the negative side, we also show that β is constrained and may not be chosen arbitrarily close to 1.
2. M OMENTUM H ELPS E SCAPE S ADDLE P OINTS : AN E MPIRICAL V IEW
Let us now establish, empirically, the clear beneﬁt of stochastic momentum on the problem of saddle-point escape.
We construct two stochastic optimization tasks, and each exhibits at least one signiﬁcant saddle point.
The two objectives are as follows.
[Equation]
Problem (3) of these was considered by (Staib et al.
; Reddi et al.
(2018a))
and represents a very straightforward non-convex optimization challenge, with an embedded saddle given by the matrix
[Figure] H := diag ([1 , − 0 .
1]) , and stochastic gaussian perturbations given by b i ∼ N (0 , diag ([0 . 1 , 0 .
; the small variance in the second component provides lower noise in the escape direction.
Here we have set n = 10 .
Observe that the origin is in the neighborhood of saddle points and has objective value zero.
SGD and SGD with momentum are initialized at the origin in the experiment so that they have to escape saddle points before the convergence.
The second objective (4) appears in the phase retrieval problem, that has real applications in physical sciences (Cand ´ es et al.
Shechtman et al.
In phase retrieval 3 , one wants to ﬁnd an unknown w ∗ ∈ R d with access to but a few samples
y i
= ( a (cid:62) i w ∗ ) 2 ; the design vector a
i is known a priori.
Here we have sampled w ∗ ∼ N (0 , I d /d ) and a i ∼ N (0 , I d ) with d = 10 and n = 200 .
The empirical ﬁndings, displayed in Figure 2, are quite stark: for both objectives, convergence is signiﬁcantly accelerated by larger choices of β .
In the ﬁrst objective (Figure 4a), we see each optimization trajectory entering a saddle point region, apparent from the “ﬂat” progress, yet we observe that large-momentum trajectories escape the saddle much more quickly than those with smaller momentum.
A similar affect appears in Figure 4b.
To the best of our knowledge, this is the ﬁrst reported empirical ﬁnding that establishes the dramatic speed up of stochastic momentum for ﬁnding an optimal solution in phase retrieval.
2. R ELATED WORKS .
Heavy ball method: The heavy ball method was originally proposed by Polyak (1964).
It has been observed that this algorithm, even in the deterministic setting, provides no convergence speedup over standard gradient descent, except in some highly structure cases such as convex quadratic objectives where an “accelerated” rate is possible (Lessard et al.
; Goh (2017); Ghadimi et al.
; Sun et al.
; Loizou & Richt´arik (2017; 2018); Gadat et al.
; Yang et al.
; Kidambi et al.
; Can et al.
We provide a comprehensive survey of the related works about heavy ball method in Appendix A.
Reaching a second order stationary point: As we mentioned earlier, there are many works aim at reaching a second order stationary point.
We classify them into two categories: specialized algorithms and simple GD/SGD variants.
Specialized algorithms are those designed to exploit the negative curvature explicitly and escape saddle points faster than the ones without the explicit exploitation (e.g. Carmon et al.
Agarwal et al.
(2017); Allen-Zhu & Li (2018);
Xu et al.
Simple
GD/SGD variants are those with minimal tweaks of standard GD/SGD or their variants (e.g. Ge et al.
; Levy (2016);
Fang et al.
; Jin et al.
; Daneshmand et al.
Staib et al.
Our work belongs to this category.
In this category, perhaps the pioneer works are (Ge et al.
and (Jin et al.
Jin et al.
(2017) show that explicitly adding isotropic noise in each iteration guarantees that GD escapes saddle points and ﬁnds a second order stationary point with high probability.
Following (Jin et al.
, Daneshmand et al.
(2018) assume that stochastic gradient inherently has a component to escape.
Speciﬁcally, they make assumption of the
Correlated Negative Curvature (CNC) for stochastic gradient g t so that E t [ (cid:104) g t , v t (cid:105) 2 ] ≥ γ > 0 .
The assumption allows the algorithm to avoid the procedure of perturbing the updates by adding isotropic noise.
Our work is motivated by (Daneshmand et al.
but assumes CNC for the stochastic momentum m t instead.
In Appendix A, we compare the results of our work with the related works.
M AIN R ESULTS
We assume that the gradient ∇ f is L -Lipschitz; that is, f is L -smooth.
Further, we assume that the
Hessian ∇ 2 f is ρ -Lipschitz.
These two properties ensure that (cid:107)∇ f ( w ) −∇ f ( w (cid:48) ) (cid:107) ≤ L (cid:107) w − w (cid:48) (cid:107) and that (cid:107)∇ 2 f ( w ) − ∇ 2 f ( w (cid:48) )
(cid:107) ≤ ρ (cid:107) w − w (cid:48)
(cid:107) , ∀ w, w (cid:48) .
The L -Lipschitz gradient assumption implies that | f ( w (cid:48) ) − f ( w ) −(cid:104)∇ f ( w ) , w (cid:48) − w (cid:105)| ≤ L 2 (cid:107) w − w (cid:48) (cid:107) 2 , ∀ w, w (cid:48) , while the ρ -Lipschitz Hessian assumption implies that | f ( w (cid:48) ) − f ( w ) −(cid:104)∇ f ( w ) , w (cid:48) − w (cid:105)− ( w (cid:48) − w ) (cid:62) ∇ 2 f ( w )( w (cid:48) − w ) | ≤ ρ 6 (cid:107) w − w (cid:48)
(cid:107) 3 , ∀ w, w (cid:48) .
Furthermore, we assume that the stochastic gradient has bounded noise (cid:107)∇ f ( w ) − ∇ f ( w ; ξ ) (cid:107) 2 ≤ σand that the norm of stochastic momentum is bounded so that (cid:107) m t (cid:107) ≤ c m .
We denote Π i M i as the matrix product of matrices { M i } and we use σ max ( M )
= (cid:107) M (cid:107) 2 :=
sup x (cid:54) =0 (cid:104) x,Mx (cid:105) (cid:104) x,x (cid:105) to denote the spectral norm of the matrix M .
3. R EQUIRED P ROPERTIES WITH E MPIRICAL V ALIDATION
Our analysis of stochastic momentum relies on three properties of the stochastic momentum dynamic.
These properties are somewhat unusual, but we argue they should hold in natural settings, and later we aim to demonstrate that they hold empirically in a couple of standard problems of interest.
Deﬁnition 2.
We say that SGD with stochastic momentum satisﬁes A lmost P ositively A ligned with
G radient (APAG) 4 if we have
[Equation]
We say that SGD with stochastic momentum satisﬁes A lmost P ositively C orrelated with G radient (APCG) with parameter τ if ∃ c (cid:48) > 0 such that,
[Equation]
where the PSD matrix M t is deﬁned as
[Equation]
with
[Equation]
for any integer 1 ≤ k ≤ τ − 1 , and η is any step size chosen that guarantees each G s,t is PSD.
Deﬁnition 3.
We say that the SGD with momentum exhibits Gr adient A lignment or C urvature
E xploitation (GrACE) if ∃ c h ≥ 0 such that
[Equation]
APAG requires that the momentum term m t must, in expectation, not be signiﬁcantly misaligned with the gradient ∇ f ( w t ) .
This is a very natural condition when one sees that the momentum term is acting as a biased estimate of the gradient of the deterministic f .
APAG demands that the bias can not be too large relative to the size of ∇ f ( w t ) .
Indeed this property is only needed in our analysis when the gradient is large (i.e. (cid:107)∇ f ( w t ) (cid:107) ≥ (cid:15) ) as it guarantees that the algorithm makes progress; our analysis does not require APAG holds when gradient is small.
APCG is a related property, but requires that the current momentum term m t is almost positively correlated with the the gradient ∇ f ( w t ) , but measured in the Mahalanobis norm induced by M t . It
[Figure] properties in experiments. an interesting observation is that, for the may appear to be an unusual object, but one can view the PSD matrix M t as measuring something about the local curvature of the function with respect to the trajectory of the SGD with momentum dynamic.
We will show that this property holds empirically on two natural problems for a reasonable constant c (cid:48) .
APCG is only needed in our analysis when the update is in a saddle region with signiﬁcant negative curvature, (cid:107)∇ f ( w ) (cid:107) ≤ (cid:15) and λ min ( ∇ 2 f ( w )) ≤ − (cid:15) .
Our analysis does not require APCG holds when the gradient is large or the update is at an ( (cid:15), (cid:15) )
-second order
stationary point.
For GrACE, the ﬁrst term on l.h.s of (7) measures the alignment between stochastic momentum m t and the gradient ∇ f ( w t ) , while the second term on l.h.s measures the curvature exploitation.
The ﬁrst term is small (or even negative) when the stochastic momentum m t is aligned with the gradient ∇ f ( w t ) , while the second term is small (or even negative) when the stochastic momentum m t can exploit a negative curvature (i.e. the subspace of eigenvectors that corresponds to the negative eigenvalues of the Hessian ∇ 2 f ( w t ) if exists).
Overall, a small sum of the two terms (and, consequently, a small c h ) allows one to bound the function value of the next iterate (see Lemma 8).
On Figure 3, we report some quantities related to APAG and APCG as well as the gradient norm when solving the previously discussed problems (3) and (4) using SGD with momentum.
We also report a quantity regarding GrACE on Figure 4 in the appendix.
3. C ONVERGENCE RESULTS
The high level idea of our analysis follows as a similar template to (Jin et al.
; Daneshmand et al.
; Staib et al.
Our proof is structured into three cases: either (a) (cid:107)∇ f ( w ) (cid:107) ≥ (cid:15) , or (b) (cid:107)∇ f ( w ) (cid:107) ≤ (cid:15) and λ min ( ∇ 2 f ( w )) ≤ − (cid:15) , or otherwise (c) (cid:107)∇ f ( w ) (cid:107) ≤ (cid:15) and λ min ( ∇ 2 f ( w )) ≥ − (cid:15) , meaning we have arrived in a second-order stationary region.
The precise algorithm we analyze is
Algorithm 2, which identical to Algorithm 1 except that we boost the step size to a larger value r on occasion.
We will show that the algorithm makes progress in cases (a) and (b).
In case (c), when the goal has already been met, further execution of the algorithm only weakly hurts progress.
Ultimately, we prove that a second order stationary point is arrived at with high probability.
While our proof borrows tools from (Daneshmand et al.
; Staib et al.
, much of the momentum analysis is entirely novel to our knowledge.
Theorem 1.
Assume that the stochastic momentum satisﬁes CNC.
Set 5 r =
O
( (cid:15) 2 ) ,
η = O ( (cid:15) 5 ) , and
T thred =
[Equation]
log(
[Equation]
) δγ(cid:15) ) = O ((1 − β ) log( Lc
[Equation] (1 − β ) δγ(cid:15) ) (cid:15) − 6 ) for some constant c > 0 . If
SGD with momentum (Algorithm 2) has APAG property when gradient is large ( (cid:107)∇ f ( w ) (cid:107) ≥ (cid:15) ), APCG T thred property when it enters a region of saddle points that exhibits a negative curvature ( (cid:107)∇ f ( w ) (cid:107) ≤ (cid:15) and λ min ( ∇ 2 f ( w )) ≤ − (cid:15) ), and GrACE property throughout the iterations, then it reaches an ( (cid:15), (cid:15) ) second order stationary point in T = 2 T thred ( f ( w 0 ) − min w f ( w )) / ( δ F thred ) =
O ((1 − β ) log( Lc m σ 2 ρc (cid:48)
c h (1 − β ) δγ(cid:15) ) (cid:15) − 10 ) iterations with high probability 1 − δ , where F thred =
O ( (cid:15) 4 ) .
The theorem implies the advantage of using stochastic momentum for SGD.
Higher β leads to reaching a second order stationary point faster.
As we will show in the following, this is due to that higher β enables escaping the saddle points faster.
In Subsection 3.2.1, we provide some key details of the proof of Theorem 1.
The interested reader can read a high-level sketch of the proof, as well as the detailed version, in Appendix G.
Remark 1: (constraints on β )
We also need some minor constraints on β so that β cannot be too close to 1.
They are 1) L (1 − β ) 3 > 1 , 2) σ 2 (1 − β ) 3 > 1 , 3) c (cid:48)
(1 − β ) 2 > 1 , 4) η ≤ 1 − βL , 5) η ≤ 1 − β(cid:15) , and 6) T thred ≥ 1 + 2 β 1 − β .
Please see Appendix E.1 for the details and discussions.
Remark 2: (escaping saddle points)
Note that Algorithm 2 reduces to CNC-SGD of Daneshmand et al.
(2018) when β
=
(i.e. without momentum).
Therefore, let us compare the results.
We show that the escape time of Algorithm 2 is T thred :=
˜ O (cid:0) (1 − β ) η(cid:15) (cid:1) (see Appendix E.3.3, especially (81-82)).
On the other hand, for CNC-SGD, based on Table 3 in their paper, is T thred =
˜ O (cid:0) 1 η(cid:15) (cid:1) .
One can clearly see that T thred of our result has a dependency 1 − β , which makes it smaller than that of
Daneshmand et al.
(2018) for any same η and consequently demonstrates escaping saddle point faster with momentum.
Remark 3: (ﬁnding a second order stationary point)
Denote (cid:96) a number such that ∀ t, (cid:107) g t (cid:107) ≤ (cid:96) .
In
Appendix G.3, we show that in the high momentum regime where (1 − β ) <<
ρ 2 (cid:96) 10 c 9 m c 2 h c (cid:48) , Algorithmis strictly better than CNC-SGD of Daneshmand et al.
(2018), which means that a higher momentum can help ﬁnd a second order stationary point faster.
Empirically, we ﬁnd out that c (cid:48)
≈
(Figure 3) and c h
≈
(Figure 4) in the phase retrieval problem, so the condition is easily satisﬁed for a wide range of β .
3.2. E SCAPING SADDLE POINTS
In this subsection, we analyze the process of escaping saddle points by SGD with momentum.
Denote t 0 any time such that ( t 0 mod T thred ) = 0 .
Suppose that it enters the region exhibiting a small gradient but a large negative eigenvalue of the Hessian (i.e. (cid:107)∇ f ( w t 0 ) (cid:107) ≤ (cid:15) and λ min ( ∇ 2 f ( w t 0 )) ≤ − (cid:15) ).
We want to show that it takes at most T thred iterations to escape the region and whenever it escapes, the function value decreases at least by F thred =
O ( (cid:15) 4 ) on expectation, where the precise expression of F thred will be determined later in
Appendix E.
The technique that we use is proving by contradiction.
Assume that the function value on expectation does not decrease at least F thred in T thred iterations.
Then, we get an upper bound of the expected distance E t
[ (cid:107) w t 0 + T thred − w t 0 (cid:107) 2 ] ≤ C upper .
Yet, by leveraging the negative curvature, we also show a lower bound of the form E t
[ (cid:107) w t 0 + T thred − w t 0 (cid:107) 2 ] ≥ C lower .
The analysis will show that the lower bound is larger than the upper bound (namely, C lower > C upper ), which leads to the contradiction and concludes that the function value must decrease at least F thred in T thred iterations on expectation.
Since
T thred =
O ((1 − β ) log( 1 (1 − β ) (cid:15) )
(cid:15) 6 ) , the dependency on β suggests that larger β can leads to smaller
T thred , which implies that larger momentum helps in escaping saddle points faster.
Lemma 1 below provides an upper bound of the expected distance.
The proof is in Appendix C.
Lemma 1.
Denote t 0 any time such that ( t 0 mod T thred ) = 0 .
Suppose that E t
[ f ( w t 0 ) − f ( w t 0 + t )] ≤ F thred for any 0 ≤ t ≤ T thred .
Then, E t
[ (cid:107) w t 0 + t − w t 0 (cid:107) 2 ] ≤ C upper ,t := 8 ηt (cid:0) F
[Table]
[Equation]
We see that C upper,t in Lemma 1 is monotone increasing with t , so we can deﬁne C upper :=
C upper , T thred .
Now let us switch to obtaining the lower bound of E t
[ (cid:107) w t 0 + T thred − w t 0 (cid:107) 2 ] .
The key to get the lower bound comes from the recursive dynamics of SGD with momentum.
Lemma 2.
Denote t 0 any time such that ( t 0 mod T thred ) = 0 .
Let us deﬁne a quadratic approximation at w t 0 , Q ( w ) :
= f ( w t 0 ) + (cid:104) w − w t 0 , ∇ f ( w t 0 ) (cid:105) + 12 ( w − w t 0 ) (cid:62) H ( w − w t 0 ) , where
H := ∇ 2 f ( w t 0 ) .
Also, deﬁne G s :
= ( I − η (cid:80) sk =1 β s − k H ) .
Then we can write w t 0 + t − w t 0 exactly using the following decomposition.
[Equation]
The proof of Lemma 2 is in Appendix D. Furthermore, we will use the quantities q v,t − 1 , q m,t − 1 , q q,t − 1 , q w,t − 1 , q ξ,t − 1 as deﬁned above throughout the analysis.
Lemma 3.
Following the notations of Lemma 2, we have that
[Equation]
We are going to show that the dominant term in the lower bound of E t
[ (cid:107) w t 0 + t − w t 0 (cid:107) 2 ] is E [-t 
-][Equation]
[Equation] q [Equation] 2 ] , which is the critical component for ensuring that the lower bound is larger than the upper bound of the expected distance.
Lemma 4.
Denote θ j :=
(cid:80) jk
=1 β j − k
= (cid:80) jk =
1 β k − 1 and λ := − λ min ( H ) .
Following the conditions and notations in Lemma 1 and Lemma 2, we have that
[Equation]
Proof.
We know that λ min ( H ) ≤ − (cid:15) < 0 .
Let v be the eigenvector of the Hessian H with unit norm that corresponds to λ min ( H ) so that Hv = λ min ( H ) v .
We have ( I − ηH ) v = v − ηλ min ( H ) v =
[Equation]
1 − ηλ min ( H )) v. Then,
[Equation]
where ( a ) is because v is with unit norm, ( b ) is by Cauchy–Schwarz inequality, ( c ) , ( d ) are by the deﬁnitions [Equation] and ( e ) is by the CNC assumption so that[-E t 
-]
[Equation]
Observe that the lower bound in (8) is monotone increasing with t and the momentum parameter β .
Moreover, it actually grows exponentially in t .
To get the contradiction, we have to show that the lower bound is larger than the upper bound.
By Lemma 1 and Lemma 3, it sufﬁces to prove the following lemma.
We provide its proof in Appendix E.
Lemma 5.
Let F thred =
O ( (cid:15) 4 ) and η 2 T thred ≤ r 2 .
By following the conditions and notations in
Theorem 1, Lemma 1 and Lemma 2, we conclude that if SGD with momentum (Algorithm 2) has the
APCG property, then we have that C lower :=
E t
[ (cid:107) q v, T thred − 1 (cid:107) 2 ]+2 η E t
[ (cid:104) q v, T thred − 1 , q m, T thred − 1 + q q, T thred − 1 + q w, T thred − 1 + q ξ, T thred − 1 (cid:105) ] > C upper .
C ONCLUSION
In this paper, we identify three properties that guarantee SGD with momentum in reaching a secondorder stationary point faster by a higher momentum, which justiﬁes the practice of using a large value of momentum parameter β .
We show that a greater momentum leads to escaping strict saddle points faster due to that SGD with momentum recursively enlarges the projection to an escape direction.
However, how to make sure that SGD with momentum has the three properties is not very clear.
It would be interesting to identify conditions that guarantee SGD with momentum to have the properties.
Perhaps a good starting point is understanding why the properties hold in phase retrieval.
We believe that our results shed light on understanding the recent success of SGD with momentum in non-convex optimization and deep learning. A CKNOWLEDGMENTS
We gratefully acknowledge ﬁnancial support from NSF IIS awards 1910077 and 1453304.
A L ITERATURE SURVEY
Heavy ball method: The heavy ball method was originally proposed by Polyak (1964).
It has been observed that this algorithm, even in the deterministic setting, provides no convergence speedup over standard gradient descent, except in some highly structure cases such as convex quadratic objectives where an “accelerated” rate is possible (Lessard et al.
; Goh (2017)).
In recent years, some works make some efforts in analyzing heavy ball method for other classes of optimization problems besides the quadratic functions.
For example, Ghadimi et al.
(2015) prove an O (1 /T ) ergodic convergence rate when the problem is smooth convex, while Sun et al.
(2019) provide a non-ergodic convergence rate for certain classes of convex problems.
Ochs et al.
(2014) combine the technique of forward-backward splitting with heavy ball method for a speciﬁc class of nonconvex optimization problem.
For stochastic heavy ball method, Loizou & Richt´arik (2017) analyze a class of linear regression problems and shows a linear convergence rate of stochastic momentum, in which the linear regression problems actually belongs to the case of strongly convex quadratic functions.
Other works includes (Gadat et al.
(2016)), which shows almost sure convergence to the critical points by stochastic heavy ball for general non-convex coercive functions.
Yet, the result does not show any advantage of stochastic heavy ball over other optimization algorithms like SGD.
Can et al.
(2019) show an accelerated linear convergence to a stationary distribution under Wasserstein distance for strongly convex quadratic functions by SGD with stochastic heavy ball momentum.
Yang et al.
(2018) provide a uniﬁed analysis of stochastic heavy ball momentum and Nesterov’s momentum for smooth non-convex objective functions.
They show that the expected gradient norm converges at rate
O (1 / √ t ) .
Yet, the rate is not better than that of the standard SGD.
We are also aware of the works (Ghadimi & Lan (2016; 2013)), which propose some variants of stochastic accelerated algorithms with ﬁrst order stationary point guarantees.
Yet, the framework in (Ghadimi & Lan (2016; 2013)) does not capture the stochastic heavy ball momentum used in practice.
There is also a negative result about the heavy ball momentum.
Kidambi et al.
(2018) show that for a speciﬁc strongly convex and [Figure] [Figure] strongly smooth problem, SGD with heavy ball momentum fails to achieving the best convergence rate while some algorithms can.
Reaching a second order stationary point: As we mentioned earlier, there are many works aim at reaching a second order stationary point.
We classify them into two categories: specialized algorithms and simple GD/SGD variants.
Specialized algorithms are those designed to exploit the negative curvature explicitly and escape saddle points faster than the ones without the explicit exploitation (e.g. Carmon et al.
Agarwal et al.
(2017); Allen-Zhu & Li (2018);
Xu et al.
Simple
GD/SGD variants are those with minimal tweaks of standard GD/SGD or their variants (e.g. Ge et al.
; Levy (2016);
Fang et al.
; Jin et al.
; Daneshmand et al.
Staib et al.
Our work belongs to this category.
In this category, perhaps the pioneer works are (Ge et al.
and (Jin et al.
Jin et al.
(2017) show that explicitly adding isotropic noise in each iteration guarantees that GD escapes saddle points and ﬁnds a second order stationary point with high probability.
Following (Jin et al.
, Daneshmand et al.
(2018) assume that stochastic gradient inherently has a component to escape.
Speciﬁcally, they make assumption of the
Correlated Negative Curvature (CNC) for stochastic gradient g t so that E t [ (cid:104) g t , v t (cid:105) 2 ] ≥ γ > 0 .
The assumption allows the algorithm to avoid the procedure of perturbing the updates by adding isotropic noise.
Our work is motivated by (Daneshmand et al.
but assumes CNC for the stochastic momentum m t instead.
Very recently, Jin et al.
(2019) consider perturbing the update of SGD and provide a second order guarantee.
Staib et al.
(2019) consider a variant of RMSProp (Tieleman & Hinton (2012)), in which the gradient g t is multiplied by a preconditioning matrix G t and the update is w t
+1 =
w t − G − 1 / 2 t g t .
The work shows that the algorithm can help in escaping saddle points faster compared to the standard SGD under certain conditions.
Fang et al.
(2019) propose average-SGD, in which a sufﬁx averaging scheme is conducted for the updates.
They also assume an inherent property of stochastic gradients that allows SGD to escape saddle points.
We summarize the iteration complexity results of the related works for simple SGD variants on
Table 1.
The readers can see that the iteration complexity of (Fang et al.
and (Jin et al.
(2019)) are better than (Daneshmand et al.
; Staib et al.
and our result.
So, we want to explain the results and clarify the differences.
First, we focus on explaining why the popular algorithm, SGD with heavy ball momentum, works well in practice, which is without the sufﬁx averaging scheme used in (Fang et al.
and is without the explicit perturbation used in (Jin et al.
Speciﬁcally, we focus on studying the effect of stochastic heavy ball momentum and showing the advantage of using it.
Furthermore, our analysis framework is built on the work of (Daneshmand et al.
We believe that, based on the insight in our work, one can also show the advantage of stochastic momentum by modifying the assumptions and algorithms in (Fang et al.
or (Jin et al.
and consequently get a better dependency on (cid:15) [Table]
B L EMMA 6, 7, AND
In the following, Lemma 7 says that under the APAG property, when the gradient norm is large, on expectation SGD with momentum decreases the function value by a constant and consequently makes progress.
On the other hand, Lemma 8 upper-bounds the increase of function value of the next iterate (if happens) by leveraging the GrACE property.
Lemma 6.
If SGD with momentum has the APAG property, then, considering the update step w t +1 =
w t − ηm t , we have that E t [
[Equation]
+
[Equation]
Proof.
By the L -smoothness assumption,
[Equation]
Taking the expectation on both sides.
We have
[Equation]
where we use the APAG property in the last inequality.
Lemma 7.
Assume that the step size η satisﬁes η ≤ {+(cid:15) 2 
+}
Lc 2 m .
If SGD with momentum has the
APAG property, then, considering the update step w t +1 =
w t − ηm t , we have that E t
[ f ( w t +1 )] ≤ f ( w t ) − η 4 (cid:15) 2 when (cid:107)∇ f ( w t ) (cid:107) ≥ (cid:15) .
Proof.
E t [ f ( w t +1 )
− f ( w t )]
Lemma 6 ≤
[Equation]
where the last inequality is due to the constraint of η .
Lemma 8.
If SGD with momentum has the GrACE property, then, considering the update step w t +1 =
w t − ηm t , we have that E t
[ f ( w t +1 )] ≤ f ( w t ) + η 2 c h + ρη 3 6 c 3 m .
Proof.
Consider the update rule w t
+1 =
w t − ηm t , where m t represents the stochastic momentum and η is the step size.
By ρ -Lipschitzness of Hessian, we have f ( w t +1 ) ≤ f ( w t ) − η
(cid:104)∇ f [Equation] w t ) , g t
[Equation]
+ η (cid:104)∇ f ( w t ) , g t [Equation] m t [Equation]
+ η [Equation] (cid:62) t ∇
[Equation]
f ( w t [Equation] t + [Equation]
3 6 (cid:107) m t (cid:107) 3 .
Taking the conditional expectation, one has
E t [
[Equation]
t +1 )]
[Equation]
[Equation]
[Equation]
C P ROOF OF L EMMA
Lemma 1 Denote t
any time such that ( t 0 mod T thred ) = 0 .
Suppose that E t
[ f ( w t 0 ) − f ( w t 0 + t )] ≤ F thred for any 0 ≤ t ≤ T thred .
Then,
[Equation]
Proof.
Recall that the update is w t {+
+}
+1 =
w t 0 − rm t 0 , and w t 0 + t = w t 0 + t − 1 − ηm t 0 + t − 1 , for t > 1 .
We have that
[Equation]
where the ﬁrst inequality is by the triangle inequality and the second one is due to the assumption that (cid:107) m t (cid:107) ≤ c m for any t .
Now let us denote
• α s := (cid:80) t − 1 − s j = β j •
A t − 1 := (cid:80) t − 1 s =
α s
and let us rewrite g t = ∇ f ( w t ) + ξ t , where ξ t is the zero-mean noise.
We have that
[Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
[Equation] [Table] [Equation] [Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation] [Table] [Equation] [Table]
[Equation]
[Table]
[Equation]
[Equation]
{+To proceed, we need to upper bound E t 
 [4 η 2 (cid:107) (cid:80) t − 1 s =
+}
α s ∇ f ( w t 0 + s ) (cid:107) 2 ] .
We have that
E t 0 [4 η 2 (cid:107)
[Equation]
where ( a ) is by Jensen’s inequality, ( b ) is by max s α s ≤ 1 1 − β , and ( c ) is by A t − 1 ≤ t 1 − β .
Now let us switch to bound the other term.
[Equation]
where ( a ) is because E t 0 [ ξ (cid:62) t
+ i ξ t 0 + j ] = 0 for i (cid:54) =
j , ( b ) is by that (cid:107) ξ t (cid:107) 2 ≤ σ 2 and max t α t ≤ 1 1 − β .
Combining (14), (15), (16), (17),
[Equation]
{+Now we need to bound E t 
+}
[ (cid:80)
[Equation]
f
[Equation]
w t[-
-]
[Equation]
s )
[Equation]
By using ρ -Lipschitzness of Hessian, we have that
[Equation]
By adding η (cid:104)∇ f ( w t 0 + s − 1 ) , g t
+ s − 1 (cid:105) on both sides, we have
[Equation]
Taking conditional expectation on both sides leads to
[Equation]
[Equation]
where E {+t 
+}
+ s − 1 [ η
[Equation]
f ( w
[Equation]
[Equation]
∇ 2 f
[Equation]
{+w t 
+}
{++ s − 1 ) m t 
+}
+ s − 1 ] ≤ η 2 c h by the GrACE property.
We have that for t 0 ≤ t 0 + s −
[Equation]
Summing the above inequality from s = 2 , 3 , . .
. , t leads to
E t
[ t [Equation]
[Equation]
where ( a ) is by the assumption (made for proving by contradiction) that E t
[ f ( w t 0 ) − f ( w t 0 + s )] ≤
F thred for any 0 ≤ s ≤ T thred .
By (21) with s = 1 and η = r , we have
[Equation]
By (23) and (24), we know that
[Equation]
where ( a ) is by the constraint that η 2 t ≤ r 2 for 0 ≤ t ≤ T thred and ( b ) is by the constraint that r ≥ η .
By combining (25) and (18)
[Equation]
[Equation]
D P ROOF OF L EMMA 2 AND L EMMA
Lemma 2 Denote t 0 any time such that ( t 0 mod T thred ) = 0 .
Let us deﬁne a quadratic approximation at w t 0 , Q ( w ) :
= f ( w t 0 ) + (cid:104) w − w t 0 , ∇ f ( w t 0 ) (cid:105) + 12 ( w − w t 0 ) (cid:62) H ( w − w t 0 ) , where
H := ∇ 2 f ( w t 0 ) .
Also, deﬁne G s :
= ( I − η (cid:80) sk =1 β s − k H ) and
• q v,t − 1 :=
(cid:0) Π t − 1 j =1 G j (cid:1)(cid:0) − rm t 0 (cid:1) .
• q m,t − 1 := − (cid:80) t − 1 s =1 (cid:0) Π t − 1 j = s +1 G j (cid:1) β s m t 0 .
• q q,t − 1 := − (cid:80) t − 1 s =1 (cid:0) Π t − 1 j
=
s
+
G j (cid:1) (cid:80) sk =1 β s − k (cid:0) ∇ f ( w t 0 + k ) − ∇ Q ( w t 0 + s ) (cid:1) .
• q w,t − 1 := − (cid:80) t − 1 s =1 (cid:0) Π t − 1 j = s +1 G j (cid:1) (cid:80) sk =1 β s − k ∇ f ( w t 0 ) .
• q ξ,t − 1 := − (cid:80) t − 1 s =1 (cid:0) Π t − 1 j = s +1 G j (cid:1) (cid:80) s k =1 β
s − k ξ t 0 + k .
Then, w t 0 + t − w t 0 = q v,t − 1 + ηq m,t − 1 + ηq q,t − 1 + ηq w,t − 1 + ηq ξ,t − 1 .
Notations:
Denote t 0 any time such that ( t 0 mod T thred ) = 0 .
Let us deﬁne a quadratic approximation at
[Equation]
t
[Equation]
[Equation]
[Equation]
where H := ∇ 2 f ( w t 0 ) .
Also, we denote
[Equation]
Proof.
First, we rewrite m t 0 + j for any j ≥ 1 as follows.
[Equation]
We have that
[Equation]
[Equation]
where ( a ) is by using (29) with j = t − 1 , ( b ) is by subtracting and adding back the same term, and ( c ) is by ∇ Q ( w t
+ t − 1 ) = ∇ f ( w t 0 )
+ H ( w t 0 + t − 1 − w
[Equation]
To continue, by using the nations in (28)
[Equation]
we can rewrite (30) as
[Equation]
Recursively expanding (31) leads to
[Equation]
where ( a ) we use the notation that Π t − 1 j = s G j :=
G s × G s +1 × . .
G t − 1 and the notation that Π t − 1 j = t G j = 1 and ( b ) is by the update rule.
By using the deﬁnitions of { q (cid:63),t − 1 } in the lemma statement, we complete the proof.
Lemma
Following the notations of Lemma 2, we have that
[Equation]
[Equation]
Proof.
Following the proof of Lemma 2, we have
[Equation]
Therefore, by using (cid:107) a + b (cid:107) 2 ≥ (cid:107) a (cid:107) 2 + 2 (cid:104) a, b (cid:105) ,
E t
[ (cid:107) w t 0 + t − w t 0 (cid:107) 2 ] ≥ E t
[ (cid:107) q v,t − 1 (cid:107) 2 ] + 2 η E t
[ (cid:104) q v,t − 1 , q m,t − 1 + q q,t − 1 + q w,t − 1 + q ξ,t − 1 (cid:105) ] .
E P ROOF OF L EMMA
Lemma 5 Let F thred =
O ( (cid:15) 4 ) and η 2 T thred ≤ r 2 .
By following the conditions and notations in
Theorem 1, Lemma 1 and Lemma
[Table]
we conclude that if SGD with momentum (Algorithm 2) has the
APCG property, then we have that
[Equation]
[Table]
W.l.o.g, we assume that c m , L , σ 2 , c (cid:48) , c h , and ρ are not less than one and that (cid:15) ≤ 1 .
E. S OME CONSTRAINTS ON β .
We require that parameter β is not too close to 1 so that the following holds,
• 1) L (1 − β ) 3 > 1 .
• 2) σ 2 (1 − β ) 3 > 1 .
• 3) c (cid:48) (1 − β ) 2 > 1 . • 4)
η ≤ 1 − βL .
η ≤ 1 − β(cid:15) .
• 6) T thred ≥ c (1 − β ) η(cid:15) [Table]
[Table]
σ 2 ρc
[Table]
c h (1 − β ) δγ(cid:15) ) ≥ 1 + 2 β 1 − β .
The constraints upper-bound the value of β .
That is, β cannot be too close to 1.
We note that the β dependence on L , σ , and c (cid:48) are only artiﬁcial.
We use these constraints in our proofs but they are mostly artefacts of the analysis.
For example, if a function is L -smooth, and L < 1 , then it is also 1 -smooth, so we can assume without loss of generality that L > 1 .
Similarly, the dependence on σ is not highly relevant, since we can always increase the variance of the stochastic gradient, for example by adding an O (1) gaussian perturbation.
E. S OME LEMMAS
To prove Lemma 5, we need a series of lemmas with the choices of parameters on Table 3.
Upper bounding E t
[ (cid:107) q q,t − 1 (cid:107) ] :
Lemma 9.
Following the conditions in Lemma 1 and Lemma 2, we have
[Equation]
Proof.
[Equation]
[Table]
(37) where ( a ) , ( c ) , ( d ) is by triangle inequality, ( b ) is by the fact that (cid:107) Ax (cid:107) 2 ≤ (cid:107) A (cid:107) 2 (cid:107) x (cid:107) 2 for any matrix
A and vector x .
Now that we have an upper bound of (cid:107)∇ f ( w t 0 + k ) − ∇ f ( w t 0 + s ) (cid:107)
[Equation]
where ( a ) is by the assumption of L-Lipschitz gradient and ( b ) is by applying the triangle inequality ( s − k ) times and that (cid:107) w t − w t − 1 (cid:107) ≤ η (cid:107) m t − 1 (cid:107) ≤ ηc m , for any t .
We can also derive an upper bound of E t
[ (cid:107)∇ f ( w t 0 + s
[Equation]
∇ Q
[Equation]
w
[Equation]
[Equation]
[Equation]
Above, ( a ) is by the fact that if a function f ( · ) has ρ Lipschitz Hessian, then
[Equation]
[Equation]
(c.f. Lemma 1.2.4 in (Nesterov (2013))) and using the deﬁnition that
[Equation] [Equation]
Combing (37), (38) [Equation] (39), we have that
[Equation]
where on the last line we use the notation that
[Equation]
[Equation]
To continue
[Equation]
let us analyze (cid:107)
[Equation]
Π
[Equation]
− 1 j = s +1 G j
[Equation]
2 ﬁrst
[Equation]
Above, we use the notation that θ j := (cid:80) jk
=
β j − k .
For (a)
[Equation]
it is due to that λ := − λ min ( H ) , λ max ( H )
≤ L
[Equation]
and the choice of η so that 1 ≥ ηL 1 − β , or equivalently,
[Equation]
For ( b )
[Equation]
it is due to that θ j ≥ 1 for any j and λ ≥ (cid:15) .
Therefore, we can upper-bound the ﬁrst term on r
[Equation]
h
[Equation]
s of (42) as
[Equation]
where ( a ) is by that fact that (cid:80) ∞ k =1 β k k ≤ β [Equation]
[Equation] 2 for any 0 ≤ β < 1 , ( b ) is by using (44) , and ( c ) is by using that (cid:80) ∞ s =1 ( 1 1+ η(cid:15) )
s ≤ 1 η(cid:15) .
Now let us switch to bound (cid:80) t − 1 s =1 (cid:107) (cid:0) Π t − 1 j = s +1 G j (cid:1) (cid:107)
(cid:80) sk =1 β s − k ρ 2 ( ν s + ν ) on (42).
We have that
[Equation]
where ( a ) is by the fact that (cid:80) sk =1 β s − k ≤ 1 / (1 − β ) , ( b ) is by (44), ( c ) is by using that (cid:80) ∞ s =1 ( 1 1+ η(cid:15) )
s ≤ 1 η(cid:15) , ( d ) is by (cid:80) ∞
k =1 z k k ≤ z (1 − z ) 2 for any | z | ≤ 1 and substituting z = 1 1+ η(cid:15) , which leads to (cid:80)
[Equation]
in which the last inequality is by chosen the step size η so that η(cid:15) ≤ 1 .
By combining (42)
[Equation] [Equation]
and (47)
[Equation]
we have that
[Equation]
(48) which completes the proof.
Upper bounding (cid:107) q
[Equation]
Lemma 10.
Following the conditions in Lemma 1 and Lemma 2, we have
[Equation]
Proof.
[Equation]
where the last inequality is because η is chosen so that 1 ≥ ηL 1 − β
and the fact that λ max ( H ) ≤ L .
Lower bounding E {+t 
+}
[2 η
[Equation]
Lemma 11.
Following the conditions in Lemma 1 and Lemma 2, we have
[Equation]
Proof.
By the results of Lemma 9 and Lemma
[Equation]
Lower bounding
[Equation]
Lemma 12.
Following the conditions in Lemma 1 and Lemma 2, we have
[Equation]
Proof.
[Equation]
where ( a ) holds for some coefﬁcients α k , ( b ) is by the tower rule, ( c ) is because q v,t − 1 is measureable with t 0 , and ( d ) is by the zero mean assumption of ξ ’s.
Lower bounding E t
[Equation]
Lemma 13.
Following the conditions in Lemma 1 and Lemma 2, we have
[Equation]
Proof.
[Equation]
where ( a ) is by deﬁning the matrix B := (cid:0)
Π t − 1 j =1 G j (cid:1) (cid:62) (cid:0) (cid:80)
[Equation]
− 1 s =1 (cid:0) Π t − 1 j = s +1 G j (cid:1) β s (cid:1) .
For (b), notice that the matrix B is symmetric positive semideﬁnite.
To see that the matrix B is symmetric positive semideﬁnite, observe that each G j :=
( I − η (cid:80) jk
=1 β j − k H ) can be written in the form of G j =
UD j U (cid:62) for some orthonormal matrix U and a diagonal matrix D j .
Therefore, the matrix product (cid:0) Π t − 1 j =1 G j (cid:1) (cid:62) (cid:0) Π t − 1 j = s +
G j (cid:1)
=
U (Π t − 1 j =1 D j )(Π t − 1 j = s +1 D j ) U (cid:62) is symmetric positive semideﬁnite as long as each G j is.
So, ( b ) is by the property of a matrix being symmetric positive semideﬁnite.
Lower
[Equation]
Lemma 14.
Following the conditions in Lemma 1 and Lemma
[Equation]
if SGD with momentum has the APCG property
[Equation]
then
[Equation]
Proof.
Deﬁne
[Equation]
[Table]
[Equation]
where ( a ) is by the APCG property.
We also have that
[Equation]
where (a) and (b) is by (44).
Substituting the result back to (58)
[Equation]
we get
[Equation] [Equation]
Using the fact that (cid:107)∇ f ( w t 0 ) (cid:107) ≤ (cid:15) completes the proof.
E. P ROOF OF L EMMA
Recall that the strategy is proving by contradiction.
Assume that the function value does not decrease at least F thred in T thred iterations on expectation.
Then, we can get an upper bound of the expected distance E t
[ (cid:107) w t 0 + T thred − w t 0 (cid:107) 2 ] ≤ C upper but, by leveraging the negative curvature, we can also show a lower bound of the form E t
[ (cid:107) w t 0 + T thred − w t 0 (cid:107) 2 ] ≥ C lower .
The strategy is showing that the lower bound is larger than the upper bound, which leads to the contradiction and concludes that the function value must decrease at least F thred in T thred iterations on expectation.
To get the contradiction, according to Lemma 1 and Lemma 3, we need to show that
[Equation]
>
[Equation]
Yet, by Lemma 13 and Lemma 12, we have that η E t
[ (cid:104) q v, T thred − 1 , q m, T thred − 1 (cid:105) ] ≥ 0 and η E t
[ (cid:104) q v, T thred − 1 , q ξ, T thred − 1 (cid:105) ]
= 0 .
So, it sufﬁces to prove that
[Equation]
and it sufﬁces to show that
• 14 E t
[ (cid:107) q v, T thred − 1 (cid:107) 2 ] + 2 η E t
[ (cid:104) q v, T thred − 1 , q q, T thred − 1 (cid:105) ] ≥ 0 .
• 14 E t
[ (cid:107) q v, T thred − 1 (cid:107) 2 ] + 2 η E t
[ (cid:104) q v, T thred − 1 , q w, T thred − 1 (cid:105) ] ≥ 0 .
• 14 E t
[ (cid:107) q v, T thred − 1 (cid:107) 2 ] ≥ C upper .
[Table]
By Lemma 4 and Lemma
[Table]
we have that
[Table]
[Equation]
To show that the above is nonnegative, it sufﬁces to show that
[Equation]
64) and
[Equation]
and
[Equation]
[Equation]
Now w.l.o.g, we assume that c m , L , σ 2 , c (cid:48) , and ρ are not less than one and that (cid:15) ≤ 1 .
By using the values of parameters on Table 3, we have the following results; a sufﬁcient condition of (64) is that
[Equation]
A sufﬁcient condition of (65) is that
[Equation]
[Equation]
and
[Equation]
69) and
[Equation]
A sufﬁcient condition of (66) is that
[Equation]
71) and a sufﬁcient condition for the above (71), by the assumption that both σ 2 ≥ 1 and c m ≥ 1 , is
[Equation]
Now let us verify if (67), (68), (69), (70), (72) are satisﬁed.
For (67), using the constraint of c η on Table 3, we have that 1 c η ≥ c 5 m ρL 2 σ 2 c (cid:48)
c
h c. Using this inequality, it sufﬁces to let c r ≥ c 0 (cid:15) 2 c 3 m ρLσ 2 c (cid:48) c h (1 − β ) 2 for getting (67), which holds by using the constraint that c (cid:48) (1 − β ) 2 > 1 and (cid:15) ≤ 1 .
For (68), using the constraint of c F on Table 3, we have that 1 c F ≥ c 4 m [Equation] Lσ
[Equation]
Using this inequality, it sufﬁces to let c r ≥
[Equation]
3 , which holds by using the constraint that σ 2 (1 − β ) 3 > 1 .
For (69), it needs
[Equation]
r , which hold by using the constraint that σ 2 (1 − β ) 3 > 1 .
For (70), it sufﬁces to let (1 − β )
[Equation]
2 m ρ ≥
[Equation]
h ≥ c r which holds by using the constraint that σ 2 (1 − β ) 3 > 1 .
For (72), it sufﬁces to let
[Equation]
r , which holds by using the constraint that L (1 − β ) 3 > 1 and (cid:15) ≤ 1 .
Therefore, by choosing the parameter values as Table 3, we can guarantee that
[Equation]
E.3.
P ROVING THAT 14 E
[Equation]
[Equation]
By Lemma 4 and Lemma 14, we have that
[Equation]
To show that the above is nonnegative, it sufﬁces to show that
[Equation]
A sufﬁcient condition is c r c η ≥ 8 (cid:15) 4 c (cid:48) 1 − β .
Using the constraint of c η on Table 3, we have that 1 c η ≥ c 5 m [Equation] 2 σ
[Equation]
c (cid:48)
[Equation]
So, it sufﬁces to let c r ≥
[Equation] [Equation]
β ) , which holds by using the constraint that
L (1 − β ) 3 > 1 (so that L (1 − β ) > 1 ) and (cid:15) ≤ 1 .
E.3. P ROVING THAT 14 E
[Equation]
From Lemma 4 and Lemma 1, we need to show that
[Equation]
We know that 14 (cid:0) Π T
[Equation] [Equation]
=
[Equation]
[Equation]
1 + ηθ j
[Equation]
) (cid:1) 2 r 2 γ .
It sufﬁces to show that
[Equation]
Note that the left hand side is exponentially growing in T thred .
We can choose the number of iterations T thred large enough to get the desired result.
Speciﬁcally, we claim that T thred ≥ c (1 − β ) [Equation] log( [Equation]
σ
[Equation]
c h [Equation] [Equation] ) δγ(cid:15) ) for some constant c > 0 .
To see this, let us ﬁrst apply log on both sides of (76),
[Equation]
where we denote a := 4 η (cid:0) F thred + [Equation] [Equation]
+
[Equation]
3 r [Equation] 3 m (cid:1) (
[Equation]
and b := 4 r 2 σ 2 (1 − β ) 2 + 2 η 2 (cid:0) β 1 − β (cid:1) 2 c 2 m + r 2 c 2 m .
To proceed, we are going to use the inequality log(1 + x ) ≥ x 2 , for x ∈ [0 , ∼ 2 . 51] .
We have that
[Equation]
[Equation]
as guaranteed by the constraint of η .
So,
[Equation]
where ( a ) is by using the inequality log(1 + x ) ≥ x 2 with x = ηθ j (cid:15) ≤ 1 and ( b ) is by making
T thred [Equation] [Equation] ≥ β
[Equation]
2 , which is equivalent to the condition that
[Equation]
Now let us substitute the result of (79) back to (77).
We have that
[Equation]
which is what we need to show.
By choosing T thred large enough,
[Equation]
for some constant c > 0 , we can guarantee that the above inequality (81) holds.
F P ROOF OF L EMMA
Lemma 15 (Daneshmand et al.
Let us deﬁne the event
Υ k :=
{(cid:107)∇ f ( w k T thred ) (cid:107) ≥ (cid:15) or λ min ( ∇ 2 f ( w k T thred )) ≤ − (cid:15) } .
The complement is Υ ck :
= {(cid:107)∇ f ( w k T thred ) (cid:107) ≤ (cid:15) and λ min ( ∇ 2 f ( w k T thred )) ≥ − (cid:15) } , which suggests that w k T thred is an ( (cid:15), (cid:15) )
-second order
stationary points.
Suppose that
[Equation]
Set T = 2 T thred (cid:0) f
[Equation]
− min
[Equation]
f
[Equation]
(cid:1) / ( δ ∆) .
We return w uniformly randomly from w 0 , w T thred , w 2 T thred , . . .
, w k T thred , . .
. , w K T thred , where K := (cid:98) T/ T thred (cid:99) .
Then, with probability at least 1 − δ , we will have chosen a w k where Υ k did not occur.
Proof.
Let P k be the probability that Υ k occurs.
[Equation]
[Equation]
Summing over all K , we have
K +
K (cid:88) k
=
E [ f ( w ( k +1) T thred ) − f ( w k T thred )] ≤ ∆
K +
K (cid:88) k
= ( δ/ 2 − P k ) ⇒
K +
K (cid:88) k
=
P k ≤ δ/
+ f ( w 0 ) − min w f ( w ) ( K + 1)∆ ≤ δ ⇒
K +
K (cid:88) k
= (1 − P k ) ≥ 1 − δ.
G P ROOF OF T HEOREM
Theorem 1 Assume that the stochastic momentum satisﬁes CNC.
Set r = O ( (cid:15) 2 ) ,
η = O ( (cid:15) 5 ) , and
T thred =
c (1 − β ) η(cid:15) log( Lc m σ 2 ρc (cid:48) c h (1 − β ) δγ(cid:15) )
=
O ((1 − β ) log( Lc m σ 2 ρc (cid:48) c h (1 − β ) δγ(cid:15) ) (cid:15) − 6 ) for some constant c > 0 . If
SGD with momentum (Algorithm 2) has APAG property when gradient is large ( (cid:107)∇ f ( w ) (cid:107) ≥ (cid:15) ), APCG T thred property when it enters a region of saddle points that exhibits a negative curvature ( (cid:107)∇ f ( w ) (cid:107) ≤ (cid:15) and λ min ( ∇ 2 f ( w )) ≤ − (cid:15) ), and GrACE property throughout the iterations, then it reaches an ( (cid:15), (cid:15) ) second order stationary point in T = 2 T thred ( f ( w 0 ) − min w f ( w )) / ( δ F thred ) =
O ((1 − β ) log( Lc m σ 2 ρc (cid:48)
c h (1 − β ) δγ(cid:15) )
(cid:15) − 10 ) iterations with high probability 1 − δ , where F thred = O ( (cid:15) 4 ) .
G. P ROOF SKETCH OF T HEOREM
In this subsection, we provide a sketch of the proof of Theorem 1.
The complete proof is available in
Appendix G. Our proof uses a lemma in (Daneshmand et al.
(2018)), which is Lemma 15 below.
The lemma guarantees that uniformly sampling a w from { w k T thred } , k = 0 , 1 , 2 , . . .
, (cid:98) T/ T thred (cid:99) gives an ( (cid:15), (cid:15) )
-second order
stationary point with high probability.
We replicate the proof of Lemma 15 in
Appendix F.
Lemma 15.
(Daneshmand et al.
Let us deﬁne the event
Υ k :=
{(cid:107)∇ f ( w k T thred ) (cid:107) ≥ (cid:15) or λ min ( ∇ 2 f ( w k T thred )) ≤ − (cid:15) } .
The complement is Υ ck :
= {(cid:107)∇ f ( w k T thred ) (cid:107) ≤ (cid:15) and λ min ( ∇ 2 f ( w k T thred )) ≥ − (cid:15) } , which suggests that w k T thred is an ( (cid:15), (cid:15) )
-second order
stationary points.
Suppose that
[Equation]
[Equation]
Set T = 2 T thred (cid:0) f ( w 0 ) − min w f ( w ) (cid:1) / ( δ ∆) .
We return w uniformly randomly from w 0 , w T thred , w 2 T thred , . . .
, w k T thred , . .
. , w K T thred , where K := (cid:98) T/ T thred (cid:99) .
Then, with probability at least 1 − δ , we will have chosen a w k where Υ k did not occur.
To use the result of Lemma 15, we need to let the conditions in (86) be satisﬁed.
We can bound E [ f ( w ( k +1)
T thred ) − f ( w k T thred ) | Υ k ] ≤ −F thred , based on the analysis of the large gradient norm regime (Lemma 7) and the analysis for the scenario when the update is with small gradient norm but a large negative curvature is available (Subsection 3.2.1).
For the other condition, E [ f ( w ( k +1) T thred ) − f ( w k T thred ) | Υ c k ] ≤ δ F thred, it requires that the expected amortized increase of function value due to taking the large step size r is limited (i.e. bounded by δ F thred) when w k T thred is a second order stationary point.
By having the conditions satisﬁed, we can apply Lemma 15 and ﬁnish the proof of the theorem.
G. F ULL PROOF OF T HEOREM
Proof.
Our proof is based on Lemma 15.
So, let us consider the events in Lemma 15, Υ k
[Equation] ≤ − (cid:15) } .
We ﬁrst show that E [ f ( w ( k +1)
T thred ) − f ( w k T thred ) | Υ k ] ≤ F thred .
When (cid:107)∇ f ( w k T thred ) (cid:107) ≥ (cid:15) :
Consider that Υ k is the case that (cid:107)∇ f ( w k T thred ) (cid:107) ≥ (cid:15) .
Denote t 0 :=
k T thred in the following.
We have that
[Equation]
(87) where ( a ) is by using Lemma 6 with step size r , ( b ) is by using Lemma 8, ( c ) is due to the constraint that η 2 T thred ≤ r 2 , ( d ) is by the choice of r , ( e ) is by (cid:107)∇ f ( w t ) (cid:107) ≥ (cid:15) , ( f ) is by the choice of r so that r ≤ (cid:15) 2 4( Lc 2 m + c h ) , and ( g ) is by
[Equation]
thred .
When (cid:107)∇ f ( w k T thred ) (cid:107) ≤ (cid:15) and λ min ( ∇ 2 f ( w k T thred )) ≤ − (cid:15) :
The scenario that Υ k is the case that (cid:107)∇ f ( w k T thred ) (cid:107) ≤ (cid:15) and λ min ( ∇ 2 f ( w k T thred )) ≤ − (cid:15) has been analyzed in Appendix E, which guarantees that E [ f ( w t 0 + T thred ) − f ( w t 0 )] ≤ −F thred under the setting.
When (cid:107)∇ f ( w k T thred ) (cid:107) ≤ (cid:15) and λ min ( ∇ 2 f ( w k T thred )) ≥ − (cid:15) :
Now let us switch to show that E [ f ( w ( k +1)
T thred ) − f ( w k T thred ) | Υ ck ] ≤ δ F thred.
Recall that Υ ck means that (cid:107)∇ f ( w k T thred ) (cid:107) ≤ (cid:15) and λ min ( ∇ 2 f ( w k T thred )) ≥ − (cid:15) .
Denote t 0 :=
k T thred in the following.
We have that
[Equation] [Equation]
where ( a ) is by using Lemma 8 with step size r , ( b ) is by using Lemma 8 with step step size η , ( c ) is by setting η 2 T thred ≤ r 2 and η ≤ r , ( d ) is by the choice of r so that 8 r 2 c h ≤ δ F thred .
Now we are ready to use Lemma 15, since both the conditions are satisﬁed.
According to the lemma and the choices of parameters value on Table 3, we can set T = 2 T thred (cid:0) f ( w 0 ) − min w f ( w ) (cid:1) /
( δ F thred ) =
O (
[Equation]
1 − β ) log
[Equation]
) (cid:15) − 10 ) , which will return a w that is an ( (cid:15), (cid:15) ) second order stationary point.
Thus, we have completed the proof.
G. C OMPARISON TO D ANESHMAND ET AL . (2018)
Theorem 2 in Daneshmand et al.
(2018) states that, for CNC-SGD to ﬁnd an ( (cid:15), ρ 1 / 2 (cid:15) ) stationary point, the total number of iterations is T
= O ( (cid:96)
10 L 3 δ 4 γ 4 log 2 ( (cid:96)L(cid:15)δγ ) (cid:15) − 10 ) , where (cid:96) is the bound of the stochastic gradient norm (cid:107) g t (cid:107) ≤ (cid:96) which can be viewed as the counterpart of c m in our paper.
By translating their result for ﬁnding an ( (cid:15), (cid:15) ) stationary point, it is T
= O (
[Equation] [Table] [Equation] [Table]
[Equation]
[Table]
[Equation]
ρ(cid:96)L(cid:15)δγ ) (cid:15) − 10 ) .
On the other hand, using the parameters value on Table 3, we have that T = 2 T thred
[Equation]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
for Algorithm 2.
Before making a comparison, we note that their result does not have a dependency on the variance of stochastic gradient (i.e. σ 2 ), which is because they assume that the variance is also bounded by the constant (cid:96) (can be seen from (86) in the supplementary of their paper where the variance terms (cid:107) ζ i (cid:107) are bounded by (cid:96) ).
Following their treatment, if we assume that σ 2 ≤ c m , then on (71) we can instead replace ( σ 2 + 3 c 2 m ) with 4 c 2 m and on (72) it becomes 1 ≥ 576 c 3 m ρc r (cid:15) (1 − β ) 3 .
This will remove all the parameters’ dependency on σ 2 .
Now by comparing ˜ O
((1 − β ) c
9 m c 2 h c (cid:48) ·
[Table]
(cid:15) − 10 ) of ours and
T =
˜ O ( ρ 2 (cid:96) 10 · ρ 3 L 3 δ 4 γ 4 (cid:15) − 10 ) of Daneshmand et al.
(2018), we see that in the high momentum regime where (1 − β ) <<
ρ 2 (cid:96) 10 c 9 m c 2 h c (cid:48) , Algorithm 2 is strictly better than that of Daneshmand et al.
(2018), which means that a higher momentum can help to ﬁnd a second order stationary point faster.