Scheduled Restart Momentum for Accelerated Stochastic Gradient Descent

1 INTRODUCTION . Training many machine learning ( ML ) models reduces to solving the following finite-sum optimization problem min w f ( w ) : = min w 1 N N∑ i=1 fi ( w ) : = min w 1 N N∑ i=1 L ( g ( xi , w ) , yi ) , w ∈ Rd , ( 1 ) where { xi , yi } Ni=1 are the training samples and L is the loss function , e.g. , cross-entropy loss for a classification task , that measure the discrepancy between the ground-truth label yi and the prediction by the model g ( · , w ) , parametrized by w. The problem ( 1 ) is known as empirical risk minimization ( ERM ) . In many applications , f ( w ) is non-convex , and g ( · , w ) is chosen among deep neural networks ( DNNs ) due to their preeminent performance across various tasks . These deep models are heavily overparametrized and require large amounts of training data . Thus , both N and the dimension of w can scale up to millions or even billions . These complications pose serious computational challenges . One of the simplest algorithms to solve ( 1 ) is gradient descent ( GD ) , which updates w according to : wk+1 = wk − sk 1 N N∑ i=1 ∇fi ( wk ) , ( 2 ) where sk > 0 is the step size at the k-th iteration . Computing ∇f ( wk ) on the entire training set is memory intensive and often prohibitive for devices with limited random access memory ( RAM ) such as graphics processing units ( GPUs ) used for deep learning ( DL ) . In practice , we sample a subset of the training set , of size m with m N , to approximate ∇f ( wk ) by the mini-batch gradient 1/m ∑m j=1∇fij ( wk ) , resulting in the ( mini-batch ) -stochastic gradient descent ( SGD ) . SGD and its Under review as a conference paper at ICLR 2021 accelerated variants are among the most used optimization algorithms in ML . These gradient-based algorithms have low computational complexity , and they are easy to parallelize , making them suitable for large scale and high dimensional problems ( Zinkevich et al. , 2010 ; Zhang et al. , 2015 ) . Nevertheless , GD and SGD have issues with slow convergence , especially when the problem is ill-conditioned . There are two common techniques to accelerate GD and SGD : adaptive step size ( Duchi et al. , 2011 ; Hinton et al . ; Zeiler , 2012 ) and momentum ( Polyak , 1964 ) . The integration of both adaptive step size and momentum with SGD leads to Adam ( Kingma & Ba , 2014 ) , one of the most used optimizers for training DNNs . Many recent developments have improved Adam ( Reddi et al. , 2019 ; Dozat , 2016 ; Loshchilov & Hutter , 2018 ; Liu et al. , 2020 ) . GD with constant momentum leverages the previous step to accelerate GD according to : vk+1 = wk − sk∇f ( wk ) ; wk+1 = vk+1 + µ ( vk+1 − vk ) , ( 3 ) where µ > 0 is a constant . A similar acceleration can be achieved by the heavy-ball ( HB ) method ( Polyak , 1964 ) . The momentum update in both ( 3 ) and HB have the same convergence rate of O ( 1/k ) as that of GD for convex smooth optimization . A breakthrough due to Nesterov ( 1983 ; 2018 ) replaces µ with ( k − 1 ) / ( k + 2 ) , which is known as the Nesterov accelerated gradient ( NAG ) with time-varying momentum . For simplicity , we denote this method as NAG below . NAG accelerates the convergence rate to O ( 1/k2 ) , which is optimal for convex and smooth loss functions ( Nesterov , 1983 ; 2018 ) . NAG can also speed up the process of escaping from saddle points ( Jin et al. , 2017 ) . In practice , NAG momentum can accelerate GD for nonconvex optimization , especially when the underlying problem is poorly conditioned ( Goh , 2017 ) . However , NAG accumulates error and causes instability when the gradient is inexact ( Devolder et al. , 2014 ; Assran & Rabbat , 2020 ) . In many DL applications , constant momentum achieves state-of-the-art result . For instance , training DNNs for image classification . Since NAG momentum achieves a much better convergence rate than constant momentum with exact gradient for general convex optimization , we consider the following question : Can we leverage NAG with a time-varying momentum parameter to accelerate SGD in training DNNs and improve the test accuracy of the trained models ? Contributions . We answer the above question by proposing the first algorithm that integrates scheduled restart NAG momentum with plain SGD . Here , we restart the momentum , which is orthogonal to the learning rate restart ( Loshchilov & Hutter , 2016 ) . We name the resulting algorithm scheduled restart SGD ( SRSGD ) . Theoretically , we prove the error accumulation of Nesterov accelerated SGD ( NASGD ) and the convergence of SRSGD . The major practical benefits of SRSGD are fourfold : • SRSGD remarkably speeds up DNN training . For image classification , SRSGD significantly reduces the number of training epochs while preserving or even improving the network ’ s accuracy . In particular , on CIFAR10/100 , the number of training epochs is reduced by half with SRSGD , while on ImageNet the reduction in training epochs is also remarkable . • DNNs trained by SRSGD generalize significantly better than the current benchmark optimizers . The improvement becomes more significant as the network grows deeper as shown in Fig . 1 . • SRSGD reduces overfitting in training very deep networks such as ResNet-200 for ImageNet classification , enabling the accuracy to keep increasing with depth . • SRSGD is straightforward to implement and only requires changes in a few lines of the SGD code . There is also no additional computational or memory overhead . We focus on image classification with DNNs , in which SGD with constant momentum is the choice . Related Work . Momentum has long been used to accelerate SGD . SGD with scheduled momentum and a good initialization can handle the curvature issues in training DNNs and enable the trained models to generalize well ( Sutskever et al. , 2013 ) . Kingma & Ba ( 2014 ) and Dozat ( 2016 ) integrated momentum with adaptive step size to accelerate SGD . In this work , we study the time-varying momentum version of NAG with restart for stochastic optimization . Adaptive and scheduled restart have been used to accelerate NAG with the exact gradient ( Nemirovskii & Nesterov , 1985 ; Nesterov , 2013 ; Iouditski & Nesterov , 2014 ; Lin & Xiao , 2014 ; Renegar , 2014 ; Freund & Lu , 2018 ; Roulet et al. , 2015 ; O ’ donoghue & Candes , 2015 ; Giselsson & Boyd , 2014 ; Su et al. , 2014 ) . These studies of restart NAG momentum are for convex optimization with the exact gradient . Restart techniques have also been used for stochastic optimization ( Kulunchakov & Mairal , 2019 ) . In particular , Aybat et al . ( 2019 ) developed a multistage variant of NAG with momentum restart between stages . Our work focuses on developing NAG-based optimization for training DNNs . Many efforts have also been Under review as a conference paper at ICLR 2021 devoted to studying the non-acceleration issues of SGD with HB and NAG momentum ( Kidambi et al. , 2018 ; Liu & Belkin , 2020 ) , as well as accelerating first-order algorithms with noise-corrupted gradients ( Cohen et al. , 2018 ; Aybat et al. , 2018 ; Lan , 2012 ) . Ghadimi & Lan ( 2013 ; 2016 ) provides analysis for the general stochastic gradient-based optimization algorithms. . Organization . In Section 2 , we review and discuss momentum for accelerating GD for convex smooth optimization . In Section 3 , we present the SRSGD algorithm and its theoretical guarantees . In Section 4 , we verify the efficacy of the proposed SRSGD in training DNNs for image classification on CIFAR and ImageNet . In Section 4.3 , we perform empirical analysis of SRSGD . We end with some concluding remarks . Technical proofs , some experimental details , and more results in training LSTMs ( Hochreiter & Schmidhuber , 1997 ) and WGANs ( Arjovsky et al. , 2017 ; Gulrajani et al. , 2017 ) are provided in the Appendix . Notation . We denote scalars and vectors by lower case and lower case bold face letters , respectively , and matrices by upper case bold face letters . For a vector x = ( x1 , · · · , xd ) ∈ Rd , we denote its ` p norm ( p ≥ 1 ) by ‖x‖p = ( ∑d i=1 |xi|p ) 1/p . For a matrix A , we use ‖A‖p to denote its induced norm by the vector ` p norm . Given two sequences { an } and { bn } , we write an = O ( bn ) if there exists a positive constant s.t . an ≤ Cbn . We denote the interval a to b ( included ) as ( a , b ] . For a function f ( w ) : Rd → R , we denote its gradient as∇f ( w ) and its Hessian as∇2f ( w ) . 2 REVIEW : MOMENTUM IN GRADIENT DESCENT . GD . GD ( 2 ) is a popular approach to solve ( 1 ) , which dates back to Cauchy ( 1847 ) . If f ( w ) is convex and L-smooth ( i.e. , ‖∇2f ( w ) ‖2 ≤ L ) , then GD converges with rate O ( 1/k ) by letting sk ≡ 1/L ( we use this sk in all the discussion below ) , which is independent of the dimension of w. HB . HB ( 4 ) ( Polyak , 1964 ) accelerates GD by using the historical information , which gives wk+1 = wk − sk∇f ( wk ) + µ ( wk −wk−1 ) , µ > 0 . ( 4 ) We can also accelerate GD by using the Nesterov/lookahead momentum , which leads to ( 3 ) . Both ( 3 ) and ( 4 ) have a convergence rate of O ( 1/k ) for convex smooth optimization . Recently , several variants of ( 3 ) have been proposed for DL , e.g. , ( Sutskever et al. , 2013 ) and ( Bengio et al. , 2013 ) . NAG . NAG ( Nesterov , 1983 ; 2018 ; Beck & Teboulle , 2009 ) replaces µ with ( tk − 1 ) /tk+1 , where tk+1 = ( 1 + √ 1 + 4t2k ) /2 with t0 = 1 . NAG iterates as following vk+1 = wk − sk∇f ( wk ) ; wk+1 = vk+1 + tk − 1 tk+1 ( vk+1 − vk ) . ( 5 ) NAG achieves a convergence rate O ( 1/k2 ) with the step size sk = 1/L . Remark 1 . Su et al . ( 2014 ) showed that ( k− 1 ) / ( k+ 2 ) is the asymptotic limit of ( tk − 1 ) /tk+1 . In the following presentation of NAG with restart , for the ease of notation , we will replace the momentum coefficient ( tk − 1 ) /tk+1 with ( k − 1 ) / ( k + 2 ) . Adaptive Restart NAG ( ARNAG ) . The sequences , { f ( wk ) − f ( w∗ ) } where w∗ is the minimum of f ( w ) , generated by GD and GD with constant momentum ( GD + Momentum , which follows ( 3 ) ) converge monotonically to zero . However , that sequence generated by NAG oscillates , as illustrated in Fig . 2 ( a ) when f ( w ) is a quadratic function . O ’ donoghue & Candes ( 2015 ) proposed ARNAG Under review as a conference paper at ICLR 2021 ( 6 ) , which restart the time-varying momentum of NAG according to the change of function values , to alleviate this oscillatory phenomenon . ARNAG iterates as following vk+1 = wk − sk∇f ( wk ) ; wk+1 = vk+1 + m ( k ) − 1 m ( k ) + 2 ( vk+1 − vk ) , ( 6 ) where m ( 1 ) = 1 ; m ( k + 1 ) = m ( k ) + 1 if f ( wk+1 ) ≤ f ( wk ) , and m ( k + 1 ) = 1 otherwise . Scheduled Restart NAG ( SRNAG ) . SR is another strategy to restart the time-varying momentum of NAG . We first divide the total iterations ( 0 , T ] ( integers only ) into a few intervals { Ii } mi=1 = ( Ti−1 , Ti ] , such that ( 0 , T ] = ⋃m i=1 Ii . In each Ii we restart the momentum after every Fi iterations . The update rule is then given by : vk+1 = wk − sk∇f ( wk ) ; wk+1 = vk+1 + ( k mod Fi ) ( k mod Fi ) + 3 ( vk+1 − vk ) . ( 7 ) Both AR and SR accelerate NAG to linear convergence for convex problems with the PolyakLojasiewicz ( PL ) condition ( Roulet & d ’ Aspremont , 2017 ) . Case Study – Quadratic Function . Consider the following quadratic optimization ( Hardt , 2014 ) min x f ( x ) = 1 2 xTLx− xT b , ( 8 ) where L ∈ Rd×d is the Laplacian of a cycle graph , and b is a d-dimensional vector whose first entry is 1 and all the other entries are 0 . Note that f ( x ) is convex with Lipschitz constant 4 . In particular , we set d = 1K ( 1K : = 103 ) . We run T = 50K iterations with step size 1/4 . In SRNAG , we restart , i.e. , we set the momentum to 0 , after every 1K iterations . Fig . 2 ( a ) shows that GD + Momentum as in ( 3 ) converges faster than GD , while NAG speeds up GD + Momentum dramatically and converges to the minimum in an oscillatory fashion . Both AR and SR accelerate NAG significantly .

Summary: This paper proposes to restart the momentum parameter in SGD (with Nesterov's momentum) according to some carefully chosen schedules in training deep neural network, which is named as SRSGD. Two different restarting schedules are proposed: linear schedule and exponential schedule. The strong point of this paper is its extensive experimental evaluations, which justify that SRSGD significantly improves the convergence speed and generalization over standard momentum SGD. The empirical analysis also sheds some light on the parameter tuning and interpretation of SRSGD.