SGD at the Edge of Stability: Decoding the Stochastic Sharpness Gap

TL;DR

Why does training with smaller batches often lead to "flatter" and better-generalizing models? This paper provides a rigorous mathematical answer. While standard Gradient Descent orbits the "Edge of Stability" (where sharpness equals $2/\eta$), SGD settles into a flatter regime. The authors prove that gradient noise acts as a hidden hand, interacting with the loss landscape's 3rd-order geometry to push the model toward lower curvature.

Background: The Edge of Stability (EoS)

In classical optimization, we are taught that the learning rate $\eta$ must be smaller than $2/{\lambda }_{max}$ to ensure convergence. However, modern neural networks defy this. They undergo Progressive Sharpening until they hit the threshold ${\lambda }_{max}\approx 2/\eta$, and then they stay there, oscillating but still decreasing the loss. This is the Edge of Stability.

While this was well-understood for full-batch GD, mini-batch SGD showed a mysterious "gap": the sharpness was always significantly lower than $2/\eta$.

The Core Insight: Noise as a Suppressor

The authors argue that the "stochasticity" in SGD isn't just random jitter—it’s an active force for flatness. By projecting gradient noise onto the top Hessian eigenvector (the direction of maximum curvature), they show that this noise injects energy into the system's oscillations.

Because the loss landscape isn't perfectly quadratic (it has cubic terms), these noise-inflated oscillations trigger a "restoring force" that drives the sharpness down.

The Formal Mechanism

The authors define a Stochastic Coupling Theorem. They show that the SGD trajectory can be decomposed into:

1. A constrained "ideal" trajectory.
2. A stochastic oscillatory component.

In the figure above, note how different batch sizes reach different plateaus. The dashed line (Full Batch) is the highest, while smaller batches (b=50) settle much lower.

The Sharpness Gap Formula

The crowning achievement of the paper is a closed-form solution for the equilibrium sharpness:

$S_{eq}\approx \frac{2}{\eta }-\frac{\eta \beta {\sigma }_u^2}{4\alpha }$

Where:

• ${\sigma }_u^2$: Variance of the gradient noise along the top eigenvector.
• $\beta$: Self-stabilization strength (related to the 3rd derivative).
• $\alpha$: Progressive sharpening rate.

This formula reveals that the gap is proportional to the noise and inversely proportional to the batch size. As $bo\infty$ (Full Batch), the noise ${\sigma }^{2}o0$, and we recover the standard EoS of $2/\eta$.

Experimental Proof: The 1/b Scaling

The authors tested this on various architectures (MLPs, CNNs, ResNets). The data was remarkably consistent with their theory.

• Batch Sharpness Saturation: They discovered that while the Full-Batch sharpness is suppressed, the Batch Sharpness (curvature calculated on the specific mini-batch) actually hits the $2/\eta$ limit.
• Power Law: The gap between GD and SGD sharpness follows a perfect $1/b$ slope on a log-log scale.

Figure 2: The sharpness gap $\Delta S$ follows the theoretical $1/b$ reference slope almost perfectly.

Conclusion and Future Impact

This research bridges the gap between empirical "tricks" (like using small batches for better generalization) and rigorous optimization theory. It identifies Batch Sharpness as the true metric that governs stability in stochastic settings.

Takeaway for Practitioners: The "flatness" benefit of SGD is not a magical property—it is a result of how noise interacts with the high-order curvature of your loss function. This suggests that in tasks where EoS dynamics don't emerge (like standard Cross-Entropy on very small datasets), the batch-size benefit might vanish, as seen in the authors' CNN+CE experiments.