Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability
By Jeremy M. Cohen, Simran Kaur, Yuanzhi Li, J. Kolter, Ameet Talwalkar. International Conference on Learning Representations, 2021
455 citations recorded by the source corpus.
We empirically demonstrate that full-batch gradient descent on neural network training objectives typically operates in a regime we call the Edge of Stability. In this regime, the maximum eigenvalue of the training loss Hessian hovers just above the numerical value $2 / \text{(step size)}$, and the training loss behaves non-monotonically over short timescales, yet consistently decreases over long timescales. Since this behavior is inconsistent with several widespread presumptions in the field of optimization, our findings raise questions as to whether these presumptions are relevant to neural network training. We hope that our findings will inspire future efforts aimed at rigorously understanding optimization at the Edge of Stability. Code is available at https://github.com/locuslab/edge-of-stability.
View Paper: https://www.semanticscholar.org/paper/026bb8a1066f50ddc8797e1341353603149a8cb8
How researchers see this paper
The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.
- Yoshua Bengio
- Shuangfei Zhai
- Jiatao Gu
- Percy Liang
- Yizhe Zhang
- Chunyuan Li
- Tatiana Likhomanenko
- Kyunghyun Cho
Citations in the corpus
- Adaptive Gradient Methods at the Edge of Stability (2022, background): “Prior work has not discussed preconditioned gradient descent, but in the next section we confirm that the results of [8] carry over to the non-adaptive preconditioned setting.”
- Stabilizing Transformer Training by Preventing Attention Entropy Collapse (2023, methodology): “Eigenvalues of the Hessian As properties of the Hessian have been successfully used to gain an understanding of stability of the learning process (Ghorbani et al., 2019; Yao et al., 2020; Cohen et al., 2021; 2022; Gilmer et al., 2021), we will also use them in our analysis.”
- LoRA+: Efficient Low Rank Adaptation of Large Models (2024, background): “Notice also that optimal learning rates ( η A , η B ) are generally close to the edge of stability, a well-known behaviour in training dynamics of deep networks (Cohen et al., 2021).”