Edge of Stochastic Stability (EoSS)
- EoSS is a regime in stochastic optimization characterized by the stabilization of batch sharpness near 2/η due to the interplay of noise, learning rate, and batch size.
- It generalizes the deterministic Edge of Stability from full-batch gradient descent to stochastic settings, leading to distinctive two-phase dynamics including a progressive sharpening phase and a stabilization plateau.
- Empirical studies and theoretical analyses show that EoSS influences implicit regularization and generalization by imposing a batch-size dependent sharpness gap that favors flatter minima.
The Edge of Stochastic Stability (EoSS) is a fundamental regime in the optimization dynamics of stochastic algorithms, particularly in the training of deep neural networks using stochastic gradient descent (SGD) and its variants. EoSS generalizes the deterministic Edge of Stability (EoS) observed in full-batch gradient descent (GD) to the stochastic regime, where curvature and implicit regularization properties emerge from the interplay of step size, noise, and batch size. EoSS encompasses a new class of dynamic behaviors, sharpness diagnostics, and regularization effects that govern convergence and generalization across architectures and optimization algorithms.
1. Formal Foundations and Definitions
EoSS is typically defined as the regime where a stochastic optimizer self-organizes so that the average directional curvature encountered by the updates—the so-called Batch Sharpness—stabilizes near a threshold set by the optimizer’s learning rate and, in the case of momentum or zeroth-order strategies, other hyperparameters. For standard mini-batch SGD with learning rate , Batch Sharpness is defined as
where is the loss on mini-batch , and the expectation is over batch sampling. The EoSS regime is characterized by
after an initial progressive sharpening phase (Andreyev et al., 2024). The same concept applies to related diagnostics, including Batch Sharpness in the gradient direction and full-batch sharpness, with explicit stochastic mechanisms for the difference between them.
2. Classical Edge of Stability and Its Stochastic Generalization
In full-batch GD, the largest Hessian eigenvalue grows until it hovers near —the classical edge of stability. Beyond this threshold, updates become unstable. In stochastic settings, however, two key changes occur:
- The mean (full-batch) sharpness stabilizes strictly below , with the gap increasing as batch size decreases (Liao et al., 22 Apr 2026).
- Batch Sharpness, an average over per-batch Hessians or mini-batch gradients, rather than the global Hessian, stabilizes at or near —the operative diagnostic of the stochastic regime (Andreyev et al., 2024, Andreyev et al., 15 Apr 2026).
This fundamental difference is due to gradient and Hessian noise, which alter the curvature landscape encountered by stochastic updates and enforce a new, noise-dependent stability constraint.
3. Dynamics, Diagnostics, and Sharpness Plateaus
Stochastic algorithms in the EoSS regime display a two-phase dynamical structure (Emmanouilidis et al., 29 Jun 2026):
- Progressive Sharpening: The effective curvature (as measured by Batch Sharpness or MiniBS) grows rapidly early in training.
- EoSS Plateau: After sharpening, the stochastic optimizer self-stabilizes such that Batch Sharpness remains near 0, while the full-batch sharpness is suppressed by a batch-size dependent gap.
For SGD with momentum, two noise-dependent plateaus for Batch Sharpness arise (Andreyev et al., 15 Apr 2026):
| Regime | Batch Sharpness Plateau | Effective Curvature Constraint |
|---|---|---|
| Small batch | 1 | Amplified noise, favors flat minima |
| Large batch | 2 | Recovers classical heavy-ball stability |
Noise in the mini-batch gradients amplifies curvature constraints at small batches and relaxes them at large batch sizes or full-batch. The transition between plateaus is continuous and batch-size dependent.
4. Theoretical Mechanisms: Self-Stabilization and Sharpness Gap
The emergence of EoSS is explained by stochastic self-stabilization mechanisms (Liao et al., 22 Apr 2026). The interplay between gradient noise and the third-order structure of the loss introduces an enhanced cubic restoring force in the sharpness dynamics, yielding a suppressed equilibrium sharpness. Define the sharpness gap 3 as the difference between the deterministic and stochastic plateaus: 4 where 5 is the progressive sharpening rate, 6 is the self-stabilization strength, and 7 the variance of gradient noise projected along the top eigenvector (Liao et al., 22 Apr 2026). This formula predicts that 8, so smaller batch sizes yield flatter minima due to increased noise regularization.
The mathematical distinction between full-batch and stochastic stability thresholds is as follows (Agarwala et al., 2024, Song et al., 16 Apr 2026):
- Full-batch: Stability set by 9 of the Hessian.
- Stochastic: Mean-square stability and convergence are governed by either the average mini-batch curvature (Batch Sharpness) or, in high dimensions, alternative quantities such as the trace of the Neural Tangent Kernel (NTK).
5. Implicit Regularization and Practical Implications
EoSS induces a trace-based regularization effect unique to stochastic (and especially zeroth-order) methods (Song et al., 16 Apr 2026). Unlike first-order methods, where large step sizes induce implicit regularization on the largest eigenvalue, stochastic and zeroth-order algorithms bias training toward regions where the spectrum is flatter or the trace of the Hessian is minimized. This has several consequences:
- Smaller batch sizes and larger learning rates regularize toward flatter minima, aligning with empirical observations of improved generalization (Andreyev et al., 2024, Andreyev et al., 15 Apr 2026, Liao et al., 22 Apr 2026).
- In zeroth-order settings, stability is governed by the Hessian trace rather than the top eigenvalue, with large step sizes primarily regularizing 0 (Song et al., 16 Apr 2026).
- The suppression of full-batch sharpness from the 1 threshold explains the batch-size-dependent sharpness gap and provides a mechanistic account for the performance benefits of noisy stochastic training (Liao et al., 22 Apr 2026).
6. Empirical Evidence and Experimental Protocols
Systematic experimental findings confirm the EoSS phenomenology across feed-forward, convolutional, and vision transformer architectures (Andreyev et al., 2024, Andreyev et al., 15 Apr 2026, Song et al., 16 Apr 2026):
- Mini-batch SGD: Batch Sharpness rises during early training and plateaus at 2, while full-batch sharpness sits strictly below and converges slowly.
- Learning rate and batch size jumps: Interventions induce “catapult” effects—spikes in curvature or sharpness—followed by rapid self-stabilization around the adjusted threshold.
- Scaling of the sharpness gap: Measured 3 versus batch size gives approximate power-law scaling with exponents close to 4, matching theoretical predictions (Liao et al., 22 Apr 2026).
- Extensions to momentum and zeroth-order optimizers: Appropriate curvature diagnostics reveal similar stochastic plateaus and sensitivity to hyperparameters.
7. Extensions: Zeroth-Order Methods and High-Dimensional Regimes
For zeroth-order (ZO) optimization, the EoSS regime is defined through mean-square stability of the dynamics, analyzed via recursion over the entire Hessian spectrum (Song et al., 16 Apr 2026). The stability threshold 5 is characterized by
6
with practical upper and lower bounds given by the Hessian trace and largest eigenvalue. In high-dimensional models, analogous results have been established through the noise kernel norm and trace-of-NTK approximations, indicating that, in small-batch regimes, stability is controlled by the collective NTK trace rather than just the top eigenmode (Agarwala et al., 2024).
References
- (Andreyev et al., 2024) "Edge of Stochastic Stability: Revisiting the Edge of Stability for SGD."
- (Andreyev et al., 15 Apr 2026) "Momentum Further Constrains Sharpness at the Edge of Stochastic Stability."
- (Liao et al., 22 Apr 2026) "SGD at the Edge of Stability: The Stochastic Sharpness Gap."
- (Emmanouilidis et al., 29 Jun 2026) "SGD at the Edge of Stability: Stochastic Stabilization with Large Learning Rates."
- (Agarwala et al., 2024) "High dimensional analysis reveals conservative sharpening and a stochastic edge of stability."
- (Song et al., 16 Apr 2026) "Zeroth-Order Optimization at the Edge of Stability."