Table of Contents
Fetching ...

Adan: Adaptive Nesterov Momentum Algorithm for Faster Optimizing Deep Models

Xingyu Xie, Pan Zhou, Huan Li, Zhouchen Lin, Shuicheng Yan

TL;DR

Adan introduces a novel adaptive optimization approach by employing a Nesterov momentum estimation (NME) that computes the gradient at the current point while maintaining the look-ahead advantages of Nesterov acceleration. The method combines NME with adaptive first- and second-moment updates, decoupled weight decay, and a restart mechanism, achieving faster convergence and robustness across diverse architectures. Theoretical results show matching lower bounds for stochastic gradient complexity under Lipschitz conditions, with improved rates under Lipschitz Hessian, and extensive experiments demonstrate consistent superiority over SoTA optimizers on CNNs, ViTs, LSTMs, BERT, GPT-2, PPO, and GNNs. The work significantly reduces the need to tailor optimizers per architecture, and the public code release facilitates broad adoption in vision, language, and RL applications.

Abstract

In deep learning, different kinds of deep networks typically need different optimizers, which have to be chosen after multiple trials, making the training process inefficient. To relieve this issue and consistently improve the model training speed across deep networks, we propose the ADAptive Nesterov momentum algorithm, Adan for short. Adan first reformulates the vanilla Nesterov acceleration to develop a new Nesterov momentum estimation (NME) method, which avoids the extra overhead of computing gradient at the extrapolation point. Then, Adan adopts NME to estimate the gradient's first- and second-order moments in adaptive gradient algorithms for convergence acceleration. Besides, we prove that Adan finds an $ε$-approximate first-order stationary point within $\mathcal{O}(ε^{-3.5})$ stochastic gradient complexity on the non-convex stochastic problems (e.g., deep learning problems), matching the best-known lower bound. Extensive experimental results show that Adan consistently surpasses the corresponding SoTA optimizers on vision, language, and RL tasks and sets new SoTAs for many popular networks and frameworks, e.g., ResNet, ConvNext, ViT, Swin, MAE, DETR, GPT-2, Transformer-XL, and BERT. More surprisingly, Adan can use half of the training cost (epochs) of SoTA optimizers to achieve higher or comparable performance on ViT, GPT-2, 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, and has been used in multiple popular deep learning frameworks or projects.

Adan: Adaptive Nesterov Momentum Algorithm for Faster Optimizing Deep Models

TL;DR

Adan introduces a novel adaptive optimization approach by employing a Nesterov momentum estimation (NME) that computes the gradient at the current point while maintaining the look-ahead advantages of Nesterov acceleration. The method combines NME with adaptive first- and second-moment updates, decoupled weight decay, and a restart mechanism, achieving faster convergence and robustness across diverse architectures. Theoretical results show matching lower bounds for stochastic gradient complexity under Lipschitz conditions, with improved rates under Lipschitz Hessian, and extensive experiments demonstrate consistent superiority over SoTA optimizers on CNNs, ViTs, LSTMs, BERT, GPT-2, PPO, and GNNs. The work significantly reduces the need to tailor optimizers per architecture, and the public code release facilitates broad adoption in vision, language, and RL applications.

Abstract

In deep learning, different kinds of deep networks typically need different optimizers, which have to be chosen after multiple trials, making the training process inefficient. To relieve this issue and consistently improve the model training speed across deep networks, we propose the ADAptive Nesterov momentum algorithm, Adan for short. Adan first reformulates the vanilla Nesterov acceleration to develop a new Nesterov momentum estimation (NME) method, which avoids the extra overhead of computing gradient at the extrapolation point. Then, Adan adopts NME to estimate the gradient's first- and second-order moments in adaptive gradient algorithms for convergence acceleration. Besides, we prove that Adan finds an -approximate first-order stationary point within stochastic gradient complexity on the non-convex stochastic problems (e.g., deep learning problems), matching the best-known lower bound. Extensive experimental results show that Adan consistently surpasses the corresponding SoTA optimizers on vision, language, and RL tasks and sets new SoTAs for many popular networks and frameworks, e.g., ResNet, ConvNext, ViT, Swin, MAE, DETR, GPT-2, Transformer-XL, and BERT. More surprisingly, Adan can use half of the training cost (epochs) of SoTA optimizers to achieve higher or comparable performance on ViT, GPT-2, 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, and has been used in multiple popular deep learning frameworks or projects.
Paper Structure (45 sections, 11 theorems, 114 equations, 3 figures, 22 tables, 2 algorithms)

This paper contains 45 sections, 11 theorems, 114 equations, 3 figures, 22 tables, 2 algorithms.

Key Result

Lemma 1

Assume $\mathop{\mathrm{\mathop{\mathbb{E}}}}\nolimits(\bm{\xi}_k) = \bm{0}$, $\operatorname{Cov}(\bm{\xi}_i, \bm{\xi}_j) =0$ for any $k,i,j>0$, $\Bar{\bm{\theta}}_k$ and $\Bar{\mathbf{m}}_k$ be the iterate and momentum of the vanilla AGD in Eqn. eq:AGD-I, respectively. Let $\bm{\theta}_{k+1}\colone where $\mathbf{g}_k'\coloneqq{\mathbf{g}}_k + \qty(1-{\color{orange}\beta_1})\qty( {\mathbf{g}}_k -

Figures (3)

  • Figure 1: Training and test curves of various optimizers on ImageNet. The different magnitude of training and test loss is due to data argumentation. Best viewed in 2$\times$-sized color pdf file.
  • Figure 2: Comparison of PPO and our PPO-Adan on several RL games simulated by MuJoCo. Here PPO-Adan simply replaces the Adam optimizer in PPO with our Adan and does not change others. Best viewed in 2$\times$-sized color pdf file.
  • Figure 3: Effects of momentum coefficients $(\beta_1,\beta_2,\beta_3)$ to top-1 accuracy (%) of Adan on ViT-B under MAE training framework (800 pretraining and 100 fine-tuning epochs on ImageNet).

Theorems & Definitions (21)

  • Lemma 1
  • Lemma 2
  • proof
  • Theorem 1
  • proof
  • Proposition 1
  • proof
  • Lemma 3
  • proof
  • Lemma 4
  • ...and 11 more