Adaptive Gradient Methods at the Edge of Stability

By Jeremy M. Cohen, Behrooz Ghorbani, Shankar Krishnan, Naman Agarwal, Sourabh Medapati, Michal Badura, Daniel Suo, David E. Cardoze, Zachary Nado, George E. Dahl and 1 more. 2022

86 citations recorded by the source corpus.

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 $\eta$ and $\beta_1 = 0.9$, this stability threshold is $38/\eta$. 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.

Subfields: Artificial Intelligence, Computational Mechanics, Statistical and Nonlinear Physics, Optimization & Training, Scaling & Learning Theory

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computational Mechanics

In Statistical and Nonlinear Physics

In Optimization & Training

In Scaling & Learning Theory

Citations in the corpus

Stable and low-precision training for large-scale vision-language models (2023, background): “[12] which offers a general and principled treatment of fast loss spikes, and which we recommend to readers.”
Stabilizing Transformer Training by Preventing Attention Entropy Collapse (2023, methodology): “non-fully adaptive) approximation of AdamW, has been shown empirically as a suitable stability threshold for the full algorithm (Cohen et al., 2022), and is given by Γ = 2 + 2β1 1− β1 1 η = 38 η , (25) where β1 = 0.”
H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models (2023, unknown): “…(Variš & Bojar, 2021; Zhang & Vaidya, 2021), a better initialization (Liu et al., 2020b;a; Xu et al., 2020; Zhu et al., 2021), customized optimizers (Cohen et al., 2022), improved normalization (Wang et al., 2022; Yang et al., 2022), weight decay (Loshchilov & Hutter, 2017), and early stopping.”