Understanding Gradient Descent on the Edge of Stability in Deep Learning

Sanjeev Arora, Zhiyuan Li, Abhishek Panigrahi

Research output: Contribution to journalConference articlepeer-review

14 Scopus citations

Abstract

Deep learning experiments by Cohen et al. (2021) using deterministic Gradient Descent (GD) revealed an Edge of Stability (EoS) phase when learning rate (LR) and sharpness (i.e., the largest eigenvalue of Hessian) no longer behave as in traditional optimization. Sharpness stabilizes around 2/LR and loss goes up and down across iterations, yet still with an overall downward trend. The current paper mathematically analyzes a new mechanism of implicit regularization in the EoS phase, whereby GD updates due to non-smooth loss landscape turn out to evolve along some deterministic flow on the manifold of minimum loss. This is in contrast to many previous results about implicit bias either relying on infinitesimal updates or noise in gradient. Formally, for any smooth function L with certain regularity condition, this effect is demonstrated for (1) Normalized GD, i.e., GD with a varying LR ηt = k∇L(ηx(t))k and loss L; (2) GD with constant LR and loss pL − minx L(x). Both provably enter the Edge of Stability, with the associated flow on the manifold minimizing λ1(∇2L). The above theoretical results have been corroborated by an experimental study.

Original languageEnglish (US)
Pages (from-to)948-1024
Number of pages77
JournalProceedings of Machine Learning Research
Volume162
StatePublished - 2022
Event39th International Conference on Machine Learning, ICML 2022 - Baltimore, United States
Duration: Jul 17 2022Jul 23 2022

All Science Journal Classification (ASJC) codes

  • Artificial Intelligence
  • Software
  • Control and Systems Engineering
  • Statistics and Probability

Fingerprint

Dive into the research topics of 'Understanding Gradient Descent on the Edge of Stability in Deep Learning'. Together they form a unique fingerprint.

Cite this