Win: Weight-Decay-Integrated Nesterov Acceleration for Adaptive Gradient Algorithms
Abstract
Training deep networks on increasingly large-scale datasets is computationally challenging.
In this work, we explore the problem of “ how to accelerate the convergence of adaptive gradient algorithms in a general manner ”, and aim at providing practical insights to boost the training efﬁciency.
To this end, we propose an effective Weight-decay-Integrated Nesterov acceleration (Win) for adaptive algorithms to enhance their convergence speed.
Taking AdamW and Adam as examples, we minimize a dynamical loss per iteration which combines the vanilla training loss and a dynamic regularizer inspired by proximal point method (PPM) to improve the convexity of the problem.
Then we respectively use the ﬁrst- and second-order Taylor approximations of vanilla loss to update the variable twice while ﬁxing the above dynamic regularization brought by PPM.
In this way, we arrive at our Win acceleration (like Nesterov acceleration) for AdamW and Adam that uses a conservative step and a reckless step to update twice and then linearly combines these two updates for acceleration.
Next, we extend this Win acceleration to LAMB and SGD.
Our transparent acceleration derivation could provide insights for other accelerated methods and their integration into adaptive algorithms.
Besides, we prove the convergence of Win-accelerated adaptive algorithms and justify their convergence superiority over their non-accelerated counterparts by taking AdamW and Adam as examples.
Experimental results testify the faster convergence speed and superior performance of our Win-accelerated AdamW, Adam,
LAMB and SGD over their vanilla counterparts on vision classiﬁcation tasks and language modeling tasks with CNN and Transformer backbones.
Introduction
Deep neural networks (DNNs) are effective to model realistic data and have been successfully applied to various applications, e g. image classiﬁcation [1–10] and speech recognition [11–14].
Typically, their training models can be formulated as the following nonconvex optimization problem:
[Equation]
where z is the model parameters; sample ⇣ is drawn from a data distribution D ; the loss f is differentiable;  is a constant.
To solve problem (1), SGD [15, 16] uses its compositional structure to efﬁciently estimate gradient via minibatch data, and has become a dominant algorithm to train
DNNs.
However, on sparse data or ill-conditioned problems, SGD suffers from slow convergence speed [17, 18], as it scales the gradient uniformly in all parameter coordinate and ignores the data or problem properties on each coordinate.
To resolve this issue, recent work has proposed a variety of adaptive methods, e.g. Adam [17] and AdamW [19], that scale each gradient coordinate according to the current geometry curvature of the loss F ( z ) .
This coordinate-wise scaling greatly accelerates the optimization convergence and helps them, e.g. Adam and AdamW, become much more popular in DNN training.
Unfortunately, along with the increasing scale of both datasets and models, efﬁcient DNN training even with SGD or adaptive algorithms has become very challenging.
In this work, we are particularly interested in the problem of “ how to accelerate the convergence of adaptive algorithms in a general manner ” because of their dominant popularity across many DNNs.
Heavy ball acceleration [20] and
Nesterov acceleration [21] are widely used in SGD but are rarely studied in adaptive algorithms.
Contributions:
In this work, based on a recent Nesterov-type acceleration formulation [22] and proximal point method (PPM) [23], we propose a new Weight-decay-Integrated Nesterov acceleration (Win) to accelerate adaptive algorithms.
By taking AdamW and Adam as examples, we follow
PPM spirit and minimize a dynamically regularized loss which combines vanillas loss and a dynamical regularization, and independently approximate the vanilla loss by its ﬁrst- and second-order
Taylor expansions to update the variable twice while ﬁxing the above dynamic regularization.
As a result, we achieve at our Win acceleration, a Nesterov-alike acceleration, for AdamW and Adam that uses a conservative step and a reckless step to update twice and then linearly combines these two updates for acceleration.
Then we extend Win acceleration to LAMB [24] and SGD.
This transparent acceleration derivation may motivate other accelerations and provide examples to introduce other accelerations into adaptive algorithms.
Moreover, we analyze the convergence of Win-accelerated adaptive algorithms to justify their convergence superiority by using AdamW & Adam as examples.
Finally, experimental results on both vision classiﬁcation tasks and language modeling tasks show that our Win-accelerated algorithms, i.e. accelerated AdamW, Adam, LAMB and SGD, can accelerate the convergence speed and also improve the performance of their corresponding non-accelerated counterparts by a remarkable margin on both CNN and transformer architectures.
Weight-decay-Integrated Nesterov Acceleration
In deterministic optimization, one widely used optimization-stabilizing and acceleration approach is proximal point method (PPM) [23, 25].
At the k -th iteration, PPM optimizes an ` 2 -regularized loss
F ( z )+ 12 ⌘ k k z  z k  1 k 22 instead of the vanilla loss F ( z ) .
This change enhances the convexity of the problem, accelerating and also stabilizing optimization [26, 27].
To make the ` 2 -regularized problem solvable iteratively, PPM approximates F ( z ) by its ﬁrst- or second-order Taylor expansion to get a close-form solution.
At below, we borrow the idea in PPM to induce a Weight-decay-Integrated
Nesterov acceleration (Win) for adaptive algorithms by using AdamW and Adam as examples.
Win-Accelerated AdamW and Adam.
To begin with, following most adaptive gradient algorithms, e.g. Adam and AdamW, we estimate the ﬁrst- and second-order moments m k and v k of gradient as g
[Equation]
where m 0 = g 0 , v
= g 20 ,  1 2 [0 , 1] and  2 2 [0 , 1] .
For brevity, with a small scaler ⌫ > 0 , we deﬁne
[Equation]
Then following PPM spirit, at the k -th iteration, we minimize a regularized loss F ( x )+ 12 ⌘ k x  x k k 2 s k .
Here we use the regularizer k x  x k k 2 s k instead of the ` 2 -regularization k x  x k k 22 , since 1) this new regularization allows us to handle adaptive algorithms as shown below Eqn.
(4), and 2) it also helps increase the problem convexity to speed up the convergence.
To make problem solvable iteratively, we approximate F ( z ) by its ﬁrst-order Taylor expansion at the point z k and update x k + as
[Equation]
where k x k s k
= p
h x , s k
⇤ x i with element-wise product ⇤ , m k is used to approximate the full gradient r F ( z k ) for Taylor expansion.
We add a small regularization  2 k x k 2 s k , as 1) it can improve the generalization in practice [19, 28]; 2) it allows us to derive Adam (  = 0 ) and AdamW (  > 0 ).
If  = 0 , the updating (4) becomes the exact Adam algorithm.
If  > 0 , the updating (4) can approximate the updating rule x k +1 =(1  ⌘ k ) x k  ⌘ k u k of AdamW.
This is because
[Equation]
Input: initialization x 0 =
z
= 0 , step size { ( ⌘ k , ¯ ⌘ k ) }
Tk =0 , moment parameters {  1 ,  2 } .
1 while k < T do
[Equation]
=
[Equation]
b
[Equation]
for AdamW and Adam
[Equation]
=
[Equation]
for LAMB
[Equation]
with ⌧
[Equation]
8 end while consider ⌘ k is small in practice, we approximate (1 + ⌘ k )  1 = 1  ⌘ k + O (  2 ⌘ 2 k ) and thus1+ ⌘ k ( x k  ⌘
[Equation]
u
[Equation] [Equation]
[Equation]
⌘
[Equation]
O (
[Equation]
⌘
[Equation]
)] x k  [ ⌘ k  O ( ⌘ 2 k )+ O (  3 ⌘ 3 k )] u k which becomes AdamW by ignoring O ( ⌘ 2 k ) and O ( ⌘ 3 k ) .
This is one reason that we adopt the regularizer k x  x k k 2 s k in (4) instead of the ` 2 -regularization in PPM, as we can ﬂexibly derive Adam and AdamW.
Similarly, we minimize a regularized loss F ( z ) + 1 2 ⌘ k k z  x t +
k 2 s k , and further approximate F ( z ) by its second-order approximation F ( z k )
+
h m k , z
z
k
i + 12¯ ⌘ k k z  z k k 2 s k :
z k +1 = argmin
[Equation]
where ⌧ k = 1 ⌘ k +¯ ⌘ k + ⌘ k ¯ ⌘ k , m k can approximate r F ( x k ) as shown in Theorem 1 in Appendix B.
For more clear, we introduce a virtual sequence y k
+
=
z k  ¯ ⌘ k u k in Win, and rewrite (5) as
[Equation]
See detailed steps in Algorithm 1.
Interestingly, Win acceleration is similar to Nesterov-type acceleration, since they both use a conservative step ⌘ k and a reckless step ¯ ⌘ k to update x k +1 and y k + respectively, and then linearly combine x k +
and y k +1 to obtain z k +1 .
Our Win-acceleration is quit simple and efﬁcient, as our accelerated AdamW/Adam only adds an extra simple algorithmic step, i.e. the 7th step in Algorithm 1, on vanilla AdamW/Adam.
Moreover, for the only extra hyper-parameter, the reckless step ¯ ⌘ k , in Algorithm 1 over AdamW/Adam, we always set it as ¯ ⌘ k =2 ⌘ k , which works well in our all experiments.
Extension to LAMB and SGD.
Here we generalize Win acceleration to LAMB [24] and SGD [15].
For LAMB, it scales the update u k of AdamW in Eqn.
(3) so that u k is at the same magnitude of the network weight x k .
That is, it changes the update rule x k +
=
(1  ⌘ k ) x k  ⌘ k m k / s k in
AdamW to x k
+
[Equation]
2 ( r k +  x k ) where r k = m k / s k .
This modiﬁcation is to avoid too large or small update, improving optimization efﬁciency.
To extend Win acceleration to
LAMB, we inherit this scaling spirit, and scale the update u k in (3) to the following one:
[Equation]
Next, we can follow Eqn.
(4) and (5) to update, and summarize detailed steps in Algorithm 1.
For SGD, applying Win acceleration to it is quite direct.
Speciﬁcally, the only algorithmic difference between SGD and AdamW on the ` 2 -regularized problems is that SGD has no second-order moment v k while AdamW has.
So we can borrow the acceleration framework of AdamW to accelerate SGD by setting s k = 1 2 R d in Eqn.
(3), (4) and (5), and obtain WIN-accelerated SGD:
[Equation]
where  0 1 2 [0 , 1] is dampening parameter.
Here we slightly modify the moment m k to accord with the one used in Nesterov-accelerated SGD ( e.g. SGD-M in Pytorch).
Convergence Analysis.
Theorem 1 in Appendix B analyzes the convergence of Win-accelerated adaptive algorithms to justify their convergence superiority by using AdamW & Adam as examples.
[Table] [Table]
Experiments
For clarity
[Table]
we call our accelerated algorithm “X-Win”, where “X” denotes vanilla optimizers.
In all experiments
[Table]
our accelerated algorithms, e.g. AdamW-Win, always use the default hyper-parameters of vanilla optimizers, e.g. moment parameters  1 and  2 in AdamW; and set ¯ ⌘ k =2 ⌘ k .
Results on ResNets and ViTs.
Table 1 reports accuracy of ResNets under the setting in [29], and Table 2 gives the performance of ViT [2] and PoolFormer [30].
Our accelerated algorithms always outperform their corresponding non-accelerated version.
On ResNet, LAMB-Win achieves remarkable improvement over the ofﬁcial optimizer LAMB for this setting; SGD-Win also surpasses heavy-ball accelerated SGD (SGD-H) and Nesterov accelerated SGD (SGD-M).
On ViTs, our accelerated algorithms consistently outperform the corresponding non-accelerated counterparts.
Fig.
1 shows the faster faster convergence behaviors of our accelerated algorithms over non-accelerated counterparts which could beneﬁt their better performance under the same computational cost.
[Table] [Table]
Results on Transformer-XL.
Table 3 shows that under different training steps on WikiText-103 dataset, our accelerated Adam-Win always achieves lower test PPL than the ofﬁcial Adam optimizer of Transformer-XLbase, and improves 1.5 average test PPL over Adam.
Conclusion
In this work, we adopt proximal point method to derive a weight-decay-integrated Nesterov acceleration for AdamW and Adam, and extend it to LAMB and SGD.
Moreover, we prove the convergence of our accelerated algorithms, i.e. accelerated AdamW, Adam and SGD, and observe the superiority of the accelerated Adam-type algorithm over the vanilla ones in terms of stochastic gradient complexity.
Finally, experimental results validate the advantages of our accelerated algorithms.