SGD at the Edge of Stability: Stochastic Stabilization with Large Learning Rates

2026-06-29

AI Summary

This work investigates the convergence behavior of stochastic gradient descent (SGD) operating near the edge of stability under large learning rates, addressing a gap in existing theory that lacks rigorous analysis in stochastic settings. Focusing on linear classifiers and two-layer neural networks trained with multiclass cross-entropy loss, the study integrates tools from stochastic optimization theory, dynamical systems analysis, and expected loss control to uncover an alternating mechanism between curvature-driven oscillations and stable descent. The paper establishes the first rigorous convergence guarantees for large-learning-rate SGD, demonstrating that its intrinsic stochasticity induces self-stabilization and enables optimal convergence rates within a fixed number of iterations. These theoretical findings are corroborated by experiments, highlighting the efficacy and superiority of large-step SGD.

Abstract

Modern deep learning has been shown to operate at the edge of stability, routinely using learning rates far larger than those justified by classical optimization theory. Most prior analyses of the edge of stability phenomenon focus on deterministic gradient descent, leaving the stochastic setting largely unexplored. In this work, we provide sharp convergence guarantees for Stochastic Gradient Descent (SGD) applied to the multiclass cross-entropy loss, for both linear classifiers and two-layer neural networks. We show that the stochasticity of SGD may cause the dynamics to alternate between an edge-of-stability regime that is dominated by curvature-driven oscillations, and a stable regime in which the expected loss decreases at a controlled rate. Despite that, we prove that SGD self-stabilizes the dynamics, ensuring that the iterates return to stability in a fixed number of iterations and allowing convergence in the best-iterate sense even with large learning rates. Experiments validate our theoretical findings and illustrate the benefits of SGD in the large-stepsize regime.

Problem

• edge of stability
• stochastic gradient descent
• large learning rates
• convergence
• stochastic optimization

Innovation

• Stochastic Gradient Descent
• Edge of Stability
• Large Learning Rates
• Self-stabilization
• Convergence Guarantee