Understanding AdamW through Proximal Methods and Scale-Freeness

1 INTRODUCTION . Recent years have seen a surge of interest in applying deep neural networks ( LeCun et al. , 2015 ) to a myriad of areas . While Stochastic Gradient Descent ( SGD ) ( Robbins & Monro , 1951 ) remains the dominant method for optimizing such models , its performance depends crucially on the step size hyperparameter which controls how far the algorithm proceeds along the negative stochastic gradient in each step to update the model parameters . To alleviate this problem , people have developed a fruitful line of research on adaptive gradient methods ( e.g . Duchi et al. , 2010a ; McMahan & Streeter , 2010 ; Tieleman & Hinton , 2012 ; Zeiler , 2012 ; Luo et al. , 2018 ; Zhou et al. , 2018 ) . These methods provide mechanisms to automatically set stepsizes , and have been shown to greatly reduce the tuning effort while maintaining good performance . Among those adaptive algorithms , one of the most widely used is Adam ( Kingma & Ba , 2015 ) which achieves good results across a variety of problems even by simply adopting the default hyperparameter setting . In practice , to improve the generalization ability , Adam is typically combined with a ℓ2 regularization which adds the squared ℓ2 norm of the model weights on top of the loss function ( which we will call Adam-ℓ2 hereafter ) . This technique is usually referred to as weight decay because when using SGD , the ℓ2 regularization works by first shrinking the model weights by a constant factor in addition to moving along the negative gradient direction in each step . However , as pointed out in Loshchilov & Hutter ( 2019 ) , for Adam , there is no fixed regularization that achieves this same effect . To address this , they provide a method called AdamW that decouples the gradient of the ℓ2 regularization from the update of Adam and directly decays the weights . The two algorithms are shown in Algorithm 1 . Although AdamW frequently outperforms Adam-ℓ2 , the approach is primarily motivated empirically without a clear understanding of why it works so well . Recently , however , Bjorck et al . ( 2020 ) applied AdamW in Natural Language Processing ( NLP ) and Reinforcement Learning problems and found no improvement of performance over sufficiently tuned Adam-ℓ2 . Considering the huge popularity of AdamW ( Kuen et al. , 2019 ; Lifchitz et al. , 2019 ; Carion et al. , 2020 ) , we investigate when and why AdamW has a significant improvement over Adam-ℓ2 . In this paper , we focus on understanding the training and testing dynamics of the AdamW update in contrast to Adam-ℓ2 . We consider this contrast from the lens of optimization theory rather than directly investigating generalization over multiple epochs . First , we unveil the surprising connection between AdamW and proximal updates . In particular , we show that AdamW is an approximation of the latter and we also confirm such similarity with an empirical study . Noticing that AdamW and the proximal update are both scale-free while Adam-ℓ2 is not , we derive a theorem showing that scale-free optimizers enjoy an automatic acceleration with respect to the condition number on certain cases . This gives AdamW a concrete theoretical advantage in training over Adam-ℓ2 . Next , we empirically identify the scenario of training very deep neural networks with batch normalization ( BN ) switched off as a case in which AdamW substantially outperforms Adam-ℓ2 in both testing and training . Note that the setting of removing BN is not our invention : indeed , there is already active research in this ( De & Smith , 2020 ; Zhang et al. , 2019 ) . The reason is that BN has many disadvantages ( Brock et al. , 2021 ) including added memory overhead ( Bulò et al. , 2018 ) and training time ( Gitman & Ginsburg , 2017 ) , and a discrepancy between training and inferencing ( Singh & Shrivastava , 2019 ) . BN has also been found to be not suitable for many cases including distributed computing with a small minibatch per GPU ( Wu & He , 2018 ; Goyal et al. , 2017 ) , sequential modeling tasks ( Ba et al. , 2016 ) , and contrastive learning algorithms ( Chen et al. , 2020 ) . Moreover , there are already SOTA architectures that do not use BN including the Vision transformer ( Dosovitskiy et al. , 2021 ) and the BERT model ( Devlin et al. , 2019 ) . For such settings of removing BN , we observe that the magnitudes of the coordinates of the updates during training are much more concentrated about a fixed value for AdamW than for Adam-ℓ2 , which is an expected property of scale-free algorithms . Further , as depth increases , we expect a greater diversity of gradient scalings , a scenario that should favor scale-free updates . Our experiments support this hypothesis : deeper networks have more dramatic differences between the distributions of update scales between Adam-ℓ2 and AdamW , and larger accuracy advantages for AdamW . Specifically , the contributions of this paper are : 1 . We show that AdamW can be seen as an approximation of the proximal updates , which utilize the closed-form proximal mapping of the regularizer instead of only its gradient . 2 . We point out the scale-freeness property enjoyed by AdamW and show the advantage of such a property on concrete problem classes . 3 . We find a scenario where AdamW is significantly better than Adam-ℓ2 in both training and testing performance and report an empirical observation of the correlation between such advantage and the scale-freeness property of AdamW . The rest of this paper is organized as follows : In Section 2 we discuss the relevant literature . The connection between AdamW and the proximal updates as well as its scale-freeness are explained in Section 3 . We then report the empirical observations in Section 4 . Finally , we conclude with a discussion of the results and point out some potential future directions . 2 RELATED WORK . By enforcing the magnitude of the model weights to be small , weight decay has long been a standard technique to improve the generalization ability in machine learning ( Krogh & Hertz , 1991 ; Bos & Chug , 1996 ) and is still widely employed in training modern deep neural networks ( Devlin et al. , 2019 ; Tan & Le , 2019 ) . Here , we do not attempt to explain the generalization ability of AdamW . Rather , we assume that the regularization and the topology of the network guarantee good generalization performance . Instead , we study algorithms from the perspective of convergence rate . The use of proximal updates in the batch optimization literature dates back at least to 1965 ( Moreau , 1965 ; Martinet , 1970 ; Rockafellar , 1976 ; Parikh & Boyd , 2014 ) , and more recently used even in the stochastic setting ( Toulis & Airoldi , 2017 ; Asi & Duchi , 2019 ) . We are not aware of any previous paper pointing out the connection between AdamW and proximal updates . The scale-free property was first proposed in the online learning field ( Orabona & Pál , 2018 ) . There , they do not need to know a priori the Lipschitz constant bounding the gradient norms while still able to achieve the optimal rates . To the best of our knowledge , scale-freeness has not been explored as an explanation for the efficiency of deep learning optimization algorithms . Algorithm 1 Adam with L2 regularization ( Adam-ℓ2 ) and Adam with decoupled weight decay ( AdamW ) Loshchilov & Hutter ( 2017 ) ( Note that these two are exactly the same when λ = 0 , namely no weight decay . ) 1 : Given α , β1 , β2 , ϵ , λ ∈ R , { ηt } t≥0 . All operations on vectors are element-wise . 2 : Initialize : x0 ∈ Rd , first moment vector m0 ← 0 , second moment vector v0 ← 0 3 : for t = 1 , 2 , . . . , T do 4 : Compute the stochastic gradient∇ft ( xt−1 ) evaluated on a mini-batch of samples 5 : gt ← ∇ft ( xt−1 ) +λxt−1 6 : mt ← β1mt−1 + ( 1− β1 ) gt , vt ← β2vt−1 + ( 1− β2 ) g2t 7 : m̂t ←mt/ ( 1− βt1 ) , v̂t ← vt/ ( 1− βt2 ) 8 : xt ← xt−1 −ηtλxt−1 −ηtαm̂t/ ( √ v̂t + ϵ ) 9 : end for 3 THEORETICAL INSIGHTS ON THE MERITS OF ADAMW . AdamW and Proximal Updates : Here , we show that AdamW approximates a proximal update with squared ℓ2 regularization . This provides a first theoretical motivation for AdamW . Consider that we want to minimize the objective function F ( x ) = λ2 ∥x∥ 2 2 + f ( x ) , ( 1 ) where λ > 0 and f ( x ) : Rd → R is a function bounded from below . We could use a stochastic optimization algorithm that updates in the following fashion xt = xt−1 −Mtpt , ( 2 ) where Mt is a generic matrix containing the learning rates and pt denotes the update direction . Specifically , we consider Mt = ηtId where ηt is a learning rate schedule , e.g. , the constant one or the cosine annealing ( Loshchilov & Hutter , 2017 ) . This update covers many cases : 1. pt = gt gives us the vanilla SGD ; 2. pt = gt√∑t i=1 g 2 i+ϵ gives the AdaGrad algorithm ( Duchi et al. , 2011 ) ; 3. pt = αm̂t/ ( √ v̂t + ϵ ) recovers the Adam algorithm ( see Algorithm 1 ) . Note that in the above we use gt to denote the stochastic gradient of the entire objective function : ∇ft ( xt−1 ) + λxt−1 ( if the regularizer is not present λ = 0 ) , with m̂t and v̂t both updated using gt . This update rule ( 2 ) is given by the following online mirror descent update ( Nemirovsky & Yudin , 1983 ; Warmuth & Jagota , 1997 ; Beck & Teboulle , 2003 ) : xt = argmin x∈Rd λ 2 ∥xt−1∥ 2 2 + f ( xt−1 ) + p ⊤ t ( x− xt−1 ) + 12 ( x− xt−1 ) ⊤M−1t ( x− xt−1 ) . This approximates minimizing a first-order Taylor approximation of F centered in xt−1 plus a term that measures the distance between the xt and xt−1 according to the matrix M−1t . The approximation becomes exact when pt = ∇f ( xt−1 ) + λxt−1 . However , this is not the only way to construct first-order updates for the objective function ( 1 ) . An alternative route is to linearize only f and to keep the squared ℓ2 norm in its function form : xt = argmin x∈Rd λ 2 ∥x∥ 2 2 + f ( xt−1 ) + p ⊤ t ( x− xt−1 ) + 12 ( x− xt−1 ) ⊤M−1t ( x− xt−1 ) , This update rule is using the proximal operator ( Moreau , 1965 ; Parikh & Boyd , 2014 ) of 12∥ · ∥ 2 2 with respect to the norm ∥ · ∥M−1t . It is intuitive why this would be a better update : We directly minimize the squared ℓ2 norm instead of approximating it . From the first-order optimality condition , we have xt = ( Id + λMt ) −1 ( xt−1 −Mtpt ) . ( 3 ) When λ = 0 , the update in ( 2 ) and this one coincide . Yet , when λ ̸= 0 , they are no longer the same . We now show how the update in ( 3 ) generalizes the one in AdamW . The update of AdamW is xt = ( 1− ληt ) xt−1 − ηtαm̂t/ ( √ v̂t + ϵ ) . ( 4 ) On the other hand , using Mt = ηtId and pt = αm̂t/ ( √ v̂t + ϵ ) in ( 3 ) gives : xt = ( 1 + ληt ) −1 ( xt−1 − ηtαm̂t/ ( √ v̂t + ϵ ) ) , ( 5 ) which we will call AdamProx afterward . Its first-order Taylor approximation around ηt = 0 is xt ≈ ( 1− ληt ) xt−1 − ηtαm̂t/ ( √ v̂t + ϵ ) , which is exactly the AdamW update ( 4 ) as in the original paper . The careful reader might notice that the approximation of AdamW to AdamProx becomes less accurate when ηt becomes too large , and so be concerned whether this approximation is practical at all . Fortunately , in practice , ηt is never large enough for this to be an issue . Most practical learning rate schedules , e.g. , cosine , exponential , polynomial , and step decay , all decrease from η0 = 1 , or some even smaller value . Thus , the remainder term of this approximation is O ( λη2t ) ≤ O ( λη20 ) which we should always expect to be small as both λ and ηt are small . Consequently , we can expect AdamW and AdamProx to perform similarly for learning rate schedules ηt commonly employed in practice , and we will indeed confirm this empirically in Section 4.3 . While perhaps not widely known , the theory of proximal methods is a deep and beautiful branch of optimization ( see Parikh & Boyd ( 2014 ) for a survey ) . On the other hand , although AdamW is a very popular practical algorithm ( e.g. , for training BERT ( Devlin et al. , 2019 ) and vision transformer-based architectures ( Dosovitskiy et al. , 2021 ) ) , it is still unclear how it works so well . So , this connection opens the door to new ways to design optimization algorithms . For example , the proximal update gives an immediate theoretical advantage : the convergence rate , at least in the convex case , will depend on ∥∇f ( xt ) ∥2Mt rather than on ∥∇f ( xt ) + λxt∥ 2 Mt Duchi et al . ( 2010b ) . AdamW is Scale-Free An optimization algorithm is said to be scale-free if its iterates do not change when one multiplies any coordinate of all the gradients by a positive constant ( Orabona & Pál , 2018 ) . It turns out that the update ( 4 ) of AdamW and the update ( 5 ) of AdamProx are both scale-free when ϵ = 0 . This is evident for AdamW , since the scaling factor for any coordinate of the gradient is kept in both m̂t and √ v̂t and will be canceled out when dividing them . In contrast , for Adam-ℓ2 , the addition of the weight decay vector to the gradient ( see Line 5 of Algorithm 1 ) destroys this property . Note that in practical applications ϵ is very small but not zero , yet we empirically verify in Section 4.2 that it is small enough to still approximately guarantee the scale-free property . We want to emphasize the comparison between Adam-ℓ2 and AdamW : once Adam-ℓ2 adopts nonzero λ , it loses the scale-freeness property ; in contrast , AdamW still enjoys this for arbitrary λ . The same applies to any AdaGrad-type and Adam-type algorithms that employ weight decay strategy in the same way as Adam-ℓ2 by simply adding the gradient of the ℓ2 regularizer directly to the gradient of f ( as implemented in Tensorflow and Pytorch ) . Such algorithms are scale-free only when they do not use weight decay . In fact , this is one of our main observations : scale-freeness helps optimizers in deep learning and AdamW preserves this scale-freeness property even with an ℓ2 regularizer . We stress that the scale-freeness is an important but largely overlooked property of an optimization algorithm . It has already been utilized to explain the success of AdaGrad ( Orabona & Pál , 2018 ) . Recently , Agarwal et al . ( 2020 ) also advocates for setting the ϵ in the denominator of AdaGrad to be 0 thus making the update indeed scale-free , and they also provide an NLP experiment where such choice yields the best performance . Now , we show the effect of scale-freeness from a theoretical point of view . Consider a twice continuously differentiable function f . It is well-known that the convergence rate of many optimization algorithms , like gradient descent , on minimizing f depends on the condition number κ ( ∇2f ( x ) ) , i.e. , the ratio of its largest eigenvalue to its smallest eigenvalue of the Hessian ( see , e.g . Nesterov , 2004 ) . It turns out that scale-freeness can effectively reduce the effects of the condition number , as detailed by the following theorem whose proof can be found in the appendix . Theorem 1 . Let f be a twice continuously differentiable function such that x∗ ∈ argmin f ( x ) . Then , let f̃Λ be the family of functions such that x∗ ∈ argmin f̃Λ ( x ) , and ∇2f̃Λ ( x ) = Λ∇2f ( x ) , where Λ = diag ( λ1 , . . . , λd ) ⪰ 0 . Consider a scale-free algorithm whose convergence rate can be bounded depending on the condition number of the objective function . Then , the convergence of such algorithm used to minimize f will only depend on the smallest condition number among all functions f̃Λ . To give an example of when this is advantageous , consider when∇2f ( x ) is a diagonal matrix : ∇2f ( x ) = diag ( g1 ( x ) , g2 ( x ) , . . . , gd ( x ) ) . Assume µ ≤ µi ≤ gi ( x ) ≤ Mi ≤ M for i ∈ { 1 , . . . , d } . Denote j = argmaxi Mi/µi . Choose λi s.t . µj ≤ λiµi ≤ λigi ( x ) ≤ λiMi ≤ Mj and we have that Λ∇2f ( x ) has a condition number κ′ = Mj/µj . This gives scale-free algorithms a big advantage when maxi Mi/µi ≪M/µ . Specifically : ( a ) Condition number 1 . ( b ) Condition number 100000. strongly convex functions due to scale-freeness . Note that the folklore justification for such improvements is that the learning rate of AdaGrad approximates the inverse of the Hessian matrix , but this is incorrect : AdaGrad does not compute Hessians and there is no reason to believe it approximates them in general . More importantly , another scenario demonstrating the advantage of scale-freeness is training deep neural networks . Neural networks are known to suffer from the notorious problem of vanishing/exploding gradients ( Bengio et al. , 1994 ; Glorot & Bengio , 2010 ; Pascanu et al. , 2013 ) . This problem leads to the gradient scales being very different across layers , especially between the first and the last layers . The problem is particularly severe when the model is not equipped with normalization mechanisms like Batch Normalization ( BN ) ( Ioffe & Szegedy , 2015 ) . In such cases , when using a non-scale-free optimization algorithm ( e.g . SGD ) , the first layers and the last layers will proceed at very different speeds , whereas a scale-free algorithm ensures that each layer is updated at a similar pace .

Summary: This work first re-interprets AdamW as an approximation of a proximal gradient method, which takes advantage of the closed-form proximal mapping of the regularizer instead of only utilizing its gradient information as in Adam-l2.  Then it further shows that AdamW is “scale-freeness” algorithm which may enjoys faster convergence speed.  In this way, this work may explain why AdamW generalizes better than Adam-L2.