Curvature, Conditioning, and Preconditioning

Different directions tolerate very different movements, and the tolerance is not a metaphor. Run full-batch gradient descent and the largest eigenvalue of the training-loss Hessian climbs until it hovers just above 2/(step size), then stays there. Cohen, Kaur, Li, Kolter and Talwalkar measured that in 2021 and named the regime the Edge of Stability. So one scalar learning rate does more than feel uncomfortable in the steep direction. It sets a ceiling on how sharp that direction is allowed to get.

In a simple quadratic problem, the ratio between the largest and the smallest relevant curvature decides how hard one fixed step size will be. Deep networks are not fixed quadratics. The intuition survives contact with them anyway, and in the full-batch case the link between the step size and the steepest direction has been measured rather than assumed. A 2021 paper named that regime the Edge of Stability: “In this regime, the maximum eigenvalue of the training loss Hessian hovers just above the numerical value 2 / (step size), and the training loss behaves non-monotonically over short timescales, yet consistently decreases over long timescales.” The sharpest curvature does not sit still waiting to be measured. It rises until it meets the largest value the chosen step size can survive. Then it stays there. An independent group analysed the same regime in 2022 and reported that “sharpness stabilizes around 2/LR and loss goes up and down across iterations, yet still with an overall downward trend”.

Two things follow that a reader can use. The threshold is a number you can compute from your own configuration — 2/(step size) — not an impression. And a loss curve that rises across short windows while falling over long ones is the signature of that regime, not automatic evidence of a broken run.

Poor conditioning otherwise shows up as oscillation in some directions while others barely move. Parameterization, normalization, data scale and architecture all shape it. Two optimizers were built for exactly this problem. K-FAC came from Martens and Grosse in 2015; it approximates large blocks of the Fisher matrix as Kronecker products of two smaller matrices, and is “only several times more expensive to compute than the plain stochastic gradient”. Shampoo followed in 2018. It “maintains a set of preconditioning matrices, each of which operates on a single dimension, contracting over the remaining dimensions”, at a runtime per step its authors report as comparable to “simple gradient methods such as SGD, AdaGrad, and Adam”.

Each method uses a different approximation, and several of these rows now carry a measured price rather than an adjective.

Changing the coordinate system before optimization is the cheapest intervention, and the one with the largest published effect on step count. In 2015 Ioffe and Szegedy inserted normalization into a state-of-the-art image classification model: “Batch Normalization achieves the same accuracy with 14 times fewer training steps, and beats the original model by a significant margin.” An ensemble of batch-normalized networks reached 4.9% top-5 validation error (4.8% test error) on ImageNet; the conference copy gives that ensemble result as 4.82% top-5 test error. The cost written in the table — altered model behavior — is real. But the benefit is 14 times fewer steps on a named model, not merely "simpler gradient scales".

Momentum accumulates persistent directions and damps some oscillation. It works from temporal history alone and carries no coordinate matrix. A diagonal adaptive optimizer uses coordinate-wise gradient magnitudes as a scale, which suits heterogeneous coordinates, at the price of extra state and a basis-dependent notion of geometry.

The second-order row is the one usually dismissed on memory and computation. It has now been tested head-to-head. MLCommons announced the results of the inaugural AlgoPerf: Training Algorithms benchmark competition on 1 August 2024: 18 submissions from 10 teams, fixed hardware, multiple workloads, a $25,000 prize. The external-tuning winner was Distributed Shampoo, from a Meta team. It scored 0.78 and trained models about 28% faster than the baseline. The self-tuning winner, Schedule Free AdamW, scored 0.85 for an 8% speedup — a different ruleset, with the tuning budget written into the rules. The competition's own analysis states the conclusion plainly: “The winning submission in the external tuning ruleset, using Distributed Shampoo, demonstrates the effectiveness of non-diagonal preconditioning over popular methods like Adam, even when compared on wall-clock runtime.” The cost of richer curvature is still memory and computation. What is no longer safe is the assumption that the cost goes unrepaid in wall-clock time.

These observations motivate a conditioning diagnosis. They do not prove one. Two of them have been followed all the way to a measurement, and it is worth seeing what the measurement looked like.

The full Hessian spectrum of an ImageNet-scale network stopped being a thought experiment in 2019, when Ghorbani and colleagues built a scalable estimator for it. A change of scale could now be tested against the spectrum itself: “In non-batch normalized networks, we observe the rapid appearance of large isolated eigenvalues in the spectrum, along with a surprising concentration of the gradient in the corresponding eigenspaces. In a batch normalized network, these two effects are almost absent.” The clue and the cure sit in the same picture. Outlier eigenvalues, gradient piling into their eigenspaces, and both effects almost gone once the coordinates are normalized.

Layerwise imbalance has a documented case too. LARS trains ResNet on ImageNet in minutes using layerwise adaptive learning rates, and it performs poorly on attention models like BERT. That failure is what a “layerwise adaptive large batch optimization technique called LAMB” was built for, published in 2020: “In particular, for BERT training, our optimizer enables use of very large batch sizes of 32868 without any degradation of performance. By increasing the batch size to the memory limit of a TPUv3 Pod, BERT training time can be reduced from 3 days to just 76 minutes.” Note both halves. Per-layer rescaling was the fix. And a per-layer rule that worked on ResNet was the thing that failed first on BERT. A conditioning story has to be tested on the model in front of you.

Use controlled changes that target scale rather than everything at once. Two of the five steps below now have a published reference standard to imitate.

Step 2, normalizing inputs, is the intervention Ioffe and Szegedy measured in training steps — 14 times fewer to the same accuracy. It is also the intervention measured in the Hessian spectrum, where the large isolated eigenvalues are almost absent once normalization is present. "Did normalization help?" therefore has a step-count answer and a spectral answer. They can disagree.

Step 3, comparing optimizers fairly, is what AlgoPerf formalized: fixed hardware, multiple workloads, separate external-tuning and self-tuning rulesets, and scoring on wall-clock runtime. The scores of 0.78 for Distributed Shampoo and 0.85 for Schedule Free AdamW mean something only because the tuning budget was written into the rules. A comparison that leaves tuning effort unstated measures the tuner, not the optimizer.

If you can afford one measurement beyond loss curves, take the largest Hessian eigenvalue and hold it against 2/(step size). The Edge of Stability result tells you where it is expected to sit. That turns a vague sense of instability into a prediction that can fail.