Escaping Saddle Points Faster with Stochastic Momentum

1 INTRODUCTION . 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 . ( 2012 ) ; He et al . ( 2016 ) ; Cubuk et al . ( 2018 ) ; 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 . ( 2017 ) ; Loshchilov & Hutter ( 2019 ) ; Wilson et al . ( 2017 ) ) . 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 justification 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 1Heavy ball momentum is the default choice of momentum method in PyTorch and Tensorflow , instead of Nesterov ’ s momentum . See the manual pages https : //pytorch.org/docs/stable/_modules/ torch/optim/sgd.html and https : //www.tensorflow.org/api_docs/python/tf/ keras/optimizers/SGD . Algorithm 1 : SGD with stochastic heavy ball momentum 1 : Required : Step size parameter η and momentum parameter β . 2 : Init : w0 ∈ Rd and m−1 = 0 ∈ Rd . 3 : for t = 0 to T do 4 : Given current iterate wt , obtain stochastic gradient gt : = ∇f ( wt ; ξt ) . 5 : Update stochastic momentum mt : = βmt−1 + gt . 6 : Update iterate wt+1 : = wt − ηmt . 7 : end for packages such as PyTorch and Tensorflow . 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 benefit 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 finding a second-order stationary point for smooth non-convex optimization by SGD with stochastic heavy ball momentum . Specifically , we consider the stochastic nonconvex optimization problem , minw∈Rd 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 ( , ) -second-order stationary point w satisfies ‖∇f ( w ) ‖ ≤ and ∇2f ( w ) − I . ( 1 ) Obtaining a second order guarantee has emerged as a desired goal in the nonconvex optimization community . Since finding 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 . ( 2017 ) ; Daneshmand et al . ( 2018 ) ; Du et al . ( 2017 ) ; Fang et al . ( 2018 ; 2019 ) ; Ge et al . ( 2015 ) ; Jin et al . ( 2017 ; 2019 ) ; Kohler & Lucchi ( 2017 ) ; Lei et al . ( 2017 ) ; Lee et al . ( 2019 ) ; Levy ( 2016 ) ; Mokhtari et al . ( 2018 ) ; Nesterov & Polyak ( 2006 ) ; Reddi et al . ( 2018a ) ; Staib et al . ( 2019 ) ; Tripuraneni et al . ( 2018 ) ; Xu et al . ( 2018 ) ) . 2 We follow these related works for the goal and aim at showing the benefit of the use of the momentum in reaching an ( , ) -second-order stationary point . We introduce a required condition , akin to a model assumption made in ( Daneshmand et al . ( 2018 ) ) , that ensures the dynamic procedure in Algorithm 2 produces updates with suitable correlation with the negative curvature directions of the function f . Definition 1 . Assume , at some time t , that the Hessian Ht = ∇2f ( wt ) has some eigenvalue smaller than − and ‖∇f ( wt ) ‖ ≤ . Let vt be the eigenvector corresponding to the smallest eigenvalue of ∇2f ( wt ) . The stochastic momentum mt satisfies Correlated Negative Curvature ( CNC ) at t with parameter γ > 0 if Et [ 〈mt , vt〉2 ] ≥ γ . ( 2 ) 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 ) , defined in the later section , then it takes T = O ( ( 1−β ) log ( 1/ ( 1−β ) ) −10 ) iterations to return an ( , ) second order stationary point . Alternatively , one can obtain an ( , √ ) second order stationary point in T = O ( ( 1 − β ) log ( 1/ ( 1 − β ) ) −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 . 2We apologize that the list is far from exhaustive . ( 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 Et [ · ] to represent conditional expectation E [ ·|w1 , w2 , . . . , wt ] , which is about fixing the randomness upto but not including t and notice that wt was determined at t− 1 . 2 BACKGROUND . 2.1 A THOUGHT EXPERIMENT . Let us provide some high-level intuition about the benefit of stochastic momentum with respect to escaping saddle points . In an iterative update scheme , at some time t0 the parameters wt0 can enter a saddle point region , that is a place where Hessian∇2f ( wt0 ) has a non-trivial negative eigenvalue , say λmin ( ∇2f ( wt0 ) ) ≤ − , and the gradient ∇f ( wt0 ) is small in norm , say ‖∇f ( wt0 ) ‖ ≤ . 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 . ( 2017 ) ; Lee et al . ( 2019 ) ) for additional details . On the other hand , if the iterates were to move in one particular direction , namely along vt0 the direction of the smallest eigenvector of ∇2f ( wt0 ) , then a fast escape is guaranteed under certain constraints on the step size η ; see e.g . ( Carmon et al . ( 2018 ) ) . 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 gt0 is strongly non-orthogonal to vt0 the direction of large negative curvature . This suffices 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 mt0 is strongly non-orthogonal to vt0 ; more precisely , Et0 [ 〈mt0 , vt0〉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 amplified in successive iterations through the momentum update when β is close to 1 . Assume that at some wt0 we have mt0 which possesses significant correlation with the negative curvature direction vt0 , then on successive rounds mt0+1 is quite close to βmt0 , mt0+2 is quite close to β 2mt0 , 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 first question in the affirmative , 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.2 MOMENTUM HELPS ESCAPE SADDLE POINTS : AN EMPIRICAL VIEW . Let us now establish , empirically , the clear benefit of stochastic momentum on the problem of saddle-point escape . We construct two stochastic optimization tasks , and each exhibits at least one significant saddle point . The two objectives are as follows . min w f ( w ) : = 1 n n∑ i=1 ( 1 2 w > Hw + b > i w + ‖w‖1010 ) , ( 3 ) min w f ( w ) : = 1 n n∑ i=1 ( ( a > i w ) 2 − y ) 2 . ( 4 ) Problem ( 3 ) of these was considered by ( Staib et al . ( 2019 ) ; Reddi et al . ( 2018a ) ) and represents a very straightforward non-convex optimization challenge , with an embedded saddle given by the matrix H : = diag ( [ 1 , −0.1 ] ) , and stochastic gaussian perturbations given by bi ∼ N ( 0 , diag ( [ 0.1 , 0.001 ] ) ) ; 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 ( Candés et al . ( 2013 ) ; Shechtman et al . ( 2015 ) ) . In phase retrieval3 , one wants to find an unknown w∗ ∈ Rd with access to but a few samples yi = ( a > i w ∗ ) 2 ; the design vector ai is known a priori . Here we have sampled w∗ ∼ N ( 0 , Id/d ) and ai ∼ N ( 0 , Id ) with d = 10 and n = 200 . The empirical findings , displayed in Figure 2 , are quite stark : for both objectives , convergence is significantly accelerated by larger choices of β . In the first objective ( Figure 4a ) , we see each optimization trajectory entering a saddle point region , apparent from the “ flat ” 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 first reported empirical finding that establishes the dramatic speed up of stochastic momentum for finding an optimal solution in phase retrieval .

Summary: This paper makes an interesting theoretical contribution; namely, that SGD with momentum (and with a slight modification to the step-size rule) is guaranteed to quickly converge to a second-order stationary point, implying it quickly escapes saddle points. SGD with momentum is widely used in the practice of deep learning, but a theoretical analysis has remained largely elusive. This paper sheds light theoretical properties justifying its use for deep learning.