Stochastic Normalized Gradient Descent with Momentum for Large Batch Training

1 INTRODUCTION . In machine learning , we often need to solve the following empirical risk minimization problem : min w∈Rd F ( w ) = 1 n n∑ i=1 fi ( w ) , ( 1 ) where w ∈ Rd denotes the model parameter , n denotes the number of training samples , fi ( w ) denotes the loss on the ith training sample . The problem in ( 1 ) can be used to formulate a broad family of machine learning models , such as logistic regression and deep learning models . Stochastic gradient descent ( SGD ) Robbins & Monro ( 1951 ) and its variants have been the dominating optimization methods for solving ( 1 ) . SGD and its variants are iterative methods . In the tth iteration , these methods randomly choose a subset ( also called mini-batch ) It ⊂ { 1 , 2 , . . . , n } and compute the stochastic mini-batch gradient 1/B ∑ i∈It ∇fi ( wt ) for updating the model parameter , where B = |It| is the batch size . Existing works Li et al . ( 2014b ) ; Yu et al . ( 2019a ) have proved that with the batch size of B , SGD and its momentum variant , called momentum SGD ( MSGD ) , achieve a O ( 1/ √ TB ) convergence rate for smooth non-convex problems , where T is total number of model parameter updates . With the population of multi-core systems and the easy implementation for data parallelism , many distributed variants of SGD have been proposed , including parallel SGD Li et al . ( 2014a ) , decentralized SGD Lian et al . ( 2017 ) , local SGD Yu et al . ( 2019b ) ; Lin et al . ( 2020 ) , local momentum SGD Yu et al . ( 2019a ) and so on . Theoretical results show that all these methods can achieve aO ( 1/ √ TKb ) convergence rate for smooth non-convex problems . Here , b is the batch size on each worker and K is the number of workers . By setting Kb = B , we can observe that the convergence rate of these distributed methods is consistent with that of sequential methods . In distributed settings , a small number of model parameter updates T implies a small synchronize cost and communication cost . Hence , a small T can further speed up the training process . Based on theO ( 1/ √ TKb ) convergence rate , we can find that if we adopt a larger b , the T will be smaller . Hence , large batch training can reduce the number of communication rounds in distributed training . Another benefit of adopting large batch training is to better utilize the computational power of current multi-core systems like GPUs You et al . ( 2017 ) . Hence , large batch training has recently attracted more and more attention in machine learning . Unfortunately , empirical results LeCun et al . ( 2012 ) ; Keskar et al . ( 2017 ) show that existing SGD methods with a large batch size will lead to a drop of generalization accuracy on deep learning models . Figure 1 shows the comparison of training loss and test accuracy between MSGD with a small batch size and MSGD with a large batch size . We can find that large batch training does degrade both training loss and test accuracy . Many works try to explain this phenomenon Keskar et al . ( 2017 ) ; Hoffer et al . ( 2017 ) . They observe that SGD with a small batch size typically makes the model parameter converge to a flatten minimum while SGD with a large batch size typically makes the model parameter fall into the region of a sharp minimum . And usually , a flatten minimum can achieve better generalization ability than a sharp minimum . Hence , large batch training has also become a challenging topic . Recently , many methods have been proposed for improving the performance of SGD with a large batch size . The work in Goyal et al . ( 2017 ) ; You et al . ( 2020 ) proposes many tricks like warmup , momentum correction and linearly scaling the learning rate , for large batch training . The work in You et al . ( 2017 ) observes that the norms of gradient at different layers of deep neural networks are widely different and the authors propose the layer-wise adaptive rate scaling method ( LARS ) . The work in Ginsburg et al . ( 2019 ) also proposes a similar method that updates the model parameter in a layer-wise way . Most of these methods lack theoretical evidence to explain why they can adopt a large batch size . Although the work in You et al . ( 2020 ) proposes some theoretical explanations for LARS , the implementation is still not consistent with its theorems in which both of the momentum coefficient and weight decay are set as zeros . In this paper , we propose a novel method , called stochastic normalized gradient descent with momentum ( SNGM ) , for large batch training . SNGM combines normalized gradient Nesterov ( 2004 ) ; Hazan et al . ( 2015 ) ; Wilson et al . ( 2019 ) and Polyak ’ s momentum technique Polyak ( 1964 ) together . The main contributions of this paper are outlined as follows : • We theoretically prove that compared to MSGD which is one of the most widely used variants of SGD , SNGM can adopt a larger batch size to converge to the -stationary point with the same computation complexity ( total number of gradient computation ) . That is to say , SNGM needs a smaller number of parameter update , and hence has faster training speed than MSGD . • For a relaxed smooth objective function ( see Definition 2 ) , we theoretically show that SNGM can achieve an -stationary point with a computation complexity of O ( 1/ 4 ) . To the best of our knowledge , this is the first work that analyzes the computation complexity of stochastic optimization methods for a relaxed smooth objective function . • Empirical results on deep learning also show that SNGM can achieve the state-of-the-art accuracy with a large batch size . 2 PRELIMINARIES . In this paper , we use ‖·‖ to denote the Euclidean norm , use w∗ to denote one of the optimal solutions of ( 1 ) , i.e. , w∗ ∈ argminw F ( w ) . We call w an -stationary point of F ( w ) if ‖∇F ( w ) ‖ ≤ . The computation complexity of an algorithm is the total number of its gradient computation . We also give the following assumption and definitions : Assumption 1 ( σ-bounded variance ) For any w , E‖∇fi ( w ) −∇F ( w ) ‖2 ≤ σ2 ( σ > 0 ) . Definition 1 ( Smoothness ) A function φ ( · ) is L-smooth ( L > 0 ) if for any u , w , φ ( u ) ≤ φ ( w ) +∇φ ( w ) > ( u−w ) + L 2 ‖u−w‖2 . L is called smoothness constant in this paper . Definition 2 ( Relaxed smoothness Zhang et al . ( 2020 ) ) A function φ ( · ) is ( L , λ ) -smooth ( L ≥ 0 , λ ≥ 0 ) if φ ( · ) is twice differentiable and for any w , ‖∇2φ ( w ) ‖ ≤ L+ λ‖∇φ ( w ) ‖ , where∇2φ ( w ) denotes the Hessian matrix of φ ( w ) . From the above definition , we can observe that if a function φ ( w ) is ( L , 0 ) -smooth , then it is a classical L-smooth function Nesterov ( 2004 ) . For a ( L , λ ) -smooth function , we have the following property Zhang et al . ( 2020 ) : Lemma 1 If φ ( · ) is ( L , λ ) -smooth , then for any u , w , α such that ‖u−w‖ ≤ α , we have ‖∇φ ( u ) ‖ ≤ ( Lα+ ‖∇φ ( w ) ‖ ) eλα . All the proofs of lemmas and corollaries of this paper are put in the supplementary . 3 RELATIONSHIP BETWEEN SMOOTHNESS CONSTANT AND BATCH SIZE . In this section , we deeply analyze the convergence property of MSGD to find the relationship between smoothness constant and batch size , which provides insightful hint for designing our new method SNGM . MSGD can be written as follows : vt+1 = βvt + gt , ( 2 ) wt+1 = wt − ηvt+1 , ( 3 ) where gt = 1/B ∑ i∈It ∇fi ( wt ) is a stochastic mini-batch gradient with a batch size of B , and vt+1 is the Polyak ’ s momentum Polyak ( 1964 ) . We aim to find how large the batch size can be without loss of performance . The convergence rate of MSGD with the batch size B for L-smooth functions can be derived from the work in Yu et al . ( 2019a ) . That is to say , when η ≤ ( 1− β ) 2/ ( ( 1 + β ) L ) , we obtain 1 T T−1∑ t=0 E‖∇F ( wt ) ‖2 ≤ 2 ( 1− β ) [ F ( w0 ) − F ( w∗ ) ] ηT + Lησ2 ( 1− β ) 2B + 4L2η2σ2 ( 1− β ) 2 , =O ( B ηC ) +O ( η B ) +O ( η2 ) , ( 4 ) where C = TB denotes the computation complexity ( total number of gradient computation ) . According to Corollary 1 in Yu et al . ( 2019a ) , we set η = √ B/ √ T = B/ √ C and obtain that 1 T T−1∑ t=0 E‖∇F ( wt ) ‖ ≤ √ O ( 1√ C ) +O ( B 2 C ) . ( 5 ) Algorithm 1 SNGM Initialization : η > 0 , β ∈ [ 0 , 1 ) , B > 0 , T > 0 , u0 = 0 , w0 ; for t = 0 , 1 , . . . , T − 1 do Randomly choose B function indices , denoted as It ; Compute a mini-batch gradient gt = 1B ∑ i∈It ∇fi ( wt ) ; ut+1 = βut + gt ‖gt‖ ; wt+1 = wt − ηut+1 ; end for Since η ≤ ( 1− β ) 2/ ( ( 1+ β ) L ) is necessary for ( 4 ) , we firstly obtain that B ≤ O ( √ C/L ) . Furthermore , according to the right term of ( 5 ) , we have to set B such that B2/C ≤ 1/ √ C , i.e. , B ≤ C1/4 , for O ( 1/ 4 ) computation complexity guarantee . Hence in MSGD , we have to set the batch size satisfying B ≤ O ( min { √ C L , C1/4 } ) . ( 6 ) We can observe that a larger L leads to a smaller batch size in MSGD . If B does not satisfy ( 6 ) , MSGD will get higher computation complexity . In fact , to the best of our knowledge , among all the existing convergence analysis of SGD and its variants on both convex and non-convex problems , we can observe three necessary conditions for the O ( 1/ 4 ) computation complexity guarantee Li et al . ( 2014b ; a ) ; Lian et al . ( 2017 ) ; Yu et al . ( 2019b ; a ) : ( a ) the objective function is L-smooth ; ( b ) the learning rate η is less than O ( 1/L ) ; ( c ) the batch size B is proportional to the learning rate η . One direct corollary is that the batch size is limited by the smooth constant L , i.e. , B ≤ O ( 1/L ) . Hence , we can not increase the batch size casually in these SGD based methods . Otherwise , it may slow down the convergence rate and we need to compute more gradients , which is consistent with the observations in Hoffer et al . ( 2017 ) .

Summary: This paper proposes a new stochastic normalized gradient descent method with momentum (SNGM) for large batch training. They prove that unlike mometum SGD (MSGD), SNGM can adopt larger batch size to converge to the epsilon-stationary point with the same computation complexity (total number of gradient computation). The paper shows that SNGM with large batches is comparable to MSGD with small batches for training ResNet on CIFAR10 and ImageNet. The paper also shows that SNGM outperforms LARS on CIFAR10.