Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
Arseniy Andreyev * 1 2 Advikar Ananthkumar * 3 Marc Walden * 3 Tomaso Poggio 4 Pierfrancesco Beneventano 4
arXiv:2604.14108v1 [cs.LG] 15 Apr 2026
Abstract
instead, Andreyev & Beneventano (2024) propose the Edge of Stochastic Stability (E O SS) and identify a directional mini-batch curvature statistic (Batch Sharpness, Definition 3.1) that saturates at 2/η for vanilla SGD.
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.
Where does momentum live relative to the stochastic edge? Momentum (Polyak heavy-ball, SGDM) and Nesterov acceleration (SGDN) are standard in deep learning and often essential for fast, stable training. Yet the “edge” picture is incomplete for momentum methods with minibatch gradients: even in deterministic quadratics, momentum and Nesterov have different stability regions, and in the stochastic regime the relevant instability certificate is not obvious. For full-batch GD with momentum, Cohen et al. (2021) established cleanly that the stability threshold increases monotonically in β, namely as 2+2β η . We observe a consistent but surprising phenomenon in the mini-batch case. In small-batch settings, this trend reverses: SGDM stabilization decreases monotonically in β. Moreover, SGD with momentum hovers in regions of extremely small curvature, where common “effective learning rate” heuristics are insufficient to predict which curvature levels or solutions training will select; see Figure 1. Following these observations, this paper asks:
1. Introduction
Question 1: Does SGD with momentum or Nesterov acceleration self-organize at an instability boundary?
Optimization at the edge of stability. A growing body of evidence suggests that modern deep-network training with constant (or piecewise-constant) step size operates in a regime of controlled instability. In full-batch training, the top Hessian eigenvalue λmax often sharpens until it hovers near a deterministic stability boundary (the Edge of Stability, E O S) (Xing et al., 2018; Jastrz˛ebski et al., 2019; 2020; Cohen et al., 2021; 2024). In mini-batch training, the full-batch sharpness λmax can fail to diagnose stability;
Moreover, a central question is Question 2: If so, what boundary is actually being saturated?
Main empirical finding: momentum splits E O SS into two batch regimes. Across architectures and hyperparameters (Appendix H), the Batch Sharpness statistic progressively sharpens and then plateaus, but the plateau level depends sharply on batch size:
1
Meta. 2 Work done at Princeton University, Princeton, NJ, USA. 3 Harvard University, Cambridge, MA, US 4 Massachusetts Institute of Technology, Cambridge, MA, US.. Correspondence to: Arseniy Andreyev <[email protected]>, Pierfrancesco Beneventano <[email protected]>. Code available on github at https://github.com/arseniqum/edge-ofstochastic-stability.
BSplateau ≈ 1
2(1 − β) η
(1)
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
Figure 1. λ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 and varying β. The stabilization level of Batch Sharpness (Definition 3.1) inverts its monotonicity in β.
Figure 2. EoSS phenomenon using SGDM (left) and SGDN (right). MLPs on an 8k subset of CIFAR-10 under different step sizes η and with β = 0.9. Batch Sharpness stabilizes around the 2(1 − β)/η = 1/(5η) threshold, shown by the dotted lines.
in the small-batch (noise-dominated) regime and ( 2(1+β) (SGDM), η BSplateau ≈ 2(1+β) (SGDN) η(1+2β)
Contributions.
More precisely, we establish:
• Batch-size–dependent E O SS under momentum. We show that SGDM and SGDN exhibit E O SS-like selforganization, but the saturated curvature level depends fundamentally on batch size and cannot be explained by a single “momentum-corrected” threshold. Empirically, we identify a noise-dominated small-batch plateau at 2(1 − β)/η and a large-batch plateau approaching the deterministic momentum stability thresholds identified by Cohen et al. (2021) (with different large-batch limits for SGDM and SGDN), together with a broad intermediate regime interpolating between these two limits (Section 3.1). • Qualitative flip: small-batch momentum biases toward flatter regions. The small-batch plateau at 2(1 − β)/η implies that, in sufficiently small batches, momentum does not relax the operative curvature constraint relative to vanilla SGD; it tightens it, biasing training toward flatter regions. This contrasts with the large-batch regime, where momentum recovers the classical stabilizing picture and permits sharper curvature (Section 3.1). • Intervention evidence for an E O SS-like regime of instability. We show that destabilizing interventions produce catapults precisely when they push Batch Sharpness above its operating plateau, while stabilizing inter-
(2)
in the large-batch (deterministic) regime. 1 The small-batch plateau is strictly lower than the vanilla-SGD threshold 2/η and therefore reveals a qualitative “flip”: with small batches, momentum enforces stricter curvature constraints and biases training toward flatter regions. Interventions certify an instability-adjacent regime. To distinguish “mere plateaus” from genuine stability constraints, we use checkpoint interventions (Andreyev & Beneventano, 2024): small destabilizing changes to hyperparameters (e.g. η ↑, b ↓, or β ↑ in the small-batch regime) trigger catapult dynamics, followed by re-stabilization near the new plateau. This provides operational evidence that SGDM (and analogously SGDN) trains near an active stochastic stability boundary. 1 Here and below, statements that momentum permits higher large-batch curvature than vanilla SGD apply literally to SGDM; for SGDN, the deterministic plateau still exceeds the small-batch value 2(1 − β)/η, but remains below the vanilla-SGD threshold 2/η. This distinction does not affect the qualitative regime change: small batches tighten the operative curvature constraint, whereas large batches recover the optimizer-specific deterministic threshold.
2
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
ventions reopen progressive sharpening. These intervention responses provide operational evidence that SGDM and SGDN train in an E O SS-like regime of instability, with Batch Sharpness as the operative diagnostic, rather than merely exhibiting descriptive curvature plateaus (Section 3.3). • Mechanism via mean-square stability. We provide a linear mean-square stability analysis showing that in the noise-dominated regime the instability threshold is governed by the effective step size ηeff = η/(1 − β). This explains the small-batch threshold 2(1−β)/η and is consistent with the observed interpolation across batch sizes. Importantly, this reduction concerns stability rather than full trajectory equivalence.2 (Section 4).
2. Preliminaries and Related Work 2.1. Notation and Optimizers We analyze mini-batch SGD with Polyak momentum or Nesterov acceleration on relatively simple vision classification tasks trained with MSE (see Section 6 for limitations). Precisely, let θt ∈ Rd denote the model parameters at iteration t, let D = {xi }ni=1 be the dataset, and ℓ(θ; xi ) be the loss on a sample (xi ). We define the empirical risk
Figure 3. Stabilization levels of Batch Sharpness and λmax across varying batch sizes for an MLP trained with SGDM (top) and SGDN (bottom) at η = 0.005 and β = 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.
n
L(θ) =
1X ℓ(θ; xi ). n i=1
(3)
2.2. The Value of Momentum
At each iteration, a mini-batch Bt of size b is sampled uniformly at random, and the stochastic gradient is 1 X gt = ∇θ ℓ(θt ; xi ). (4) b
The added value of momentum. Polyak heavy-ball momentum and Nesterov acceleration are ubiquitous in modern deep learning—often as explicit buffers (SGDM/SGDN) or implicitly inside adaptive methods—and are frequently key to fast and stable training in practice (Krizhevsky et al., 2012; Sutskever et al., 2013; Gitman et al., 2019; Fu et al., 2023). A large body of work has proposed complementary explanations for why momentum helps (see Appendix A for further related work): (i) stability enlargement / effective step-size rescaling, where momentum can enlarge the usable learning-rate range, and in some regimes SGDM can be related to SGD after matching an effective step size (Fu et al., 2023; Wang et al., 2024; Gitman et al., 2019); (ii) temporal filtering and noise-shaping: the momentum buffer is an exponential moving average of stochastic gradients, motivating stationary-distribution, SDE, and modified-equation analyses (Mandt et al., 2017; Li et al., 2019; Gitman et al., 2019); (iii) inertial/underdamped geometry, where momentum is viewed as a discretization of a second-order flow and can therefore produce different transient dynamics than vanilla gradient descent (Su et al., 2014; Wibisono et al., 2016; Shi et al., 2022; Wilson et al., 2021); (iv) solution selection and generalization, where momentum can preserve or change implicit bias depending on regime and can affect generalization through implicit-regularization or stability
i∈Bt
We use the heavy-ball (HB) momentum formulation standard in deep learning libraries, rather than an EMA-style update. The two algorithms considered are: • SGD with Polyak Momentum (HB or SGDM): vt+1 = βvt + gt , θt+1 = θt − ηvt+1 ,
(5)
with momentum β ∈ [0, 1), learning rate η > 0, and v0 = 0. • SGD with Nesterov Acceleration (NAG or SGDN): vt+1 = βvt + gt (θt − βηvt ) ,
(6) θt+1 = θt − ηvt+1 . P 1 Let LB (θ) = |B| i∈B ℓ(θ; xi ) be the mini-batch loss for a batch B ⊆ D of size b drawn from the mini-batch sampling distribution Pb . Define the mini-batch gradient gB (θ) = ∇LB (θ) and mini-batch Hessian HB (θ) = ∇2 LB (θ). 2
Empirically, even when ηeff is matched so that curvature statistics stabilize similarly, SGD and SGDM remain separated in parameter/function space; see Appendix F.
3
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
η ↑, b ↓, or β ↑). If training is near the active stability boundary, the perturbed run exhibits a catapult (transient excursion + loss spike), whereas stabilizing perturbations do not. Note that this mechanism complements that of Cohen et al. (2021), who perturb the step size in the opposite direction to witness recovered stability and progressive sharpening. We use both mechanisms to establish that both SGDM and SGDN train at the Edge of Stochastic Stability.
mechanisms (Wang et al., 2022; Jelassi & Li, 2022; Ghosh et al., 2023; Ramezani-Kebrya et al., 2024; Lyu, 2025). These perspectives motivate a wider central question: What are the effects of momentum and acceleration on the training dynamics? What is known (and what remains unclear). Recent work has also clarified that the effect of momentum is strongly regime dependent: (1) small learning-rate regimes can make momentum nearly redundant after matching effective learning rates, with SGD and SGDM often following closely tracking trajectories, or at least closely matching trajectory statistics, and showing limited additional gains (Fu et al., 2023; Wang et al., 2024); (2) deterministic large learning-rate regimes highlight momentum’s stabilizing role, where it can substantially enlarge the range of usable learning rates and becomes most helpful near (or beyond) an instability boundary (Cohen et al., 2021); (3) stochastic regimes admit diffusion/modified-equation limits in which momentum changes how noise is accumulated over time and thus can influence exploration, stationary behavior, and sometimes implicit bias/generalization (Liu et al., 2018; Li et al., 2019; Jelassi & Li, 2022; Ramezani-Kebrya et al., 2024). At the same time, much of the theoretical momentum literature either (a) analyzes convergence/implicit bias under assumptions consistent with a stable descent-type regime, or (b) studies continuous-time limits that abstract away the batch-dependent curvature actually seen by minibatch methods.
We further test whether instability criteria (Definition A.1) known to govern optimizers without momentum or acceleration saturate here as well, and if so, at which level. On top of λmax , we track Definition 3.1 (Batch Sharpness). Assume the batches are drawn from the mini-batch sampling distribution Pb . The Batch Sharpness at θ is gB (θ)⊤ HB (θ) gB (θ) BS(θ) := EB∼Pb . (7) ∥gB (θ)∥22 Andreyev & Beneventano (2024) showed that BS > (2 + ε)/η is a sharp sufficient instability certificate for SGD, making BS the natural mini-batch analogue of the E O S curvature threshold. Further discussion appears in Appendix A. Section 2 introduced the instability-centric viewpoint and the empirical diagnostics that we use throughout (including Batch Sharpness). Here we focus on a narrower question: how do these dynamics change once momentum is introduced? Concretely, we study SGD with Polyak (heavyball) momentum (SGDM) and with Nesterov acceleration (SGDN), as defined in Section 2.1.
E O S vs. mini-batch: why the full-batch λmax signal breaks. In full-batch GD, training typically self-organizes near the quadratic stability boundary λmax (∇2 L(θt )) ≈ 2/η and enters the oscillatory “central-flow” regime (Xing et al., 2018; Jastrz˛ebski et al., 2019; 2020; Cohen et al., 2021; 2024). For deterministic heavy-ball momentum (β), the analogous linear boundary on quadratics is λmax ≈ 2(1 + β)/η. In mini-batch training, however, this full-batch λmax diagnostic does not carry over cleanly: λmax (∇2 L(θt )) can plateau far below 2/η (and may not stabilize), while loss oscillations are ubiquitous and by themselves do not diagnose instability (Cohen et al., 2021; Andreyev & Beneventano, 2024).
The stability thresholds for full-batch gradient descent with momentum have been characterized in the classical optimization literature. For Heavy-Ball momentum, the stability boundary is λmax < 2(1 + β)/η (Polyak, 1964), while 2(1+β) for Nesterov acceleration it is λmax < η(1+2β) (Nesterov, 1988). Cohen et al. (2021) demonstrated that when using full-batch GD with momentum, λmax indeed reaches and remains near these thresholds during training. However, it remains unclear whether an analogous regime of instability governs training when momentum is combined with mini-batch stochastic gradients. In this section, we establish that SGD with momentum (SGDM) and Nesterov acceleration (SGDN) train at the Edge of Stochastic Stability, and we characterize how this regime depends fundamentally on batch size.
3. SGDM and SGDN Typically Operate at the Edge of Stochastic Stability
3.1. Batch-size–dependent curvature plateau under momentum
The mechanism: instability criteria + perturbations. Following Andreyev & Beneventano (2024), we treat an “edge” as saturation of a computable one-sided instability certificate for the local (quadratic) dynamics, and we test for this regime via checkpoint perturbations: restart from a checkpoint θt and apply a small destabilizing change (e.g.
We begin by examining within-run dynamics of Batch Sharpness for SGDM and SGDN. Across all settings we tested (architectures, activations, and hyperparameter sweeps; see Appendix H), Batch Sharpness and λmax increase during 4
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
Figure 4. Dynamics of curvature statistics for SGDM with β = 0.5. Top row: MLP; bottom row: CNN. Columns correspond to batch sizes b ∈ {4, 64, 256}. Batch Sharpness and λ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 − β)/η, the middle column lies in transition, and the right column approaches the large-batch level 2(1 + β)/η.
the early stage of training and then plateau. The primary difference from vanilla SGD is that the plateau level depends sharply on the batch size.
Between these extremes, there is a broad transition region in which the plateau interpolates between the small-batch value and the large-batch value; see the middle column of Figure 4. Notably, much of practical CIFAR-10 training lies specifically in this intermediate regime (He et al., 2015; Zagoruyko & Komodakis, 2017; Masters & Luschi, 2018). The trend is monotone in b: larger batches yield systematically higher plateau levels (Figure 3). A second, more qualitative observation is that SGDM tends to transition later than SGDN: for fixed (η, β), SGDN often reaches its large-batch (deterministic) plateau at smaller b, i.e., it tends to have a smaller critical batch size (Figure 3).
Two regimes with a transition. Empirically, two plateau regimes are consistently observed: • Small-batch regime. For sufficiently small3 b, the Batch Sharpness plateau is approximately BSplateau ≈
2(1 − β) , η
(8)
for both SGDM and SGDN in our experiments (see Figure 4). Qualitatively, this means that the effect of momentum in small-batch SGD is opposite to its effect in full-batch GD: momentum now leads to more restrictive curvature levels.
Batch Sharpness as an indicator, not a certificate. An important distinction from the vanilla SGD case must be emphasized. Andreyev & Beneventano (2024) established that for vanilla SGD, Batch Sharpness serves as an instability criterion—crossing 2/η is sufficient to guarantee divergence on the quadratic approximation. For SGDM/SGDN, we do not have an analogous theoretical result: Batch Sharpness does not necessarily govern the stability of the momentum dynamics in the same direct sense, as we further discuss in Section 4. Instead, Batch Sharpness here functions as an empirical indicator of a particular dynamical regime. Still, as discussed in Section 2, the precise way to establish whether the dynamics are in a regime of instability is through perturbation experiments, which we present in Section 3.3.
• Large-batch regime. For large batch sizes we recover the full-batch behavior: Batch Sharpness and 2(1+β) λmax stabilize4 at 2(1+β) for SGDM and at η(1+2β) η for Nesterov’s Accelerated Gradient (NAG), see Figure 3 and the right column of Figure 4. In this regime, momentum plays its classical role of allowing training in regions of higher curvature than its non-momentum counterparts. 3
This is dataset-size-dependent, but for the 8k subset of CIFAR10, “small” in this context means b ≲ 16. 4 Notice that Batch Sharpness sometimes stabilizes slightly below that threshold, consistent with (Andreyev & Beneventano, 2024); this can be explained by the fact that the full-batch gradient contains the self-stabilizing component of (Damian et al., 2023) in addition to the highest-eigenvector component.
3.2. Consequences for λmax . Although we avoid interpreting Batch Sharpness as the stability quantity for momentum, it is still informative to track 5
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
to hyperparameters should trigger characteristic catapults, i.e., abrupt loss spikes followed by restabilization. We now show that SGDM and SGDN exhibit precisely this behavior. Destabilizing perturbations trigger catapults. We consider three classes of mid-training interventions that lower the effective stability threshold: • Increasing step size η → η ′ > η: this directly reduces the threshold from 2(1 − β)/η (small-batch) or 2(1 + β)/η (large-batch) to the corresponding value at η ′ . • Increasing momentum β → β ′ > β: in the smallbatch regime, this tightens the constraint from 2(1 − β)/η to 2(1 − β ′ )/η; in the large-batch regime, the effect is reversed. • Decreasing batch size b → b′ < b: this increases the value of Batch Sharpness, thus putting it above its stabilization threshold.
Figure 5. Within-run dynamics for an MLP with batch size b = 4. The SGDM run uses learning rate η = 0.001 with momentum β = 0.9, while the SGD run uses learning rate η = 0.01, chosen to match the effective step size.
how full-batch sharpness behaves alongside it. Empirically, as in the case of vanilla SGD, stabilization of Batch Sharpness induces a corresponding stabilization of the full-batch top eigenvalue λmax ; see Figure 4 and Appendix H. Because Batch Sharpness stabilizes at 2(1 − β)/η in the small-batch regime, which is strictly lower than the vanilla SGD threshold of 2/η, the full-batch eigenvalue λmax is suppressed to even lower values than in vanilla SGD. This implies that momentum with small batches biases training toward flatter regions than either vanilla SGD or momentum with large batches.
In each case, when the intervention causes the new threshold to fall below the current value of Batch Sharpness, we observe a catapult: a sharp spike in the training loss accompanied by a transient excursion in the curvature statistics (Figure 6). After the catapult, Batch Sharpness re-stabilizes around the new, lower threshold. This is the signature behavior of training at an instability boundary. Stabilizing perturbations restart progressive sharpening. Conversely, although not crucial for establishing a regime of instability, interventions that raise the effective threshold—decreasing η, decreasing β, or increasing b—do not trigger catapults. Instead, these perturbations open a gap between the current Batch Sharpness and the new, higher threshold. This gap permits a renewed phase of progressive sharpening: Batch Sharpness gradually increases until it again approaches the updated threshold (Figure 10).
Matching stabilization levels. In the small-batch regime, η empirically, SGDM (η, β, b) and vanilla SGD( 1−β , b) reach approximately the same λmax stabilization level (Figure 5). This suggests that the stabilization level of Batch Sharpness is the primary determinant of where the full-batch eigenvalue λmax settles, with Section 4 proposing a mechanism behind this behavior. Importantly, we observe in Appendix F that the two trajectories do not track each other; they merely share the same instability threshold.
Batch Sharpness, not λmax , governs the transition. A key observation from these experiments is that the catapult/progressive-sharpening dynamics are predicted by Batch Sharpness, not by λmax . Specifically, when we change the batch size mid-training, λmax does not change instantaneously—the full-batch loss landscape is unaffected by the choice of batch size. Yet we observe either a catapult or renewed progressive sharpening depending on whether Batch Sharpness crosses or falls below its stabilization level. This mirrors the findings of Andreyev & Beneventano (2024) for vanilla SGD and provides strong evidence that Batch Sharpness controls the stability of SGDM/SGDN dynamics.
Two effects of increasing batch size. As batch size increases, two concurrent effects raise the stabilization level of λmax . First, the stabilization threshold for Batch Sharpness itself increases from 2(1−β)/η toward 2(1+β)/η. Second, the gap between Batch Sharpness and λmax decreases and can even flip sign as batch size grows. Both effects push λmax to stabilize at progressively higher values (Figures 3 and 4). 3.3. Showing Instability Through Interventions The stabilization of Batch Sharpness at batch-sizedependent plateaus is suggestive of an E O SS-like regime, but does not by itself establish that SGDM and SGDN operate at an instability boundary. Following the discussion in Section 2, the definitive diagnostic for such a regime is the intervention experiment: if training self-organizes near an instability threshold, then small destabilizing perturbations
Instability without a complete theory. We emphasize an important distinction from the vanilla SGD case. For SGD without momentum, Andreyev & Beneventano (2024) established that Batch Sharpness crossing 2/η is a valid instability criterion: on the quadratic approximation, exceeding this 6
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
Figure 6. Within-run EoSS dynamics for an MLP under destabilizing interventions at step 75k with batch size b = 16, learning rate η = 0.004, and momentum β = 0.9. Left: destabilizing momentum intervention, increasing β to 0.95. Middle: destabilizing learning rate intervention, increasing η to 0.0067. Right: destabilizing batch size intervention, decreasing b to 8. Top: Batch Sharpness and λmax . Bottom: Training loss.
xt := θt − θ⋆ and Hi := ∇2 ℓi (θ⋆ ), the mini-batch Hes⋆ b t := 1 P sian is H j∈Bt Hj . Near θ , the SGDM update (5) b linearizes to
threshold guarantees divergence, see their Theorem 1. For SGDM and SGDN, we do not have an analogous theoretical result. Nevertheless, the empirical evidence confirms that SGDM and SGDN train at the edge of stochastic stability, potentially governed by Batch Sharpness:
b t xt−1 , vt = βvt−1 + H
xt = xt−1 − ηvt .
(9)
4.1. Warm-Up: The One-Dimensional Case
• Batch Sharpness undergoes progressive sharpening and saturates at batch-size–dependent plateaus. • Destabilizing interventions that push Batch Sharpness above its plateau trigger catapults. • Stabilizing interventions that lower Batch Sharpness below its plateau restart progressive sharpening. • These transitions occur precisely when Batch Sharpness crosses its stabilization level, independent of λmax .
We begin with a sketch of the one-dimensional case, with the details in Appendix B and C for SGDM and SGDN, respectively. When d = 1, or along any direction where the b t } commute, the system (9) reduces mini-batch Hessians {H to a scalar second-order recursion with random curvature ht having mean a := E[ht ] and variance σb2 := Var(ht ). Analyzing mean-square stability via the induced recursion on (E[x2t ], E[xt xt−1 ], E[x2t−1 ]), the dominant eigenvalue of the mean-square operator admits the expansion 2 2 η a η η2 2 λ⋆ (η) = 1 − 2a + σ +O . 1−β (1 − β)2 b (1 − β)3
The catapult-and-restabilization pattern is the operational signature of an instability boundary, even in the absence of a closed-form divergence theorem. Developing such a theory for momentum methods, including identifying the precise instability criterion and proving that it is saturated under progressive sharpening, remains an important open question.
In the noise-dominated regime where σb2 ≫ a2 /(1 − β), this recovers the same leading-order mean-square stability condition as vanilla SGD with step size ηeff = η/(1 − β), as outlined in Wu et al. (2018).
4. The Instability Threshold Interpolation between regimes. In d = 1, we can derive the full stability boundary to showcase interpolation between the two limits:
To understand the batch-size-dependent stability thresholds observed in Section 3, we perform a linear stability analysis following Wu et al. (2018); Ma & Ying (2021). We focus this formal development on SGDM, as the more commonly used momentum method. The key finding is that in the noisedominated (small-batch) regime, the linearized mean-square stability of SGDM(η, β) admits a leading-order reduction to that of vanilla SGD with effective step size ηeff := η/(1 − β).
1 ηmax
=
a σb2 + . 2(1 + β) 2a(1 − β) | {z } | {z } deterministic
(10)
stochastic
When σb2 → 0 (large batch), this recovers the classical heavy-ball threshold ηmax = 2(1 + β)/a. When σb2 ≫ a2 (small batch), the stochastic term dominates, yielding the effective threshold 2(1 − β)/η for the curvature. In between, this yields the interpolation observed empirically in Section 3.1.
Setup: Linearization near a Minimizer. Let θ⋆ be an interpolating5 minimizer with ∇ℓi (θ⋆ ) = 0 for all i. Writing 5 Relaxing the interpolation assumption to a general local minimum affects the steady-state variance of the iterates but leaves the
fundamental stability criteria unchanged.
7
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
4.2. Extension to Multiple Dimensions
Adam, and SGD in the E O S literature (Cohen et al., 2021; 2022; 2024; Andreyev & Beneventano, 2024). The resulting implications include:
The one-dimensional intuition extends to the general case. The mean-square dynamics of the augmented state (xt , vt ) are governed by a 4d2 × 4d2 linear operator whose spectrum splits into fast modes (contracting at rate β) and d2 slow modes near unity. The following theorem, proved in Appendix D, characterizes the leading-order slow dynamics. Define the Kronecker moments and drift operator: b t ], H̄ := E[H
bt ⊗ H b t ], G := E[H
• Descent-lemma proof templates fail at (Eo)SS. If training self-organizes near an instability boundary, uniform-smoothness arguments enforcing step-by-step monotone descent are typically not informative; this caveat extends to momentum methods once they exhibit EoSS-like plateaus. • What “large step size” means with momentum. The relevant finite-η comparator is direction-aware mini-batch curvature (Batch Sharpness), not the fullbatch λmax . Under momentum, the operative plateau level is batch- and method-dependent (e.g., small-batch BSplateau ≈ 2(1 − β)/η, while large-batch plateaus approach classical full-batch momentum thresholds). In particular, this reverses the full-batch intuition that momentum simply allows larger step sizes. • Stabilization becomes distributional. Under progressive sharpening, “where the dynamics stabilizes” can depend on the distribution of mini-batch Hessians {HB } (and, with momentum, plausibly also on timecouplings induced by the velocity state), not only on the mean/full-batch Hessian. • Limits of “GD + noise” and diffusion/SDE surrogates. Modeling momentum SGD by generic noise injection or diffusion limits can discard the discretetime, batch-dependent curvature constraints that govern EoSS-like behavior; faithful surrogates should preserve the mini-batch sampling structure relevant to the governing curvature statistics.
K := H̄ ⊗ Id + Id ⊗ H̄
Theorem 4.1. In the noise-dominated regime and for sufficiently small η, the leading-order mean-square stability condition for SGDM(η, β) is: 2 ρ Id2 − ηeff K + ηeff G < 1, (11) Notice that (11) coincides with the vanilla SGD stability condition of Ma & Ying (2021), Equation (31), evaluated at step size ηeff . We leave the details of the proof to Appendix D. Crucially, theoretical links between SGDM stability and SGD with a modified step size already exist in the literature, see e.g. Yuan et al. (2016)6 Our operator-centric approach recovers, to leading order, the vanilla-SGD stability criterion (rather than merely a sufficient condition); it also leaves open the possibility of explaining the empirically observed similarity of quantities like λmax between SGDM and SGD through the explicit dependence on H̄ in the constraint. 4.3. Connection to Batch Sharpness Theorem 4.1 therefore suggests that small-batch SGDM should exhibit a stability threshold of 2/ηeff = 2(1 − β)/η, matching the empirical findings about Batch Sharpness stabilization in Section 3.1. In the noise-dominated small-batch regime, this link is indirect but theoretical: Theorem 4.1 reduces SGDM mean-square stability to vanilla SGD with step size ηeff , and Andreyev & Beneventano (2024) prove that for vanilla SGD Batch Sharpness is the governing instability criterion with threshold 2/ηeff . Thus, in this regime, our theory supports the Batch Sharpness threshold through reduction to SGD. What is still missing is a direct momentum-specific bridge from the operator condition (11) to Batch Sharpness itself, especially outside the small-batch reduction.
On top of this, we established that momentum does not merely shift the edge of stability—it makes the edge itself batch-dependent, and in the stochastic regime it can enforce strictly flatter solutions than vanilla SGD. Momentum is not uniformly stabilizing. In deterministic (large-batch) regimes, momentum recovers its classical stability enlargement and permits training near sharper curvature thresholds. In contrast, in noise-dominated (smallbatch) regimes, momentum acts as an effective step-size amplifier, tightening the operative curvature constraint to BSplateau ≈ 2(1 − β)/η and biasing the search toward flatter regions.
5. Implications Instability viewpoint on hyperparameter coupling. In the small-batch regime, ηeff = η/(1 − β) emerges as the leading-order stability parameter, so changing β without a compensating change in η can move the dynamics across the stochastic edge. This provides a stability-centric rationale for coupling rules that approximately keep η/(1 − β) fixed when varying momentum under noisy training.
We showed that SGDM and SGDN train neural networks at the Edge of Stochastic Stability and therefore inherit the implications previously established for GD, full-batch 6
Although they assume smallness of the step size and prove strong approximation of SGDM and SGD trajectories, which does not hold for the “large” step sizes considered here; see Appendix F.
8
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
vanilla SGD with effective step size ηeff , and then invoke Andreyev & Beneventano (2024), where Batch Sharpness is proved to govern instability for SGD. Thus the present theory justifies the role of Batch Sharpness indirectly in that regime; a direct momentum-specific bridge, especially beyond the small-batch reduction and for SGDN, remains open.
Toward adaptive instability control. Interventions show that catapults occur when Batch Sharpness is pushed above its operating plateau, suggesting that Batch Sharpness tracking could serve as a practical instability monitor for tuning η, β, or b to stay near—but not beyond—the stochastic edge.
6. Conclusion We studied how momentum interacts with instability-limited training in modern neural networks. We show that SGD with Polyak momentum and Nesterov acceleration operates at an Edge of Stochastic Stability (E O SS), characterized by the saturation of a batch-dependent directional curvature rather than by the full-batch Hessian spectrum.
6.2. Future Work Future work includes (i) finding a valid instability criterion for SGDM and SGDN, (ii) understanding the implications for learning and performance, (iii) identifying the sources of progressive sharpening and (iv) the self-stabilization mechanism, and (v) extending this line of analysis toward the stabilization levels of Adam, AdamW, and Muon (Jordan, 2024).
Our main finding is that momentum induces two distinct instability regimes. In the large-batch (near-deterministic) regime, SGDM and SGDN recover classical full-batch behavior, stabilizing at sharper curvatures consistent with known momentum-dependent stability thresholds. In contrast, in the small-batch regime, momentum tightens the effective stability constraint: Batch Sharpness stabilizes at 2(1 − β)/η, biasing training toward significantly flatter regions than vanilla SGD. This regime reversal explains why momentum can simultaneously accelerate training while increasing implicit regularization under stochastic gradients. Methodologically, we adopt an intervention-based diagnostic that makes instability empirically testable: small destabilizing hyperparameter changes trigger catapult-like excursions only when training is genuinely stability-limited. This mechanism allows us to distinguish curvature-driven instability from noise-induced oscillations and to identify the operative stability boundary in mini-batch training. Overall, our results place momentum-based optimizers within the broader Edge-of-Stability framework, clarify when classical momentum intuitions apply, and highlight Batch Sharpness as a central geometric quantity governing stochastic optimization dynamics. 6.1. Limitations Our main empirical limitation is that we only test this on CIFAR-10 and SVHN using relatively small models (up to ResNet18). This limitation is shared by much of the Edge-of-Stability literature. Indeed, computing distributional quantities (such as moments) of the per-sample Hessians {Hi (θ)}i∈D,θ∈Rd is computationally infeasible for large models d ≫ 1 and datasets n ≫ 1; see (Andreyev & Beneventano, 2024) for a discussion. A second limitation is theoretical. Our analysis does not directly prove from the momentum dynamics themselves that Batch Sharpness governs stabilization. Rather, in the noise-dominated small-batch regime, we reduce SGDM to 9
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
Impact Statement
Cutkosky, A. and Orabona, F. Momentum-based variance reduction in non-convex SGD. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
This work clarifies the practical operating regime of momentum-based stochastic gradient methods in modern deep learning. Our findings have broader implications for optimizers and hyperparameter tuning, helping clarify the impact of momentum on optimization and generalization. We characterize deep learning optimization as shaped jointly by momentum and mini-batch noise, help reconcile disparate empirical observations, and provide a principled framework for understanding the dynamics of momentum in real-world training.
Damian, A., Nichani, E., and Lee, J. D. Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability, April 2023. URL http://arxiv.org/abs/ 2209.15594. arXiv:2209.15594 [cs, math, stat]. Fu, J., Wang, B., Zhang, H., Zhang, Z., Chen, W., and Zheng, N. When and why momentum accelerates SGD: An empirical study, 2023. URL https://arxiv.org/ abs/2306.09000.
Agarwala, A. and Pennington, J. High dimensional analysis reveals conservative sharpening and a stochastic edge of stability. arXiv preprint arXiv:2404.19261, 2024.
Ghosh, A., Lyu, H., Zhang, X., and Wang, R. Implicit regularization in heavy-ball momentum accelerated stochastic gradient descent. In International Conference on Learning Representations (ICLR), 2023. URL https: //arxiv.org/abs/2302.00849.
Ahn, K., Zhang, J., and Sra, S. Understanding the unstable convergence of gradient descent. In Proceedings of the 39th International Conference on Machine Learning, June 2022. URL https://proceedings.mlr. press/v162/ahn22a.html.
Gitman, I., Lang, H., Zhang, P., and Xiao, L. Understanding the role of momentum in stochastic gradient methods. In Advances in Neural Information Processing Systems (NeurIPS), 2019. URL https://arxiv.org/abs/ 1910.13962.
References
Granziol, D., Zohren, S., and Roberts, S. Learning rates as a function of batch size: A random matrix theory approach to neural network training, 2021. URL https: //arxiv.org/abs/2006.09092.
Andreyev, A. and Beneventano, P. Edge of Stochastic Stability: Revisiting the Edge of Stability for SGD. December 2024. doi: 10.48550/arXiv.2412.20553. URL http:// arxiv.org/abs/2412.20553. arXiv:2412.20553.
He, K., Zhang, X., Ren, S., and Sun, J. Deep Residual Learning for Image Recognition. arXiv:1512.03385 [cs], December 2015. URL http://arxiv.org/abs/ 1512.03385. arXiv:1512.03385.
Chen, L. and Bruna, J. Beyond the edge of stability via two-step gradient updates. In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J. (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 4330–4391. PMLR, 23–29 Jul 2023. URL https://proceedings.mlr.press/ v202/chen23b.html.
Jastrz˛ebski, S., Kenton, Z., Ballas, N., Fischer, A., Bengio, Y., and Storkey, A. On the Relation Between the Sharpest Directions of DNN Loss and the SGD Step Length, December 2019. URL http://arxiv.org/ abs/1807.05031. arXiv:1807.05031 [stat].
Cohen, J. M., Kaur, S., Li, Y., Kolter, J. Z., and Talwalkar, A. Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability. arXiv:2103.00065 [cs, stat], June 2021. URL http://arxiv.org/abs/2103. 00065. arXiv:2103.00065. Cohen, J. M., Ghorbani, B., Krishnan, S., Agarwal, N., Medapati, S., Badura, M., Suo, D., Cardoze, D., Nado, Z., Dahl, G. E., and Gilmer, J. Adaptive Gradient Methods at the Edge of Stability, July 2022. URL http://arxiv. org/abs/2207.14484. arXiv:2207.14484 [cs]. Cohen, J. M., Damian, A., Talwalkar, A., Kolter, Z., and Lee, J. D. Understanding Optimization in Deep Learning with Central Flows, October 2024. URL http://arxiv. org/abs/2410.24206. arXiv:2410.24206. 10
Jastrz˛ebski, S., Szymczak, M., Fort, S., Arpit, D., Tabor, J., Cho, K., and Geras, K. The BreakEven Point on Optimization Trajectories of Deep Neural Networks. arXiv:2002.09572 [cs, stat], February 2020. URL http://arxiv.org/abs/2002. 09572. arXiv:2002.09572. Jelassi, S. and Li, Y. Towards understanding how momentum improves generalization in deep learning. In Proceedings of the 39th International Conference on Machine Learning (ICML), Proceedings of Machine Learning Research, 2022. URL https://arxiv.org/ abs/2207.05931. Jordan, K. Muon: An optimizer for hidden layers in neural networks, dec 2024. URL https://kellerjordan. github.io/posts/muon/. Blog post.
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
Kidambi, R., Netrapalli, P., Jain, P., and Kakade, S. M. On the insufficiency of existing momentum schemes for stochastic optimization, 2018. URL https://arxiv. org/abs/1803.05591.
Ma, J. and Yarats, D. Quasi-hyperbolic momentum and adam for deep learning, 2018. URL https://arxiv. org/abs/1810.06801.
Liu, Y., Gao, Y., and Yin, W. An improved analysis of stochastic gradient descent with momentum. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
Ramezani-Kebrya, A., Antonakopoulos, K., Cevher, V., Khisti, A., and Liang, B. On the generalization of stochastic gradient descent with momentum. Journal of Machine Learning Research, 25(22):1–56, 2024. URL https: //jmlr.org/papers/v25/22-0068.html.
Mandt, S., Hoffman, M. D., and Blei, D. M. Stochastic gradient descent as approximate Bayesian inference. Krizhevsky, A., Sutskever, I., and Hinton, G. E. ImageNet Journal of Machine Learning Research, 18(134):1–35, Classification with Deep Convolutional Neural Networks. 2017. URL https://www.jmlr.org/papers/ In Advances in Neural Information Processing Systems, v18/16-511.html. volume 25. Curran Associates, Inc., 2012. URL https: //papers.nips.cc/paper/2012/hash/ Masters, D. and Luschi, C. Revisiting Small Batch Training c399862d3b9d6b76c8436e924a68c45b-Abstract. for Deep Neural Networks, April 2018. URL http:// html. arxiv.org/abs/1804.07612. arXiv:1804.07612. Lee, S. and Jang, C. A new characterization of Mitliagkas, I., Zhang, C., Hadjis, S., and Ré, C. Asynthe edge of stability based on a sharpness meachrony begets momentum, with an application to deep sure aware of batch gradient distribution. In Inlearning, 2016. URL https://arxiv.org/abs/ ternational Conference on Learning Representations, 1605.09774. 2023. URL https://api.semanticscholar. org/CorpusID:259298833. Mulayoff, R. and Michaeli, T. Exact mean square linear stability analysis for sgd, 2024. URL https://arxiv. Li, Q., Tai, C., and E, W. Dynamics of stochastic gradient org/abs/2306.07850. algorithms, 2015. URL https://arxiv.org/abs/ 1511.06251. Nesterov, Y. On an approach to the construction of optimal methods of minimization of smooth convex funcLi, Q., Tai, C., and E, W. Stochastic modified equations tions. Ekonomika i Mateaticheskie Metody, 24(3):509– and adaptive stochastic gradient algorithms. In Proceed517, 1988. ings of the 34th International Conference on Machine Paquette, C. and Paquette, E. Dynamics of stochastic Learning (ICML), 2017. URL https://arxiv.org/ momentum methods on large-scale quadratic models. abs/1612.06277. In Advances in Neural Information Processing Systems Li, Q., Tai, C., and E, W. Stochastic modified equations I: (NeurIPS), 2021. URL https://arxiv.org/abs/ Mathematical foundations. Journal of Machine Learning 2104.03485. Research, 2019. URL https://www.jmlr.org/ Polyak, B. Some methods of speeding up the convergence papers/v20/17-526.html. of iteration methods. USSR Computational Mathematics Liu, T., Chen, Z., Zhou, E., and Zhao, T. A diffusion apand Mathematical Physics, 4(5):1–17, 1964. ISSN proximation theory of momentum SGD in nonconvex op0041-5553. doi: 10.1016/0041-5553(64)90137-5. timization, 2018. URL https://arxiv.org/abs/ URL https://www.sciencedirect.com/ 1802.05155. science/article/pii/0041555364901375.
Lyu, B. Effects of momentum in implicit bias of gradient flow for diagonal linear networks. In Proceedings of the AAAI Conference on Artificial Intelligence (AAAI), 2025. URL https://ojs.aaai.org/ index.php/AAAI/article/view/34118.
Shi, B., Du, S., Jordan, M. I., and Su, W. Understanding the acceleration phenomenon via high-resolution differential equations. Mathematical Programming, 2022. doi: 10.1007/s10107-021-01681-8. URL https://doi. org/10.1007/s10107-021-01681-8.
Ma, C. and Ying, L. On linear stability of SGD and input-smoothness of neural networks. In Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems, 2021. URL https://openreview.net/forum? id=yAvCV6NwWQ.
Su, W., Boyd, S., and Candès, E. A differential equation for modeling Nesterov’s accelerated gradient method: Theory and insights. In Advances in Neural Information Processing Systems (NeurIPS), 2014. URL https:// arxiv.org/abs/1503.01243. 11
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
Sutskever, I., Martens, J., Dahl, G. E., and Hinton, G. E. On the importance of initialization and momentum in deep learning. In Proceedings of the 30th International Conference on Machine Learning (ICML), volume 28 of Proceedings of Machine Learning Research, pp. 1139–1147, 2013. URL https://proceedings.mlr.press/ v28/sutskever13.html. Wang, B., Meng, Q., Zhang, H., Sun, R., Chen, W., Ma, Z.-M., and Liu, T.-Y. Does momentum change the implicit regularization on separable data? In Advances in Neural Information Processing Systems (NeurIPS), 2022. URL https://openreview.net/forum? id=i-8uqlurj1f. Wang, R., Malladi, S., Wang, T., Lyu, K., and Li, Z. The marginal value of momentum for small learning rate SGD. In International Conference on Learning Representations (ICLR), 2024. URL https://arxiv.org/ abs/2307.15196. Wibisono, A., Wilson, A. C., and Jordan, M. I. A variational perspective on accelerated methods in optimization, 2016. URL https://arxiv.org/abs/1603.04245. Wilson, A. C., Recht, B., and Jordan, M. I. A Lyapunov analysis of accelerated methods in optimization. Journal of Machine Learning Research, 2021. Wu, L. and Su, W. J. The Implicit Regularization of Dynamical Stability in Stochastic Gradient Descent, June 2023. URL http://arxiv.org/abs/2305. 17490. arXiv:2305.17490 [stat]. Wu, L., Ma, C., and E, W. How SGD Selects the Global Minima in Over-parameterized Learning: A Dynamical Stability Perspective. In Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018. URL https://proceedings.neurips. cc/paper_files/paper/2018/hash/ 6651526b6fb8f29a00507de6a49ce30f-Abstract. html. Wu, L., Wang, M., and Su, W. The alignment property of sgd noise and how it helps select flat minima: A stability analysis, 2022. Xing, C., Arpit, D., Tsirigotis, C., and Bengio, Y. A Walk with SGD, May 2018. URL http://arxiv.org/ abs/1802.08770. arXiv:1802.08770 [cs, stat]. Yuan, K., Ying, B., and Sayed, A. H. On the influence of momentum acceleration on online learning, 2016. URL http://arxiv.org/abs/1603.04136. Zagoruyko, S. and Komodakis, N. Wide Residual Networks. arXiv:1605.07146 [cs], June 2017. URL http:// arxiv.org/abs/1605.07146. arXiv:1605.07146. 12
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
Appendix A. Further Related Work The added value of momentum. Polyak heavy-ball momentum and Nesterov-style acceleration are ubiquitous in modern deep learning and are often practically necessary for stable and fast training (Sutskever et al., 2013; Gitman et al., 2019; Fu et al., 2023; Wang et al., 2024). A large body of work has tried to isolate what momentum contributes beyond learning-rate tuning, and several complementary explanations have emerged. A first theme is stability enlargement and step-size rescaling: relative to vanilla SGD, adding momentum can enlarge the admissible learning-rate range by stabilizing the dynamics, and in certain regimes its net effect is well approximated by an effective learning rate ηeff ≈ γ/(1 − β), so that SGD and SGDM can exhibit closely tracking trajectories (or statistics of trajectories) after matching ηeff (Fu et al., 2023; Wang et al., 2024); related sufficient conditions can also be expressed in terms of global smoothness-type constants (Yuan et al., 2016). A second theme is temporal filtering and noise shaping: the momentum buffer is an exponential moving average of past Pk−1 stochastic gradients, mk = j=0 β k−1−j gj , so SGDM can be viewed as applying a linear time-invariant filter to gradient noise, motivating stationary-distribution, SDE-limit, and stochastic modified-equation analyses (Gitman et al., 2019; Li et al., 2015; 2017; 2019; Mandt et al., 2017). A third theme is inertial (underdamped) dynamics and geometry: momentum induces second-order behavior that can change transient exploration in narrow valleys and ill-conditioned directions relative to overdamped SGD; this viewpoint is often formalized via continuous-time limits and Lyapunov analyses (Su et al., 2014; Wibisono et al., 2016; Shi et al., 2022; Wilson et al., 2021; Liu et al., 2018). Beyond optimization speed, momentum can influence solution selection and generalization through implicit regularization and stability-based bounds (Wang et al., 2022; Jelassi & Li, 2022; Ghosh et al., 2023; Lyu, 2025; Ramezani-Kebrya et al., 2024), and it also interacts with systems/implementation effects (e.g., asynchrony), where staleness can induce implicit momentum and motivate “negative momentum” corrections (Mitliagkas et al., 2016). When does momentum help? A hyperparameter-dependent view. Recent empirical and theoretical work suggests that the practical value of momentum is strongly regime-dependent. When the learning rate that yields good performance is already small—as is common in small/medium batch training and many fine-tuning setups—SGD and SGDM often have closely tracking trajectories after matching effective learning rates, and momentum can deliver only marginal gains in both optimization and generalization (Fu et al., 2023; Wang et al., 2024). In contrast, once one pushes toward larger step sizes, plain SGD often encounters an instability threshold first; in this high-step regime, momentum can matter substantially by enlarging the stability region and delaying or mitigating instability (Fu et al., 2023; Cohen et al., 2021). In genuinely noisy regimes, momentum also changes the limiting stochastic dynamics (e.g., via underdamped diffusion limits and stochastic modified equations), affecting transient exploration and stationary behavior (Liu et al., 2018; Li et al., 2015; 2017; 2019; Mandt et al., 2017; Gitman et al., 2019). The picture for implicit bias and generalization is correspondingly mixed: in some separable linear settings momentum preserves the max-margin implicit bias (Wang et al., 2022), whereas in other regimes—including feature-learning settings, certain linear-network parameterizations, and multi-epoch stability analyses—momentum can provably or empirically change the selected solution and its test performance (Jelassi & Li, 2022; Ghosh et al., 2023; Lyu, 2025; Ramezani-Kebrya et al., 2024). Pointers into the momentum literature (deep-learning oriented). On the empirical and algorithmic side, momentum schedules have long been motivated by their role in training large neural networks (Sutskever et al., 2013), with subsequent work proposing unifying parameterizations and guidelines (e.g., QHM/QHAdam) (Ma & Yarats, 2018) and characterizing the regimes in which momentum accelerates or becomes largely redundant (Fu et al., 2023; Wang et al., 2024). A complementary body of work uses dynamical-systems perspectives—ODE limits, variational formulations, “high-resolution” differential equations, and Lyapunov analyses—to explain how accelerated discretizations reshape stability and transient behavior (Su et al., 2014; Wibisono et al., 2016; Shi et al., 2022; Wilson et al., 2021). In stochastic regimes, diffusion approximations, stationary-law analyses, and stochastic modified-equation techniques provide a principled language for how momentum filters noise and induces underdamped dynamics (Liu et al., 2018; Li et al., 2015; 2017; 2019; Gitman et al., 2019; Mandt et al., 2017). Nonasymptotic convergence theory clarifies both limitations of vanilla momentum and circumstances where modified momentum schemes (often coupled with variance reduction) achieve provable gains (Kidambi et al., 2018; Liu et al., 2020; Cutkosky & Orabona, 2019; Paquette & Paquette, 2021). Finally, a growing literature studies momentum’s effect on implicit bias and generalization in overparameterized learning (Jelassi & Li, 2022; Wang et al., 2022; Ghosh et al., 2023; Lyu, 2025; Ramezani-Kebrya et al., 2024), and systems work highlights how asynchrony can generate implicit momentum and how “negative momentum” can correct for staleness (Mitliagkas et al., 2016). 13
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
A.1. Further Explanation: Edge of Stability: deterministic picture and why mini-batch is different Progressive sharpening and E O S. A sequence of works (Jastrz˛ebski et al., 2019; 2020; Cohen et al., 2021) pinpoints a rapid early-time change in local curvature during training: along GD/SGD trajectories, the top Hessian eigenvalue λmax typically exhibits a short initial dip, followed by a sustained increase. Jastrz˛ebski et al. (2020) further report a sharp transition that ends this “progressive sharpening” phase. Subsequent evidence suggests that the timing of this transition is closely tied to optimizer stability (and can differ across algorithms even on the same objective) (Jastrz˛ebski et al., 2019; 2020; Cohen et al., 2021; 2022). For full-batch methods, Cohen et al. (2021; 2022) relate the transition to the corresponding instability threshold; for instance, under GD, λmax settles into oscillations around 2/η, whereas for full-batch Adam (standard hyperparameters) the top eigenvalue of the preconditioned Hessian hovers around 38/η. In MSE experiments, much of the remaining optimization appears to take place in this near-threshold regime, which in turn largely sets the λmax of the final iterate (Cohen et al., 2021; 2022). Chen & Bruna (2023); Lee & Jang (2023) explain why λmax can slightly exceed 2/η in practice: gradient nonlinearity introduces higher-order corrections that shift the effective boundary. The mechanism sustaining E O S-style dynamics has also been studied; Damian et al. (2023) argue that, under empirically supported alignment assumptions, third-order derivatives can provide a global stabilizing effect even when the local linearization is unstable. Full-batch E O S as a stability boundary (Cohen et al.). A useful mental model is to view optimization as “surfing” the boundary of stability: the algorithm pushes toward higher curvature (which accelerates loss decrease) until it reaches the largest curvature that still permits sustained progress. Concretely, for full-batch gradient descent (GD) with constant step size η, θt+1 = θt − η∇L(θt ), (12) the local quadratic model L(θ) ≈ L(θt ) + g ⊤ e + 21 e⊤ He predicts linear stability only if η λmax (H) < 2. Cohen et al. observed that, in deep networks, training often enters a regime where the full-batch sharpness λmax (∇2 L(θt )) increases (“progressive sharpening”) until it hovers near the quadratic stability threshold, λmax (∇2 L(θt )) ≈ 2/η, and training continues while remaining close to this boundary (the “Edge of Stability”, EOS). Momentum near the deterministic edge. For momentum methods (e.g. Polyak heavy-ball or Nesterov acceleration), classical linear analysis on quadratics predicts that the stability boundary shifts compared to GD, i.e. the maximal stable curvature increases by a β-dependent factor. Empirically, in deterministic or very-large-batch regimes this can manifest as an E O S-like plateau at a sharper level than 2/η (often summarized heuristically as a “∼ (1 + β)” increase in the plateau sharpness), consistent with the intuition that momentum can tolerate and traverse higher-curvature directions when noise is small. Mini-batch training: why full-batch sharpness alone can be misleading. Mini-batch optimizers do not evolve on a single fixed landscape: each step uses a sampled loss LB (with gradient gB and Hessian HB ), so stability is governed by the distribution of mini-batch geometries rather than only the full-batch Hessian. Practically, this means that the full-batch λmax (∇2 L(θt )) need not stabilize near 2/η in mini-batch regimes, even when training exhibits sharp transitions and intermittent “catapult”-like excursions. Moreover, stochastic training can exhibit persistent loss oscillations for reasons other than instability-limited dynamics, so the onset of oscillations is not, by itself, a reliable mini-batch analogue of the full-batch E O S signature. These differences motivate stability notions and diagnostics that directly target the stochastic dynamics. SGD stability analyses and a direction-aware mini-batch analogue of E O S. Several papers (Wu et al., 2018; Ma & Ying, 2021; Granziol et al., 2021; Wu et al., 2022; Mulayoff & Michaeli, 2024) analyze constant-step-size SGD on quadratic objectives by studying linear stability of the parameter second moment, yielding sharp criteria of the form EB∼Pb (I − ηH(LB ))⊗2 ≶ 1. These conditions can be interpreted as Lyapunov-style stability tests (e.g., for the squared distance to the minimizer) and clarify how sampling-induced randomness alters stability relative to full-batch GD. In deep networks, directly evaluating such d2 -dimensional operators is typically impractical, so a parallel literature explores more tractable diagnostics and proxies, including scalar curvature summaries such as trace-style Hessian/NTK quantities (Wu & Su, 2023; Agarwala & Pennington, 2024) and empirical characterizations of oscillatory regimes (Cohen et al., 2021; Xing et al., 2018; Lee & Jang, 2023; Ahn et al., 2022). A complementary and especially simple direction-aware statistic is the curvature along the 14
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
stochastic update direction, captured by the Batch Sharpness: gB (θ)⊤ HB (θ) gB (θ) BS(θ) := EB∼Pb , ∥gB (θ)∥2
gB (θ) = ∇LB (θ), HB (θ) = ∇2 LB (θ).
(13)
BS(θ) is the expected Rayleigh quotient of the mini-batch Hessian along the mini-batch gradient direction; informally, it is the expected curvature “felt” along the stochastic update. Under a local quadratic approximation, for vanilla SGD the threshold BS(θ) > 2/η (by any fixed margin) is sufficient to trigger a catapult-like excursion, suggesting a mini-batch edge-of-stochastic-stability condition E O SS:
BS(θt ) progressively sharpens and then hovers near 2/η.
(14)
For momentum methods, the exact closed-form instability criterion is more delicate; nevertheless, direction-aware curvature tracking (e.g., BS(θt )) together with a targeted intervention test (below) provides a practical way to diagnose whether training is genuinely constrained by an instability boundary. A.2. More on the Mechanism: Instability Criteria and an intervention From local descent models to instability criteria (what an “edge” means stochastically). In stochastic optimization, insisting on a single global descent lemma is often too blunt; a common alternative is a local viewpoint: fix a checkpoint θt and a neighborhood Ut where a local approximation (e.g., a quadratic model) is believed to be informative. Within such a neighborhood, one can formalize the notion of a stability boundary through an algorithm-dependent instability criterion. Definition A.1 (Instability criterion). Consider a training algorithm (a discrete-time dynamical system) (θt )t≥0 on a parameter space Θ ⊆ Rd with fixed hyperparameters h (e.g. learning rate, batch size). Let U ⊆ Θ be an open set (typically, a region where a local approximation of the loss is trusted), and let f : U → R and c ∈ R. We say that f is a valid instability criterion with threshold c for the algorithm on U if f (θ0 ) > c =⇒ (θt )t≥0 leaves every compact subset of U . Equivalently: for any compact K ⊂ U containing θ0 , there exists a finite time T such that P(θT ∈ / K | θ0 ) > 0. We say that f is saturated at θ if f (θ) is (approximately) equal to c; in practice, up to an O(η · poly(log(η))) tolerance. Stochastic systems can admit multiple valid instability criteria (depending on which local model and which Lyapunov function one uses), so the “edge” is best understood operationally: training self-organizes so that some valid criterion saturates, f (θt ; h) ≈ c(h), under progressive sharpening. A plateau in a diagnostic plot can therefore be suggestive, but is not, by itself, conclusive evidence that training is truly instability-limited. A checkpoint perturbation test for “training at the edge”. An intervention-based mechanism can make “training at the edge” more directly falsifiable. Take a checkpoint θt from a baseline run with hyperparameters h0 , and restart training from θt under a small destabilizing perturbation of h0 (e.g. η ↑ or b ↓, with other settings fixed). If the baseline dynamics is genuinely stability-limited at θt , then this small perturbation typically triggers a rapid transient runaway (a “catapult”): a large excursion accompanied by a sharp loss spike, followed (in many cases) by re-stabilization and saturation at a new level consistent with the perturbed hyperparameters. Conversely, stabilizing perturbations (e.g. η ↓ or b ↑) should suppress catapults and can reopen a “gap” that allows renewed progressive sharpening. This intervention test is useful precisely because it distinguishes quantities that merely appear to plateau from quantities that are actually governing local instability: if a diagnostic seems to saturate but small destabilizing perturbations do not trigger catapults, then that diagnostic is unlikely to be a valid instability criterion for the relevant stochastic dynamics.
15
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
B. Proof of the 1D Warm-Up Case for SGDM This Appendix provides the proofs for the claims for the one-dimensional case in Section 4.1. We consider the one-dimensional linearized SGDM: xt =
1 + β − ηht xt−1 − βxt−2 ,
(15)
where (ht )t≥1 are i.i.d., independent of (xt−1 , xt−2 ), with σb2 := Var(ht ). Pn n−b σ 2 with σ 2 := n1 i=1 a2i − a2 .) (Under without-replacement mini-batching, σb2 = b(n−1) a := E[ht ],
Closed recursion for second moments.
Let αt := 1 + β − ηht and define the second-moment state E[x2t ] wt := E[xt xt−1 ] . E[x2t−1 ]
Thus, x2t = αt2 x2t−1 − 2βαt xt−1 xt−2 + β 2 x2t−2 ,
xt xt−1 = αt x2t−1 − βxt−1 xt−2 .
Combined with α1 := E[αt ] = 1 + β − ηa,
α2 := E[αt2 ] = (1 + β − ηa)2 + η 2 σb2 = α12 + η 2 σb2 ,
we get α2 R(η) := α1 1
wt = R(η) wt−1 ,
−2βα1 −β 0
β2 0 . 0
(16)
Mean-square linear stability is thus equivalent to ρ(R(η)) < 1 Perturbative expansion.
Characteristic polynomial of ρ(R(η)): pη (λ) = λ3 + β − α2 λ2 + 2βα12 − βα2 − β 2 λ − β 3 .
(17)
At η = 0, α1 = 1 + β and α2 = (1 + β)2 , and (17) factorizes as p0 (λ) = (λ − 1)(λ − β)(λ − β 2 ), so the only eigenvalue on the unit circle is λ = 1. By continuity, for small η there is a unique eigenvalue λ⋆ (η) with λ⋆ (0) = 1 governing the stability boundary. To expand λ⋆ (η), set the ansatz λ⋆ (η) = 1 + c1 η + c2 η 2 + O(η 3 ), and impose pη (λ⋆ (η)) ≡ 0. Expanding (17) in η and matching coefficients yields c1 = −
2a , 1−β
c2 =
Therefore, 2a λ⋆ (η) = 1 − η+ 1−β
σb2 1 − 3β 2 + a . (1 − β)2 (1 − β)3
σb2 1 − 3β 2 2 + a η + O(η 3 ). (1 − β)2 (1 − β)3
2
a In the noise-dominated regime where 1−β ≪ σb2 (so the a2 -term is lower order relative to σb2 ), this reduces to 2 λ⋆ (η) = 1 − 2a ηeff + σb2 ηeff + o(η 2 ),
ηeff :=
η , 1−β
which is the same leading-order stability condition as in the Wu et al. (2018) analysis, with stepsize ηeff . 16
(18)
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
Exact interpolation formula for ηmax . The convenience of 1D case is that we compute the exact interpolation case. Since the eigenvalues are {1, β, β 2 } at η = 0 and β, β 2 < 1, the first loss of stability occurs when the dominant mode crosses at λ = 1. The Jury conditions for (17) reduce to the single active inequality pη (1) > 0. Evaluating (17) at λ = 1 and simplifying gives pη (1) = η 2a(1 + β)(1 − β) − η (1 − β)a2 + (1 + β)σb2 . Thus pη (1) > 0 holds iff 0 < η < ηmax , where ηmax =
2a(1 + β) 2a(1 + β)(1 − β) = . 2 1+β 2 2 (1 − β)a + (1 + β)σb σb a2 + 1−β
(19)
a σb2 + . 2(1 + β) 2a(1 − β)
(20)
For convenience, we can take the reciprocal 1 ηmax
=
The deterministic limit σb2 → 0 yields ηmax = 2(1 + β)/a, while the noise-dominated limit σb2 ≫ a2 yields ηmax ≈
2a(1 − β) σb2
⇐⇒
ηeff,max :=
ηmax 2a ≈ 2, 1−β σb
matching the 1D SGD noise-dominated threshold with the effective stepsize ηeff = η/(1 − β).
C. Proof of the 1D Warm-Up Case for SGDN This Appendix provides the proofs for the claims for the one-dimensional case of SGD with Nesterov momentum (SGDN) in the sense of (6). This is essentially repeating Appendix B, but for SGDN, showcasing the same stability threshold. We consider the 1D quadratic linearization near a global minimizer x⋆ = 0, for which the (mini-batch) gradient evaluated at any point y ∈ R takes the form gt (y) = ht y, where (ht )t≥0 are i.i.d., independent of the past iterates, with σb2 := Var(ht ). Pn n−b σ 2 with σ 2 := n1 i=1 a2i − a2 .) (Under without-replacement mini-batching, σb2 = b(n−1) a := E[ht ],
Deriving the 2-step recursion Let xt := θt − θ⋆ = θt . The NAG update is vt+1 = βvt + gt (θt − βηvt ) = βvt + ht (xt − βηvt ),
xt+1 = xt − ηvt+1 .
Hence vt+1 = ht xt + β(1 − ηht )vt ,
xt+1 = xt − ηht xt − βη(1 − ηht )vt = (1 − ηht ) xt − βηvt .
Using xt = xt−1 − ηvt , i.e. vt = (xt−1 − xt )/η, we obtain xt − βηvt = xt − β(xt−1 − xt ) = (1 + β)xt − βxt−1 . Therefore the 1D linearized NAG dynamics reduces to the random-coefficient 2-step recursion xt+1 =
1 − ηht (1 + β)xt − βxt−1 .
Closed recursion for second moments Let rt := 1 − ηht , and define the second-moment state E[x2t ] wt := E[xt xt−1 ] . E[x2t−1 ]
17
(21)
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
From (21) we have x2t+1 = rt2 (1 + β)2 x2t − 2β(1 + β)xt xt−1 + β 2 x2t−1 , xt+1 xt = rt (1 + β)x2t − βxt xt−1 . Define the first two moments of rt : q := E[rt2 ] = E (1 − ηht )2 = (1 − ηa)2 + η 2 σb2 = p2 + η 2 σb2 .
p := E[rt ] = 1 − ηa, We get
wt+1 = RNAG (η) wt ,
(1 + β)2 q RNAG (η) := (1 + β)p 1
−2β(1 + β)q −βp 0
β2q 0 . 0
(22)
Mean-square linear stability is thus equivalent to ρ(RNAG (η)) < 1. Perturbative expansion of the dominant eigenvalue Let pη (λ) := det(λI − RNAG (η)). A direct determinant calculation gives the characteristic polynomial pη (λ) = λ3 + βp − (1 + β)2 q λ2 + βq (1 + β)2 p − β λ − β 3 pq.
(23)
At η = 0, p = q = 1 and (23) factorizes as p0 (λ) = (λ − 1)(λ − β)(λ − β 2 ), so the only eigenvalue on the unit circle is λ = 1. By continuity, for small η there is a unique eigenvalue λ⋆ (η) with λ⋆ (0) = 1 governing the stability boundary. To expand λ⋆ (η), set the ansatz λ⋆ (η) = 1 + c1 η + c2 η 2 + O(η 3 ), and impose pη (λ⋆ (η)) ≡ 0. Expanding (23) in η and matching coefficients yields c1 = −
2a , 1−β
c2 =
σb2 1 − β − 2β 2 2 + a . (1 − β)2 (1 − β)3
Therefore, 2a λ⋆ (η) = 1 − η+ 1−β
σb2 1 − β − 2β 2 2 2 + a η + O(η 3 ). (1 − β)2 (1 − β)3
2
a In the noise-dominated regime where 1−β ≪ σb2 (so the a2 -term is lower order relative to σb2 ), (24) reduces to 2 λ⋆ (η) = 1 − 2a ηeff + σb2 ηeff + o(η 2 ),
ηeff :=
η , 1−β
which matches the leading-order stability condition of 1D SGD with stepsize ηeff , just like in the case of SGDM.
18
(24)
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
D. Proofs for the Multi-dimensional Case In this appendix, we provide the proof that the stability of SGDM(η, β) reduces to that of SGD(ηeff ) in the multi-dimensional case with non-commuting mini-batch Hessians. b t and assume {Ht }t≥1 are i.i.d. with finite second moment and indepenProof of Theorem 4.1. Throughout, write Ht := H dent of (xt−1 , vt−1 ). Define K := H̄ ⊗ I + I ⊗ H̄,
H̄ := E[Ht ],
G := E[Ht ⊗ Ht ].
Consider the linearized SGDM recursion (9), xt = xt−1 − ηvt .
vt = βvt−1 + Ht xt−1 , x Introduce the augmented state zt := t ∈ R2d , so that vt zt = At zt−1 ,
Second-moment operator.
At :=
I − ηHt Ht
−ηβI . βI
(25)
Let Σt := E[zt zt⊤ ] and mt := vec(Σt ). Vectorizing, we get: mt = THB (η, β)mt−1 ,
where
THB (η, β) := E[At ⊗ At ].
(26)
Mean-square stability is then equivalent to ρ(THB (η, β)) < 1. Partition mt into the four (d2 )-blocks mxx t := E[xt ⊗ xt ],
mxv t := E[xt ⊗ vt ],
vv vx xv and set ut := mxx t and yt := mt , mt , mt
⊤
mvx t := E[vt ⊗ xt ],
mvv t := E[vt ⊗ vt ],
2
∈ R3d . Then (26) can be written in block form
ut M11 (η) = yt M21 (η)
M12 (η) ut−1 . M22 (η) yt−1
(27)
A direct expansion of THB (η, β) = E[At ⊗ At ] yields M11 (η) = E[(I − ηHt ) ⊗ (I − ηHt )] = Id2 − ηK + η 2 G, and M12 (η) = O(η), M21 (η) = O(1), while βId2 0 M22 (0) = β(H̄ ⊗ I)
0 βId2 β(I ⊗ H̄)
0 0 , β 2 Id2
ρ M22 (0) = |β| < 1.
(28)
(29)
Hence for η sufficiently small, ρ(M22 (η)) ≤ |β| + O(η) < 1. Schur-complement reduction. Since M22 (η) is strictly stable, the eigenvalues of the full operator in a neighborhood of 1 are governed by the Schur complement: −1 Tslow (η, β) := M11 (η) + M12 (η) I3d2 − M22 (η) M21 (η),
(30)
in the sense that ρ THB (η, β) < 1
⇐⇒
ρ Tslow (η, β) < 1,
and the remaining spectrum of THB (η, β) stays uniformly bounded by |β| + O(η). 19
(31)
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
Small-η expansion.
Expanding (30) around η = 0 and retaining terms up to O(η 2 ) gives Tslow (η, β) = Id2 −
η2 η G + R(η, β), K+ 1−β (1 − β)2
(32)
where R(η, β) collects the (deterministic) curvature-squared contributions and higher-order terms; a crude norm bound of the correct scaling is η2 ∥H̄∥2 . (33) ∥R(η, β)∥ ≲ (1 − β)3 Define ηeff := η/(1 − β). Then (32) reads 2 Tslow (η, β) = Id2 − ηeff K + ηeff G + R(η, β).
(34)
Conclusion. By (31)–(34), mean-square stability is controlled by ρ(Tslow (η, β)). In the noise-dominated regime (small 2 batch) where the multiplicative term dominates the remainder, i.e. ηeff ∥G∥ ≫ ∥R(η, β)∥, we obtain the advertised condition 2 ρ Id2 − ηeff K + ηeff G < 1, which matches the SGD stability operator of Ma & Ying (2021, Eq. (31)) with step size ηeff .
20
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
E. Interventions Mid-Training We present additional intervention experiments that complement our main points by illustrating the robustness of the batch sharpness response. E.1. Destabilizing Intervention We first consider interventions in which a single hyperparameter is modified mid-training so as to lower the effective stability threshold (by increasing the learning rate, increasing the momentum, or decreasing the batch size). Figure 6, Figure 7, Figure 8, and Figure 9 show destabilizing intervention experiments across varying batch-size regimes and intervention timings (i.e, both during and after the progressive sharpening phase). The batch sharpness trajectory of the intervention run (red) closely follows that of the baseline run (blue) prior to the intervention and rapidly transitions to track that of the destabilized run (purple) after the intervention. When the intervention causes the effective stability threshold to drop below the current batch sharpness level, training exhibits a catapult: a sharp increase in loss and curvature statistics followed by restabilization near the new threshold. These results confirm that batch sharpness responds immediately and predictably to destabilizing hyperparameter changes and supports the interpretation that training dynamically tracks a hyperparameter-dependent stability boundary.
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, momentum β = 0.9, and batch size b = 16. Left: destabilizing momentum intervention, increasing β to 0.95. Middle: destabilizing learning-rate intervention, increasing η to 0.0067. Right: destabilizing batch-size intervention, decreasing batch size b to 8. Top: Batch Sharpness and λ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, momentum β = 0.9, and batch size b = 512. Left: destabilizing momentum intervention, increasing β to 0.95. Middle: destabilizing learning-rate intervention, increasing η to 0.0067. Right: destabilizing batch-size intervention, decreasing batch size b to 256. Top: Batch Sharpness and λmax . Bottom: Training loss.
21
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
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, momentum β = 0.5, and batch size b = 16384. Left: destabilizing momentum intervention, increasing β to 0.52. Middle: destabilizing learning-rate intervention, increasing η to 0.035. Right: destabilizing batch-size intervention, decreasing b to 12288. Top: Batch Sharpness and λmax . Bottom: Training loss.
E.2. Stabilizing Intervention We next consider stabilizing interventions by decreasing the learning rate, decreasing the momentum, or increasing the batch size across varying batch sizes and timings in Figure 10, Figure 11, and Figure 12. In contrast to destabilizing interventions, stabilizing interventions induce an immediate decrease in the training loss instead of a catapult. In addition, the intervention modifies the long-term evolution of Batch Sharpness by permitting further progressive sharpening toward a higher plateau. For small batches and for interventions both during and after the progressive sharpening phase, the intervention run transitions away from the baseline trajectory and gradually approaches the trajectory of the stabilized run trained from initialization with the modified hyperparameters. This indicates that raising the stability threshold reopens a sharpening phase that had previously saturated. Nevertheless, for stabilizing interventions at intermediate batch sizes (Figure 12), the batch sharpness remains close to that of the baseline run rather than approaching the stabilized run. Under this hyperparameter regime, we do not observe a continuation of progressive sharpening, as the network has already effectively converged at the time of intervention. This suggests that once the learning dynamics have saturated, raising the stability threshold alone is insufficient to reinitiate sharpening.
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 η = 0.004, and momentum β = 0.9. Left: stabilizing momentum intervention, decreasing β to 0.875. Middle: stabilizing learning-rate intervention, decreasing η to 0.003. Right: stabilizing batch-size intervention, increasing batch size b to 32. Top: Batch Sharpness and λmax . Bottom: Training loss.
22
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
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, momentum β = 0.9, and batch size b = 16. Left: stabilizing momentum intervention, decreasing β to 0.875. Middle: stabilizing learning-rate intervention, decreasing η to 0.003. Right: stabilizing batch-size intervention, increasing batch size b to 32. Top: Batch Sharpness and λ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, momentum β = 0.9, and batch size b = 512. Left: stabilizing momentum intervention, decreasing β to 0.875. Middle: stabilizing learning-rate intervention, decreasing η to 0.003. Right: stabilizing batch-size intervention, increasing batch size b to 768. Top: Batch Sharpness and λmax . Bottom: Training loss.
23
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
F. Distance in the Parameter Space In Figure 13, we evaluate whether SGDM and vanilla SGD with matching stabilization levels of Batch Sharpness and λmax (as in Figure 5) follow comparable parameter trajectories. To compute the L2 distance in the left plot, we apply a fixed JL projection to reduce the weights to a 5,000-dimensional subspace, allowing us to compute trajectory distances efficiently while maintaining geometric fidelity. On the right plot, we use the notion of a test prediction distance, defined as the Frobenius norm between the networks’ output logits on CIFAR-10’s held-out test set of size 10,000. In addition to measuring distance at matching training steps, we introduce the notion of true distance: the minimum distance from each SGD step to any point along the entire SGDM trajectory. This metric is motivated by the observation that even if two trajectories diverge at matching step counts, they may still traverse the same regions of parameter space at different rates. Figure 13 uses the distance from initialization primarily as a baseline to provide context for the separation between SGD and SGDM trajectories. While both runs move a similar total distance through parameter space, the distance between them is of a comparable order of magnitude to their distance from initialization. This lack of point-by-point proximity suggests that matching Batch Sharpness stabilization levels does not ensure a strong approximation between the two methods. This divergence is equally present in function space, where test prediction distances indicate that the networks learn meaningfully different input-output mappings. These results confirm that SGD and SGDM explore geometrically distinct regions of the landscape, though this observation does not rule out the possibility of a weak approximation (as in Wang et al. (2024)) in which the statistical properties of the trajectories might still align.
η , b) for η = 0.001, β = 0.9, and b = 4. Figure 13. Distance metrics comparing training trajectories of SGDM(η, β, b) and SGD( 1−β 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 L2 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.
G. Stabilization vs Batch Size Ablations This section shows how curvature-related quantities stabilize as a function of batch size for a fixed total sample budget for training runs on a CNN and MLP. For each optimizer and learning-rate configuration, we plot the plateau values of Batch Sharpness and λmax obtained from within-run dynamics across increasing batch sizes. These plots illustrate the transition from the small-batch regime to the large-batch regime, and we overlay the corresponding theoretical stability thresholds implied by momentum and learning rate. Together, they make explicit how optimizer choice reshapes the location and sharpness of this transition; in particular, Nesterov typically reaches the critical batch size at much smaller batch sizes than Polyak momentum, consistent with its effectively reduced stability threshold. 24
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
Figure 14. MLP, η = 0.01, β = 0.9
Figure 15. MLP, η = 0.005, β = 0.9
Figure 16. MLP, η = 0.004, β = 0.9
Figure 17. MLP, η = 0.003, β = 0.9
Figure 18. MLP, η = 0.001, β = 0.9
Figure 20. MLP, η = 0.002, β = 0.9, Nesterov
Figure 19. MLP, η = 0.005, β = 0.9, Nesterov
25
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
Figure 21. CNN, η = 0.005, β = 0.9
Figure 22. CNN, η = 0.003, β = 0.9
Figure 23. CNN, η = 0.001, β = 0.9
Figure 24. CNN, η = 0.01, β = 0.8
26
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
H. Further Experiments This appendix reports additional within-run dynamics grouped by ablation family. For each configuration we show batch sharpness, λmax , and training loss as a function of optimization steps, with subfigures corresponding to different batch sizes. Unless otherwise noted, all within-run dynamics are shown for batch sizes B ∈ {2, 4, 6, 8, 16, 32, 64, 128, 256, 8192}. Notice how for smaller batch sizes we do not always have convergence—as illustrated in the flatness of the loss. This happens due to excessive noise for selected batch and step size. H.1. Momentum ablations (MLP, η = 0.01)
(a) B = 2
(b) B = 4
(c) B = 6
(d) B = 8
(e) B = 16
(f) B = 32
(g) B = 64
(h) B = 128
(i) B = 256
(j) B = 8192 Figure A1. Within-run dynamics for an MLP trained with η = 0.01 and momentum β = 0.5 across batch sizes.
27
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
(a) B = 2
(b) B = 4
(c) B = 6
(d) B = 8
(e) B = 16
(f) B = 32
(g) B = 64
(h) B = 128
(i) B = 256
(j) B = 8192 Figure A2. Within-run dynamics for MLP with η = 0.01 and momentum β = 0.9 across batch sizes.
H.1.1. N OTE ON SMALLNESS OF λmax As a concrete illustration, consider an MLP trained with learning rate η = 0.01, momentum β = 0.9, and batch size B = 16. λ2max η 2 For this run, the final observed curvature satisfies λmax ≈ 12.76. Evaluating quantity (1−β) 3 yields a value of approximately 16.28.
28
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
H.2. Learning-rate ablations (MLP, β = 0.5)
(a) B = 2
(b) B = 4
(c) B = 6
(d) B = 8
(e) B = 16
(f) B = 32
(g) B = 64
(h) B = 128
(i) B = 256
(j) B = 8192 Figure A3. Within-run dynamics for MLP with η = 0.005 and momentum β = 0.5 across batch sizes.
29
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
H.3. Architecture ablations MLP We reuse Fig. A1 as the MLP reference. CNN: R E LU, η = 0.02, β = 0.5
(a) B = 2
(b) B = 4
(c) B = 6
(d) B = 8
(e) B = 16
(f) B = 32
(g) B = 64
(h) B = 128
(i) B = 256
(j) B = 8192 Figure A4. Within-run dynamics for CNN+ReLU with η = 0.02 and β = 0.5.
30
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
CNN: S I LU, η = 0.01, β = 0.5
(a) B = 2
(b) B = 4
(c) B = 6
(d) B = 8
(e) B = 16
(f) B = 32
(g) B = 64
(h) B = 128
(i) B = 256
(j) B = 8192 Figure A5. Within-run dynamics for CNN+SiLU with η = 0.01 and β = 0.5.
31
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
CNN: S I LU, η = 0.01, β = 0.8
(a) B = 2
(b) B = 4
(c) B = 6
(d) B = 8
(e) B = 16
(f) B = 32
(g) B = 64
(h) B = 128
(i) B = 256
(j) B = 8192 Figure A6. Within-run dynamics for CNN+SiLU with η = 0.01 and β = 0.8.
32
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
CNN: S I LU, η = 0.02, β = 0.5
(a) B = 2
(b) B = 4
(c) B = 6
(d) B = 8
(e) B = 16
(f) B = 32
(g) B = 64
(h) B = 128
(i) B = 256
(j) B = 8192 Figure A7. Within-run dynamics for CNN+SiLU with η = 0.02 and β = 0.5.
33
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
CNN: S I LU, η = 0.02, β = 0.8
(a) B = 2
(b) B = 4
(c) B = 6
(d) B = 8
(e) B = 16
(f) B = 32
(g) B = 64
(h) B = 128
(i) B = 256
(j) B = 8192 Figure A8. Within-run dynamics for CNN+SiLU with η = 0.02 and β = 0.8.
34
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
R ES N ET: R E LU, β = 0.5, η = 0.02
(a) B = 2
(b) B = 4
(c) B = 6
(d) B = 8
(e) B = 16
(f) B = 32
(g) B = 64
(h) B = 128
(i) B = 256
(j) B = 8192 Figure A9. Within-run dynamics for ResNet+ReLU with η = 0.02 and β = 0.5.
35
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
H.4. Activation ablations This section isolates activation function effects in an MLP at fixed optimizer settings: β = 0.5 and η = 0.01. MLP: R E LU, β = 0.5, η = 0.01
(a) B = 2
(b) B = 4
(c) B = 6
(d) B = 8
(e) B = 16
(f) B = 32
(g) B = 64
(h) B = 128
(i) B = 256
(j) B = 8192 Figure A10. Within-run dynamics for ReLU MLP (η = 0.01, β = 0.5) across batch sizes.
36
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
MLP: S I LU, β = 0.5, η = 0.01
(a) B = 2
(b) B = 4
(c) B = 6
(d) B = 8
(e) B = 16
(f) B = 32
(g) B = 64
(h) B = 128
(i) B = 256
(j) B = 8192 Figure A11. Within-run dynamics for SiLU MLP (η = 0.01, β = 0.5) across batch sizes.
37
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
V I T, β = 0.5, η = 0.007
(a) B = 2
(b) B = 4
(c) B = 6
(d) B = 8
(e) B = 16
(f) B = 32
(g) B = 64
(h) B = 128
(i) B = 256
(j) B = 8192 Figure A12. Within-run dynamics for a ViT (η = 0.007, β = 0.5) across batch sizes using the MSE loss.
38
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
CNN, β = 0.9, η = 0.001
(a) B = 2
(b) B = 4
(c) B = 6
(d) B = 8
(e) B = 16
(f) B = 32
(g) B = 64
(h) B = 128
(i) B = 256
(j) B = 8192 Figure A13. Within-run dynamics for a CNN (η = 0.001, β = 0.9) trained on the SST dataset across batch sizes using the MSE loss.
39
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
MLP, β = 0.75, η = 0.008
(a) B = 2
(b) B = 4
(c) B = 6
(d) B = 8
(e) B = 16
(f) B = 32
(g) B = 64
(h) B = 128
(i) B = 256
(j) B = 8192 Figure A14. Within-run dynamics for an MLP (η = 0.008, β = 0.75) trained on the SST dataset across batch sizes using the MSE loss.
40