Non-Euclidean Gradient Descent Operates at the Edge of Stability

TL;DR

This paper extends the understanding of the Edge of Stability phenomenon in gradient descent to non-Euclidean norms, providing a unified, geometry-aware spectral measure that explains training dynamics across various optimizers.

Contribution

It introduces a generalized sharpness measure applicable to arbitrary norms, extending the EoS analysis beyond Euclidean settings and encompassing multiple optimization methods.

Findings

Non-Euclidean GD exhibits EoS behavior similar to Euclidean GD.

The generalized sharpness measure captures optimizer dynamics across different geometries.

Experiments confirm progressive sharpening and oscillations around the stability threshold.

Abstract

The Edge of Stability (EoS) is a phenomenon where the sharpness (largest eigenvalue) of the Hessian converges to $2/\eta$ during training with gradient descent (GD) with a step-size $\eta$. 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/\eta$. Practically, our framework provides a single, geometry-aware spectral measure that works across optimizers.