Adan: Adaptive Nesterov Momentum Algorithm for Faster Optimizing Deep Models
Abstract
Adaptive gradient algorithms [1–4] combine the moving average idea with heavy ball acceleration to estimate accurate ﬁrst- and second-order moments of gradient for accelerating convergence.
But Nesterov acceleration which converges faster than heavy ball acceleration in theory [5] and also in many empirical cases [6] is much less investigated under the adaptive gradient setting.
In this work, we propose the ADAptive Nesterov momentum algorithm (Adan) to speed up the training of deep neural networks.
Adan ﬁrst reformulates the vanilla Nesterov acceleration to develop a new Nesterov momentum estimation (NME) method that avoids the extra computation and memory overhead of computing gradient at the extrapolation point.
Then Adan adopts NME to estimate the ﬁrst- and second-order gradient moments in adaptive gradient algorithms for convergence acceleration.
Besides, we prove that
Adan ﬁnds an ✏ -approximate stationary point within O  ✏  stochastic gradient complexity on the non-convex stochastic problems, matching the best-known lower bound.
Extensive experimental results show that Adan surpasses the corresponding
SoTA optimizers for vision, language, and RL tasks and sets new SoTAs for many popular networks and frameworks, e.g. ResNet [7], ConvNext [8], ViT [9],
Swin [10], MAE [11], Transformer-XL [12] and BERT [13].
More surprisingly,
Adan can use half of the training cost (epochs) of SoTA optimizers to achieve higher or comparable performance on ViT, ResNet, MAE, etc , and also shows great tolerance to a large range of minibatch size, e.g. from 1k to 32k. Code is released at https://github. com/sail-sg/Adan .
Introduction
Deep neural networks (DNNs) have made remarkable success in many ﬁelds, e.g. computer vision [7, 8, 14–16] and natural language processing [17, 18].
A noticeable part of such success is contributed by the stochastic gradient based optimizers which ﬁnd satisfactory solutions with high efﬁciency.
Starting from AdaGrad [19] and RMSProp [20], adaptive gradient algorithms [1–3, 19–25] have gained wide attention with faster convergence speed.
They adjust the learning rate for each gradient coordinate according to the current geometry curvature of the loss objective.
Indeed, Adam [1] and
AdamW [3], as two representatives which often offer fast convergence speed across many DNN frameworks, have become the default choice to train CNNs and ViTs [9], respectively.
However, none of the above optimizers can always stay undefeated among all its competitors across different network architectures and application settings.
For instance, for vanilla ResNet, SGD often achieves better generalization performance than adaptive gradient algorithms such as Adam, whereas on vision transformers (ViTs) [9, 10, 26], SGD often fails and AdamW is the dominant optimizer with higher and more stable performance.
Moreover, these commonly used optimizers usually fail for large-batch training, but which is a default setting of the prevalent distributed training.
Although there is some performance degradation, we still tend to choose the large-batch setting for large-scale deep learning training tasks due to the unaffordable training time.
Though some methods, e.g. LARS [27] and LAMB [4], have been proposed to handle large batch sizes, their performance often varies signiﬁcantly across batch sizes.
This performance inconsistency increases the training cost and engineering burden, since one usually has to try various optimizers for different architectures.
When we rethink the current adaptive gradient algorithms, we ﬁnd that they mainly combine the moving average idea with the heavy ball acceleration technique to estimate the ﬁrst- and second-order moments of the gradient [1–4].
However, previous studies [5, 28, 29] have revealed that Nesterov acceleration can theoretically achieve a faster convergence speed than heavy ball acceleration, as it uses gradient at an extrapolation point of the current solution and sees a slight “future".
Moreover, a recent work [6] has shown the potential of Nesterov acceleration for large-batch training [30].
Thus we are inspired to consider efﬁciently integrating Nesterov acceleration with adaptive algorithms.
Contributions: 1) We propose an efﬁcient dnn optimizer, named Adan, to train DNNs.
Adan develops a Nesterov momentum estimation method to estimate stable and accurate ﬁrst- and secondorder gradient moments in adaptive algorithms for acceleration.
Adan enjoys provably faster convergence speed than previous adaptive algorithms, e.g. Adam. 3) Empirically, Adan shows superior performance over the SoTA deep optimizers across vision, language, and RL tasks.
Methodology
In this work, we study the following regularized nonconvex optimization problem:
[Equation]
where loss f ( · , · ) is differentiable and possibly nonconvex, data ⇣ is drawn from an unknown distribution D , ✓ is learnable parameters, and k · k is the classical ` 2 norm.
Here we consider the square ` 2 regularizer as it can improve generalization performance and is widely used in practice [3].
2.1 Preliminaries
Taking a deeper step into Adam, one can easily observe that the key moving average idea in Adam is similar to the classical (stochastic) heavy-ball acceleration (HBA) technique [31]:
[Equation]
where g k is the minibatch gradient g k :
= E ⇣ ⇠ D [ r f ( ✓ k , ⇣ )]
+ ⇠ k , ⇠ k is the gradient noise, and the scalar constant ⌘ is the base learning rate.
In addition to HBA, Nesterov’s accelerated (stochastic) gradient descent (AGD) [5, 28, 29] is another popular acceleration technique in the optimization community:
[Equation]
Unlike HBA, AGD uses the gradient at the extrapolation point ✓ 0 k = ✓ k  (1   1 )( ✓ k  ✓ k  1 ) .
Hence AGD sees a slight “future” to converge faster.
Indeed, AGD theoretically converges faster than
HBA and achieves optimal convergence rate on the general smooth convex problems [5].
Meanwhile, since the over-parameterized DNNs have been observed/proved to have many convex-alike local basins [32–39], AGD seems more suitable than HBA for DNNs.
For large-batch training, the recent work [6] shows that AGD has the potential to achieve comparable performance to some speciﬁcally designed optimizers, e.g. LARS and LAMB.
With its advantage in convergence and large-batch training, we consider applying AGD to improve adaptive algorithms.
2.2 Adaptive Nesterov Momentum Algorithm
Main iteration.
We temporarily set  = 0 in Eqn. As aforementioned, AGD computes gradient at an extrapolation point ✓ 0 k instead of the current iterate ✓ k , which however brings extra computation and memory overhead for computing ✓ 0 k and preserving both ✓ k and ✓ 0 k . To solve the issue, we reformulate AGD (2) into its equivalent (See Lemma 1 in Appendix) but more DNN-efﬁcient version:
Reformulated AGD: m k = (1   1 ) m k  1 + [ g k + (1   1 )( g k  g k  1 )] ,
[Equation]
[Equation]
[Equation]
while where g k = E ⇣ ⇠ D [ r f ( ✓ k , ⇣ )] + ⇠ k .
The main idea here is that we maintain ( ✓ k  ⌘ (1   1 ) m k  1 ) rather than ✓ k in vanilla AGD per iteration, as there is no difference between them when the algorithm converges.
Like Adam, by regarding g 0 k = g k + (1   1 )( g k  g k  1 ) as the current stochastic gradient and movingly averaging g 0 k to estimate the ﬁrst- and second-moments of gradient, we obtain
Vanilla Adan:
[Equation]
The main difference of Adan with Adam-type methods is that as compared in Eqn.
(3), the moment m k of Adan is the average of { g t + (1   1 )( g t  g t  1 ) } kt =1 while those of Adam-type are the average of { g t } kt =1 .
So is their second-order term n k .
[Equation]
The ﬁrst-order moment m k =
P kt
[Equation] [Equation]
[Equation]
g
[Equation]
1 )] consists of two terms, i.e., gra- dient term g t and gradient difference term ( g t  g t  1 ) , which actually have different physic meanings.
So we further decouple them for greater ﬂexibility and also better trade-off between them:
[Equation]
where c 0 k,t =  2 (1   2 ) ( k  t ) for t > 0 , c 0 k,t = (1   2 ) k for t = 0 , and m k and v k are deﬁned as
[Equation]
This change for a ﬂexible estimation does not impair convergence speed.
As we show in Theorem(see Sec.
C in Appendix), the complexity of Adan under this change matches the lower complexity bound.
We do not separate the gradients and their difference in the second-order moment n k , since
E ( n k ) contains the correlation term
Cov( g k , g k  1 ) 6 = 0 which may have statistical signiﬁcance.
Decay Weight by Proximation.
As observed in AdamW, decoupling the optimization objective and ` 2 regularizer can largely improve the generalization performance.
Intuitively, at each iteration, we minimize a ﬁrst-order approximation of F ( · ) at the point ✓ k :
[Equation]
where k x k 2 p n k :=
h x , p n k + "  x i and ¯ m k :=
m k + (1   2 ) v k is the ﬁrst-order derivative of
F ( · ) in some sense.
Interestingly, follow the idea of proximal gradient descent [40, 41], we decouple the ` 2 regularizer from F ( · ) and only linearize the loss function f ( · ) :
[Equation]
From the perspective of the proximal operator, we ﬁnd that the optimization objective of Decoupled Regularization methods ( e.g. Adan and AdamW) at the k -th iteration is changed from the
[Table]
[Table]
vanilla “static" function F ( · ) in (1) to a “dynamic" function F k ( · ) , which adaptively regularizes the coordinates with larger gradient square terms more.
We summarize our Adan in Algorithm 1.
Convergence Analysis: As shown in Theorems 1 in Appendix.
C, the convergence speed of Adan matches the best-known theoretical lower bound for non-convex stochastic optimization problems.
This conclusion is still valid when it also uses the decoupled weight decay.
Experimental Results
[Table]
For all the tested vision tasks, NLP, and RL tasks, we only replace the default optimizer with our Adan, and do not make other changes, e.g. network architectures and data augmentation.
Vision Results.
1) supervised settings: we report the results on CNN-type architectures and ViTs in Tables 1 and 2, respectively.
2) self-supervised settings: we follow the MAE training framework to pretrain and ﬁne-tune ViT-B and ViT-L, and report results in Table 3.
All these results show that in most cases, Adan can use half of the training cost (epochs) of SoTA optimizers to achieve higher or comparable performance on ViT, ResNet, MAE, etc .
[Table]
NLP Results.
1) supervised settings: we investigate the performance of Adan on Transformer-XL, and report the results in
Table 4. 2) self-supervised settings: we use Adan to train BERT from scratch, and report the results in Table 5.
For all NLP tasks,
Adan achieves higher performance than the default SoTAs, and suppress Adam within half training steps on Transformer-XL.
[Table]
[Figure]
We replace the default Adam optimizer in PPO [44] (one of the most popular policy gradient method), and do not make other change in PPO.
Fig.
1 shows that on representative MuJoCo games, PPO-Adan achieves much higher rewards than PPO with Adam as its optimizer.
More Extra Results .
Due to space limitation, we defer more extra results on vision, NLP and RL tasks ( e [Figure] g. results with large batch size, loss curve, ablation study, etc .) to Appendix.