Momentum Further Constrains Sharpness at the Edge of Stochastic Stability

Arseniy Andreyev; Advikar Ananthkumar; Marc Walden; Tomaso Poggio; Pierfrancesco Beneventano

cs.LGApr 15, 2026

Momentum Further Constrains Sharpness at the Edge of Stochastic Stability

Authors: Arseniy Andreyev, Advikar Ananthkumar, Marc Walden, Tomaso Poggio, Pierfrancesco Beneventano

Organizations: Meta. · Work done at Princeton University, Princeton, NJ, USA. · Harvard University, Cambridge, MA, US · Massachusetts Institute of Technology, Cambridge, MA, US

Abstract

Recent work suggests that (stochastic) gradient descent self-organizes near an instability boundary, shaping both optimization and the solutions found. Momentum and mini-batch gradients are widely used in practical deep learning optimization, but it remains unclear whether they operate in a comparable regime of instability. We demonstrate that SGD with momentum exhibits an Edge of Stochastic Stability (EoSS)-like regime with batch-size-dependent behavior that cannot be explained by a single momentum-adjusted stability threshold. Batch Sharpness (the expected directional mini-batch curvature) stabilizes in two distinct regimes: at small batch sizes it converges to a lower plateau $2(1-β)/η$ , reflecting amplification of stochastic fluctuations by momentum and favoring flatter regions than vanilla SGD; at large batch sizes it converges to a higher plateau $2(1+β)/η$ , where momentum recovers its classical stabilizing effect and favors sharper regions consistent with full-batch dynamics. We further show that this aligns with linear stability thresholds and discuss the implications for hyperparameter tuning and coupling.

Figures & tables

Panel of Figure 1 — Figure 1: $\lambda_{\max}$ under full-batch GD with momentum (left) and mini-batch SGD with momentum (right). MLP on an 8k subset of CIFAR-10 for fixed step size $\eta=.004$ and varying $\beta$ . The stabilization level of Batch Sharpness (Definition 3.1 ) inverts its monotonicity in $\beta$ .

Figure 2: EoSS phenomenon using SGDM (left) and SGDN (right). MLPs on an 8k subset of CIFAR-10 under different step sizes $\eta$ and with $\beta=0.9$ . Batch Sharpness stabilizes around the $2(1-\beta)/\eta=1/(5\eta)$ threshold, shown by the dotted lines.

Panel of Figure 3 — Figure 3: Stabilization levels of Batch Sharpness and $\lambda_{\max}$ across varying batch sizes for an MLP trained with SGDM (top) and SGDN (bottom) at $\eta=0.005$ and $\beta=0.9$ . The critical batch size, defined heuristically as the threshold at which training dynamics enter the large-batch regime, is marked for each optimizer. Notably, SGDN reaches this regime at a batch size almost an order of magnitude smaller than SGDM.

Figure 4: Dynamics of curvature statistics for SGDM with $\beta=0.5$ . Top row: MLP; bottom row: CNN. Columns correspond to batch sizes $b\in\{4,64,256\}$ . Batch Sharpness and $\lambda_{\max}$ rise and then plateau, with larger batches yielding higher plateau levels. For Batch Sharpness, the left column is near the small-batch level $2(1-\beta)/\eta$ , the middle column lies in transition, and the right column approaches the large-batch level $2(1+\beta)/\eta$ .

Figure 5: Within-run dynamics for an MLP with batch size $b=4$ . The SGDM run uses learning rate $\eta=0.001$ with momentum $\beta=0.9$ , while the SGD run uses learning rate $\eta=0.01$ , chosen to match the effective step size.

Panel of Figure 6 — Figure 6: Within-run EoSS dynamics for an MLP under destabilizing interventions at step 75k with batch size $b=16$ , learning rate $\eta=0.004$ , and momentum $\beta=0.9$ . Left: destabilizing momentum intervention, increasing $\beta$ to $0.95$ . Middle: destabilizing learning rate intervention, increasing $\eta$ to $0.0067$ . Right: destabilizing batch size intervention, decreasing $b$ to $8$ . Top: Batch Sharpness and $\lambda_{\max}$ . Bottom: Training loss.

Appendix figures & tables23 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 7 — Figure 7: Within-run EoSS dynamics for early destabilizing interventions during the progressive sharpening phase at step 10k on an MLP with baseline learning rate $\eta=0.004$ , momentum $\beta=0.9$ , and batch size $b=16$ . Left: destabilizing momentum intervention, increasing $\beta$ to $0.95$ . Middle: destabilizing learning-rate intervention, increasing $\eta$ to $0.0067$ . Right: destabilizing batch-size intervention, decreasing batch size $b$ to $8$ . Top: Batch Sharpness and $\lambda_{\max}$ . Bottom: Training loss.

Panel of Figure 8 — Figure 8: Within-run EoSS dynamics for destabilizing interventions at an intermediate batch size at step 30k on an MLP with baseline learning rate $\eta=0.004$ , momentum $\beta=0.9$ , and batch size $b=512$ . Left: destabilizing momentum intervention, increasing $\beta$ to $0.95$ . Middle: destabilizing learning-rate intervention, increasing $\eta$ to $0.0067$ . Right: destabilizing batch-size intervention, decreasing batch size $b$ to $256$ . Top: Batch Sharpness and $\lambda_{\max}$ . Bottom: Training loss.

Panel of Figure 9 — Figure 9: Within-run EoSS dynamics for destabilizing interventions at high batch size at step 50k on an MLP with baseline learning rate $\eta=0.03$ , momentum $\beta=0.5$ , and batch size $b=16384$ . Left: destabilizing momentum intervention, increasing $\beta$ to $0.52$ . Middle: destabilizing learning-rate intervention, increasing $\eta$ to $0.035$ . Right: destabilizing batch-size intervention, decreasing $b$ to $12288$ . Top: Batch Sharpness and $\lambda_{\max}$ . Bottom: Training loss.

Panel of Figure 10 — Figure 10: Within-run EoSS dynamics for stabilizing interventions with low batch sizes at step 150k on an MLP with batch size $b=16$ , learning rate $\eta=0.004$ , and momentum $\beta=0.9$ . Left: stabilizing momentum intervention, decreasing $\beta$ to $0.875$ . Middle: stabilizing learning-rate intervention, decreasing $\eta$ to $0.003$ . Right: stabilizing batch-size intervention, increasing batch size $b$ to $32$ . Top: Batch Sharpness and $\lambda_{\max}$ . Bottom: Training loss.

Panel of Figure 11 — Figure 11: Within-run EoSS dynamics for early stabilizing interventions during the progressive sharpening phase at step 10k on an MLP with baseline learning rate $\eta=0.004$ , momentum $\beta=0.9$ , and batch size $b=16$ . Left: stabilizing momentum intervention, decreasing $\beta$ to $0.875$ . Middle: stabilizing learning-rate intervention, decreasing $\eta$ to $0.003$ . Right: stabilizing batch-size intervention, increasing batch size $b$ to $32$ . Top: Batch Sharpness and $\lambda_{\max}$ . Bottom: Training loss.

Panel of Figure 12 — Figure 12: Within-run EoSS dynamics for stabilizing interventions with intermediate batch sizes at step 75k on an MLP with baseline learning rate $\eta=0.004$ , momentum $\beta=0.9$ , and batch size $b=512$ . Left: stabilizing momentum intervention, decreasing $\beta$ to $0.875$ . Middle: stabilizing learning-rate intervention, decreasing $\eta$ to $0.003$ . Right: stabilizing batch-size intervention, increasing batch size $b$ to $768$ . Top: Batch Sharpness and $\lambda_{\max}$ . Bottom: Training loss.

Figure 13: Distance metrics comparing training trajectories of $\text{SGDM}(\eta,\beta,b)$ and $\text{SGD}(\frac{\eta}{1-\beta},b)$ for $\eta=0.001$ , $\beta=0.9$ , and $b=4$ . The hyperparameters are chosen as in Figure 5 , so that Batch Sharpness stabilizes at 200 for both runs. Each panel shows four curves: the distance of the SGD weights from initialization (orange), the distance of the SGDM weights from initialization (green), the distance between the two runs at matching steps (red), and the true distance between the two trajectories. The left plot uses $L_{2}$ distance in a weight space projected down to 5,000 dimensions. The right plot uses test prediction distance, measuring functional rather than weight-space divergence.

Panel of Figure 16 — Figure A1 : Within-run dynamics for an MLP trained with $\eta=0.01$ and momentum $\beta=0.5$ across batch sizes.

Panel of Figure 17 — Figure A2 : Within-run dynamics for MLP with $\eta=0.01$ and momentum $\beta=0.9$ across batch sizes.

Panel of Figure 18 — Figure A3 : Within-run dynamics for MLP with $\eta=0.005$ and momentum $\beta=0.5$ across batch sizes.

Panel of Figure 19 — Figure A4 : Within-run dynamics for CNN+ReLU with $\eta=0.02$ and $\beta=0.5$ .

Panel of Figure 20 — Figure A5 : Within-run dynamics for CNN+SiLU with $\eta=0.01$ and $\beta=0.5$ .

Panel of Figure 21 — Figure A6 : Within-run dynamics for CNN+SiLU with $\eta=0.01$ and $\beta=0.8$ .

Panel of Figure 22 — Figure A7 : Within-run dynamics for CNN+SiLU with $\eta=0.02$ and $\beta=0.5$ .

Panel of Figure 23 — Figure A8 : Within-run dynamics for CNN+SiLU with $\eta=0.02$ and $\beta=0.8$ .

Panel of Figure 24 — Figure A9 : Within-run dynamics for ResNet+ReLU with $\eta=0.02$ and $\beta=0.5$ .

Panel of Figure 25 — Figure A10 : Within-run dynamics for ReLU MLP ( $\eta=0.01$ , $\beta=0.5$ ) across batch sizes.

Panel of Figure 26 — Figure A11 : Within-run dynamics for SiLU MLP ( $\eta=0.01$ , $\beta=0.5$ ) across batch sizes.

Panel of Figure 27 — Figure A12 : Within-run dynamics for a ViT ( $\eta=0.007$ , $\beta=0.5$ ) across batch sizes using the MSE loss.

Panel of Figure 28 — Figure A13 : Within-run dynamics for a CNN ( $\eta=0.001$ , $\beta=0.9$ ) trained on the SST dataset across batch sizes using the MSE loss.

Panel of Figure 29 — Figure A14 : Within-run dynamics for an MLP ( $\eta=0.008$ , $\beta=0.75$ ) trained on the SST dataset across batch sizes using the MSE loss.

Stochastic Gradient Descent

Stochastic Gradient Descent (SGD) uses noisy, data-sampled gradients to train models and solve large-scale optimization problems.

See all 214 papers

SGD at the Edge of Stability: The Stochastic Sharpness Gap
Apr 22, 2026Fangshuo Liao, Afroditi Kolomvaki, Anastasios KyrillidisEdge of Stability Neural Network Optimization
Rice CS · Rice CS, Ken Kennedy Institute
SGD at the Edge of Stability: Stochastic Stabilization with Large Learning Rates
Jun 29, 2026Konstantinos Emmanouilidis, Lachlan MacDonald, Salma Tarmoun +1Edge of Stability Stochastic Optimization Convergence
Department of CIS1 · IDEAS4 · University of Pennsylvania +2
Dynamics of Stochastic Momentum with Sparse Updates in High Dimensions
May 27, 2026Katie Everett, Elliot PaquetteHigh-Dimensional Asymptotics Momentum Methods
Google DeepMind & MIT · McGill University & Mila

Apr 22, 2026J/K move · Enter open · S save

SGD at the Edge of Stability: The Stochastic Sharpness Gap

Fangshuo Liao, Afroditi Kolomvaki, Anastasios Kyrillidis

Rice CS · Rice CS, Ken Kennedy Institute

Edge of Stability Neural Network Optimization Neural Network Training Dynamics Gradient Descent Dynamics Stochastic Gradient Descent

Open paper AskPDF arXiv

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.

Abstract

Stochastic Gradient Descent

Explore similar work

SGD at the Edge of Stability: The Stochastic Sharpness Gap

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

Dynamics of Stochastic Momentum with Sparse Updates in High Dimensions