Non-Euclidean Gradient Descent Operates at the Edge of Stability
This work provides a unified theoretical framework for understanding the Edge of Stability across a broader range of optimization methods, which is significant for researchers studying the dynamics of deep learning optimization.
This paper investigates the Edge of Stability (EoS) phenomenon in gradient descent, where the Hessian's largest eigenvalue converges to 2/η. The authors extend the EoS interpretation to non-Euclidean norms, defining a generalized sharpness measure that encompasses various optimizers, and demonstrate through experiments on neural networks that non-Euclidean GD also exhibits progressive sharpening and oscillations around or above the 2/η threshold.
The Edge of Stability (EoS) is a phenomenon where the sharpness (largest eigenvalue) of the Hessian converges to $2/η$ during training with gradient descent (GD) with a step-size $η$. Despite (apparently) violating classical smoothness assumptions, EoS has been widely observed in deep learning, but its theoretical foundations remain incomplete. We provide an interpretation of EoS through the lens of Directional Smoothness Mishkin et al. [2024]. This interpretation naturally extends to non-Euclidean norms, which we use to define generalized sharpness under an arbitrary norm. Our generalized sharpness measure includes previously studied vanilla GD and preconditioned GD as special cases, as well as methods for which EoS has not been studied, such as $\ell_{\infty}$-descent, Block CD, Spectral GD, and Muon without momentum. Through experiments on neural networks, we show that non-Euclidean GD with our generalized sharpness also exhibits progressive sharpening followed by oscillations around or above the threshold $2/η$. Practically, our framework provides a single, geometry-aware spectral measure that works across optimizers.