Quickly Finding a Benign Region via Heavy Ball Momentum in Non-Convex Optimization

1 INTRODUCTION . Poylak ’ s Heavy Ball method ( Polyak ( 1964 ) ) has been very popular in modern non-convex optimization and deep learning , and the stochastic version ( a.k.a . SGD with momentum ) has become the de facto algorithm for training neural nets . Many empirical results show that the algorithm is better than the standard SGD in deep learning ( see e.g . Hoffer et al . ( 2017 ) ; Loshchilov & Hutter ( 2019 ) ; Wilson et al . ( 2017 ) ; Sutskever et al . ( 2013 ) ) , but there are almost no corresponding mathematical results that show a benefit relative to the more standard ( stochastic ) gradient descent . Despite its popularity , we still have a very poor justification theoretically for its success in non-convex optimization tasks , and Kidambi et al . ( 2018 ) were able to establish a negative result , showing that Heavy Ball momentum can not outperform other methods in certain problems . Furthermore , even for convex problems it appears that strongly convex , smooth , and twice differentiable functions ( e.g . strongly convex quadratic functions ) are one of just a handful examples for which a provable speedup over standard gradient descent can be shown ( e.g ( Lessard et al. , 2016 ; Goh , 2017 ; Ghadimi et al. , 2015 ; Gitman et al. , 2019 ; Loizou & Richtárik , 2017 ; 2018 ; Gadat et al. , 2016 ; Scieur & Pedregosa , 2020 ; Sun et al. , 2019 ; Yang et al. , 2018a ; Can et al. , 2019 ; Liu et al. , 2020 ; Sebbouh et al. , 2020 ; Flammarion & Bach , 2015 ) ) . There are even some negative results when the function is strongly convex but not twice differentiable . That is , Heavy Ball momentum might lead to a divergence in convex optimization ( see e.g . ( Ghadimi et al. , 2015 ; Lessard et al. , 2016 ) ) . The algorithm ’ s apparent success in modern non-convex optimization has remained quite mysterious . In this paper , we identify two non-convex optimization problems for which the use of Heavy Ball method has a provable advantage over vanilla gradient descent . The first problem is phase retrieval . It has some useful applications in physical science such as microscopy or astronomy ( see e.g . ( Candés et al. , 2013 ) , ( Fannjiang & Strohmer , 2020 ) , and ( Shechtman et al. , 2015 ) ) . The objective is minw∈Rd f ( w ) : = 1 4n ∑n i=1 ( ( x > i w ) 2 − yi ) 2 , ( 1 ) where xi ∈ Rd is the design vector and yi = ( x > i w∗ ) 2 is the label of sample i . The goal is to recover w∗ up to the sign that is not recoverable ( Candés et al. , 2013 ) . Under the Gaussian design setting ( i.e . xi ∼ N ( 0 , Id ) ) , it is known that the empirical risk minimizer ( 1 ) is w∗ or −w∗ , as long as the number of samples n exceeds the order of the dimension d ( see e.g . Bandeira et al . ( 2014 ) ) . Therefore , solving ( 1 ) allows one to recover the desired vectorw∗ ∈ Rd up to the sign . Unfortunately the problem is non-convex which limits our ability to efficiently find a minimizer . For this problem , there are many specialized algorithms that aim at achieving a better computational complexity and/or sample complexity to recoverw∗ modulo the unrecoverable sign ( e.g . ( Cai et al. , 2016 ; Candés & Li , 2014 ; Candés et al. , 2015 ; 2013 ; Chen & Candés , 2017 ; Duchi & Ruan , 2018 ; Ma et al. , 2017 ; 2018 ; Netrapalli et al. , 2013 ; Qu et al. , 2017 ; Tu et al. , 2016 ; Wang et al. , 2017a ; b ; Yang et al. , 2018b ; Zhang et al. , 2017a ; b ; Zheng & Lafferty , 2015 ) ) . Our goal is not about providing a state-of-the-art algorithm for solving ( 1 ) . Instead , we treat this problem as a starting point of understanding Heavy Ball momentum in non-convex optimization and hope for getting some insights on why Heavy Ball ( Algorithm 1 and 2 ) can be faster than the vanilla gradient descent in non-convex optimization and deep learning in practice . If we want to understand why Heavy Ball momentum leads to acceleration for a complicated non-convex problem , we should first understand it in the simplest possible setting . We provably show that Heavy Ball recovers the desired vector w∗ , up to a sign flip , given a random isotropic initialization . Our analysis divides the execution of the algorithm into two stages . In the first stage , the ratio of the projection of the current iterate wt on w∗ to the projection of wt on the perpendicular component keeps growing , which makes the iterate eventually enter a benign region which is strongly convex , smooth , twice differentiable , and contains a global optimal point . Therefore , in the second stage , Heavy Ball has a linear convergence rate . Furthermore , up to a value , a larger value of the momentum parameter has a faster linear convergence than the vanilla gradient descent in the second stage . Yet , most importantly , we show that Heavy Ball momentum also has an important role in reducing the number of iterations in the first stage , which is when the iterate might be in a non-convex region . We show that the higher the momentum parameter β , the fewer the iterations spent in the first stage ( see also Figure 1 ) . Namely , momentum helps the iterate to enter a benign region faster . Consequently , using a non-zero momentum parameter leads to a speedup over the standard gradient descent ( β = 0 ) . Therefore , our result shows a provable acceleration relative to the vanilla gradient descent , for computing a global optimal solution in non-convex optimization . The second of these is solving a class of cubic-regularized problems , minw f ( w ) : = 1 2w > Aw + b > w + ρ3‖w‖ 3 , ( 2 ) where the matrixA ∈ Rd×d is symmetric and possibly indefinite . Problem ( 2 ) is a sub-routine of the Nesterov-Polyak cubic-regularized method ( Nesterov & Polyak ( 2006 ) ) , which aims to minimize a non-convex objective F ( · ) by iteratively solving wt+1 = arg minw∈Rd { ∇F ( wt ) > ( w − wt ) + 12 ( w − wt ) > ∇2F ( w ) ( w − wt ) + ρ3‖w − wt‖ 3 } , ( 3 ) With some additional post-processing , the iterate wt converges to an ( g , h ) second order stationary point , defined as { w : ‖∇f ( w ) ‖ ≤ g and ∇2f ( w ) − hId } for any small g , h > 0 . However , their algorithm needs to compute a matrix inverse to solve ( 2 ) , which is computationally expensive when the dimension is high . A very recent result due to Carmon & Duchi ( 2019 ) shows that vanilla gradient descent approximately finds the global minimum of ( 2 ) under mild conditions , which only needs a Hessian-vector product and can be computed in the same computational complexity as computing gradients ( Pearlmutter , 1994 ) , and hence is computationally cheaper than the matrix inversion of the Hessian . Our result shows that , similar to the case of phase retrieval , the use of Heavy Ball momentum helps the iterate to enter a benign region of ( 3 ) that contains a global optimal solution faster , compared to vanilla gradient descent . For certain non-convex problems , e.g . dictionary learning ( Sun et al. , 2015 ) , matrix completion ( Chi et al. , 2019 ) , robust PCA ( Ge et al. , 2017 ) , and learning a neural network ( Ge et al . ( 2019 ) ; Bai & Lee ( 2020 ) ) , where it suffices to find a second-order stationary point , our result consequently might have application . To summarize , our theoretical results of the two non-convex problems provably show the benefit of using Heavy Ball momentum . Compared to the vanilla gradient descent , the use of momentum helps to accelerate the optimization process . The key to showing the acceleration in getting into benign regions of these problems is a family of simple dynamics due to Heavy Ball momentum . We will argue that the simple dynamics are not restricted to the two main problems considered in this paper . Specifically , the dynamics also naturally arise when solving the problem of top eigenvector computation ( Golub & Loan , 1996 ) and the problem of saddle points escape ( e.g . ( Jin et al. , 2017 ; Wang et al. , 2020 ) ) , which might imply the broad applicability of the dynamics for analyzing Heavy Ball in non-convex optimization . Algorithm 1 : Heavy Ball method ( Polyak , 1964 ) ( Equivalent version 1 ) 1 : Required : step size η and momentum parameter β ∈ [ 0 , 1 ] . 2 : Init : w0 = w−1 ∈ Rd 3 : for t = 0 to T do 4 : Update iterate wt+1 = wt− η∇f ( wt ) +β ( wt−wt−1 ) . 5 : end for Algorithm 2 : Heavy Ball method ( Polyak , 1964 ) ( Equivalent version 2 ) 1 : Required : step size η and momentum parameter β ∈ [ 0 , 1 ] . 2 : Init : w0 ∈ Rd and m−1 = 0 . 3 : for t = 0 to T do 4 : Update momentum mt : = βmt−1 +∇f ( wt ) . 5 : Update iterate wt+1 : = wt − ηmt . 6 : end for 2 MORE RELATED WORKS . Heavy Ball ( HB ) : HB has two exactly equivalent presentations in the literature ( see Algorithm 1 and 2 ) . Given the same initialization , both algorithms generate the same sequence of { wt } . In Algorithm 2 , we note that the momentum mt can be written as mt = ∑t s=0 β t−s∇f ( ws ) and can be viewed as a weighted sum of gradients . As we described in the opening paragraph , there is little theory of showing a provable acceleration of the method in non-convex optimization . The only exception that we are aware of is ( Wang et al. , 2020 ) . They show that HB momentum can help to escape saddle points faster and find a second-order stationary point faster for smooth non-convex optimization . They also observed that stochastic HB solves ( 1 ) and that using higher values of the momentum parameter β leads to faster convergence . However , while their work focused on the stochastic setting , their main result required some assumptions on the statistical properties of the sequence of observed gradients ; it is not clear whether these would hold in general . In appendix A , we provide a more detailed literature review of HB . To summarize , current results in the literature imply that we are still very far from understanding deterministic HB in non-convex optimization , let alone understanding the success of stochastic HB in deep learning . Hence , this work aims to make progress on a simple question : can we give a precise advantage argument for the acceleration effect of Heavy Ball in the deterministic setting ? Phase retrieval : The optimization landscape of problem ( 1 ) and its variants has been studied by ( Davis et al. , 2018 ; Soltanolkotabi , 2014 ; Sun et al. , 2016 ; White et al. , 2016 ) , which shows that as long as the number of samples is sufficiently large , it has no spurious local optima . We note that the problem can also be viewed as a special case of matrix sensing ( e.g . Li et al . ( 2018 ) ; Gunasekar et al . ( 2017 ) ; Li & Lin ( 2020 ) ; Li et al . ( 2019 ) ; Gidel et al . ( 2019 ) ; You et al . ( 2020 ) ) ; in Appendix A , we provide a brief summary of matrix sensing . For solving phase retrieval , Mannellia et al . ( 2020 ) study gradient flow , while Chen et al . ( 2018 ) show that the standard gradient descent with a random initialization like Gaussian initialization solves ( 1 ) and recovers w∗ up to the sign . Tan & Vershynin ( 2019 ) show that online gradient descent with a simple random initialization can converge to a global optimal point in an online setting where fresh samples are required for each step . In this paper , we show that Heavy Ball converges even faster than the vanilla gradient descent . Zhou et al . ( 2016 ) propose leveraging Nesterov ’ s momentum to solve phase retrieval . However , their approach requires delicate and computationally expensive initialization like spectral initialization so that the initial point is already within the neighborhood of a minimizer . Similarly , Xiong et al . ( 2018 ; 2020 ) show local convergence of Nesterov ’ s momentem and Heavy Ball momentum for phase retrieval , but require the initial point to be in the neighborhood of an optimal point . Jin et al . ( 2018 ) propose an algorithm that uses Nesterov ’ s momentum together as a subroutine with perturbation for finding a second-order stationary point , which could be applied for solving phase retrieval . Compared to ( Zhou et al. , 2016 ; Jin et al. , 2018 ; Xiong et al. , 2018 ; 2020 ) , we consider directly applying gradient descent with Heavy Ball momentum ( i.e . HB method ) to the objective function with simple random initialization , e.g . Gaussian initialization , which is what people do in practice and is what we want to understand . The goals of the works are different . Finally , we note that there are some efforts in integrating the technique of generative models and phase retrieval , which could help the task of image recovery ( e.g . Hand et al . ( 2018 ) ) . Phase retrieval might also be a good entry point of understanding some observations in optimization and neural net training ( e.g . Mannellia et al . ( 2020 ) ) . f ( w1 ) ≥ 0.5 } , i.e . if the relative change of objective value compared to the initial value has been increased to 50 % . Algorithm 3 and Algorithm 4 in Appendix H describe the procedures . All the lines are obtained by initializing the iterate at the same pointw0 ∼ N ( 0 , Id/ ( 10000d ) ) and using the same step size η = 5× 10−4 . Here we set w∗ = e1 and sample xi ∼ N ( 0 , Id ) with dimension d = 10 and number of samples n = 200 . We see that the higher the momentum parameter β , the faster the algorithm enters the linear convergence regime . ( a ) : Objective value ( 1 ) vs. iteration t. We see that the higher the momentum parameter β , the faster the algorithm enters the linear convergence regime . ( b ) : The size of projection of wt on w∗ over iterations ( i.e . |w‖t | vs. t ) , which is non-decreasing throughout the iterations until reaching an optimal point ( here , ‖w∗‖ = 1 ) . ( c ) : The size of the perpendicular component over iterations ( i.e . ‖w⊥t ‖ vs. t ) , which is increasing in the beginning and then it is decreasing towards zero after some point . We see that the slope of the curve corresponding to a larger momentum parameter β is steeper than that of a smaller one , which confirms Lemma 1 and Lemma 3 .

Summary: This paper analyzes Polyak momentum in the deterministic case for two simple but important non-convex problems: phase retrieval and finding the cubic-regularized newton step. It is shown that in both cases the problems posses a “benign region” in which the objective “looks” a bit convex. In the phase retrieval case, assuming that the data is generated from a normal distribution, it is shown that having a non-zero momentum parameter in fact allows the momentum method to find the benign region faster, after which standard convex analysis can take over to find the optimal point.