Beyond Pointwise Sharpness: Understanding Generalization via the Sharpness Dimension

TL;DR

Why do neural networks generalize even when the training process is unstable and chaotic? This paper argues that at the Edge of Stability (EoS), the optimizer doesn't settle at a single point but explores a fractal attractor. The authors introduce the Sharpness Dimension (SD)—a measure of this attractor's complexity—and prove that this dimension, rather than the total number of parameters, dictates the generalization error.

Problem & Motivation: The Chaos at the Edge

Classical optimization theory suggests that for Gradient Descent to converge, the learning rate $\eta$ must be less than $2/{\lambda }_{max}$ (where ${\lambda }_{max}$ is the largest eigenvalue of the Hessian). However, modern practitioners often use large learning rates where ${\lambda }_{max}$ frequently exceeds $2/\eta$. This is the Edge of Stability.

In this regime, training is a wild ride: trajectories oscillate and exhibit deterministic chaos. Pointwise metrics like "Hessian Trace" or "Flatness" at a single final iterate become insufficient because the model is never truly "still." The authors propose that we should instead look at the entire set of solutions the optimizer visits in the long run.

Methodology: The Geometry of Attractors

The researchers treat optimization as a Random Dynamical System (RDS). Instead of a single minimum, they focus on the Pullback Random Attractor—a structured "snapshot" of where the model's weights end up after infinite training under a specific noise realization.

The Sharpness Dimension (SD)

They define a new spectral metric based on the Jacobian of the update step. While the first singular value ${\lambda }_1>0$ indicates expansion (instability), the global system remains dissipative (contractive overall). The Sharpness Dimension is the "sweet spot" volume index where the expansion and contraction balance out.

Figure 1: Conceptual view of the RDS. At EoS, the dynamics are confined to a lower-dimensional attractor (SD < d), providing a principled reason for generalization despite overparameterization.

Generalization Bounds

The paper's core theoretical contribution is a new generalization bound: ${\mathcal{G}}_s\le \mathcal{O}\left(\sqrt{\frac{{\dim }_S\mathcal{A}}{n}}\right)$ This formula directly replaces the ambient parameter dimension $d$ with the fractal Sharpness Dimension ${\dim }_S\mathcal{A}$. It provides a mathematical guarantee that if the attractor is thin and "fractal-like," the model will generalize well, even if it has millions of parameters.

Experimental Evidence: From MLPs to GPT-2

The authors validated their theory using both small-scale MLPs (where exact Hessians can be computed) and large-scale GPT-2 models (using Stochastic Lanczos Quadrature to estimate the spectrum).

Key Insights:

1. Superior Correlation: In the EoS regime, SD showed significantly higher Kendall's correlation with the actual generalization gap compared to traditional "Sharpness" or "Hessian Trace."
2. Tracking Grokking: The measure accurately captures "Grokking"—the sudden generalization after long periods of overfitting. During the transition, SD reflects the shift in the spectral structure of the attractor.
3. Transformer Stability: Experiments on GPT-2 showed that for different optimizers (SGD, AdamW), the RDS-based quantities remained consistently informative, whereas classical sharpness often failed for AdamW.

Figure 2: Correlation matrices showing that the proposed Sharpness Dimension (SD) is a much better predictor of generalization (green region) than traditional trajectory-based tools like PH-Dim.

Critical Analysis & Conclusion

Takeaway

The paper provides a rigorous bridge between chaos theory and deep learning. It suggests that the "Edge of Stability" is not just a numerical curiosity but a mechanism that naturally regularizes models by forcing them onto a lower-dimensional solution manifold.

Limitations

• Computational Cost: Even with SLQ, estimating the full Hessian spectrum for 1.5B+ parameter models remains expensive.
• Regularity Assumptions: The theory requires $C^2$ continuity of the loss, which may be challenged by certain non-smooth activations or discrete components.

Future Work

The next frontier is extending this to adaptive optimizers like Adam more deeply and understanding how specific architecture choices (like Attention or LayerNorm) influence the fractal dimension of the resulting attractor.