The Edge of Stability

Classical optimisation says that gradient descent on a smooth function with learning rate \(\eta\) converges only if the largest Hessian eigenvalue stays below \(2/\eta\). Above that threshold the iterates oscillate and diverge along the sharpest direction. Every convergence proof for gradient descent assumes you are below it.

Neural network training is not below it. It goes above it, on purpose, and stays there.

What was measured

Cohen et al. ran full-batch gradient descent on standard architectures and simply tracked \(\lambda_{\max}\), the top eigenvalue of the training loss Hessian, throughout training (Cohen et al., 2021, Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability, arXiv:2103.00065, ICLR 2021). Two phases appeared reliably.

Progressive sharpening. Early in training, \(\lambda_{\max}\) rises steadily. Nothing in the objective asks for this; the trajectory drifts into progressively sharper regions of the landscape on its own.

The edge. Once \(\lambda_{\max}\) reaches \(2/\eta\), it stops rising and hovers just above it, for the rest of training. The loss stops decreasing monotonically and starts behaving non-monotonically over a handful of steps, while still decreasing reliably over hundreds. Training does not diverge.

The key relation is worth stating precisely: the stability threshold depends on the learning rate you chose, so

is a self-organised outcome. Halve the learning rate and the network finds a region twice as sharp. Curvature is not a property of the problem that the optimiser must respect; it is something the optimiser negotiates.

Why it does not blow up

The naive picture says an unstable direction grows exponentially. It does grow, but the growth is self-limiting, and the mechanism is third-order. As the iterate oscillates along the sharpest eigendirection, it moves into a region where \(\lambda_{\max}\) is lower, because sharpness itself varies across the landscape. The oscillation reduces the curvature that caused it, which damps the oscillation, which lets sharpening resume. Damian et al. call this self-stabilisation and derive it as an implicit constraint on \(\lambda_{\max}\) arising from the negative gradient of sharpness (Damian et al., 2023, Self-Stabilization, arXiv:2209.15594, ICLR 2023).

The effect is a genuine implicit regulariser. Training at the edge biases the trajectory toward flatter regions, without any explicit sharpness penalty. This is one reason large learning rates generalise better than small ones at matched final training loss, and one reason sharpness-aware minimisation, which penalises curvature explicitly, gives smaller gains than its motivation suggests: gradient descent was already doing a version of it.

What this changes in practice

Warmup has a mechanistic account. At initialisation, \(\lambda_{\max}\) can already exceed \(2/\eta\) for the target learning rate. Starting small keeps you below threshold while the network sharpens into a region where the full rate is survivable. That is a sharper story than "it helps stability", and it predicts that the required warmup length scales with initial curvature rather than with dataset size. See warmup and why it helps.

Non-monotone loss is not a bug. A loss curve that ratchets down while wobbling over a five-step window is the expected signature of edge-of-stability training. Alerting on short-window loss increases produces false alarms; alerting on a rising trend over hundreds of steps does not. Distinguish this from real loss spikes, which are orders of magnitude larger and correlate with attention logit growth.

Learning rate and curvature are one knob. Because \(\lambda_{\max}\) tracks \(2/\eta\), the learning rate is choosing the sharpness of the solution, not just the step size. That reframes learning rate schedules: the decay phase of a WSD schedule is a controlled walk into sharper, lower-loss regions.

When it breaks

The original result is for full-batch gradient descent on the squared or cross-entropy loss. Stochastic mini-batch training shows a related but distinct phenomenon, often called the edge of stochastic stability, where the controlled quantity involves batch-gradient noise rather than \(\lambda_{\max}\) alone, and the effective threshold is lower.

Adaptive optimisers complicate it further. Adam preconditions the gradient, so the relevant operator is not the raw Hessian but the Hessian in the preconditioned geometry, and the observed threshold moves. The qualitative story survives, which is why the practical advice above holds for Adam-trained transformers, but the clean \(2/\eta\) number does not transfer directly. Treat it as the right mental model and the wrong constant.