SGD at the Edge of Stability: The Stochastic Sharpness Gap

2026-04-22

AI Summary

Existing theory struggles to explain why the sharpness in stochastic gradient descent (SGD) remains stably below $2/\eta$ and decreases as the batch size shrinks. This work proposes a stochastic self-stabilization theory that extends deterministic self-stabilization mechanisms to SGD, revealing that gradient noise enhances the cubic sharpness suppression effect, thereby lowering the sharpness equilibrium. By analyzing stochastic dynamics along the top Hessian eigenvector direction, modeling projected trajectories, and establishing a coupling theorem, the authors derive—for the first time—a closed-form expression for the sharpness gap: $\Delta S = \eta\beta\sigma_u^2/(4\alpha)$. This result quantitatively characterizes the relationship between batch size and solution flatness, showing that smaller batches yield flatter minima and naturally recover full-batch gradient descent behavior in the deterministic limit.

Technology Category

• Machine Learning: Learning Theory
• Reasoning under Uncertainty: Stochastic Optimization
• Search and Optimization: Non-convex Optimization

Application Category

• Security and Privacy: Large-scale security measurements
• Graph Algorithms and Modeling for the Web: Algorithms and analysis for incomplete, noisy, or partially observed Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

When training neural networks with full-batch gradient descent (GD) and step size $η$, the largest eigenvalue of the Hessian -- the sharpness $S(\boldsymbolθ)$ -- rises to $2/η$ and hovers there, a phenomenon termed the Edge of Stability (EoS). \citet{damian2023selfstab} showed that this behavior is explained by a self-stabilization mechanism driven by third-order structure of the loss, and that GD implicitly follows projected gradient descent (PGD) on the constraint $ S(\boldsymbolθ)\leq 2/η$. For mini-batch stochastic gradient descent (SGD), the sharpness stabilizes below $2/η$, with the gap widening as the batch size decreases; yet no theoretical explanation exists for this suppression. We introduce stochastic self-stabilization, extending the self-stabilization framework to SGD. Our key insight is that gradient noise injects variance into the oscillatory dynamics along the top Hessian eigenvector, strengthening the cubic sharpness-reducing force and shifting the equilibrium below $2/η$. Following the approach of \citet{damian2023selfstab}, we define stochastic predicted dynamics relative to a moving projected gradient descent trajectory and prove a stochastic coupling theorem that bounds the deviation of SGD from these predictions. We derive a closed-form equilibrium sharpness gap: $ΔS = ηβσ_{\boldsymbol{u}}^{2}/(4α)$, where $α$ is the progressive sharpening rate, $β$ is the self-stabilization strength, and $σ_{ \boldsymbol{u}}^{2}$ is the gradient noise variance projected onto the top eigenvector. This formula predicts that smaller batch sizes yield flatter solutions and recovers GD when the batch equals the full dataset.

Problem

• stochastic gradient descent
• sharpness
• Edge of Stability
• gradient noise
• Hessian eigenvalue

Innovation

• stochastic self-stabilization
• Edge of Stability
• sharpness gap
• gradient noise
• projected gradient descent