U NDERSTANDING A DAM W THROUGH P ROXIMAL M ETHODS AND S CALE -F REENESS
Anonymous authors
Paper under double-blind review
A BSTRACT
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 ℓ 2 regularizer (referred to as
Adam- ℓ 2 ).
However, even better performance can be obtained with AdamW, which decouples the gradient of the regularizer from the update rule of Adam- ℓ 2 .
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- ℓ 2 .
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- ℓ 2 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.
I NTRODUCTION
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 switched off as a case in which AdamW substantially outperforms Adam- ℓ 2 in both testing and training.
For such settings, 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:
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.
We point out the scale-freeness property enjoyed by AdamW and show the advantage of such a property on concrete problem classes.
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.
R ELATED W ORK
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.
T HEORETICAL I NSIGHTS ON THE M ERITS OF A DAM W
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
[Equation]
[Table]
[Table] [Table]
[Equation]
[Equation]
[Equation] [Equation] [Equation] where λ > 0 and f ( x ) : R d → R is a function bounded from below.
We could use a stochastic optimization algorithm that updates in the following fashion
[Equation]
where M t is a generic matrix containing the learning rates and p t denotes the update direction.
Specifically, we consider M t =
η t I d where η t is a learning rate schedule, e.g., the constant one or the cosine annealing (Loshchilov & Hutter, 2017).
This update covers many cases:
p t =
g t gives us the vanilla SGD; 2.
p t =
g t √ (cid:80) ti =
1 g 2 i + ϵ gives the AdaGrad algorithm (Duchi et al., 2011); 3.
p t =
α ˆ m t / ( √ ˆ v t + ϵ ) recovers the Adam algorithm (see Algorithm 1).
Note that in the above we use g t to denote the stochastic gradient of the entire objective function: ∇ f t ( x t − 1 ) + λ x t − 1 (if the regularizer is not present λ = 0 ), with ˆ m t and ˆ v t both updated using g t .
This update rule (2) is given by the following online mirror descent update (Nemirovsky & Yudin, 1983; Warmuth & Jagota, 1997; Beck & Teboulle, 2003):
[Equation]
This approximates minimizing a first-order Taylor approximation of F centered in x t − 1 plus a term that measures the distance between the x t and x t − 1 according to the matrix M − 1 t .
The approximation becomes exact when p t = ∇ f ( x t − 1 ) + λ x t − 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:
[Equation]
This update rule is using the proximal operator (Moreau, 1965; Parikh & Boyd, 2014) of 12 ∥ · ∥ 22 with respect to the norm ∥ · ∥ M − 1 t .
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
[Equation]
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
[Equation]
On the other hand, using M t =
η t I d and p t =
α ˆ m t / ( √ ˆ v t + ϵ ) in (3) gives:
[Equation]
which we will call AdamProx afterward.
Its first-order Taylor approximation around η t = 0 is
[Equation]
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 ( λη 2 t ) ≤ 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 ( x t ) ∥ 2 M t rather than on ∥∇ f ( x t ) + λ x t ∥ 2 M t 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 ϵ = {+
+}
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.
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 κ ( ∇ 2 f ( 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 ∇ 2 ˜ f Λ ( x ) = Λ ∇ 2 f ( 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 ∇ 2 f ( x ) is a diagonal matrix:
[Equation]
Assume µ ≤ µ i ≤ g i ( x ) ≤ M i ≤ M for i ∈ { 1 , . . .
, d } .
Denote j = arg max
i M i /µ i .
Choose λ i s.t.
µ j ≤ λ i µ i ≤ λ i g i ( x ) ≤ λ i M i ≤ M j and we have that Λ ∇ 2 f ( x ) has a condition number κ ′ =
M j /µ j .
This gives scale-free algorithms a big advantage when max i M i /µ i ≪ M/µ .
Specifically:
[Figure]
Corollary 1.
For quadratic problems f ( x ) = 12 x ⊤ Hx + b ⊤ x + c , with H diagonal and positive definite, any scale-free algorithm will not differentiate between minimizing f and ˜ f
( x ) = 12 x ⊤ x +
H − 1 b ⊤ x + c .
As the condition number of ˜ f is 1, the operation, and most importantly, the convergence, of a scale-free algorithm will not be affected by the condition number of f at all.
Figure 1 illustrates Corollary 1. Here, we compare GD (non-scalefree) with AdamW (scale-free) on optimizing two quadratic functions with the same minimizer, but one hessian matrix being a rescaled version of the other, resulting in different condition numbers.
Even after tuning the learning rates, the figure clearly shows that the updates of AdamW (starting from the same point) and thus its convergence to the minimizer, is completely unaffected by the condition number, while GD’s updates change drastically and its performance deteriorates significantly when the condition number is large.
Moreover, note that poor training performance mostly likely will also lead to poor testing performance.
This also explains AdaGrad’s improvements over SGD in certain scenarios. As an example, in Appendix B we analyze a variant of AdaGrad with restarts and show an improved convergence on 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.
E MPIRICAL F INDINGS OF A DAM W ON D EEP L EARNING E XPERIMENTS
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.
(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 Network 1 model (He et al., 2016) of 20/44/56/110/218 layers; and for CIFAR-100, we utilize the DenseNet-BC 2 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 −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 [Figure] { 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. A DAM W VS . A DAM - ℓ 2 ON I MAGE C LASSIFICATION T ASKS , THE I NFLUENCE OF B ATCH N ORMALIZATION , AND THE C ORRELATION WITH S CALE - 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 investigated 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 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.
Thus, we consider the case where
Batch Normalization is removed as that is where AdamW and Adam- ℓ 2 will show very different patterns due to scale-freeness.
[Figure]
As a side note, another reason for considering the scenario of removing BN comes from practical considerations (Zhang et al., 2019).
For example, batch normalization is inconvenient in distributed training as it requires each machine to collect a batch of statistics to update the model and it may not be accurate when a machine does not communicate frequently with each other (Goyal et al., 2017).
Without BN, AdamW Outperforms Adam- [-(cid:96) 
-]ℓ 2 :
In fact, without BN, AdamW outperforms Adam- ℓ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 [Figure] 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 1 e − 10 will lead to divergence.
4. V ERIFYING S CALE -F REENESS : A DAM W VS . A DAM - ℓ 2 UNDER L OSS M ULTIPLICATION
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 1 e - 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 batch normalization 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, for Adam- ℓ 2 , its update is not scale-free, and we can see that its behavior varies drastically and the best initial learning rates and weight decay parameters in each setting are totally different.
4. A DAM W AND A DAM P ROX B EHAVE VERY S IMILARLY IN P RACTICE
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. [Figure]
C ONCLUSION AND F UTURE W ORKS
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 ℓ(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
[Equation]
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 M t =
η t I d as
[Equation]
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:
[Equation]
for which the proximal update is obtained by solving:
[Equation]
When M t 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
[Equation]
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.
Appendices A P ROOF OF T HEOREM
Proof. From the Fundamental Theorem of Calculus we have:
[Equation] Thus, for any function ˜ f Λ ( x ) whose Hessian is Λ ∇ 2 f ( 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 Λ . B A S CALE - FREE A LGORITHM WITH D EPENDENCY ON THE C ONDITION N UMBER Algorithm are element-wise.)
Input : #Iterations T , a set K , x 1 ∈ K , stepsize η for t = 1 . . .
T do
Receive: ∇ f ( x t )
Set: η t = η √ (cid:80) ti =1 ( ∇ f ( x i ))
Update: x t +1 =
Π K ( x t − η t ∇ f ( x t )) where Π K is the projection onto K .
end for
Output: ¯ x = 1 T (cid:80) Tt =1 x t .
Algorithm
Input : #Rounds N , x 0 ∈ R d , upper bound on ∥ x 0 − x ∗ ∥ ∞ as D ∞ , strong convexity µ , smoothness M
Set: ¯ x 0 = xfor
i = 1 . . .
N do
Run Algorithm 2 to get ¯ x i with T =
d Mµ , x 1 = ¯ x i − 1 , K = { x : ∥ x − ¯ x i − 1 ∥ 2 ∞ ≤ D 2 ∞ 4 i − 1 } [Equation] η = D ∞ / √ 2 2 i −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 :
[Equation]
Theorem 3.
For a µ strongly convex and M smooth function f , denote its unique minimizer as x ∗ ∈ R d .
Given x 0 ∈ R d , assume that ∥ x 0 − x ∗ ∥ ∞ ≤ D ∞ , then Algorithm 3 guarantees:
[Equation]
Thus, to get a x such that ∥ x − x ∗ ∥ 2 ∞ ≤ ϵ , we need at most 32 d Mµ log 4 (cid:0) D 2 ∞ /ϵ (cid:1) 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:
[Equation]
where the second inequality is by the M smoothness of f .
This gives:
[Equation]
Let ¯ x i = 1 T (cid:80)
Tt =1 x t we have by the µ -strong-convexity that:
[Equation]
Put T =
d Mµ in Equation (7) we have that ∥ ¯ x i − x ∗ ∥ 2 ∞ ≤ D 2 ∞ 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 +
=
D ∞ i 2 .
This concludes the proof.
Proof of Theorem 2.
[Equation]
[Equation]
) −
[Equation]
(cid:17)
[Equation] (cid:88)
[Equation]
(cid:88) j [Equation]
j (cid:18) [Equation]
( x t ) (cid:19)≤
T (cid:88) t = d (cid:88)
j = ( x t,j − x j ) 2 − ( x t +1 ,j − x j )2 η t,j +
T (cid:88) t = d (cid:88)
j = η t,j (cid:18)
∂f ∂x t,j ( x t ) (cid:19)≤ d (cid:88)
j =
T (cid:88) t
= ( x t,j − x j ) (cid:18)η t,j − η t − 1 ,j (cid:19) + d (cid:88) j =
T (cid:88) t = η t,j (cid:18)
∂f ∂x t,j ( x t ) (cid:19) [Table] D [Table]
2 η d (cid:88) j
=
T (cid:88) t =   (cid:118)(cid:117)(cid:117)(cid:116) t (cid:88) i = (cid:18)
∂f ∂x i,j ( x i ) (cid:19)− (cid:118)(cid:117)(cid:117)(cid:116)
t − 1 (cid:88) i = (cid:18)
∂f ∂x i,j ( x i ) (cid:19) 2   + d (cid:88)
j =
T (cid:88) t =
η 2 (cid:114)(cid:80) ti =1 (cid:16) ∂f∂x i,j ( x i ) (cid:17)(cid:18)
∂f ∂x t,j ( x t ) (cid:19)≤
[Table]
(cid:88)
j =   D 2 ∞ 2 η (cid:118)(cid:117)(cid:117)(cid:116)
[Table]
(cid:88)
[Equation]
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 √ · .
C T HE H ISTOGRAMS OF THE M AGNITUDE OF EACH U PDATE C OORDINATE DURING THE E NTIRE T RAINING P HASE
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- ℓ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 .)
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]