Generalization at the Edge of Stability

2026-04-21

AI Summary

This work investigates the mechanism underlying improved generalization of neural networks trained with large learning rates near the edge of stability. By modeling stochastic optimizers as random dynamical systems, the study reveals that optimization trajectories converge to fractal attractors of low intrinsic dimensionality. Building on the Lyapunov dimension, the authors introduce a novel metric—“sharpness dimension”—to characterize generalization performance. This measure uniquely incorporates the full spectral structure of the Hessian and its partial determinants, overcoming limitations of conventional sharpness measures that rely solely on the trace or spectral norm. The theoretical framework is validated across diverse architectures, including MLPs and Transformers, and offers a new perspective on the “grokking” phenomenon.

Technology Category

• Machine Learning: Optimization
• Search and Optimization: Learning to Search
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• Security and Privacy: Large-scale security measurements
• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• User Modeling, Personalization and Recommendation: Metrics for user behavior and evaluating success

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.

Problem

• generalization
• edge of stability
• chaotic dynamics
• neural networks
• optimization

Innovation

• sharpness dimension
• edge of stability
• random dynamical systems
• fractal attractor
• Hessian spectrum