Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability

TL;DR

This paper empirically shows that neural network training with gradient descent often occurs at the Edge of Stability, where the Hessian's maximum eigenvalue hovers just above a critical value, challenging common optimization assumptions.

Contribution

It introduces the concept of the Edge of Stability in neural network training and provides empirical evidence of this regime's prevalence and characteristics.

Findings

Hessian eigenvalue hovers just above 2 / step size

Training loss exhibits non-monotonic short-term behavior

Long-term training loss decreases despite instability signs

Abstract

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.