STORM+: Fully Adaptive SGD with Recursive Momentum for Nonconvex Optimization

1 Introduction . Over the past decade non-convex models have become principal tools in ML ( Machine Learning ) , and in data-science . This predominantly includes deep models , as well as Phase Retrieval [ Candes et al. , 2015 ] , non-negative matrix factorization [ Hoyer , 2004 ] , and matrix completion problems [ Ge et al. , 2016 ] among others . The main workhorse for training ML models is SGD ( stochastic gradient descent ) and its numerous variants . One parameter that significantly affects the SGD performance is the learning rate , which often requires a careful and costly hyper-parameter tuning . Adaptive approaches to setting the learning rate like AdaGrad [ Duchi et al. , 2011 ] and Adam [ Kingma and Ba , 2014 ] as well as non-adaptive heuristics [ Loshchilov and Hutter , 2017 , He et al. , 2019 ] are very popular in modern ML applications , yet these methods also require some tuning of hyper-parameters like momentum and the scale of the learning rate schedule . A popular SGD heuristic that has proven to be crucial in many applications is the use of momentum , i.e. , the use of a weighted average of past gradients instead of the current gradient [ Sutskever et al. , 2013 , Kingma and Ba , 2014 ] . Although adaptive approaches to setting the momentum have been investigated in the past [ Srinivasan et al. , 2018 , Hameed et al. , 2016 ] , principled and theoreticallygrounded approaches to doing so are less investigated . Another aspect that has not been extensively studied , which we take into account in this work , is the interplay between learning rate and momentum . ∗A Viterbi fellow . Corresponding author 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . In this work we explore momentum-based adaptive and parameter-free methods for stochastic non-convex optimization problems . Concretely , we focus on the setting where the objective is an expectation over smooth losses ( see Eq . ( 4 ) ) , and the goal is to find an approximate stationary point . In the general case of smooth non-convex objectives it is known that one can approach a stationary point at a rate of O ( 1/T 1/4 ) , where T is the total number of samples [ Ghadimi and Lan , 2013 ] . While this rate is optimal in the general case , it is known that one can obtain an improved rate of O ( 1/T 1/3 ) if the objective is an expectation over smooth losses [ Fang et al. , 2018 , Zhou et al. , 2018 , Cutkosky and Orabona , 2019 , Tran-Dinh et al. , 2019 ] . Besides , this rate was recently shown to be tight [ Arjevani et al. , 2019 ] . Nevertheless , most of the methods developed for this setting rely on variance reduction techniques [ Johnson and Zhang , 2013 , Zhang et al. , 2013 , Mahdavi et al. , 2013 , Wang et al. , 2013 ] , which require careful maintenance of anchor points in conjunction with appropriately selected large batchsizes . This leads to a challenging hyper-parameter tuning problem , weakening their practicality . One exception is the recent STORM algorithm of Cutkosky and Orabona [ 2019 ] . STORM does not require large batches nor anchor points ; instead , it uses a corrected momentum based gradient update that leads to implicit variance reduction , which in turn facilitates fast convergence . Unfortunately , none of the aforementioned methods ( including STORM ) is parameter-free . Indeed , the knowledge of smoothness parameter together with either the noise variance or a bound on the norm of the gradients are crucial to establish their guarantees . In this work , we essentially develop a parameter-free variant of STORM algorithm . We summarize our contributions as follows , • We present STORM+ , a parameter-free momentum based method that ensures the optimal O ( 1/T 1/3 ) rate for the expectation over smooth losses setting . Similarly to STORM , our method does not require large-batches nor anchor points . • STORM+ implicitly adapts to the variance of the gradients . Concretely , it obtains convergence rate of O ( 1/ √ T + σ1/3/T 1/3 ) , which recovers the optimal O ( 1/ √ T ) rate in the noiseless case . We also improve over STORM by shaving off a ( log T ) 3/4 factor from the 1/ √ T term . • In STORM+ we demonstrate a novel way to set the learning rate by introducing an adaptive interplay between learning rate and momentum parameters . 2 Related Work . In the context of stochastic non-convex optimization with general smooth losses , it was shown in Ghadimi and Lan [ 2013 ] that SGD with an appropriately selected learning rate can obtain a rate of O ( 1/T 1/4 ) for finding an approximate stationary point , which is known to match the respective lower bound [ Arjevani et al. , 2019 ] . While the method of Ghadimi and Lan [ 2013 ] requires knowledge of the smoothness and variance parameters , recent works have shown that adaptive methods like AdaGrad are able to obtain this bound in a parameter free manner , as well as to adapt to the variance of the problem [ Li and Orabona , 2019 , Ward et al. , 2019 , Reddi et al. , 2018 ] . These results , in a sense , explain the success of adaptive2 methods like AdaGrad [ Duchi et al. , 2011 ] , Adam [ Kingma and Ba , 2014 ] , and RMSProp [ Tieleman and Hinton , 2012 ] in handling non-convex problems . The idea of using variance reduction techniques for non-convex problems was first suggested in the context of finite sum problems by Allen-Zhu and Hazan [ 2016 ] , Reddi et al . [ 2016 ] , showing a rate of O ( 1/T 1/4 ) . This was later improved by Lei et al . [ 2017 ] to a rate of O ( 1/T 3/10 ) . The first works that have obtained the optimal O ( 1/T 1/3 ) for this setting were Fang et al . [ 2018 ] , Zhou et al . [ 2018 ] . Additionally , Fang et al . [ 2018 ] shows that the same convergence behavior applies to the more general expectation over smooth losses setting ( see Eq . ( 4 ) ) – a setting that captures finite-sum problems as a private case . 2An adaptive method is a method that updates its learning rate according to the ( noisy ) gradient feedback that it receives throughout the training process . The STORM algorithm suggested in Cutkosky and Orabona [ 2019 ] is the first algorithm to obtain the optimal O ( 1/T 1/3 ) for this setting without the need to maintain anchor points or large batches . Instead , it relies on a clever correction of the momentum by making only one extra call to the oracle , which leads to an implicit variance reduction effect . Moreover , STORM adapts to the variance of the problem by obtaining a rate of O ( ( log T ) 3/4/ √ T + σ1/3/T 1/3 ) without any prior knowledge of variance . However , it needs to know the smoothness parameter and a bound on the gradient norms to set the step size and momentum parameters . Simultaneously to the work of Cutkosky and Orabona [ 2019 ] , another paper [ Tran-Dinh et al. , 2019 ] have obtained the same optimal bound by proposing a similar update rule . Note that Tran-Dinh et al . [ 2019 ] does calculate a single anchor point , and it still requires the knowledge of the smoothness and variance parameters . 3 Setting and Preliminaries . We discuss stochastic non-convex optimization problems where the objective f : Rd 7→ R is of the following form , f ( x ) : = Eξ∼D [ f ( x ; ξ ) ] , and D is an unknown distribution from which we may draw i.i.d . samples . Our goal is to find an approximate stationary point of f , i.e . after T draws from D we should output a point x̄ ∈ Rd such that E‖∇f ( x ) ‖ ≤ Poly ( 1/T ) . We focus on first order methods , i.e. , methods that may access the gradients of f ( · , ξ ) , and make the following assumptions regarding the noisy gradients and function values . Bounded values : There exists B > 0 such that , max x , y∈Rd |f ( x ) − f ( y ) | ≤ B . ( 1 ) Bounded gradients : There exists G > 0 such that , ‖∇f ( x ; ξ ) ‖2 ≤ G2 ; ∀x ∈ Rd , ξ ∈ support { D } . ( 2 ) Bounded variance : There exists σ > 0 such that , E‖∇f ( x ; ξ ) −∇f ( x ) ‖2 ≤ σ2 ; ∀x ∈ Rd . ( 3 ) Expectation over smooth losses : There exists L > 0 such that , ‖∇f ( x ; ξ ) −∇f ( y ; ξ ) ‖ ≤ L‖x− y‖ ; ∀x , y ∈ Rd , ξ ∈ support { D } . ( 4 ) The last assumption also implies that the expected loss f ( · ) is L smooth . A property of smooth functions that we will exploit throughout the paper is the following , f ( y ) ≤ f ( x ) +∇f ( x ) > ( y − x ) + ( L/2 ) ‖y − x‖2 ; ∀x , y ∈ Rd ( 5 ) In the rest of this manuscript , ∇f ( x ; ξ ) relates to gradients with respect to x , i.e. , ∇ : = ∇x . We use ‖ · ‖ to denote the Euclidean norm , and x∗ denotes a global minima of f ( · ) , i.e. , x∗ = minx∈Rd f ( x ) . 4 Method . In this section we present STORM+ ( STochastic Recursive Momentum + ) : a parameter-free stochastic optimization method that finds approximate stationary points at an optimal rate . We describe our method in Alg . 1 and Eq . ( 8 ) , and state its guarantees in Theorem 1 . The original STORM algorithm : The original STORM template of Cutkosky and Orabona [ 2019 ] relies on an SGD-style update with a corrected momentum . Concretely , the idea is to maintain a gradient estimate dt which is a corrected weighted average of past stochastic gradients , and then update the iterates similarly to SGD , xt+1 = xt − ηtdt . ( 6 ) Standard momentum is a weighted average of past gradients , dt = at∇f ( xt , ξt ) + ( 1− at ) dt−1 ; where at ∈ [ 0 , 1 ] . Under this construction , dt is generally a biased estimate of ∇f ( xt ) . In STORM it is suggested to add a correction term , ( 1 − at ) ( ∇f ( xt , ξt ) − ∇f ( xt−1 , ξt ) ) , which leads to the following update rule ( again , at ∈ [ 0 , 1 ] ) , dt = ∇f ( xt , ξt ) + ( 1− at ) ( dt−1 −∇f ( xt−1 , ξt ) ) , ( Corrected Momentum ) The correction term plays a crucial role here : it exploits the smoothness of f ( · , ξ ) in a way that leads to a variance reduction effect . To see this effect one can inspect the error of the momentum dt compared to the exact gradient at xt , t : = dt −∇f ( xt ) . The STORM update rule induces the following error dynamics , t = ( 1− at ) t−1 + at ( ∇f ( xt , ξt ) −∇f ( xt ) ) + ( 1− at ) Zt where Zt : = ( ∇f ( xt , ξt ) − ∇f ( xt−1 , ξt ) ) − ( ∇f ( xt ) − ∇f ( xt−1 ) ) . Due to the smoothness of the objective we have ‖Zt‖ ≤ O ( ‖xt − xt−1‖ ) = O ( ηt−1‖dt−1‖ ) . Intuitively , as we approach a stationary point ( and use a small enough learning rate ) then ηt−1‖dt−1‖ decreases which in turn reduces the magnitude of Zt ’ s . Moreover , the second term in the above dynamics , at ( ∇f ( xt , ξt ) − ∇f ( xt ) ) , can be controlled by choosing a small enough momentum at . Thus , carefully controlling the learning rate and momentum parameters leads to a variance reduction effect which facilitates fast convergence . The original STORM paper [ Cutkosky and Orabona , 2019 ] makes the following choices , ηt = θ/ ( w + t∑ i=1 ‖gi‖2 ) 1/3 & at = cL2η2t−1 , ( 7 ) where we denote gt : = ∇f ( xt , ξt ) . The above choice of learning rate is inspired by AdaGrad Duchi et al . [ 2011 ] , which also sets the learning rate inversely proportional to the cumulative square norms of past gradients . Note that θ and w are constants that depend on the smoothness of the objective L , as well as on the bound on the gradients G , and c is an absolute constant independent of the problem ’ s characteristics . These choices of the constants and especially the choice of at ∝ L2η2t−1 is crucial for the analysis of the original STORM . In fact , the convergence proof for STORM breaks down unless we encode this prior knowledge into ηt and at . Next , we describe our parameter-free version . Our STORM+ algorithm : STORM+ relies on the original STORM template described in Equations ( 6 ) and ( Corrected Momentum ) , with the following parameter-free choices of learning rate and momentum parameter , ηt = 1/ ( t∑ i=1 ‖di‖2/ai+1 ) 1/3 & at = 1/ ( 1 + t−1∑ i=1 ‖gi‖2 ) 2/3 , ( 8 ) where again we denote gt : = ∇f ( xt , ξt ) . Note that in contrast to the original STORM our adaptive learning rate builds on history of estimates { d1 , . . . , dt } as well as on the momentum parameters { a1 , . . . , at+1 } . Our momentum term is similar to the adaptive choice of STORM , yet it does not require a bound on the gradients nor on the smoothness parameter , which was crucial for the original analysis . Finally , note that the above choice ensures at ∈ [ 0 , 1 ] . For completeness we present our method in Alg . 1 , where it can be seen that STORM+ is a combination of the original STORM template ( Equations ( 6 ) and ( Corrected Momentum ) ) together with the specific choices of ηt and at appearing in Eq . ( 8 ) . Note that the solution that STORM+ outputs is a point chosen uniformly at random among all iterates , which is quite standard in ( stochastic ) non-convex optimization . Notation : In Alg . 1 and throughout the rest of the paper we will employ the following notation , gt : = ∇f ( xt , ξt ) ; g̃t : = ∇f ( xt , ξt+1 ) ; ḡt : = ∇f ( xt ) . Now , we are at a position to present our main theorem regarding STORM+ ( Alg . 1 ) : Algorithm 1 STORM+ Input : # iterations T , x1 ∈ Rd 1 : Sample ξ1 and set d1 = g1 = ∇f ( x1 , ξ1 ) 2 : for t = 1 , ... , T do 3 : at+1 ← 1/ ( 1 + ∑t i=1 ‖gi‖2 ) 2/3 & ηt ← 1/ ( ∑t i=1 ‖di‖2/ai+1 ) 1/3 4 : xt+1 ← xt − ηtdt 5 : Sample ξt+1 and set gt+1 : = ∇f ( xt+1 ; ξt+1 ) , and g̃t : = ∇f ( xt ; ξt+1 ) 6 : dt+1 ← gt+1 + ( 1− at+1 ) ( dt − g̃t ) 7 : end for 8 : Choose x̄T uniformly at random from { x1 , . . . , xT } 9 : return x̄T Theorem 1 . Under the assumption in Eq . ( 1 ) , ( 2 ) , ( 3 ) and ( 4 ) in Section 3 , STORM+ ensures , E‖∇f ( x̄T ) ‖ ≤ O ( M√ T + κσ1/3 T 1/3 ) , where κ = O ( B3/4 + L3/2 ) ; M = O ( 1 + L9/4 +B9/8 +G5 + ( LG4 ) 3/2 ) , and the expectation is with respect to the randomization of the samples as well as the algorithm ’ s . Theorem 1 demonstrates that in the stochastic case STORM+ achieves the optimal O ( 1/T 1/3 ) rate for our setting . Moreover , it can be seen that STORM+ implicitly adapts to the variance of the noise ; in the noiseless case where σ = 0 , STORM+ recovers the optimal O ( 1/ √ T ) rate . We note that scaling the learning rate by some ( absolute ) constant factor may enable us to obtain better dependence on L and B .

Summary: This paper proposes an adaptive variant of the recently proposed STORM algorithm for minimizing expectations of smooth, nonconvex losses. By doing so, it removes the hyperparameters necessary for running STORM. It offers an analysis of the method and rates of convergence, but no empirical validation of the method.