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: In this paper, the authors first interpret AdamW as an approximation of proximal mapping, and then proposed their own algorithm AdamProx. They aim to unravel the connection between AdamW and AdamProx from both theoretical and empirical perspective. They delicately designed some probing experiments to verify their hypothesis.