Momentum Further Constrains Sharpness at the Edge of Stochastic Stability

Recent work suggests that (stochastic) gradient descent self-organizes near the 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 \textit{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-\beta)/\eta$, 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+\beta)/\eta$, where momentum recovers its classical stabilizing effect and favors sharper regions consistent with full-batch dynamics. We further show this aligns with linear stability thresholds and we discuss the implications on hyperparameters tuning and coupling.

Modern neural networks are trained by repeatedly nudging their parameters to reduce errors - in such way we are trying to get lower and lower on the "loss landscape". A common trick called momentum keeps a running memory of previous nudges - much like a "heavy ball" that maintains velocity, aka the history of previous pushes. Researchers often think of momentum as making training more stable and allowing the model to move through steeper regions of the error surface. We asked whether this intuition still holds when each update/nudge is computed from a small random subset of the training data, rather than from the whole dataset (which is the regime used in practice) - i.e. in the presence of random noise from the choice of this subset. We measured how sharp the error surface looks along these noisy update directions and tested what happens when we deliberately perturb the trajectory. We find that momentum has two different roles: with large subset size it behaves as expected - it makes the training more stable, with velocity counteracting big nudges, but with small subset size momentum amplifies randomness and keeps training in flatter regions than the same method without momentum. When we perturb the trajectory slightly, the error briefly spikes, as the valley becomes too sharp and the ball jumps out, and then settles again, providing evidence that training continues in a regime that is "almost unstable". These results help explain how the trajectory looks like, how choice of the trajectory (nudge size, subset size and momentum) interact, and can guide more principled tuning of neural-network training.