Adan: Adaptive Nesterov Momentum Algorithm for Faster Optimizing Deep Models

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

Source-linked AI summary

Adan: Adaptive Nesterov Momentum Algorithm for Faster Optimizing Deep Models

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

arXiv:2208.06677v5Aug 13, 2022cs.LGmath.OC

TL;DR

Deep networks often require different optimizers, making optimizer selection costly and inconsistent across architectures. Adan reformulates Nesterov acceleration into an efficient adaptive optimizer, and the paper reports lower-bound-matching convergence with broad empirical gains across vision, language, and reinforcement-learning tasks. These results include comparable or higher performance with half the training epochs in several settings, although the paper notes remaining optimizer-performance gaps on some transformers.

Problem
Different deep-network architectures often require different optimizers, creating repeated tuning effort and inconsistent performance across architectures and training settings.
Method
Adan reformulates vanilla Nesterov acceleration into a Nesterov momentum estimator and uses it to estimate adaptive first- and second-order gradient moments without extrapolation-point gradient computation.
Results
Adan consistently surpasses corresponding SoTA optimizers across vision, language, and reinforcement-learning tasks, with stochastic gradient complexity matching the cited lower bound.
Takeaways & Limitations
Adan can reduce the need to try different optimizers across architectures and can achieve higher or comparable performance with half the training epochs in several reported settings.
Takeaways & Limitations
The paper reports that some optimizers, including SGD and Adam, can achieve much lower accuracy than AdamW on transformers.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Adan addresses the difficulty of choosing optimizers across deep-network architectures by reformulating Nesterov acceleration into an efficient adaptive optimizer. It combines this mechanism with theoretical convergence guarantees and broad empirical improvements across tasks and training regimes.

Motivation: Different architectures favor different optimizers, increasing training cost and engineering effort, while common optimizers can perform inconsistently across architectures and batch sizes.The paper notes differing behavior for ResNets, vision transformers, and large-batch settings.
Motivation: Nesterov acceleration can theoretically converge faster than heavy-ball acceleration and may improve robustness and large-batch training.Nesterov uses gradients at an extrapolation point to look slightly ahead along the optimization trajectory.
Method: Adan reformulates vanilla Nesterov acceleration into Nesterov momentum estimation, avoiding gradient computation at the extrapolation point and its associated implementation overhead.The reformulation estimates the moving gradient average from the current gradient and its difference from the previous gradient.
Method: Adan uses the reformulated estimate to compute adaptive first- and second-order moments, then applies an adaptive parameter update.This preserves Nesterov acceleration while fitting the mechanism into adaptive gradient optimization.
Theory: Adan’s stochastic gradient complexity matches the lower bound under Lipschitz Hessian assumptions and is superior to the cited complexities of A-NIGT and Adam+.The paper also compares Adan favorably with prior adaptive methods under nonconvex stochastic optimization.
Experiments: Across vision, language, and reinforcement-learning tasks, Adan surpasses corresponding SoTA optimizers, often achieving higher or comparable performance with half the training epochs.The reported minibatch tolerance extends from 1k to 32k in cited ViT experiments.

2 RELATED WORK

The related work distinguishes SGD and accelerated variants from adaptive gradient algorithms, while situating Adan among broader approaches to training efficiency.

Optimizer families: SGD updates parameters along stochastic-gradient directions, while heavy-ball and Nesterov acceleration add momentum-based mechanisms for faster convergence.Nesterov computes a gradient after looking ahead to an extrapolated point.
Adaptive optimization: Adaptive methods such as AdaGrad, RMSProp, and Adam use gradient moments to precondition updates and adjust learning rates across coordinates.Variants target more accurate moments, reduced variance, or prevention of gradient collapse.
Broader efficiency research: Training-efficiency research also includes data augmentation and strategies for training compact models, beyond the optimization techniques central to this paper.Mixup is cited as enriching training data without additional data collection.

3 METHODOLOGY

Adan builds an adaptive optimizer from an efficient reformulation of Nesterov acceleration, combining gradient and gradient-difference information with separated regularization. Its design preserves convergence guarantees while supporting practical optimization updates.

Problem setting: The optimization setting is regularized nonconvex stochastic learning, where the loss may be nonconvex and parameters are trained from data sampled from an unknown distribution.The formulation covers network training and least-squares regression.
Background: Adaptive methods adjust coordinate-wise learning rates using objective curvature, whereas Adam’s moving-average mechanism resembles heavy-ball acceleration.Moving averages accumulate consistent gradient directions into larger update steps.
Motivation: Nesterov acceleration evaluates gradients at an extrapolated point, but this creates extra computation, parameter-storage, and distributed-synchronization costs.The extrapolation step complicates the optimizer workflow and can burden multi-GPU communication.
Nesterov momentum estimation: Adan replaces the extrapolated-point gradient with g′_k = g_k + (1 − β1)(g_k − g_{k−1}), an equivalent reformulation that avoids the extra cost.The reformulation maintains an equivalent convergent trajectory under the stated lemma assumptions.
Adaptive moments: Adan averages the reformulated gradient estimate to obtain adaptive first- and second-order moments for its parameter update.Its first moment combines gradient and gradient-difference terms, while the second moment retains their correlation.
Adaptive moments: Adan decouples gradient and gradient-difference contributions for flexibility without impairing convergence speed, while retaining correlation in the second-order moment.The paper states that this complexity still matches the best-known lower bound.
Regularization: Separated regularization decouples the ℓ2 regularizer from the loss and yields a coordinate-adaptive dynamic regularization objective.The resulting update connects to the AdamW-style decoupled weight-decay rule through first-order approximation.
Algorithm: Algorithm 1 includes a momentum-reset restart condition, but the experiments generally omit it except for the specified supplementary comparison.The restart can improve performance while remaining optional for the simplified implementation.

4 CONVERGENCE ANALYSIS

Adan’s convergence analysis establishes guarantees for nonconvex stochastic optimization under smoothness and gradient-oracle assumptions, with stronger guarantees under a Lipschitz-Hessian condition. The analysis also considers momentum restarts and large-batch training.

Assumptions and convergence: Under smoothness and unbiased, bounded stochastic gradients, Theorem 1 proves convergence of Adan to an ε-approximate first-order stationary point.The analysis explicitly assumes Lipschitz gradients and a stochastic gradient oracle with zero-mean noise and bounded magnitude and variance.
Comparison with prior analyses: Adan’s convergence complexity is reported as lower than those of prior adaptive optimizers, including AdaBelief, LAMB, Adam+, and A-NIGT.The comparison is presented in terms of stochastic gradient complexity under the analyzed conditions.
Restart mechanism: The analysis introduces a restart condition that resets momentum after sufficient parameter updates, helping Adan use local geometric information in a new local basin.The restart frequency is controlled by a constant R, and the extension point is defined using the current parameters, momentum, and gradient-related terms.
Assumptions and convergence: With a Lipschitz-continuous Hessian, Theorem 2 improves Adan’s stochastic gradient complexity to O(ε^-3.5), matching the corresponding lower bound.The paper states that this complexity matches the lower bound for non-variance-reduction algorithms under the additional Hessian condition.
Large-batch setting: Theorem 2 continues to hold for large batch sizes, while the analysis permits a step size of order O(1), larger than those reported for some competing optimizers.The paper connects this larger step size with faster convergence in practice for large-minibatch training.

5 EXPERIMENTAL RESULTS

Experiments evaluate Adan across vision, language, detection, segmentation, and large-batch settings, generally comparing it with official or state-of-the-art optimizers. Across these tasks, Adan is reported to converge faster and achieve higher or comparable performance, sometimes with substantially reduced training cost.

Vision and language evaluations: Adan’s experiments compare it with official or state-of-the-art optimizers across CNNs, vision transformers, self-supervised learning, detection, segmentation, and language models.The study also includes BERT, GPT-2, and Transformer-XL under their established training settings.
CNN-type architectures: Adan achieves higher or comparable ImageNet top-1 accuracy than LAMB with 200 epochs on ResNet, compared with LAMB’s 300-epoch official result.The authors also report larger improvements when training is limited to 100 epochs.
Vision transformers and MAE: Adan outperforms AdamW across ViT and Swin model sizes and improves MAE pretraining results, including 0.5% over AdamW at 300 pretraining epochs.At 800 epochs, Adan surpasses AdamW trained for 1,600 epochs by 0.2%.
Large-batch training: Adan is robust from 2k to 32k batch sizes and achieves higher performance and robustness than LAMB in large-batch ViT training.The paper contrasts this with AdamW failing beyond batch size 4,096 for ViTs.
Detection and segmentation: Adan improves detection box AP by 1.6%–1.8% over AdamW on Deformable-DETR and is evaluated against AdamW for COCO detection and segmentation.The experiments use Deformable-DETR and Mask R-CNN with ConvNext as backbone.
Language models: Adan achieves the lowest perplexity on three LSTM models, improves average perplexity by about 1.0 over the runner-up, and improves BERT performance by 1.8 average points across GLUE tasks.It also surpasses Adam on GPT-2 within half the pre-training steps and matches Adam’s Transformer-XL results using 100k versus 200k steps.

6 CONCLUSION

The paper proposes Adan to reduce the need to try different optimizers across deep network architectures. It combines reformulated Nesterov acceleration with adaptive optimization, with theoretical and empirical advantages over existing optimizers.

Adan reformulates vanilla accelerated gradient descent and uses it to estimate first- and second-order moments in adaptive optimization.
Adan’s convergence complexity matches lower bounds and is superior to that of other adaptive optimizers.
Adan consistently surpasses other optimizers across popular vision, language, and reinforcement-learning architectures and frameworks.

14 Adan: Adaptive Nesterov Momentum Algorithm

This material outlines the paper’s supplementary organization, covering convergence discussions, implementation details, experiments, ablations, and technical proofs for Adan.

The supplementary materials include additional experiments, implementation details, and technical proofs for Adan.
Convergence discussion: The convergence appendix discusses why the lower bound is Ω(ε^-3.5) rather than Ω(ε^-3.0) and compares constants in optimizer convergence bounds.
Experiments: Additional experimental materials include detailed ViT comparisons and an ablation study.
Technical proofs: The proof appendix establishes equivalence between AGD and reformulated AGD, analyzes convergence under Lipschitz-gradient assumptions, and studies faster convergence under Lipschitz-Hessian assumptions.

APPENDIX A DISCUSSION ON CONVERGENCE RESULTS

The appendix explains the assumptions behind Adan’s convergence results and why constant factors can matter substantially for large neural networks, even when accuracy dependence is already optimal.

Under Lipschitz gradient and Hessian assumptions, stochastic gradient methods querying one point per iteration have a lower bound of Ω(ε^-3.5).
Without a Lipschitz-Hessian assumption, the nonconvex complexity lower bound is Θ(ε^-4).
Most compared optimizers already achieve optimal dependence on ε, so their complexity differences arise mainly from constant factors.
Large parameter dimension and differences between ℓ2- and ℓ∞-norm bounds can make constant-factor differences span several orders of magnitude.
For ViT-small, Adan’s bound is empirically estimated to be 7.6 × 10^6 smaller than Adam’s under the widely used Lipschitz-gradient assumption.

B.1 Pre-training Results on LLMs

The supplementary experiments examine Adan in large-scale language-model pretraining and optimizer comparisons, including MoE models, ViT-S, and related training measurements.

Pre-training experiments: MoE pretraining experiments train models from scratch with multiple model sizes and data volumes on RedPajama-v2.
Reported measurements: The supplementary tables cover training loss, optimizer accuracy, GPU memory and wall duration, and multi-node LLM training speed.
MoE results: Adan consistently outperforms AdamW across all tested MoE configurations and data volumes.
ViT-S comparisons: The experiments compare Adan with Adam, SGD-M, LAMB, and AdamW on ViT-S under ImageNet training setting II.
Ablation: Adan’s momentum coefficients are evaluated by their effects on ViT-B top-1 accuracy under MAE training.

B.3 Memory and Computation Time Comparison for Single Step

Adan adds little computational or memory overhead in single- and multi-GPU training, while its performance improvements remain significant across model sizes.

Single GPU Analysis: Over 10 CNN and transformer models showed negligible timing differences despite Adan’s slightly increased computational complexity.Measurements covered 200 training iterations on a single GPU.
Single GPU Analysis: Small models showed little peak-memory difference because PyTorch rounds reserved memory blocks to larger page sizes.For example, AdamW and Adan used 768 MB and 900 MB, respectively, while PyTorch could report 1024 MB.
Multi-GPU Distributed Training: A 1.3-billion-parameter, 8-expert MoE model showed small time and memory differences across GPUs.Distributed optimizer states reduce the effect of Adan’s additional memory load on any single GPU.
Overall: Across single- and multi-GPU settings, Adan’s additional overhead was minimal while its performance enhancements were significant.The conclusion applies regardless of model size.

B.5.1 Robustness to in momentum coefficients

Adan remains robust across momentum-coefficient settings in vision and language-model experiments, generally maintaining strong accuracy or lower loss than AdamW.

MAE: 83.7%+ accuracy was achieved by Adan in most MAE momentum-coefficient settings after 800 pretraining epochs.This exceeded the official 83.6% accuracy reported for the corresponding setup.
MAE: A smaller second-order coefficient, 1.0 −β3, can improve accuracy by increasing the use of current landscape information.The passage suggests complex MAE pretraining may prefer local geometric information less strongly.
MoE language model: Adan consistently achieved lower next-token training loss than AdamW in nearly all tested MoE momentum settings.The experiment used an 8-expert architecture and 3 billion RedPajama tokens.
MoE language model: Adan showed lower sensitivity to β fluctuations, with smaller loss variance across tested momentum settings.Because the language-model epoch was one, training loss effectively represented validation loss in this study.

B.5.2 Ablation Study on Adan’s Components

Component ablations attribute Adan’s main improvement to its reformulated Nesterov acceleration, while restart and proximal weight decay provide smaller gains.

Component ablation: 0.016 training-loss reduction came from Adan’s reformulated Nesterov acceleration relative to AdamW’s heavy-ball acceleration.The ablation used an 8-expert MoE model trained on 10 billion RedPajama tokens.
Component ablation: 0.002 and 0.006 training-loss improvements came from proximal weight decay and restart strategy, respectively.Adan does not use restart by default.
Component ablation: The ablation identifies improved Nesterov acceleration as the primary contributor to Adan’s performance enhancement.This conclusion follows the larger loss reduction associated with the reformulated acceleration.
Training settings: Adan consistently improved results across the two evaluated training settings and surpassed the reported 80.4% ViT-small accuracy.The comparison concerns ImageNet top-1 accuracy under Training Settings I and II.
Restart strategy: Restart strategy slightly improved Adan’s test performance on ViT and ConvNext after 300 training epochs.The authors omitted restart in most experiments to avoid tuning its frequency.

C.3.5 Proof for Main Theorem

The main-theorem proof establishes convergence of Adan under stated assumptions by combining restart-cycle bounds with auxiliary propositions and lemmas.

Restart-cycle analysis: The proof analyzes Algorithm 2 across restart cycles, initializing each cycle from the previous cycle’s last iterate.The total iteration bound is obtained by summing inequalities over the cycles.
Restart-cycle analysis: Each restart cycle has a bounded iteration count, enabling an overall bound on Algorithm 2’s total iterations.The supplied passages state these bounds symbolically but omit their explicit orders.
Supporting results: The proof relies on auxiliary propositions and moving-average lemmas under the stated assumptions.These supporting results include propositions concerning assumptions and moving-average sequences.