[2604.19740] Generalization at the Edge of Stability
Abstract:Training modern neural networks often relies on large learning rates, operating at the edge of stability, where the optimization dynamics exhibit oscillatory and chaotic behavior. Empirically, this regime often yields improved generalization performance, yet the underlying mechanism remains poorly understood. In this work, we represent stochastic optimizers as random dynamical systems, which often converge to a fractal attractor set (rather than a point) with a smaller intrinsic dimension. Building on this connection and inspired by Lyapunov dimension theory, we introduce a novel notion of dimension, coined the `sharpness dimension', and prove a generalization bound based on this dimension. Our results show that generalization in the chaotic regime depends on the complete Hessian spectrum and the structure of its partial determinants, highlighting a complexity that cannot be captured by the trace or spectral norm considered in prior work. Experiments across various MLPs and transformers validate our theory while also providing new insights into the recently observed phenomenon of grokking.
View PDF HTML (experimental) Abstract:Training modern neural networks often relies on large learning rates, operating at the edge of stability, where the optimization dynamics exhibit oscillatory and chaotic behavior. Empirically, this regime often yields improved generalization performance, yet the underlying mechanism remains poorly understood. In this work, we represent stochastic optimizers as random dynamical systems, which often converge to a fractal attractor set (rather than a point) with a smaller intrinsic dimension. Building on this connection and inspired by Lyapunov dimension…
saved by
related reading
- Neural population geometry and optimal coding of tasks with shared latent structurenature.com
- Neural network training makes beautiful fractals | Jascha’s blogsohl-dickstein.github.io
- Understanding deep learning requires rethinking generalizationarxiv.org
- [2605.01172] A Theory of Generalization in Deep Learningarxiv.org
- Statistical Mechanics of Deep Learningganguli-gang.stanford.edu
- [2604.21691] There Will Be a Scientific Theory of Deep Learningarxiv.org
- arxiv.org/pdf/1805.08522arxiv.org
- nn-notes.pdfboris-hanin.github.io
- A Theory of Deep Learning | Elements of a Vector Spaceelonlit.com
- [1912.02178] Fantastic Generalization Measures and Where to Find Themarxiv.org
- [2007.12826] The Interpolation Phase Transition in Neural Networks: Memorization and Generalization under Lazy Trainingarxiv.org
- Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasetsmathai-iclr.github.io