📊 statistics

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

This paper establishes sharp convergence guarantees for Stochastic Gradient Descent (SGD) with large learning rates on multiclass cross-entropy loss, demonstrating that inherent stochasticity enables the algorithm to self-stabilize and alternate between edge-of-stability oscillations and controlled descent to ensure convergence.

Original authors: Konstantinos Emmanouilidis, Lachlan MacDonald, Salma Tarmoun, Rene Vidal

Published 2026-07-01

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Konstantinos Emmanouilidis, Lachlan MacDonald, Salma Tarmoun, Rene Vidal

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to roll a heavy boulder down a very bumpy, winding hill to get it to the bottom (the "perfect" solution). In the world of machine learning, this boulder is your AI model, the hill is the "loss landscape" (a map of how wrong the model is), and the bottom is the point where the model makes the fewest mistakes.

For a long time, scientists believed you had to push the boulder gently and carefully. If you pushed too hard (used a "large learning rate"), the boulder would bounce off the walls, fly over the edge, or get stuck in a loop, never reaching the bottom. This was the "classical theory."

However, in modern AI, engineers noticed something strange: if they pushed the boulder very hard, it actually got to the bottom faster and sometimes even found a better spot than the gentle pushes. This phenomenon is called the "Edge of Stability." It's like driving a car at the very limit of its traction; it wobbles and skids, but it still moves forward.

The problem is that most of the math explaining this "wobbly driving" was written for a deterministic world (where you know exactly what the road looks like). But in real life, AI uses Stochastic Gradient Descent (SGD). This means the AI doesn't see the whole road at once; it only sees a tiny, random patch of the road (a "mini-batch") before deciding which way to push. This adds noise (randomness) to the mix.

This paper asks: How does this random noise affect the "wobbly driving" at the edge of stability? Does the AI crash, or does it find a way to stabilize itself?

Here is what the authors discovered, explained through simple analogies:

1. The Two Modes of Driving

The authors found that when you use a large learning rate with random noise, the AI's journey isn't just one chaotic mess. It actually switches between two distinct modes:

The "Skidding" Phase (Edge of Stability): Imagine the boulder is being pushed so hard that it hits a bump and starts bouncing wildly. The error (how far it is from the bottom) goes up and down. It looks unstable. In this phase, the AI is dominated by the "curvature" of the hill (the bumps).
The "Gliding" Phase (Stable Regime): Eventually, the AI settles into a rhythm. Even though it's still being pushed hard, it enters a zone where, on average, it is moving closer to the bottom. The error decreases steadily.

2. The Magic of "Self-Stabilization"

The most surprising finding is how the AI handles the randomness.

In a perfect, noise-free world, if you push too hard, you might overshoot and never come back. But in the real world with random noise, the authors proved that SGD has a built-in self-correcting mechanism.

The Analogy: Imagine you are walking on a tightrope in a windy storm (the noise). Sometimes the wind blows you off the rope (the AI leaves the stable zone).
The Discovery: The authors show that even if the wind blows you off, the shape of the tightrope (the math of the cross-entropy loss) acts like a magnet. It pulls you back. You might step off the rope a few times, but you will always be pulled back onto the rope within a specific, predictable number of steps.
The Result: The AI doesn't crash. It wobbles, falls off the "stable zone" briefly due to random noise, but then immediately corrects itself and returns to a safe, stable path. This happens over and over, but as the AI gets better (the loss gets lower), it spends more and more time on the stable path and less time falling off.

3. It Works for Both Simple and Complex Models

The authors didn't just look at simple straight lines (linear classifiers). They also looked at two-layer neural networks (a basic type of deep learning model).

They proved that even with the added complexity of these networks and the random noise, the same rules apply.
They identified specific "activation functions" (the mathematical switches inside the AI) that allow this self-stabilization to happen. They showed that common tools used in modern AI (like GELU or Softplus) fit these rules perfectly.

4. The Proof is in the Pudding (Experiments)

To make sure their math wasn't just theory, they ran experiments on real data (like the MNIST handwritten digits and CIFAR-10 images).

They trained AI models with huge learning rates (much larger than what old textbooks say is safe).
The Result: The models didn't explode. Instead, the loss (error) dropped quickly. The models trained faster and achieved high accuracy, confirming that "wobbly driving" at the edge of stability is not only safe but often superior.

Summary

This paper explains why modern AI works so well even when we use "reckless" learning rates. It turns out that the randomness (noise) in the training process isn't just a nuisance; it interacts with the shape of the problem in a way that creates a self-stabilizing loop.

The AI might stumble and wobble at the edge of stability, but it has an invisible safety net that catches it and pulls it back to a stable path, allowing it to converge to a great solution faster than if it had been driven carefully and slowly.

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

Problem Statement

Modern deep learning routinely employs learning rates significantly larger than those justified by classical smooth convex optimization theory, a phenomenon known as the Edge of Stability (EoS). While the EoS behavior of deterministic Gradient Descent (GD) has been analyzed—characterized by an initial phase of curvature-driven oscillations followed by a self-stabilizing phase where the loss decreases monotonically—the stochastic setting remains largely unexplored.

In Stochastic Gradient Descent (SGD), the distinction between harmless stochastic fluctuations (due to mini-batch noise) and genuine instability (due to large stepsizes relative to curvature) is blurred. The non-monotonicity of the loss in SGD can arise from either source, rendering classical deterministic stability thresholds (based on the Hessian's largest eigenvalue) ill-defined or insufficient. Consequently, a principled understanding of how SGD behaves under arbitrarily large learning rates, particularly for the multi-class cross-entropy loss, is missing. The core question addressed is: What is the appropriate notion of stability for SGD with cross-entropy loss, and can we characterize its dynamics beyond the classical small learning rate regime?

Methodology

The authors develop a theoretical framework for analyzing SGD applied to multi-class cross-entropy loss on linear classifiers and two-layer neural networks under arbitrarily large step sizes ( $\eta > 0$ ).

Stochastic Lyapunov Stability:
Instead of relying on pointwise monotonicity of the realized loss (which is impossible to guarantee in the presence of noise), the paper adopts a stochastic Lyapunov stability viewpoint. The loss function $L(W)$ is treated as a Lyapunov function. Stability is defined not by $L(W_{t+1}) \leq L(W_t)$ , but by the condition that the expected loss decreases in conditional expectation given the filtration $\mathcal{F}_t$ :
$\mathbb{E}[L(W_{t+1}) - L(W_t) \mid \mathcal{F}_t] \leq 0$
This criterion captures the balance between curvature and noise. It reduces to the deterministic stability condition in the full-batch limit but accounts for mini-batch fluctuations in the stochastic setting.
Two-Phase Dynamics Analysis:
The analysis decouples the training dynamics into two regimes:
- Edge-of-Stability (EoS) Regime: The initial phase where iterates may experience curvature-driven oscillations, and the loss is not monotonically decreasing.
- Stable Regime: A region where the expected loss admits a one-step descent. The authors define a stable set $\mathcal{S}$ where $L(W_t) \leq \tilde{L}$ , with $\tilde{L}$ depending on the stepsize, batch size, and data margin.
Stochastic Stabilization Mechanism:
A key theoretical contribution is the proof that SGD possesses a self-stabilizing property. Even if mini-batch noise temporarily pushes iterates outside the stable set $\mathcal{S}$ , the dynamics are guaranteed to return to the stable regime within a fixed, bounded number of iterations with high probability. This return mechanism is driven by the geometry of the cross-entropy loss, where both curvature and noise magnitude decrease as the loss approaches zero.
Extension to Neural Networks:
The framework is extended to two-layer neural networks in mean-field scaling. The authors impose mild conditions on activation functions (differentiability, smoothness, and near-homogeneity) and introduce an "effective margin" $\tilde{\gamma}$ that accounts for the activation derivative's lower bound. This allows the transfer of linear convergence guarantees to the non-linear setting.

Key Contributions

Stochastic Stability Criterion: The paper proposes a novel stability definition for SGD based on conditional Lyapunov drift, which is well-suited for separable cross-entropy minimization where parameters may diverge while the loss converges to zero.
Convergence Guarantees for Large Stepsizes: The authors provide the first theoretical convergence guarantees for SGD on multi-class cross-entropy loss with arbitrarily large step sizes ( $\eta > 0$ ). They establish that the best iterate achieves a $\tilde{O}(1/t)$ convergence rate, regardless of the stepsize magnitude.
Characterization of the Stable Regime: The work proves that SGD enters the stable regime within a fixed number of iterations with high probability. Once inside, the last iterate converges at a sublinear rate.
Stochastic Stabilization Proof: The paper rigorously proves that SGD "self-stabilizes." Even if the trajectory exits the stable set due to noise, it returns to stability in a bounded number of steps. This contrasts with deterministic GD, where once the stable phase is entered, the trajectory remains there by construction.
Neural Network Analysis: The results are extended to two-layer neural networks, showing that the large-stepsize convergence guarantees hold despite model non-linearity, provided the activation functions satisfy specific "leaky" conditions.

Results

Theoretical Bounds: Theorem 2.2 establishes an upper bound on the expected loss of the best iterate for linear classifiers, showing a $\tilde{O}(1/t)$ rate. Theorem 4.1 and Theorem 5.4 provide similar guarantees for the last iterate once the dynamics enter the stable regime.
Stability Return: Theorem 4.2 (for linear models) and Theorem 5.5 (for neural networks) quantify the probability of remaining in the stable set and the maximum number of steps required to return to it if an exit occurs. These bounds depend on the batch size, stepsize, and current loss value.
Experimental Validation: Experiments on synthetic data, MNIST (two-layer NN), and CIFAR-10 (8-layer NN) validate the theoretical findings. The results demonstrate:
- Non-monotonic loss decrease at the start (EoS phase) followed by convergence in the stable regime.
- Faster convergence of the loss with larger step sizes.
- Comparable or improved test accuracy with large stepsizes compared to smaller ones.
- Reduced noise and oscillation with larger batch sizes.

Significance and Claims

The paper claims to fill a critical gap in the theoretical understanding of modern deep learning optimization. While prior work has explained the EoS phenomenon for deterministic GD, this work extends the theory to the stochastic setting, which is the standard for practical deep learning.

The authors emphasize that their results provide a principled understanding of how stochasticity alters the two-phase structure of optimization. They argue that the "self-stabilization" mechanism is intrinsic to SGD with cross-entropy loss, ensuring that the algorithm does not diverge even with large learning rates, provided the loss decreases over time. The work does not claim to solve all optimization challenges in deep learning but offers a rigorous framework for analyzing SGD in the large-stepsize regime, validating empirical observations that large learning rates are effective and stable in practice. The theoretical guarantees are specific to separable data and cross-entropy loss, and the analysis for neural networks is restricted to two-layer architectures under mean-field scaling.