Non-Euclidean Gradient Descent Operates at the Edge of Stability

TL;DR

This work provides an interpretation of EoS through the lens of Directional Smoothness Mishkin et al.

Abstract

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.

Figures (20)

• Figure 1: (Vanilla GD) Train loss, gradient norm, directional smoothness, and sharpness during training MLP ( top) and CNN ( bottom) models on CIFAR10-5k dataset with vanilla GD. Horizontal dashed lines correspond to the value $2/\eta$.
• Figure 2: ($\ell_{\infty}$-descent) Train loss, gradient norm, directional smoothness, and generalized sharpness \ref{['eq:max-inner-norm']} during training MLP on CIFAR10-5k (top) and Transformer on Tiny Shakespeare (bottom) with $\ell_{\infty}$-descent. Horizontal dashed lines correspond to the value $2/\eta$.
• Figure 3: ( Block CD) Train loss, gradient norm, directional smoothness, and generalized sharpness \ref{['eq:sharpness-12norm']} during training MLP (top) and CNN (bottom) models on CIFAR10-5k dataset with Block CD. Horizontal dashed lines correspond to the value $2/\eta$.
• Figure 4: ( Spectral GD) Train loss, gradient norm, directional smoothness, and generalized sharpness \ref{['eq:sharpness-spectral']} during training MLP (top, CIFAR10) and Transformer (bottom, Tiny Shakespeare) models with the Spectral GD. Horizontal dashed lines correspond to the value $2/\eta$.
• Figure 5: (Normalized non-Euclidean GD) Gradient norm, train loss, directional smoothness (normalized by the dual gradient norm), and generalized sharpness (normalized by the dual gradient norm) during training a CNN model with SignGD (CIFAR10-5k dataset, top line) and Muon without momentum (CIFAR10 dataset, bottom line). Horizontal dashed lines correspond to the value $2/\eta$.

Theorems & Definitions (36)

• definition 1
• definition 2
• definition 3
• definition 4
• theorem 4
• theorem 4
• lemma 4
• lemma 5
• proof
• lemma 6