AdamP: Slowing Down the Slowdown for Momentum Optimizers on Scale-invariant Weights

1 INTRODUCTION . Normalization techniques , such as batch normalization ( BN ) ( Ioffe & Szegedy , 2015 ) , layer normalization ( LN ) ( Ba et al. , 2016 ) , instance normalization ( IN ) ( Ulyanov et al. , 2016 ) , and group normalization ( GN ) ( Wu & He , 2018 ) , have become standard tools for training deep neural network models . Originally proposed to reduce the internal covariate shift ( Ioffe & Szegedy , 2015 ) , normalization methods have proven to encourage several desirable properties in deep neural networks , such as better generalization ( Santurkar et al. , 2018 ) and the scale invariance ( Hoffer et al. , 2018 ) . Prior studies have observed that the normalization-induced scale invariance of weights stabilizes the convergence for the neural network training ( Hoffer et al. , 2018 ; Arora et al. , 2019 ; Kohler et al. , 2019 ; Dukler et al. , 2020 ) . We provide a sketch of the argument here . Given weights w and an input x , we observe that the normalization makes the weights become scale-invariant : Norm ( w > x ) = Norm ( cw > x ) ∀c > 0 . ( 1 ) The resulting equivalence relation among the weights lets us consider the weights only in terms of their ` 2-normalized vectors ŵ : = w‖w‖2 on the sphere S d−1 = { v ∈ Rd : ‖v‖2 = 1 } . We refer to Sd−1 as the effective space , as opposed to the nominal space Rd where the actual optimization algorithms operate . The mismatch between these spaces results in the discrepancy between the gradient descent steps on Rd and their effective steps on Sd−1 . Specifically , for the gradient descent updates , the effective step sizes ‖∆ŵt+1‖2 : = ‖ŵt+1 − ŵt‖2 are the scaled versions of the nominal step sizes ‖∆wt+1‖2 : = ‖wt+1−wt‖2 by the factor 1‖wt‖2 ( Hoffer et al. , 2018 ) . Since ‖wt‖2 increases during training ( Soudry et al. , 2018 ; Arora et al. , 2019 ) , the effective step sizes ‖∆ŵt‖2 decrease as the optimization progresses . The automatic decrease in step sizes stabilizes the convergence of gradient descent algorithms applied on models with normalization layers : even if the nominal learning rate is set to a constant , the theoretically optimal convergence rate is guaranteed ( Arora et al. , 2019 ) . maxw w > w ? ‖w‖2‖w ? ‖2 s. Trajectories start from w0 towards the optimal solution w ? . The problem is invariant to the scale of w. Video version in the attached code . In this work , we show that the widely used momentum-based gradient descent optimizers ( e.g . SGD and Adam ( Kingma & Ba , 2015 ) ) decreases the effective step size ∆ŵt even more rapidly than the momentum-less counterparts considered in ( Arora et al. , 2019 ) . This leads to a slower effective convergence for ŵt and potentially sub-optimal model performances . We illustrate this effect on a 2D toy optimization problem in Figure 1 . Compared to “ GD ” , “ GD+momentum ” is much faster in the nominal space R2 , but the norm growth slows down the effective convergence in S1 , reducing the acceleration effect of momentum . This phenomenon is not confined to the toy setup , for example , 95.5 % and 91.8 % of the parameters of the widely-used ResNet18 and ResNet50 ( He et al. , 2016 ) are scale-invariant due to BN . The majority of deep models nowadays are trained with SGD or Adam with momentum . And yet , our paper is first to delve into the issue in the widely-used combination of scale-invariant parameters and momentum-based optimizers . We propose a simple solution to slow down the decay of effective step sizes while maintaining the step directions of the original optimizer in the effective space . At each iteration of a momentum-based gradient descent optimizer , we propose to project out the radial component ( i.e . component parallel to w ) from the update , thereby reducing the increase in the weight norm over time . Because of the scale invariance , the procedure does not alter the update direction in the effective space ; it only changes the effective step sizes . We can observe the benefit of our optimizer in the toy setting in Figure 1 . “ Ours ” suppresses the norm growth and thus slows down the effective learning rate decay , allowing the momentum-accelerated convergence in R2 to be transferred to the actual space S1 . “ Ours ” converges most quickly and achieves the best terminal objective value . We do not discourage the use of momentum-based optimizers ; momentum is often an indispensable ingredient that enables best performances by deep neural networks . Instead , we propose to use our method that helps momentum realize its full potential by letting the acceleration operate on the effective space , rather than squandering it on increasing norms to no avail . The projection algorithm is simple and readily applicable to various optimizers for deep neural networks . We apply this technique on SGD and Adam ( SGDP and AdamP , respectively ) and verify the slower decay of effective learning rates as well as the resulting performance boosts over a diverse set of practical machine learning tasks including image classification , image retrieval , object detection , robustness benchmarks , audio classification , and language modelling . As a side note , we have identified certain similarities between our approaches and Cho & Lee ( 2017 ) . Cho & Lee ( 2017 ) have considered performing the optimization steps for the scale-invariant parameters on the spherical manifold . We argue that our approaches are conceptually different , as ours operate on the ambient Euclidean space , and are more practical . See Appendix §G.1 for a more detailed argumentation based on conceptual and empirical comparisons . 2 PROBLEM . Widely-used normalization techniques ( Ioffe & Szegedy , 2015 ; Salimans & Kingma , 2016 ; Ba et al. , 2016 ; Ulyanov et al. , 2016 ; Wu & He , 2018 ) in deep networks result in the scale invariance for weights . We show that the introduction of momentum in gradient-descent ( GD ) optimizers , when applied on such scale-invariant parameters , decreases the effective learning rate much more rapidly . This phenomenon has not yet been studied in literature , despite its ubiquity . We suspect the resulting early convergence may have introduced sub-optimality in many SGD and Adam-trained models across machine learning tasks . The analysis motivates our optimizer in §3 . 2.1 NORMALIZATION LAYER AND SCALE INVARIANCE . For a tensor x ∈ Rn1×···×nr of rank r , we define the normalization operation along the axes k ∈ { 0 , 1 } { 1 , ··· , r } as Normk ( x ) = x−µk ( x ) σk ( x ) where µk , σk are the mean and standard deviation functions along the axes k , without axes reduction ( to allow broadcasted operations with x ) . Depending on k , Normk includes special cases like batch normalization ( BN ) ( Ioffe & Szegedy , 2015 ) . For a function g ( u ) , we say that g is scale invariant if g ( cu ) = g ( u ) for any c > 0 . We then observe that Norm ( · ) is scale invariant . In particular , under the context of neural networks , Norm ( w > x ) = Norm ( ( cw ) > x ) ( 2 ) for any c > 0 , leading to the scale invariance against the weights w preceding the normalization layer . The norm of such weights ‖w‖2 does not affect the forward fw ( x ) or the backward ∇wfw ( x ) computations of a neural network layer fw parameterized by w. We may represent the scale-invariant weights via their ` 2-normalized vectors ŵ : = w‖w‖2 ∈ S d−1 ( i.e . c = 1‖w‖2 ) . 2.2 NOTATIONS FOR THE OPTIMIZATION STEPS −ηpt ŵt+1ŵt wt+1 wt ‖∆w t+ 1‖2 ‖∆ŵt+1‖2 See the illustration on the right for the summary of notations describing an optimization step . We write a gradient descent ( GD ) algorithm as : wt+1 ← wt − ηpt ( 3 ) where η > 0 is the user-defined learning rate . The norm of the difference ‖∆wt+1‖2 : = ‖wt+1 − wt‖2 = η‖pt‖2 is referred to as the step size . When p = ∇wf ( w ) , equation 3 is the vanilla GD algorithm . Momentumbased variants have more complex forms for p. In this work , we study the optimization problem in terms of the ` 2 normalized weights in Sd−1 , as opposed to the nominal space Rd . As the result of equation 3 , an effective optimization step takes place in Sd−1 : ∆ŵt+1 : = ŵt+1 − ŵt . We refer to the effective step size ‖∆ŵt+1‖2 . 2.3 EFFECTIVE STEP SIZES FOR VANILLA GRADIENT DESCENT ( GD ) . We approximate the effective step sizes for the scale-invariant w under the vanilla GD algorithm . We observe that the scale invariance f ( cw ) ≡ f ( w ) leads to the orthogonality : 0 = ∂f ( cw ) ∂c = w > ∇wf ( w ) . ( 4 ) For example , the vanilla GD update step p = ∇wf ( w ) is always perpendicular to w. Based on this , we establish the effective step size for w on Sd−1 : ‖∆ŵt+1‖2 : = ∥∥∥∥ wt+1 ‖wt+1‖2 − wt‖wt‖2 ∥∥∥∥ 2 ≈ ∥∥∥∥ wt+1 ‖wt+1‖2 − wt‖wt+1‖2 ∥∥∥∥ 2 = ‖∆wt+1‖2 ‖wt+1‖2 ( 5 ) where the approximation assumes 1‖wt+1‖2 − 1 ‖wt‖2 = o ( η ) , which holds when pt ⊥ wt as in the vanilla GD . We have thus derived that the effective step size on Sd−1 is inversely proportional to the weight norm , in line with the results in Hoffer et al . ( 2018 ) . Having established the relationship between the effective step sizes and the weight norm of a scaleinvariant parameters ( s.i.p ) , we derive the formula for its growth under the vanilla GD optimization . Lemma 2.1 ( Norm growth by GD , Lemma 2.4 in Arora et al . ( 2019 ) ) . For a s.i.p . w and the vanilla GD , where pt = ∇wf ( wt ) , ‖wt+1‖22 = ‖wt‖22 + η2‖pt‖22 . ( 6 ) The lemma follows from the orthogonality in equation 4 . It follows that the norm of a scale-invariant parameter ‖w‖2 is monotonically increasing and consequently decreases the effective step size for w. Arora et al . ( 2019 ) has further shown that GD with the above adaptive step sizes converges to a stationary point at the theoretically optimal convergence rate O ( T−1/2 ) under a fixed learning rate . 2.4 RAPID DECAY OF EFFECTIVE STEP SIZES FOR MOMENTUM-BASED GD . Momentum is designed to accelerate the convergence of gradient-based optimization by letting w escape high-curvature regions and cope with small and noisy gradients . It has become an indispensable ingredient for training modern deep neural networks . A momentum update follows : wt+1 ← wt − ηpt , pt ← βpt−1 +∇wtf ( wt ) ( 7 ) for steps t ≥ 0 , where β ∈ ( 0 , 1 ) and p−1 is initialized at 0 . Note that the step direction pt and the parameter wt may not be perpendicular anymore . We show below that momentum increases the weight norm under the scale invariance , even more so than does the vanilla GD . Lemma 2.2 ( Norm growth by momentum ) . For a s.i.p . w updated via equation 7 , we have ‖wt+1‖22 = ‖wt‖22 + η2‖pt‖22 + 2η2 t−1∑ k=0 βt−k‖pk‖22 . ( 8 ) Proof is in the Appendix §A . Comparing Lemma 2.1 and 2.2 , we notice that the formulation is identical , except for the last term on the right hand side of Lemma 2.2 . This term is not only non-negative , but also is an accumulation of the past updates . This additional term results in the significantly accelerated increase of weight norms when the momentum is used . We derive a more precise asymptotic ratio of the weight norms for the GD with and without momentum below . Corollary 2.3 ( Asymptotic norm growth comparison ) . Let ‖wGDt ‖2 and ‖wGDMt ‖2 be the weight norms at step t ≥ 0 , following the recursive formula in Lemma 2.1 and 2.2 , respectively . We assume that the norms of the updates ‖pt‖2 for GD with and without momentum are identical for every t ≥ 0 . We further assume that the sum of the update norms is non-zero and bounded : 0 < ∑ t≥0 ‖pt‖22 < ∞ . Then , the asymptotic ratio between the two norms is given by : ‖wGDMt ‖22 − ‖w0‖22 ‖wGDt ‖22 − ‖w0‖22 −→ 1 + 2β 1− β as t→∞ . ( 9 ) Proof in the Appendix §A . While the identity assumption for ‖pt‖2 between GD with and without momentum is strong , the theory is designed to illustrate an approximate norm growth ratios between the algorithms . For a popular choice of β = 0.9 , the factor is as high as 1 + 2β/ ( 1− β ) = 19 . Our observations are also applicable to Nesterov momentum and momentum-based adaptive optimizers like Adam . We later verify that the momentum induce the increase in weight norms and thus rapidly reduces the effective learning rates in many realistic setups of practical relevance ( §3.2 and §4 ) .

Summary: This paper points out that momentum in GD optimizers results in a far more rapid reduction in effective step sizes for scale-invariant weights.  To solve the problem, two algorithms called SGDP and AdamP are proposed, which project the updates to tangent space of the parameter. Experiments on several tasks including image classification, language modeling, etc show the effectiveness of the proposed algorithms. The idea to study the integration of BN and momentum is interesting. The analyses and proposed algorithms provide guidance to practitioners.