CoChat
Sign Up
PreprintCited by 97Stochastic Gradient Optimization TechniquesNeural Networks and ApplicationsMachine Learning and Data Classification

The DOI matches its DataCite record.

Adaptive Gradient Methods at the Edge of Stability

Open access
Cohen, Jeremy M., Ghorbani, Behrooz, Krishnan, Shankar, Agarwal, Naman, Medapati, Sourabh, Badura, Michal, Suo, Daniel, Cardoze, David, Nado, Zachary, Dahl, George E., Gilmer, Justin·2022·arXiv
DOI 10.48550/arxiv.2207.14484arXiv 2207.14484
Open this paper in CoChat — free accountOpen PDF

Abstract

Very little is known about the training dynamics of adaptive gradient methods like Adam in deep learning. In this paper, we shed light on the behavior of these algorithms in the full-batch and sufficiently large batch settings. Specifically, we empirically demonstrate that during full-batch training, the maximum eigenvalue of the preconditioned Hessian typically equilibrates at a certain numerical value -- the stability threshold of a gradient descent algorithm. For Adam with step size $η$ and $β_1 = 0.9$, this stability threshold is $38/η$. Similar effects occur during minibatch training, especially as the batch size grows. Yet, even though adaptive methods train at the ``Adaptive Edge of Stability'' (AEoS), their behavior in this regime differs in a significant way from that of non-adaptive methods at the EoS. Whereas non-adaptive algorithms at the EoS are blocked from entering high-curvature regions of the loss landscape, adaptive gradient methods at the AEoS can keep advancing into high-curvature regions, while adapting the preconditioner to compensate. Our findings can serve as a foundation for the community's future understanding of adaptive gradient methods in deep learning.

Overview

The claims and statistics of this paper haven’t been checked yet.

Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.

A few free checks a day · papers already verified open instantly · results are public

Integrity signals

Registry

DataCite · 10.48550/arxiv.2207.14484

Cites

Statistics

Verification

Claims the paper makes

The claims and statistics of this paper haven’t been checked yet.

Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.

A few free checks a day · papers already verified open instantly · results are public

Claims other papers make about this one

Each is checked against this paper’s own text: does it say what the citing paper says it does?

  • Cohen et al. [16] define 'preconditioned sharpness' as the top eigenvalue of the preconditioned Hessian λ1(P−1t Ht), and observe that this quantity often oscillates around the stability threshold 2+2β1(1−β1)η.

    “They define 'preconditioned sharpness' to be the top eigenvalue of the preconditioned Hessian λ 1 ( P -1 t H t ) , where H t is the Hessian of the loss, and observe that this quantity often oscillates around the stability threshold 2+2 β 1 (1 -β 1 ) η , where η is the learning rate.”

    From How Memory in Optimization Algorithms Implicitly Modifies the Loss (2025) · cites this paper as [16]

  • Cohen et al. [16] notice that in a sense Adam trains at the edge of stability.

    “Cohen et al. [16] notice that in a sense Adam trains at the edge of stability.”

    From How Memory in Optimization Algorithms Implicitly Modifies the Loss (2025) · cites this paper as [16]

  • Cohen et al. [16] view Adam as momentum gradient descent with evolving preconditioner.

    “They view Adam as momentum gradient descent with evolving preconditioner”

    From How Memory in Optimization Algorithms Implicitly Modifies the Loss (2025) · cites this paper as [16]

Self-check

Percentages, sums and reported differences re-computed from the paper’s own numbers.

The self-check needs the paper’s full text.

Citations

Cites

No references could be resolved yet.

Research works printed but not linked: 44

  1. Naman Agarwal, Rohan Anil, Elad Hazan, Tomer Koren, and Cyril Zhang. Disentangling adaptive gradient methods from learning rates. arXiv preprint arXiv:2002.11803, 2020. · arXiv
  2. Kwangjun Ahn, Jingzhao Zhang, and Suvrit Sra. Understanding the unstable convergence of gradient descent. arXiv preprint arXiv:2204.01050, 2022. · arXiv
  3. Lukas Balles and Philipp Hennig. Dissecting adam: The sign, magnitude and variance of stochastic gradients. In International Conference on Machine Learning, pages 404-413. PMLR, 2018.
  4. Anas Barakat and Pascal Bianchi. Convergence and dynamical behavior of the adam algorithm for non-convex stochastic optimization. arXiv preprint arXiv:1810.02263, 2018. · arXiv
  5. Jinghui Chen, Dongruo Zhou, Yiqi Tang, Ziyan Yang, Yuan Cao, and Quanquan Gu. Closing the generalization gap of adaptive gradient methods in training deep neural networks. arXiv preprint arXiv:1806.06763, 2018. · arXiv
  6. Xiangning Chen, Cho-Jui Hsieh, and Boqing Gong. When vision transformers outperform resnets without pre-training or strong data augmentations. arXiv preprint arXiv:2106.01548, 2021. · arXiv
  7. Xiangyi Chen, Sijia Liu, Ruoyu Sun, and Mingyi Hong. On the convergence of a class of adam-type algorithms for non-convex optimization. arXiv preprint arXiv:1808.02941, 2018. · arXiv
  8. Jeremy M. Cohen, Simran Kaur, Yuanzhi Li, J. Zico Kolter, and Ameet S. Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. ArXiv, abs/2103.00065, 2021. · arXiv
  9. Alexandre Défossez, Léon Bottou, Francis Bach, and Nicolas Usunier. A simple convergence proof of adam and adagrad. arXiv preprint arXiv:2003.02395, 2020. · arXiv
  10. Timothy Dozat. Incorporating nesterov momentum into adam. ICLR 2016, workshop track, 2016.
  11. John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(61):2121-2159, 2011.
  12. Niv Giladi, Mor Shpigel Nacson, Elad Hoffer, and Daniel Soudry. At stability's edge: How to adjust hyperparameters to preserve minima selection in asynchronous training of neural networks? arXiv preprint arXiv:1909.12340, 2019. · arXiv
  13. Justin Gilmer, Behrooz Ghorbani, Ankush Garg, Sneha Kudugunta, Behnam Neyshabur, David Cardoze, George Dahl, Zachary Nado, and Orhan Firat. A loss curvature perspective on training instability in deep learning. arXiv preprint arXiv:2110.04369, 2021. · arXiv
  14. Diego Granziol. Flatness is a false friend. arXiv preprint arXiv:2006.09091, 2020. · arXiv
  15. Diego Granziol, Xingchen Wan, Samuel Albanie, and Stephen Roberts. Explaining the adaptive generalisation gap. arXiv preprint arXiv:2011.08181, 2020. · arXiv
  16. Stanislaw Jastrzebski, Devansh Arpit, Oliver Astrand, Giancarlo B Kerg, Huan Wang, Caiming Xiong, Richard Socher, Kyunghyun Cho, and Krzysztof J Geras. Catastrophic fisher explosion: Early phase fisher matrix impacts generalization. In International Conference on Machine Learning, pages 4772-4784. PMLR, 2021.
  17. Stanisław Jastrzębski, Zachary Kenton, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. On the relation between the sharpest directions of dnn loss and the sgd step length. arXiv preprint arXiv:1807.05031, 2018. · arXiv
  18. Stanislaw Jastrzebski, Maciej Szymczak, Stanislav Fort, Devansh Arpit, Jacek Tabor, Kyunghyun Cho, and Krzysztof Geras. The break-even point on optimization trajectories of deep neural networks. arXiv preprint arXiv:2002.09572, 2020. · arXiv
  19. Nitish Shirish Keskar and Richard Socher. Improving generalization performance by switching from adam to sgd. arXiv preprint arXiv:1712.07628, 2017. · arXiv
  20. Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014. · arXiv
  21. Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint arXiv:2003.02218, 2020. · arXiv
  22. Xiaoyu Li and Francesco Orabona. On the convergence of stochastic gradient descent with adaptive stepsizes. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 983-992. PMLR, 2019.
  23. Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han. Understanding the difficulty of training transformers. arXiv preprint arXiv:2004.08249, 2020. · arXiv
  24. Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv:1608.03983, 2016. · arXiv
  25. Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101, 2017. · arXiv
  26. Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. arXiv preprint arXiv:1906.05890, 2019. · arXiv
  27. Chao Ma, Lei Wu, et al. A qualitative study of the dynamic behavior for adaptive gradient algorithms. arXiv preprint arXiv:2009.06125, 2020. · arXiv
  28. Chao Ma, Lei Wu, and Lexing Ying. The multiscale structure of neural network loss functions: The effect on optimization and origin. arXiv preprint arXiv:2204.11326, 2022. · arXiv
  29. Rotem Mulayoff, Tomer Michaeli, and Daniel Soudry. The implicit bias of minima stability: A view from function space. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
  30. Kamil Nar and Shankar Sastry. Step size matters in deep learning. Advances in Neural Information Processing Systems, 31, 2018.
  31. Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. In search of the real inductive bias: On the role of implicit regularization in deep learning. arXiv preprint arXiv:1412.6614, 2014. · arXiv
  32. Martin Popel and Ondřej Bojar. Training tips for the transformer model. arXiv preprint arXiv:1804.00247, 2018. · arXiv
  33. Sashank J. Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence of adam and beyond. In International Conference on Learning Representations, 2018.
  34. Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pages 4596-4604. PMLR, 2018.
  35. Bohan Wang, Qi Meng, Wei Chen, and Tie-Yan Liu. The implicit bias for adaptive optimization algorithms on homogeneous neural networks. In International Conference on Machine Learning, pages 10849-10858. PMLR, 2021.
  36. Rachel Ward, Xiaoxia Wu, and Leon Bottou. Adagrad stepsizes: Sharp convergence over nonconvex landscapes. In International Conference on Machine Learning, pages 6677-6686. PMLR, 2019.
  37. Ashia C Wilson, Rebecca Roelofs, Mitchell Stern, Nati Srebro, and Benjamin Recht. The marginal value of adaptive gradient methods in machine learning. Advances in neural informa- tion processing systems, 30, 2017.
  38. Lei Wu, Chao Ma, and Weinan E. How sgd selects the global minima in over-parameterized learning: A dynamical stability perspective. In S. Bengio, H. Wallach, H. Larochelle, K. Grau- man, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31.
  39. Yuege Xie, Xiaoxia Wu, and Rachel Ward. Linear convergence of adaptive stochastic gradient descent. In International Conference on Artificial Intelligence and Statistics, pages 1475-1485. PMLR, 2020.
  40. Zeke Xie, Xinrui Wang, Huishuai Zhang, Issei Sato, and Masashi Sugiyama. Adai: Separating the effects of adaptive learning rate and momentum inertia. arXiv preprint arXiv:2006.15815, 2020. · arXiv
  41. Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In International Conference on Machine Learning, pages 10524-10533. PMLR, 2020.
  42. Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. arXiv preprint arXiv:1605.07146, 2016. · arXiv
  43. Dongruo Zhou, Jinghui Chen, Yuan Cao, Yiqi Tang, Ziyan Yang, and Quanquan Gu. On the convergence of adaptive gradient methods for nonconvex optimization. arXiv preprint arXiv:1808.05671, 2018. · arXiv
  44. Pan Zhou, Jiashi Feng, Chao Ma, Caiming Xiong, Steven Chu Hong Hoi, et al. Towards theoretically understanding why sgd generalizes better than adam in deep learning. Advances in Neural Information Processing Systems, 33:21285-21296, 2020.

Cited by

Sorted by citations

Metadata

Authors
Cohen, J. M., Ghorbani, B., Krishnan, S., Agarwal, N., Medapati, S., Badura, M., Suo, D., Cardoze, D., Nado, Z., Dahl, G. E., & Gilmer, J.
Affiliations
Google (United States) (US)
ORCID
Venue
arXiv
Publisher
arXiv
Published
2022
ISSN
2331-8422 (ISSN-L)
DOI
10.48550/arxiv.2207.14484
arXiv
2207.14484
OpenAlex
W4289446290
Cited by
97
References
0
Citations by year
2026: 41 · 2025: 15 · 2024: 14 · 2023: 23 · 2022: 4
Type
posted content
Language
en
Open access
Licence
creativecommons.org/licenses/by/4.0/legalcode
Locations
  • arXiv (Cornell University) (submittedVersion)
  • arXiv (Cornell University) (cc-by)
Topics
  • Physical Sciences › Computer Science › Artificial Intelligence › Stochastic Gradient Optimization Techniques
  • Physical Sciences › Computer Science › Artificial Intelligence › Neural Networks and Applications
  • Physical Sciences › Computer Science › Artificial Intelligence › Machine Learning and Data Classification
Keywords
adaptive gradient methods, Adam optimizer, edge of stability, training dynamics, largest eigenvalue, batch size scaling, loss landscape curvature, stability threshold

About these checks

Registry, retraction, Cites and statistics results are computed, not judged, with one exception: a printed reference that no index lists is read once by a model before it counts as a citation error. Claims are extracted by a model and kept only when their quote appears verbatim in the text; each cited source is then checked, with a model’s help, for whether it says what the claim relies on. Those verdicts can be wrong — read the claim’s quote and the cited source before relying on one.

Statistics and arithmetic are recomputed, not judged.

Confidence bands are shown only for fields calibrated on our annotated test set. None is calibrated yet, so no band is shown.

What counts as a claim ›How we measure confidence ›How we check this ›

Are you an author, or do you see a problem with this page? Request a review