Year: 2021

Abstract:
Adam has been widely adopted for training deep neural networks due to less hyperparameter tuning and remarkable performance. To improve generalization, Adam is typically used in tandem with a squared l2 regularizer (referred to as Adam-l2). However, even better performance can be obtained with AdamW, which decouples the gradient of the regularizer from the update rule of Adam-l2. Yet, we are still lacking a complete explanation of the advantages of AdamW. In this paper, we tackle this question from both an optimization and an empirical point of view. First, we show how to re-interpret 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. Next, we consider the property of “scale-freeness” enjoyed by AdamW and by its proximal counterpart: their updates are invariant to component-wise rescaling of the gradients. We provide empirical evidence across a wide range of deep learning experiments showing a correlation between the problems in which AdamW exhibits an advantage over Adam-l2 and the degree to which we expect the gradients of the network to exhibit multiple scales, thus motivating the hypothesis that the advantage of AdamW could be due to the scale-free updates.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RELATED WORK
text: 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

heading: 3 THEORETICAL INSIGHTS ON THE MERITS OF ADAMW
text: 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.

heading: 4 EMPIRICAL FINDINGS OF ADAMW ON DEEP LEARNING EXPERIMENTS
text: In this section, we empirically compare Adam-ℓ2 with AdamW. First (Section 4.1), we report experiments for deep neural networks on image classification tasks (CIFAR10/100). Here, AdamW enjoys a significant advantage over Adam-ℓ2 when the batch normalization is switched off on deeper neural networks. We also report the correlation between this advantage and the scale-freeness property of AdamW. Next (Section 4.2), we show that AdamW is still almost scale-free even when the ϵ used in practice is not 0, and how, contrary to AdamW, Adam-ℓ2 is not invariant of the loss function being multiplied by a positive constant. Finally (Section 4.3), we show that AdamW performs similarly to AdamProx in practice, thus supporting the observations in Section 3.
Data Normalization and Augmentation: We consider the image classification task on CIFAR10/100 datasets. Images are normalized per channel using the means and standard deviations computed from all training images. We adopt the data augmentation technique following Lee et al. (2015) (for training only): 4 pixels are padded on each side of an image and a 32× 32 crop is randomly sampled from the padded image or its horizontal flip.
Models: For the CIFAR-10 dataset, we employ the Residual Network1 model (He et al., 2016) of 20/44/56/110/218 layers; and for CIFAR-100, we utilize the DenseNet-BC2 model (Huang et al., 2017) with 100 layers and a growth-rate of 12. The loss is the cross-entropy loss.
Hyperparameter tuning: For both Adam-ℓ2 and AdamW, we set β1 = 0.9, β2 = 0.999, ϵ = 10−8 as suggested in the original Adam paper Kingma & Ba (2015). To set the initial step size α and weight decay parameter λ, we grid search over {0.00005, 0.0001, 0.0005, 0.001, 0.005} for α and {0, 0.00001, 0.00005, 0.0001, 0.0005, 0.001} for λ. Whenever the best performing hyperparameters lie in the boundary of the searching grid, we always extend the grid to ensure that the final bestperforming hyperparameters fall into the interior of the grid.
Training: For each experiment configuration (e.g. 110-layer Resnet batch normalization), we randomly select an initialization of the model to use as a fixed starting point for all optimizers and hyperparameter settings. We use a mini-batch of 128, and train 300 epochs unless otherwise specified.
4.1 ADAMW VS. ADAM-ℓ2 ON IMAGE CLASSIFICATION TASKS, THE INFLUENCE OF BATCH NORMALIZATION, AND THE CORRELATION WITH SCALE-FREENESS
With BN, Adam-ℓ2 is on par with AdamW: Recently, Bjorck et al. (2020) found that AdamW has no improvement in absolute performance over sufficiently tuned Adam-ℓ2 in some reinforcement learning experiments. We also discover the same phenomenon in several image classification tasks, see Figure 2. Indeed, the best weight decay parameter is 0 for all cases and AdamW coincides with Adam-ℓ2 in these cases. Nevertheless, AdamW does decouple the optimal choice of the weight decay parameter from the initial step size much better than Adam-ℓ2 in all cases.
Removing BN: Notice that the models used in Figure 2 all employ Batch Normalization (BN) (Ioffe & Szegedy, 2015). BN works by normalizing the input to each layer across the mini-batch to make each
1https://github.com/akamaster/pytorch_resnet_cifar10 2https://github.com/bearpaw/pytorch-classification
coordinate have zero-mean and unit-variance. Without BN, deep neural networks are known to suffer from gradient explosion and vanishing (Schoenholz et al., 2017). This means each coordinate of the gradient will have very different scales, especially between the first and last layers. For non-scale-free algorithms, the update to the network weights will also be affected and each coordinate will proceed at a different pace. In contrast, scale-free optimizers are robust to such issues as the scaling of any single coordinate will not affect the update. Also, we already mentioned in the introduction some other reasons for considering removing BN. Thus, we consider the case where BN is removed as that is where AdamW and Adam-ℓ2 will show very different patterns due to scale-freeness.
Without BN, AdamW Outperforms Adam-ℓ2: In fact, without BN, AdamW outperforms Adam-ℓ2 even when both are finely tuned, especially on relatively deep neural networks (see Figure 3). AdamW not only obtains a much better test accuracy but also trains much faster.
Correlation between AdamW’s Advantage and Scale-freeness: We also observe that the advantage of AdamW becomes more evident as the network becomes deeper. Recall that as the depth grows, without BN, the gradient explosion and vanishing problem will become more severe. This means that for the non-scale-free Adam-ℓ2, the updates of each coordinate will be dispersed on a wider range of
scales even when the same weight decay parameter is employed. In contrast, the scales of the updates of AdamW will be much more concentrated in a smaller range. This is exactly verified empirically as illustrated in the 5th & 6th columns of figures in Figure 3. There, we report the histograms of the absolute value of updates of Adam-ℓ2 vs. AdamW of all coordinates near the end of training (for their comparison during the whole training process please refer to the Appendix).
This correlation between the advantage of AdamW over Adam-ℓ2 and the different spread of corresponding update scales which is induced by the scale-freeness property of AdamW provides an empirical justification on when and why AdamW excels over Adam-ℓ2.
SGD and Scale-freeness: The reader might wonder why SGD is known to provide state-of-the-art performance on many deep learning architectures (e.g., He et al., 2016; Huang et al., 2017) without being scale-free. At first blush, this seems to contradict our claims that scale-freeness correlates with good performance. In reality, the good performance of SGD in very deep models is linked to the use of BN that normalizes the gradients. Indeed, we verified empirically that SGD fails spectacularly when BN is not used. For example, on training the 110 layer Resnet without BN using SGD with momentum and weight decay of 0.0001, even a learning rate of 1e− 10 will lead to divergence.
4.2 VERIFYING SCALE-FREENESS: ADAMW VS. ADAM-ℓ2 UNDER LOSS MULTIPLICATION
In the previous section, we elaborated on the scale-freeness property of AdamW and its correlation with AdamW’s advantage over Adam-ℓ2. However, one may notice that in practice, the ϵ factor in the AdamW update is typically small but not 0, in our case 1e-8, thus preventing it from completely scale-free. In this section, we verify that the effect of such an ϵ on the scale-freeness is negligible.
As a simple empirical verification of the scale-freeness, we consider the scenario where we multiply the loss function by a positive number. Note that any other method to test scale-freeness would be equally good. For a feed-forward neural network without BN, this means the gradient would also be scaled up by that factor. In this case, the updates of a scale-free optimization algorithm would remain exactly the same, whereas they would change for an optimization algorithm that is not scale-free.
Figure 4 shows results of the loss function being multiplied by 10 and 100 respectively on optimizing a 110-layer Resnet with BN removed. For results of the original loss see Figure 3d. We can see that AdamW has almost the same performance across the range of initial learning rates and weight decay parameters, and most importantly, the best values of these two hyperparameters remain the same. This verifies that, even when employing a (small) non-zero ϵ, AdamW is still approximately scale-free. In contrast, Adam-ℓ2 is not scale-free and we can see that its behavior varies drastically with the best initial learning rates and weight decay parameters in each setting totally different.

heading: 4.3 ADAMW AND ADAMPROX BEHAVE VERY SIMILARLY IN PRACTICE
text: In Section 3, we showed theoretically that AdamW is the first order Taylor approximation of AdamProx (update rule (5)). Beyond this theoretical argument, we also verified empirically that the approximation is good. In Figure 5, we consider the case when ηt = 1 for all t - a relatively large constant learning rate schedule. In such cases, AdamW and AdamProx still behave very similarly. This suggests that for most learning rate schedules, e.g., cosine, exponential, polynomial, and step decay, which all monotonously decrease from η0 = 1, AdamProx will remain a very good approximation to AdamW. Thus, it is reasonable to try to understand AdamW by instead understanding the more classically-linked AdamProx.

heading: 5 CONCLUSION AND FUTURE WORKS
text: In this paper, we provide explanations for the merits of AdamW from two points of view. We first show that AdamW is an approximation of the proximal updates both theoretically and empirically. We then identify the setting of training very deep neural networks without batch normalization in which AdamW substantially outperforms Adam-ℓ2 in both training and testing, and provide an empirical explanation through the scale-freeness property of AdamW. Nevertheless, there are still many directions worth exploring and we name a few here.
Update for no-square ℓ2 regularization. Instead of using the squared ℓ2 regularization, we might think to use the ℓ2 regularization, that is without the square. This is known to have better statistical properties than the squared ℓ2 (Orabona, 2014), but it is not smooth so it is harder to be optimized. However, with proximal updates, we don’t have to worry about its non-smoothness. Hence, we can consider the objective function F (x) = λ∥x∥2 + f(x) .
The corresponding prox-SGD update was derived in Duchi & Singer (2009) for scalar learning rates and it is easy to see that it generalizes to our setting for Mt = ηtId as
xt+1 = max ( 1− ληt∥xt−ηtpt∥ , 0 ) (xt − ηtpt) .
Its performance, named AdamProxL2, as shown in Figure 6, can be on a par with AdamW.
Update for ℓ1 regularization. The ℓ1 regularization is widely used for achieving the sparsity of the solution. Recently, Neyshabur (2020) proposed β-lasso, a variant of the LASSO algorithm or namely regression with ℓ1 regularization, and showed significant improvements on training fully-connected networks. Thus, we are curious to see how the variant of AdamW derived from using ℓ1 regularization instead will behave. Specifically, we consider the objective function:
F (x) = λ∥x∥1 + f(x), for which the proximal update is obtained by solving:
argmin x
1 2 [(x− (xt −Mtpt)] TM−1t [x− (xt −Mtpt)] + λ∥x∥1 .
When Mt is a diagonal matrix with positive eigenvalues, the above objective function is decomposable and we can solve for each dimension. The solution for 1-d case was derived in Duchi & Singer (2009) and it is easy to see that it generalizes to our setting as
xt+1,i = sign(xt −Mtpt)i max(|xt −Mtpt|i − λMt,i, 0) .
Distributed Training As already noted in Section 4.1, batch normalization is not very compatible with distributed training. Since AdamW outperforms Adam-ℓ2 significantly in settings with BN switched off, at least in feed-forward neural networks, we can apply AdamW in distributed training to see if it still enjoys the same merits.

heading: A PROOF OF THEOREM 1
text: Proof. From the Fundamental Theorem of Calculus we have: ∇f(x) = ∇f(x∗)+ ∫ 1 0 ∇2f(x∗+t(x−x∗))(x−x∗)dt = ∫ 1 0 ∇2f(x∗+t(x−x∗))(x−x∗)dt .
Thus, for any function f̃Λ(x) whose Hessian is Λ∇2f(x) and ∇f̃Λ(x∗) = 0, ∇f̃Λ(x) = Λ∇f(x). Now, from the definition of a scale-free algorithm, the iterates of such an algorithm do not change when one multiplies any coordinate of all the gradients by a positive constant. Thus, a scale-free algorithm optimizing f behaves the same as if it is optimizing f̃Λ.

heading: B A SCALE-FREE ALGORITHM WITH DEPENDENCY ON THE CONDITION NUMBER
text: Algorithm 2 AdaGrad (Duchi et al., 2011; McMahan & Streeter, 2010) (All operations on vectors are element-wise.)
Input: #Iterations T , a set K, x1 ∈ K, stepsize η for t = 1 . . . T do
Receive: ∇f(xt) Set: ηt =
η√∑t i=1(∇f(xi))2
Update: xt+1 = ΠK (xt − ηt∇f(xt)) where ΠK is the projection onto K. end for Output: x̄ = 1T ∑T t=1 xt.
Algorithm 3 AdaGrad with Restart
Input: #Rounds N , x0 ∈ Rd, upper bound on ∥x0 − x∗∥∞ as D∞, strong convexity µ, smoothness M Set: x̄0 = x0 for i = 1 . . . N do
Run Algorithm 2 to get x̄i with T = 32dMµ , x1 = x̄i−1, K = {x : ∥x−x̄i−1∥ 2 ∞ ≤ D2∞ 4i−1 }, η =
D∞/ √ 2
2i−1
end for Output: x̄N .
Theorem 2. Let K be a hypercube with ∥x− y∥∞ ≤ D∞ for any x,y ∈ K. For a convex function f , set η = D∞√
2 , then Algorithm 2 guarantees for any x ∈ K:
T∑ t=1 f(xt)− f(x) ≤ √√√√2dD2∞ T∑ t=1 ∥∇f(xt)∥2 . (6)
Theorem 3. For a µ strongly convex and M smooth function f , denote its unique minimizer as x∗ ∈ Rd. Given x0 ∈ Rd, assume that ∥x0 − x∗∥∞ ≤ D∞, then Algorithm 3 guarantees:
∥x̄N − x∗∥2∞ ≤ D2∞ 4N .
Thus, to get a x such that ∥x− x∗∥2∞ ≤ ϵ, we need at most 32dMµ log4 ( D2∞/ϵ ) gradient calls.
Proof of Theorem 3. Consider round i and assume K passed to Algorithm 2 is bounded w.r.t. ℓ∞ norm by D∞i . When f is µ-strongly convex and M smooth, let x = x ∗, Equation (6) becomes:
T∑ t=1 f(xt)− f(x∗) ≤ √√√√2dD2∞i T∑ t=1 ∥∇f(xt)∥2 ≤ √√√√4MdD2∞i T∑ t=1 (f(xt)− f(x∗)) ,
where the second inequality is by the M smoothness of f . This gives: T∑
t=1
f(xt)− f(x∗) ≤ 4MdD2∞i .
Let x̄i = 1T ∑T t=1 xt we have by the µ-strong-convexity that:
∥x̄i − x∗∥2∞ ≤ ∥x̄i − x∗∥2 ≤ 2 µ (f(x̄)− f(x∗)) ≤ 2 µ 1 T T∑ t=1 (f(xt)− f(x∗)) ≤ 8MdD2∞i µT . (7)
Put T = 32dMµ in Equation (7) we have that ∥x̄i − x ∗∥2∞ ≤ D2∞i 4 . Thus, after each round, the ℓ∞
distance between the update x̄i and x∗ is shrinked by half, which in turn ensures that x∗ is still inside the K passed to Algorithm 2 in the next round with D∞i+1 = D∞i 2 . This concludes the proof.
Proof of Theorem 2. T∑
t=1
f(xt)− f(x)
≤ T∑
t=1
⟨∇f(xt),xt − x⟩
= T∑ t=1 d∑ j=1 ∂f ∂xt,j (xt) ∗ (xt,j − xj)
= T∑ t=1 d∑ j=1
(xt,j − xj)2 − ( xt,j − ηt,j ∂f∂xt,j (xt)− xj )2 2ηt,j + T∑ t=1 d∑ j=1 ηt,j 2 ( ∂f ∂xt,j (xt) )2
≤ T∑
t=1 d∑ j=1 (xt,j − xj)2 − (xt+1,j − xj)2 2ηt,j + T∑ t=1 d∑ j=1 ηt,j 2 ( ∂f ∂xt,j (xt) )2
≤ d∑
j=1 T∑ t=1 (xt,j − xj)2 2 ( 1 ηt,j − 1 ηt−1,j ) + d∑ j=1 T∑ t=1 ηt,j 2 ( ∂f ∂xt,j (xt) )2
≤ D 2 ∞
2η d∑ j=1 T∑ t=1
 √√√√ t∑
i=1
( ∂f
∂xi,j (xi)
)2 − √√√√t−1∑ i=1 ( ∂f ∂xi,j (xi) )2+ d∑ j=1 T∑ t=1
η
2 √∑t i=1 ( ∂f ∂xi,j (xi)
)2 ( ∂f ∂xt,j (xt) )2
≤ d∑
j=1
D2∞ 2η √√√√ T∑ t=1 ( ∂f ∂xt,j (xt) )2 + η √√√√ T∑ t=1 ( ∂f ∂xt,j (xt) )2 =
d∑ j=1 √√√√2D2∞ T∑ t=1 ( ∂f ∂xt,j (xt) )2
≤ √√√√2dD2∞ T∑ t=1 d∑ j=1 ( ∂f ∂xt,j (xt) )2
= √√√√2dD2∞ T∑ t=1 ∥∇f(xt))∥2 .
where the first inequality is by convexity, the second one by the projection lemma as the projection onto a hypercube equals performing the projection independently for each coordinate, the fifth one by Lemma 5 in (McMahan & Streeter, 2010), and the last one by the concavity of √ ·.

heading: C THE HISTOGRAMS OF THE MAGNITUDE OF EACH UPDATE COORDINATE DURING THE ENTIRE TRAINING PHASE
text: In this section, we report the histograms of the absolute value of updates of Adam-ℓ2 vs. AdamW of all coordinates divided by α during the whole training process. From the figures shown below, we can clearly see that AdamW’s updates remain in a much more concentrated scale range than Adam-ℓ2 during the entire training. Moreover, as the depth of the network grows, Adam-ℓ2’s updates become more and more dispersed, while AdamW’s updates are still concentrated. (Note that the leftmost bin contains all values equal or less than 2−27 ≈ 10−8.1 and the rightmost bin contains all values equal to or larger than 1.)

Review summaries:
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. 

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.

The paper give an analysis to the AdamW optimizer. This paper gives a theoretical view of the AdamW optimizing algorithm and connect it to the proximal gradients methods. Authors also explore in depth the advantages of AdamW over AdaM with l2 regularizer, from the perspective of scale-freeness property. By analyzing the scale-freeness property, authors find that the training is more stable in certain scenarios with large conditional numbers. Authors conduct experiments to validate their conclusions.