SGD at the Edge of Stability: The Stochastic Sharpness Gap

arXiv:2604.21016v1 Announce Type: cross Abstract: When training neural networks with full-batch gradient descent (GD) and step size eta, the largest eigenvalue of the Hessian -- the sharpness S(oldsymbol{heta}) -- rises to 2/eta and hovers there, a phenomenon termed the Edge of Stability (EoS). itet{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(oldsymbol{heta})leq 2/eta. For mini-batch stochastic gradient descent (SGD), the sharpness stabilizes below 2/eta, 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/eta. Following the approach of itet{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: Delta S = eta eta sigma_{oldsymbol{u}}^{2}/(4alpha), where alpha is the progressive sharpening rate, eta is the self-stabilization strength, and sigma_{ oldsymbol{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.

Source: arXiv cs.AI | 2026-04-24

Related