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: In this paper, the authors provide a parameter-free version named STORM+ of the STORM algorithm. STORM+ uses the same single-loop recursive momentum (which leads to implicit variance reduction effects) and constant batchsizes as in STORM, but features a parameter-free design by constructing the learning rate \eta_t and momentum coefficient a_t using past parameters and momentums (as a comparison, STORM needs to tune these two factors manually). Theoretically, the authors show that STORM+ achieves the same near-optimal complexity as STORM but without needing to tune the parameters to ensure this guarantee.