Amortized Nesterov's Momentum: Robust and Lightweight  Momentum for Deep Learning

1 INTRODUCTION . In recent years , Gradient Descent ( GD ) ( Cauchy , 1847 ) and its variants have been widely used to solve large scale machine learning problems . Among them , Stochastic Gradient Descent ( SGD ) ( Robbins & Monro , 1951 ) , which replaces gradient with an unbiased stochastic gradient estimator , is a popular choice of optimizer especially for neural network training which requires lower precision . Sutskever et al . ( 2013 ) found that using SGD with Nesterov ’ s momentum ( Nesterov , 1983 ; 2013b ) , which was originally designed to accelerate deterministic convex optimization , achieves substantial speedups for training neural networks . This finding essentially turns SGD with Nesterov ’ s momentum into the benchmarking method of neural network design , especially for classification tasks ( He et al. , 2016b ; a ; Zagoruyko & Komodakis , 2016 ; Huang et al. , 2017 ) . It is observed that in these tasks , the momentum technique plays a key role in achieving good generalization performance . Adaptive methods ( Duchi et al. , 2011 ; Kingma & Ba , 2015 ; Tieleman & Hinton , 2012 ; Reddi et al. , 2018 ) , which are also becoming increasingly popular in the deep learning community , diagonally scale the gradient to speed up training . However , Wilson et al . ( 2017 ) show that these methods always generalize poorly compared with SGD with momentum ( both classical momentum ( Polyak , 1964 ) and Nesterov ’ s momentum ) . In this work , we introduce Amortized Nesterov ’ s Momentum , which is a special variant of Nesterov ’ s momentum . From users ’ perspective , the new momentum has only one additional integer hyper-parameter m to choose , which we call the amortization length . Learning rate and momentum parameter of this variant are strictly aligned with Nesterov ’ s momentum and by choosing m = 1 , it recovers Nesterov ’ s momentum . This paper conducts an extensive study based on both empirical evaluation and convex analysis to identify the benefits of the new variant ( or from users ’ angle , to set m apart from 1 ) . We list the advantages of Amortized Nesterov ’ s Momentum as follows : • Increasing m improves robustness1 . This is an interesting property since the new momentum not only provides acceleration , but also enhances the robustness . We provide an understanding of this property by analyzing the relation between convergence rate andm in the convex setting . • Increasing m reduces ( amortized ) iteration complexity . • A suitably chosen m boosts the convergence rate in the early stage of training and produces comparable final generalization performance . 1In this work , robustness refers to the probability of an optimizer significantly deviating from its expected performance , which can be reflected by the deviations of accuracy or loss in the training process over multiple runs that start with the same initial guess . • It is easy to tune m. The performances of the methods are stable for a wide range of m and we prove that the methods converge for any valid choice of m in the convex setting . • Ifm is not too large , the methods obtain the optimal convergence rate in general convex setting , just like Nesterov ’ s method . The new variant does have some minor drawbacks : it requires one more memory buffer , which is acceptable in most cases , and it shows some undesired behaviors when working with learning rate schedulers , which can be addressed by a small modification . Considering these pros and cons , we believe that the proposed variant can benefit many large-scale deep learning tasks . Our high level idea is simple : the stochastic Nesterov ’ s momentum can be unreliable since it is provided only by the previous stochastic iterate . The iterate potentially has large variance , which may lead to a false momentum that perturbs the training process . We thus propose to use the stochastic Nesterov ’ s momentum based on several past iterates , which provides robust acceleration . In other words , instead of immediately using an iterate to provide momentum , we put the iterate into an “ amortization plan ” and use it later . 2 PRELIMINARIES : SGD AND NESTEROV ’ S MOMENTUM . We start with a review of SGD and Nesterov ’ s momentum . We discuss some subtleties in the implementation and evaluation , which contributes to the interpretation of our methods . Notations In this paper , we use x ∈ Rd to denote the vector of model parameters . ‖·‖ and 〈· , ·〉 denote the standard Euclidean norm and inner product , respectively . Scalar multiplication for v ∈ Rd and β ∈ R is denoted as β ·v . f : Rd → R denotes the loss function to be minimized and∇f ( x ) represents the gradient of f evaluated at x . We denote the unbiased stochastic gradient estimator of ∇f ( x ) as ∇fi ( x ) with the random variable i independent of x ( e.g. , using mini-batch ) . We use x0 ∈ Rd to denote the initial guess . SGD SGD has the following simple iterative scheme , where γ ∈ R denotes the learning rate : xk+1 = xk − γ · ∇fik ( xk ) , for k ≥ 0 . Nesterov ’ s momentum The original Nesterov ’ s accelerated gradient ( with constant step ) ( Nesterov , 1983 ; 2013b ) has the following scheme2 ( y ∈ Rd , η , β ∈ R and y0 = x0 ) : yk+1 = xk − η · ∇f ( xk ) , xk+1 = yk+1 + β · ( yk+1 − yk ) , for k ≥ 0 , ( 1 ) where we call β · ( yk+1 − yk ) the momentum . By simply replacing ∇f ( xk ) with ∇fik ( xk ) , we obtain the SGD with Nesterov ’ s momentum , which is widely used in deep learning . To make this point clear , recall that the reformulation in Sutskever et al . ( 2013 ) ( scheme ( 2 ) , also the Tensorflow ( Abadi et al. , 2016 ) version ) and the PyTorch ( Paszke et al. , 2017 ) version ( scheme ( 3 ) ) have the following schemes ( v , vpt ∈ Rd and v0 = vpt0 = 0 ) : for k ≥ 0 , ( 2 ) { vk+1 = β · vk − η · ∇fik ( yk + β · vk ) , yk+1 = yk + vk+1 . ( 3 ) { vptk+1 = β · v pt k +∇fik ( xk ) , xk+1 = xk − η · ( β · vptk+1 +∇fik ( xk ) ) . Here the notations are modified based on their equivalence to scheme ( 1 ) . It can be verified that schemes ( 2 ) and ( 3 ) are equivalent to ( 1 ) through vk = β−1 · ( xk−yk ) and vptk = η−1β−1 · ( yk−xk ) , respectively ( see Defazio ( 2018 ) for other equivalent forms of scheme ( 1 ) ) . Interestingly , both PyTorch and Tensorflow3 track the values { xk } , which we refer to as M-SGD . This choice allows a consistent implementation when wrapped in a generic optimization layer ( Defazio , 2018 ) . However , the accelerated convergence rate ( in the convex case ) is built upon { yk } ( Nesterov , 2013b ) and { xk } may not possess such a theoretical improvement . We use OM-SGD to refer to the Original M-SGD that outputs { yk } . 2We exchange the notations of x and y in Nesterov ( 2013b ) . 3Tensorflow tracks the values { yk + β · vk } = { xk } . SGD and M-SGD In order to study the features of momentum , in this work , we regard momentum as an add-on to plain SGD , which corresponds to fixing the learning rates4 γ = η . From the interpretation in Allen-Zhu & Orecchia ( 2017 ) , η represents the learning rate for the gradient descent “ inside ” Nesterov ’ s method . To introduce the evaluation metrics of this paper , we report the results of training ResNet34 ( He et al. , 2016b ) on CIFAR-10 ( Krizhevsky et al. , 2009 ) ( our basic case study ) using SGD and M-SGD in Figure 1 . In this paper , all the multiple runs start with the same initial guess x0 . Figure 1a shows that Nesterov ’ s momentum hurts the convergence in the first 60 epochs but accelerates the final convergence , which verifies the importance of momentum for achieving high accuracy . Figure 1b depicts the robustness of M-SGD and SGD , which suggests that adding Nesterov ’ s momentum slightly increases the uncertainty in the training process of SGD . Train-batch loss vs. Full-batch loss In Figure 1c , train-batch loss stands for the average of batch losses forwarded in an epoch , which is commonly used to indicate the training process in deep learning . Full-batch loss is the average loss over the entire training dataset evaluated at the end of each epoch . In terms of optimizer evaluation , full-batch loss is much more informative than trainbatch loss as it reveals the robustness of an optimizer . However , full-batch loss is too expensive to evaluate and thus we only measure it on small datasets . On the other hand , test accuracy couples optimization and generalization , but since it is also evaluated at the end of the epoch , its convergence is similar to full-batch loss . Considering the basic usage of momentum in deep learning , we mainly use test accuracy to evaluate optimizers . We provide more discussion on this issue in Appendix C.2 . M-SGD vs. OM-SGD We also include OM-SGD in Figure 1a . In comparison , the final accuracies of M-SGD and OM-SGD are 94.606 % ± 0.152 % and 94.728 % ± 0.111 % with average deviations at 1.040 % and 0.634 % , respectively . This difference can be explained following the interpretation in Hinton ( 2012 ) that { xk } are the points after “ jump ” and { yk } are the points after “ correction ” . 3 AMORTIZED NESTEROV ’ S MOMENTUM . In this section , we formally introduce SGD with Amortized Nesterov ’ s Momentum ( AM1-SGD ) in Algorithm 1 with the following remarks : Options It can be verified that if m = 1 , AM1-SGD with Option I degenerates to M-SGD and Option II corresponds to OM-SGD . Just like the case for M-SGD and OM-SGD , the accelerated convergence rate is built upon Option II while Option I is easier to be implemented in a generic optimization layer5 . Intuitively , Option I is SGD with amortized momentum and Option II applies an m-iterations tail averaging on Option I . 4Ma & Yarats ( 2019 ) observed that when effective learning rates γ = η ( 1 − β ) −1 are fixed , M-SGD and SGD have similar performance . We provide a discussion on this observation in Appendix C.1 . 5To implement Option II , we can either maintain another identical network for the shifted point x̃ or temporarily change the network parameters in the evaluation phase . Algorithm 1 AM1-SGD Input : Initial guess x0 , learning rate η , momentum β , amortization length m , iteration number K. Initialize : x← x0 , x̃← x0 , x̃+ ← 0 { a running average } . 1 : for k = 0 , . . . , K − 1 do 2 : x← x− η · ∇fik ( x ) . 3 : x̃+ ← x̃+ + 1m · x . 4 : if ( k + 1 ) mod m = 0 then 5 : x← x+ β · ( x̃+ − x̃ ) . { adding amortized momentum } 6 : x̃← x̃+ , x̃+ ← 0 . 7 : end if 8 : end for Output : Option I : x , Option II : x̃ . * The symbol ‘ ← ’ denotes assignment . Efficiency We can improve the efficiency of Algorithm 1 by maintaining a running scaled momentum ṽ+ , m · ( x̃+ − x̃ ) instead of the running average x̃+ , by replacing the following steps in Algorithm 1 : Initialize : x← x0 , x̃← x0 , ṽ+ ← −m · x0 , Step 3 : ṽ+ ← ṽ+ + x . Step 5 : x← x+ ( β/m ) · ṽ+ . Step 6 : x̃← x̃+ ( 1/m ) · ṽ+ , ṽ+ ← −m · x̃ . Then , in one m-iterations loop , for each of the first m − 1 iterations , AM1-SGD requires 1 vector addition and 1 scaled vector addition . At the m-th iteration , it requires 1 vector addition , 1 scalarvector multiplication and 3 scaled vector additions . In comparison , M-SGD ( standard PyTorch ) requires 1 vector addition , 1 ( in-place ) scalar-vector multiplication and 2 scaled vector additions per iteration . Thus , as long as m > 2 , AM1-SGD has lower amortized cost than M-SGD . For memory complexity , AM1-SGD requires one more auxiliary buffer than M-SGD . Tuning m We did a parameter sweep for m in our basic case study . We plot the final and the average deviation of test accuracies over 5 runs againstm in Figure 2a . Note that m=1 corresponds to the results of M-SGD and OM-SGD , which are already given in Figure 1 . From this empirical result , m introduces a trade-off between final accuracy and robustness ( the convergence behaviors can be found in Appendix A.1 ) . Figure 2a suggests that m= 5 is a good choice for this task . For simplicity , and also as a recommended setting , we fix m=5 for the rest of experiments in this paper . A momentum that increases robustness To provide a stronger justification , we ran 20 seeds with m = 5 in Figure 2b and the detailed data are given in Figure 3 & Table 1 . The results show that the amortized momentum significantly increases the robustness . Intuitively , the gap between Option I and Option II can be understood as the effect of tail averaging . However , the large gap between Option I and SGD is somewhat mysterious : what Option I does is to inject a very large momentum6 into SGD every m iterations . It turns out that this momentum not only provides acceleration , but also helps the algorithm become more robust than SGD . This observation basically differentiates AM1-SGD from a simple interpolation in-between M-SGD and SGD . 6Amortized momentum β · ( x̃+−x̃ ) is expected to be much large than Nesterov ’ s momentum β · ( yk+1−yk ) . Learning rate scheduler issue We observed that when we use schedulers with a large decay factor and the momentum β is too large for the task ( e.g. , 0.995 for the task of this section ) , there would be a performance drop after the learning rate reduction . We believe that it is caused by the different cardinalities of iterates being averaged in x̃+ , which leads to a false momentum . This issue is resolved by restarting the algorithm after each learning rate reduction inspired by ( O ’ donoghue & Candes , 2015 ) . We include more discussion and evidence in Appendix A.4 .

Summary: The authors proposed Amortized Nesterov’s Momentum, a variant of Nesterov’s momentum that utilizes several past iterates, instead of one iterate, to provide the momentum.  The goal is to have more robust iterates, faster convergence in the early stage and higher efficiency. The authors designed two different realizations, AM1-SGD and AM2-SGD.