SGD at the Edge of Stability: Stochastic Stabilization with Large Learning Rates
- URL: http://arxiv.org/abs/2606.30930v1
- Date: Mon, 29 Jun 2026 21:32:58 GMT
- Title: SGD at the Edge of Stability: Stochastic Stabilization with Large Learning Rates
- Abstract summary: Modern deep learning has been shown to operate at the edge of stability.<n>We provide convergence guarantees for Gradient Descent (SGD) applied to the multiclass cross-entropy loss.<n> Experiments validate our theoretical findings and illustrate the benefits of SGD in the large-stepsize regime.
- Score: 2.104588729923885
- License: http://creativecommons.org/licenses/by/4.0/
- 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.
Related papers
- The Road Taken: The Role of Optimizers at the Edge of Stability [53.81268610971847]
edge of stability refers to a phenomenon in deep learning where the Hessian eigenvalues of the loss remain stable above a threshold that the classical descent predicts to be unstable.<n>Previous works formulate the edge of stability with respect to the maximum Hessian eigenvalue and the learning rate.<n>We observe that many first-order methods, including gradient descent, significantly violate the stability bound predicted by these theories by a factor as large as $times 21.1$.
arXiv Detail & Related papers (2026-08-19T01:03:34Z) - Clipping Bottleneck: Stabilizing RLVR via Stochastic Recovery of Near-Boundary Signals [83.0127582612634]
Near-boundary Rescue (NSR) is a minimal, plug-and-play modification that retains out-of-bound tokens to recover lost signals.<n>NSR substantially improves training stability and delivers consistent gains over strong baselines such as DAPO and GSPO.
arXiv Detail & Related papers (2026-05-21T16:45:31Z) - Momentum Further Constrains Sharpness at the Edge of Stochastic Stability [0.5405981353784005]
Recent work suggests (stochastic) gradient descent self-organizes near an instability boundary, shaping both optimization and the solutions found.<n>We demonstrate that SGD with momentum exhibits an Edge of Stability (EoSS)-like regime with batch-size-dependent behavior.
arXiv Detail & Related papers (2026-04-15T17:28:46Z) - GTS: Inference-Time Scaling of Latent Reasoning with a Learnable Gaussian Thought Sampler [54.10960908347221]
We model latent thought exploration as conditional sampling from learnable densities and instantiate this idea as a Gaussian Thought Sampler (GTS)<n>GTS predicts context-dependent perturbation distributions over continuous reasoning states and is trained with GRPO-style policy optimization while keeping the backbone frozen.
arXiv Detail & Related papers (2026-02-15T09:57:47Z) - High dimensional analysis reveals conservative sharpening and a stochastic edge of stability [21.12433806766051]
We show that the dynamics of the large eigenvalues of the training loss Hessian have some remarkably robust features across models and in the full batch regime.<n>There is often an early period of progressive sharpening where the large eigenvalues increase, followed by stabilization at a predictable value known as the edge of stability.
arXiv Detail & Related papers (2024-04-30T04:54:15Z) - A Precise Characterization of SGD Stability Using Loss Surface Geometry [8.942671556572073]
Descent Gradient (SGD) stands as a cornerstone optimization algorithm with proven real-world empirical successes but relatively limited theoretical understanding.
Recent research has illuminated a key factor contributing to its practical efficacy: the implicit regularization it instigates.
arXiv Detail & Related papers (2024-01-22T19:46:30Z) - Investigating the Edge of Stability Phenomenon in Reinforcement Learning [20.631461205889487]
We explore the edge of stability phenomenon in reinforcement learning (RL)
Despite significant differences to supervised learning, the edge of stability phenomenon can be present in off-policy deep RL.
Our results suggest that, while neural network structure can lead to optimisation dynamics that transfer between problem domains, certain aspects of deep RL optimisation can differentiate it from domains such as supervised learning.
arXiv Detail & Related papers (2023-07-09T15:46:27Z) - Type-II Saddles and Probabilistic Stability of Stochastic Gradient Descent [16.849376037005452]
We show that saddle points in neural networks can be divided into two types, among which the Type-II saddles are especially difficult to escape from because the gradient noise vanishes at the saddle.
The dynamics of SGD around these saddles are thus to leading order described by a random matrix product process.
We leverage to show that saddle points can be either attractive or repulsive for SGD, and its dynamics can be classified into four different phases, depending on the signal-to-noise ratio in the gradient close to the saddle.
arXiv Detail & Related papers (2023-03-23T08:17:10Z) - NAG-GS: Semi-Implicit, Accelerated and Robust Stochastic Optimizer [45.47667026025716]
We propose a novel, robust and accelerated iteration that relies on two key elements.
The convergence and stability of the obtained method, referred to as NAG-GS, are first studied extensively.
We show that NAG-arity is competitive with state-the-art methods such as momentum SGD with weight decay and AdamW for the training of machine learning models.
arXiv Detail & Related papers (2022-09-29T16:54:53Z) - Stability and Generalization Analysis of Gradient Methods for Shallow
Neural Networks [59.142826407441106]
We study the generalization behavior of shallow neural networks (SNNs) by leveraging the concept of algorithmic stability.
We consider gradient descent (GD) and gradient descent (SGD) to train SNNs, for both of which we develop consistent excess bounds.
arXiv Detail & Related papers (2022-09-19T18:48:00Z) - Training Generative Adversarial Networks by Solving Ordinary
Differential Equations [54.23691425062034]
We study the continuous-time dynamics induced by GAN training.
From this perspective, we hypothesise that instabilities in training GANs arise from the integration error.
We experimentally verify that well-known ODE solvers (such as Runge-Kutta) can stabilise training.
arXiv Detail & Related papers (2020-10-28T15:23:49Z) - Fine-Grained Analysis of Stability and Generalization for Stochastic
Gradient Descent [55.85456985750134]
We introduce a new stability measure called on-average model stability, for which we develop novel bounds controlled by the risks of SGD iterates.
This yields generalization bounds depending on the behavior of the best model, and leads to the first-ever-known fast bounds in the low-noise setting.
To our best knowledge, this gives the firstever-known stability and generalization for SGD with even non-differentiable loss functions.
arXiv Detail & Related papers (2020-06-15T06:30:19Z)
This list is automatically generated from the titles and abstracts of the papers in this site.
This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.