Generalization at the Edge of Stability

Paper: 2604.19740 Authors: Mario Tuci, Caner Korkmaz, Umut Şimşekli, Tolga Birdal (INRIA, Imperial College London) Categories: cs.LG, stat.ML

Problem

Training modern neural networks often relies on large learning rates operating at the edge of stability (EoS), where optimization dynamics exhibit oscillatory and chaotic behavior. This regime often yields improved generalization, yet the underlying mechanism remains poorly understood.

Key Insight

Modeling stochastic optimization as a random dynamical system (RDS), which often converges to a fractal attractor set (rather than a point) with a smaller intrinsic dimension.

Contribution: Sharpness Dimension (SD)

Inspired by Lyapunov dimension theory, introduces a novel notion of dimension that measures:

• RDS Sharpness: Expansion and contraction rates of the attractor
• Sharpness Dimension (SD): Effective dimensionality of the random attractor

Key Theorem

The worst-case generalization error is governed by SD rather than:

• Parameter count (d)
• Trace of Hessian
• Spectral norm

At EoS: SD<d, meaning generalization is controlled by a provably lower-dimensional attractor.

Results

1. Generalization Bound: Rigorous connection between fractal dimension and generalization error
2. Empirical Validation: Across MLPs and transformers (GPT-2)
3. Grokking Analysis: New insights into delayed and sudden generalization phenomenon

Takeaways

• Overparameterized models can generalize beyond classical complexity measures
• Large learning rates create fractal attractors with low intrinsic dimensionality
• Sharpness dimension provides principled explanation for EoS generalization