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

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

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−β)/η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+β)/η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

Figure 1: λmax⁡\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 η=.004\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 β=0.9\beta=0.9 . Batch Sharpness stabilizes around the 2(1−β)/η=1/(5η)2(1-\beta)/\eta=1/(5\eta) threshold, shown by the dotted lines.
Figure 3: Stabilization levels of Batch Sharpness and λmax⁡\lambda_{\max} across varying batch sizes for an MLP trained with SGDM (top) and SGDN (bottom) at η=0.005\eta=0.005 and β=0.9\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 β=0.5\beta=0.5 . Top row: MLP; bottom row: CNN. Columns correspond to batch sizes b∈{4,64,256}b\in\{4,64,256\} . Batch Sharpness and λmax⁡\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−β)/η2(1-\beta)/\eta , the middle column lies in transition, and the right column approaches the large-batch level 2(1+β)/η2(1+\beta)/\eta .
Figure 5: Within-run dynamics for an MLP with batch size b=4b=4 . The SGDM run uses learning rate η=0.001\eta=0.001 with momentum β=0.9\beta=0.9 , while the SGD run uses learning rate η=0.01\eta=0.01 , chosen to match the effective step size.
Figure 6: Within-run EoSS dynamics for an MLP under destabilizing interventions at step 75k with batch size b=16b=16 , learning rate η=0.004\eta=0.004 , and momentum β=0.9\beta=0.9 . Left: destabilizing momentum intervention, increasing β\beta to 0.950.95 . Middle: destabilizing learning rate intervention, increasing η\eta to 0.00670.0067 . Right: destabilizing batch size intervention, decreasing bb to 88 . Top: Batch Sharpness and λmax⁡\lambda_{\max} . Bottom: Training loss.
Appendix figures & tables23 assets

Supplementary material from the paper’s appendix.

Appendix

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 η=0.004\eta=0.004 , momentum β=0.9\beta=0.9 , and batch size b=16b=16 . Left: destabilizing momentum intervention, increasing β\beta to 0.950.95 . Middle: destabilizing learning-rate intervention, increasing η\eta to 0.00670.0067 . Right: destabilizing batch-size intervention, decreasing batch size bb to 88 . Top: Batch Sharpness and λmax⁡\lambda_{\max} . Bottom: Training loss.
Figure 8: Within-run EoSS dynamics for destabilizing interventions at an intermediate batch size at step 30k on an MLP with baseline learning rate η=0.004\eta=0.004 , momentum β=0.9\beta=0.9 , and batch size b=512b=512 . Left: destabilizing momentum intervention, increasing β\beta to 0.950.95 . Middle: destabilizing learning-rate intervention, increasing η\eta to 0.00670.0067 . Right: destabilizing batch-size intervention, decreasing batch size bb to 256256 . Top: Batch Sharpness and λmax⁡\lambda_{\max} . Bottom: Training loss.
Figure 9: Within-run EoSS dynamics for destabilizing interventions at high batch size at step 50k on an MLP with baseline learning rate η=0.03\eta=0.03 , momentum β=0.5\beta=0.5 , and batch size b=16384b=16384 . Left: destabilizing momentum intervention, increasing β\beta to 0.520.52 . Middle: destabilizing learning-rate intervention, increasing η\eta to 0.0350.035 . Right: destabilizing batch-size intervention, decreasing bb to 1228812288 . Top: Batch Sharpness and λmax⁡\lambda_{\max} . Bottom: Training loss.
Figure 10: Within-run EoSS dynamics for stabilizing interventions with low batch sizes at step 150k on an MLP with batch size b=16b=16 , learning rate η=0.004\eta=0.004 , and momentum β=0.9\beta=0.9 . Left: stabilizing momentum intervention, decreasing β\beta to 0.8750.875 . Middle: stabilizing learning-rate intervention, decreasing η\eta to 0.0030.003 . Right: stabilizing batch-size intervention, increasing batch size bb to 3232 . Top: Batch Sharpness and λmax⁡\lambda_{\max} . Bottom: Training loss.
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 η=0.004\eta=0.004 , momentum β=0.9\beta=0.9 , and batch size b=16b=16 . Left: stabilizing momentum intervention, decreasing β\beta to 0.8750.875 . Middle: stabilizing learning-rate intervention, decreasing η\eta to 0.0030.003 . Right: stabilizing batch-size intervention, increasing batch size bb to 3232 . Top: Batch Sharpness and λmax⁡\lambda_{\max} . Bottom: Training loss.
Figure 12: Within-run EoSS dynamics for stabilizing interventions with intermediate batch sizes at step 75k on an MLP with baseline learning rate η=0.004\eta=0.004 , momentum β=0.9\beta=0.9 , and batch size b=512b=512 . Left: stabilizing momentum intervention, decreasing β\beta to 0.8750.875 . Middle: stabilizing learning-rate intervention, decreasing η\eta to 0.0030.003 . Right: stabilizing batch-size intervention, increasing batch size bb to 768768 . Top: Batch Sharpness and λmax⁡\lambda_{\max} . Bottom: Training loss.
Figure 13: Distance metrics comparing training trajectories of SGDM(η,β,b)\text{SGDM}(\eta,\beta,b) and SGD(η1−β,b)\text{SGD}(\frac{\eta}{1-\beta},b) for η=0.001\eta=0.001 , β=0.9\beta=0.9 , and b=4b=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 L2L_{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.
Figure 14
Figure 15
Figure A1 : Within-run dynamics for an MLP trained with η=0.01\eta=0.01 and momentum β=0.5\beta=0.5 across batch sizes.
Figure A2 : Within-run dynamics for MLP with η=0.01\eta=0.01 and momentum β=0.9\beta=0.9 across batch sizes.
Figure A3 : Within-run dynamics for MLP with η=0.005\eta=0.005 and momentum β=0.5\beta=0.5 across batch sizes.
Figure A4 : Within-run dynamics for CNN+ReLU with η=0.02\eta=0.02 and β=0.5\beta=0.5 .
Figure A5 : Within-run dynamics for CNN+SiLU with η=0.01\eta=0.01 and β=0.5\beta=0.5 .
Figure A6 : Within-run dynamics for CNN+SiLU with η=0.01\eta=0.01 and β=0.8\beta=0.8 .
Figure A7 : Within-run dynamics for CNN+SiLU with η=0.02\eta=0.02 and β=0.5\beta=0.5 .
Figure A8 : Within-run dynamics for CNN+SiLU with η=0.02\eta=0.02 and β=0.8\beta=0.8 .
Figure A9 : Within-run dynamics for ResNet+ReLU with η=0.02\eta=0.02 and β=0.5\beta=0.5 .
Figure A10 : Within-run dynamics for ReLU MLP ( η=0.01\eta=0.01 , β=0.5\beta=0.5 ) across batch sizes.
Figure A11 : Within-run dynamics for SiLU MLP ( η=0.01\eta=0.01 , β=0.5\beta=0.5 ) across batch sizes.
Figure A12 : Within-run dynamics for a ViT ( η=0.007\eta=0.007 , β=0.5\beta=0.5 ) across batch sizes using the MSE loss.
Figure A13 : Within-run dynamics for a CNN ( η=0.001\eta=0.001 , β=0.9\beta=0.9 ) trained on the SST dataset across batch sizes using the MSE loss.
Figure A14 : Within-run dynamics for an MLP ( η=0.008\eta=0.008 , β=0.75\beta=0.75 ) trained on the SST dataset across batch sizes using the MSE loss.
More in this topic

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

Explore similar work

CardsList
Apr 22, 2026

SGD at the Edge of Stability: The Stochastic Sharpness Gap

When training neural networks with full-batch gradient descent (GD) and step size ηη, the largest eigenvalue of the Hessian -- the sharpness S(θ)S(\boldsymbolθ) -- rises to 2/η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(θ)≤2/η S(\boldsymbolθ)\leq 2/η. For mini-batch stochastic gradient descent (SGD), the sharpness stabilizes below 2/η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/η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=ηβσu2/(4α)ΔS = ηβσ_{\boldsymbol{u}}^{2}/(4α), where αα is the progressive sharpening rate, ββ is the self-stabilization strength, and σu2σ_{ \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.

Fangshuo Liao, Afroditi Kolomvaki, Anastasios Kyrillidis

Rice CS · Rice CS, Ken Kennedy Institute

Edge of StabilityNeural Network OptimizationNeural Network Training DynamicsGradient Descent DynamicsStochastic Gradient Descent
AskPDF
Jun 29, 2026

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

Modern deep learning has been shown to operate at the edge of stability, routinely using learning rates far larger than those justified by classical optimization theory. Most prior analyses of the edge of stability phenomenon focus on deterministic gradient descent, leaving the stochastic setting largely unexplored. In this work, we provide sharp convergence guarantees for Stochastic Gradient Descent (SGD) applied to the multiclass cross-entropy loss, for both linear classifiers and two-layer neural networks. We show that the stochasticity of SGD may cause the dynamics to alternate between an edge-of-stability regime that is dominated by curvature-driven oscillations, and a stable regime in which the expected loss decreases at a controlled rate. Despite that, we prove that SGD self-stabilizes the dynamics, ensuring that the iterates return to stability in a fixed number of iterations and allowing convergence in the best-iterate sense even with large learning rates. Experiments validate our theoretical findings and illustrate the benefits of SGD in the large-stepsize regime.

Konstantinos Emmanouilidis, Lachlan MacDonald, Salma Tarmoun +1

Department of CIS1 · IDEAS4 · University of Pennsylvania +2

Edge of StabilityStochastic Optimization ConvergenceNeural Network Training DynamicsStochastic Gradient Descent
AskPDF
May 27, 2026

Dynamics of Stochastic Momentum with Sparse Updates in High Dimensions

Existing theory of momentum assumes that gradients arrive at every parameter at a roughly constant rate, an assumption violated in practice by heavy-tailed data distributions and modern architectures. We theoretically analyze the dynamics of two tractable models of momentum under sparse updates: a least squares model with sparse inputs and a logistic regression model with a rare class. Both admit exact closed-form second-moment dynamics whose high-dimensional limits we characterize across three scaling exponents for sparsity, batch size, and momentum decay. The phase structure on both problems is governed by the ratio of two intrinsic timescales: a momentum retention timescale (how many active updates the buffer survives) and a learning timescale (how many active updates it takes to reduce the squared error). When learning is much slower than retention, the limit matches SGD; when learning is faster, the system is unstable; where the timescales coincide, we recover classical heavy-ball dynamics. The oscillatory dynamics occur at different momentum values for different token sparsity, creating a spectral conflict for global momentum across token frequencies.

Katie Everett, Elliot Paquette

Google DeepMind & MIT · McGill University & Mila

High-Dimensional AsymptoticsMomentum MethodsGradient Descent DynamicsStochastic Gradient Descent
AskPDF