Adapting Stepsizes by Momentumized Gradients Improves Optimization and Generalization

1 INTRODUCTION . Prevailing first-order optimization algorithms in modern machine learning can be classified into two categories . One is stochastic gradient descent ( SGD ) ( Robbins & Monro , 1951 ) , which is widely adopted due to its low memory cost and outstanding performance . SGDM ( Sutskever et al. , 2013 ) which incorporates the notion of momentum into SGD , has become the best choice for optimizer in computer vision . The drawback of SGD ( M ) is that it scales the gradient uniformly in all directions , making the training slow especially at the begining and fail to optimize complicated models well beyond Convolutional Neural Networks ( CNN ) . The other type is adaptive gradient methods . Unlike SGD , adaptive gradient optimizers adapt the stepsize ( a.k.a . learning rate ) elementwise according to the gradient values . Specifically , they scale the gradient by the square roots of some form of the running average of the squared values of the past gradients . Popular examples include AdaGrad ( Duchi et al. , 2011 ) , RMSprop ( Tijmen Tieleman , 2012 ) and Adam ( Kingma & Ba , 2015 ) etc . Adam , in particular , has become the default choice for many machine learning application areas owing to its rapid optimizing speed and outstanding ability to handle sophisticated loss curvatures . Despite their fast speed in the early training phase , adaptive gradient methods are found by studies ( Wilson et al. , 2017 ; Zhou et al. , 2020 ) to be more likely to exhibit poorer generalization ability than SGD . This is discouraging because the ultimate goal of training in many machine learning tasks is to exhibit high performance during testing phase . In recent years researchers have put many efforts to mitigate the deficiencies of adaptive gradient algorithms . AMSGrad ( Reddi et al. , 2019 ) corrects the errors in the convergence analysis of Adam and proposes a faster version . Yogi ( Reddi et al. , 2018 ) takes the effect of batch size into consideration . M-SVAG ( Balles & Hennig , 2018 ) transfers the variance adaptation mechanism from Adam to SGD . AdamW ( Loshchilov & Hutter , 2017b ) firsttime decouples weight decay from gradient descent for Adam-alike algorithms . SWATS ( Keskar & Socher , 2017 ) switches from Adam to SGD throughout the training process via a hard schedule and AdaBound ( Luo et al. , 2019 ) switches with a smooth transation by imposing dynamic bounds on stepsizes . RAdam ( Liu et al. , 2019 ) rectifies the variance of the adaptive learning rate through investigating the theory behind warmup heuristic ( Vaswani et al. , 2017 ; Popel & Bojar , 2018 ) . AdaBelief ( Zhuang et al. , 2020 ) adapts stepsizes by the belief in the observed gradients . Nevertheless , most of the above variants can only surpass ( as they claim ) Adam or SGD in limited tasks or under specifically and carefully defined scenarios . Till today , SGD and Adam are still the top options in machine learning , especially deep learning ( Schmidt et al. , 2021 ) . Conventional rules for choosing optimizers are : from task perspective , choose SGDM for vision , and Adam ( or AdamW ) for language and speech ; from model perspective , choose SGDM for Fully Connected Networks and CNNs , and Adam ( or AdamW ) for Recurrent Neural Networks ( RNN ) ( Cho et al. , 2014 ; Hochreiter & Schmidhuber , 1997b ) , Transformers ( Vaswani et al. , 2017 ) and Generative Adversarual Networks ( GAN ) ( Goodfellow et al. , 2014 ) . Based on the above observations , a natural question is : Is there a computationally efficient adaptive gradient algorithm that can converge fast and meanwhile generalize well ? . In this work , we are delighted to discover that simply replacing the gradient term in the second moment estimation term of Adam with its momentumized version can achieve this goal . Our idea comes from the origin of Adam optimizer , which is a combination of RMSprop and SGDM . RMSprop scales the current gradient by the square root of the exponential moving average ( EMA ) of the squared past gradients , and Adam replaces the raw gradient in the numerator of the update term of RMSprop with its EMA form , i.e. , with momentum . Since the momentumized gradient is a more accurate estimation of the appropriate direction to descent , we consider putting it in the second moment estimation term as well . We find such operation makes the optimizer more suitable for the general loss curvature and can theoretically converge to minima that generalize better . Extensive experiments on a broad range of tasks and models indicate that : without bells and whistles , our proposed optimizer can be as good as SGDM on vision problems and outperforms all the competitors in other tasks , meanwhile maintaining fast convergence speed . Our algorithm is efficient with no additional memory cost , and applicable to a wide range of scenarios in machine learning , especially deep learning . More importantly , AdaMomentum requires little effort in hyperparameter tuning and the default parameter setting for adaptive gradient method works well consistently in our algorithm . Notation We use t , T to symbolize the current and total iteration number in the optimization process . θ ∈ Rd denotes the model parameter and f ( θ ) ∈ R denotes the loss function . We further use θt to denote the parameter at step t and ft to denote the noisy realization of f at time t because of the mini-batch stochastic gradient mechanism . gt denotes the t-th time gradient and α denotes stepsize . mt , vt represent the EMA of the gradient and the second moment estimation term at time t of adaptive gradient methods respectively . is a small constant number added in adaptive gradient methods to refrain the denominator from being too close to zero . β1 , β2 are the decaying parameter in the EMA formulation of mt and vt correspondingly . For any vectors a , b ∈ Rd , we employ√ a , a2 , |a| , a/b , a ≥ b , a ≤ b for elementwise square root , square , absolute value , division , greater or equal to , less than or equal to respectively . For any 1 ≤ i ≤ d , θt , i denotes the i-th element of θt . Given a vector x ∈ Rd , we use ‖x‖2 to denote its l2-norm and ‖x‖∞ to denote its l∞-norm . 2 ALGORITHM Table 1 : Comparison of AdaMomentum and classic adaptive gradient methods in mt and vt in ( 1 ) . Optimizer mt vt SGD gt 1 Rprop gt g2t RMSprop gt ( 1− β2 ) ∑t i=1 β t−i 2 g 2 i Adam ( 1− β1 ) ∑t i=1 β t−i 1 gi ( 1− β2 ) ∑t i=1 β t−i 2 g 2 i Ours ( 1− β1 ) ∑t i=1 β t−i 1 gi ( 1− β2 ) ∑t i=1 β t−i 2 mi 2 Preliminaries & Motivation Omitting the debiasing operation and the damping term , the adaptive gradient methods can be generally written in the following form : θt+1 = θt − α mt√ vt . ( 1 ) Here mt , vt are called the first and second moment estimation terms . When mt = gt and vt = 1 , ( 1 ) degenerates to the vanilla SGD . Rprop ( Duchi et al. , 2011 ) is the pioneering work using the notion of adaptive learning rate , in which mt = gt and vt = g2t . Actually it is equivalent to only using the sign of gradients for different weight parameters . RMSprop ( Tijmen Tieleman , 2012 ) forces the number divided to be similar for adjacent mini-batches by incorporating momentum acceleration into vt. Adam ( Kingma & Ba , 2015 ) is built upon RMSprop in which it turns gt into momentumized version . Both RMSprop and Adam boost their performance thanks to the smoothing property of EMA using momentum . Due to the fact that momentumized gradient is a more accurate estimation than raw gradient , we deem that there is no reason to use gt in lieu of mt in second moment estimation term vt . Therefore we propose to replace the gis in vt of Adam with their momentumized versions mis , which further smooths the exponential moving average . Algorithm 1 AdaMomentum ( ours ) . All mathematical operations are element-wise . 1 : Initialization : Parameter initialization θ0 , step size α , damping term , m0 ← 0 , v0 ← 0 , t← 0 2 : while θt not converged do 3 : t← t+ 1 . Updating time step 4 : gt ← ∇θft ( θt−1 ) . Acquiring stochastic gradient at time t 5 : mt ← β1mt−1 + ( 1− β1 ) gt . EMA of gradients 6 : vt ← β2vt−1 + ( 1− β2 ) mt2 + . EMA of squared momentumized gradients 7 : m̂t ← mt/ ( 1− βt1 ) . Bias correction of first moment estimation 8 : v̂t ← vt/ ( 1− βt2 ) . Bias correction of second moment estimation 9 : θt ← θt−1 − α · m̂t/ √ v̂t . Updating parameters 10 : end while Detailed Algorithm The detailed procedure of our proposed optimizer is displayed in Algorithm 1 . There are two major modifications based on Adam , which are marked in red and blue respectively . One is that we replace the gt in vt of Adam with mt , which is the momentumized gradient . Hence we name our proposed optimzier as AdaMomentum . The other is the location of ( in Adam is added after √ · in line 10 of Alg.1 ) . We discover that moving the adding from outside the radical symbol to inside can consistently enhance performance . To the best of our knowledge , our method is the first attempt to put momentumized gradient in the second moment estimation term of adaptive gradient methods . Note that although the modifications seem simple to some degree , they can lead to siginificant changes in the performance of an adaptive gradient optimizer due to the iterative nature of optimization methods , which will also be elaborated in the following section . 3 WHY ADAMOMENTUM OVER ADAM ? . 3.1 ADAMOMENTUM IS MORE SUITABLE FOR GENERAL LOSS CURVATURE . In this section , we show that AdaMomentum can converge to ( global ) minima faster than Adam does via a 1-D example . The left part of Figure 1 is the process of optimization from a plateau to a basin area , where global optimum is assumed to exist . The right part is the zoomed-in version of the situation near the minimum , where we have some peaks and valleys . This phenomenon frequently takes place in optimization since there is only one global minimum with probably a great number of local minima surrounding ( Hochreiter & Schmidhuber , 1997a ; Keskar et al. , 2017 ) . Benefits of substituting gt with mt . We first explain how substituting mt for gt in the preconditioner vt can improve training via decomposing the trajectory of parameter point along the loss curve . 1 ) In area A , the parameter point starts to slide down the curve and |gt| begins to enlarge abruptly . So the actual stepsize α/ √ vt is small for Adam . However the absolute value of the momentumized gradient mt is small since it is the EMA of the past gradients , making α/ √ vt still large for AdaMomentum . Hence AdaMomentum can maintain higher training speed than Adam in this changing corner of the loss curve , which is what an optimal optimizer should do . 2 ) In area B , since the exponential moving average decays the impact of past gradients exponentially w.r.t . t , the magnitude of the elements of mt will gradually becomes as large as gt . 3 ) In area C , when the parameter approaches the basin , the magnitude of gt decreases , making the stepsizes of Adam increase immediately . In contrast , the stepsize of AdaMomentum is still comparatively small as |mt| is still much larger than |gt| , which is desired for an ideal optimizer . Small stepsize near optimum has benefits for convergence and stability . A more concrete illustration is given in the right part of Figure 1 . If the stepsize is too large ( e.g . in Adam ) , the weight parameter θt may rush to θ ( 2 ) t+1 and miss the global optimum . In contrast , small stepsize can guarantee the parameter to be close to the global minimum ( see θ ( 1 ) t+1 ) even if there may be tiny oscillations within the basin before the final convergence . Benefits of changing the location of . Next we elaborate why putting under the √ · is beneficial . We denote the debiased second moment estimation in AdaMomentum as v̂t and the second moment estimation term without as v̂′t . By simple calculation , we have v̂t = ( ( 1− β2 ) / ( 1− βt2 ) ) · t∑ i=1 βt−i2 m 2 i + 1− β2 , v̂′t = ( ( 1− β2 ) / ( 1− βt2 ) ) · t∑ i=1 βt−i2 m 2 i . Hence we have v̂t = v̂′t + / ( 1 − β2 ) . Then the actual stepsizes are α/ ( √ v̂′t + / ( 1− β2 ) ) and α/ ( √ v̂′t+ ) respectively . In the final stage of optimization , v̂ ′ t is very close to 0 ( because the values of gradients are near 0 ) and far less than hence the actual stepsizes can be approximately written as√ 1− β2α/ √ and α/ .As usually takes very tiny values ranging from 10−8 to 10−16 and β2 usu- ally take values that are extremely close to 1 ( usually 0.999 ) , we have √ 1− β2α/ √ α/ . Therefore we may reasonably come to the conclusion that after moving term into the radical symbol , AdaMomentum further reduces the stepsizes when the training is near minima , which contributes to enhancing convergence and stability as we have discussed above .

Summary: This paper proposed AdaMomentum, which uses the EMA (exponential moving average) of the square of EMA of gradient as sthe denominator, while Adam uses the EMA of the square of the gradient. The authors conducted experiments in CNN, Transformer, LSTM and GAN to show the superior performance of AdaMomentum. The authors also tried to provide theoretical analysis.