II · THE IDEA · ARTIFICIAL INTELLIGENCE
Sharpness-Aware Minimization (SAM)
▶ Listen · narrated
Two weight sets can reach the same training loss. One sits in a narrow ravine; the other in a broad valley. Which you find depends on the optimiser, and it changes test performance.
At a glance
- What it is
- An optimiser that prefers parameters whose neighbourhood has uniformly low loss
- Core idea
- Minimise worst-case loss inside a small ball around the weights
- Aim
- Better generalisation than ordinary empirical risk minimisation
- Rough cost
- About two gradient evaluations per parameter update
- Related
- Gradient descent often trains at the edge of stability
Think of the training loss as a landscape of hills and valleys, and the model's weights as a ball you are rolling into a valley floor. Many valley floors are equally low, so finishing with a low training loss does not tell you which valley you are in. Some valleys are narrow ravines: move a tiny step and you are halfway up a cliff. Others are wide bowls: you can walk around near the bottom and the ground stays low.
Ordinary gradient descent only cares how low the ground is exactly under the ball. Sharpness-aware minimisation also cares about the ground a short walk away in every direction. At each step it briefly pretends to stand at the worst nearby spot, checks the slope there, and uses that information to nudge the real ball toward broader bowls rather than narrow ravines. Those broader solutions tend to behave better on data the model has not seen, which is the practical point of the method.
SAM replaces the usual empirical-risk objective L(w) with a robust surrogate: minimise over w of max_{||ε||_p ≤ ρ} L(w + ε). The inner maximisation seeks the worst loss inside an ℓ_p ball of radius ρ. In the common p=2 case a first-order approximation gives ε* ≈ ρ · ∇L(w) / ||∇L(w)||_2. The outer step then descends on ∇L(w + ε*) rather than on ∇L(w), leaving w updated and the perturbed point discarded. The result is a parameter update whose direction accounts for local sharpness without forming a Hessian.
Computationally each iteration needs two backward passes. The radius ρ is a hyperparameter controlling the neighbourhood scale; too small and the method collapses toward ordinary SGD, too large and the linear approximation of the inner max degrades. Empirically the procedure finds minima with lower Hessian top eigenvalues and improved test error under matched training loss.
Independently, analyses of gradient descent on neural nets show training often proceeds at the edge of stability: sharpness rises until it approaches a threshold set by the learning rate, after which further sharpening is limited by the discrete dynamics. SAM's explicit neighbourhood penalty biases the trajectory away from the sharpest admissible configurations under that regime. Limitations include sensitivity to ρ, the fact that the dual-norm step only approximates the true inner max, and the broader caveat that flatness metrics are scale-dependent and not a full account of generalisation.
Look closer
The objective is a neighbourhood, not a point
Standard training asks only that the loss at the current weights be small. SAM asks something stricter: that the loss stay small everywhere inside a ball of fixed radius around those weights. A deep, narrow well fails that test even if its floor is low, because a short step away the loss rises sharply. A broader basin can pass it, and that is the geometry SAM is built to prefer.
Two gradients per step
The practical algorithm is a first-order approximation of that max-over-neighbourhood objective. One gradient is used to estimate the worst-case perturbation inside the ball; a second gradient, evaluated at the perturbed weights, is then applied as the actual update to the original parameters. The extra pass is the main overhead; no second-order matrices are required.
Flatness meets the edge of stability
Separate analysis of gradient descent on neural networks finds that training often settles at an edge of stability, where the sharpness of the loss sits near a threshold set by the learning rate. SAM's explicit penalty on sharp neighbourhoods pushes solutions away from the most unstable ridges that ordinary gradient descent can still occupy, which is one reason its generalisation behaviour differs.
The story
When a neural network is trained, many different parameter settings can drive the training loss to roughly the same low value. They are not equivalent. Some sit in sharp regions of the loss surface, where a small displacement in weight space sends the loss steeply upward. Others sit in flatter regions, where the loss stays low across a whole neighbourhood. Sharpness-aware minimisation is a deliberate attempt to find the second kind.
The motivating observation is simple enough to state and harder to enforce. If a solution is only good at a single point and brittle nearby, it is more likely to be an artefact of the finite training sample than a pattern that will hold on unseen data. A solution whose entire local neighbourhood is low-loss is, in that sense, more trustworthy. SAM turns that preference into an objective: instead of minimising the training loss at the current weights alone, it minimises an approximation of the maximum loss inside a small ball around them.
Formally the idea is to seek parameters that do well even after an adversary is allowed a bounded perturbation. Directly solving that inner maximisation at every step would be expensive. The method used in practice is a first-order shortcut. Compute the gradient of the loss, take a step of fixed radius in the direction that increases loss most (the dual-norm direction of that gradient), evaluate the gradient again at the perturbed point, and use that second gradient to update the original, unperturbed weights. The neighbourhood radius is a hyperparameter; the extra gradient is the main computational cost — roughly a doubling of the backward-pass work per update.
What this buys is not a lower training loss. Ordinary gradient methods are already good at driving training loss down. What changes is which of the many low-loss solutions the optimiser lands in. Solutions that sit in uniformly low neighbourhoods tend, on the evidence reported for SAM, to generalise better: the gap between training and test performance shrinks relative to standard empirical risk minimisation under comparable training budgets.
A related line of work on the dynamics of gradient descent itself finds that training on neural networks often occurs at the edge of stability. Sharpness, in that analysis, climbs until it sits near a threshold determined by the learning rate, and then training proceeds with progressive sharpening held in check by that limit. SAM's neighbourhood objective acts as an explicit pressure against the sharpest of those configurations. It does not remove the geometry of the loss surface, but it changes which parts of it are attractive under the update rule.
The method is therefore best read not as a claim that flatness is the only explanation of generalisation, but as a practical optimiser that builds a flatness preference into the training loop without requiring second-order information at every step. The neighbourhood is local, the radius is chosen rather than learned, and the inner maximisation is only approximated — yet the shift in the solutions found is large enough to matter on standard supervised benchmarks.
Why it mattered then
By the time SAM was proposed, it was already clear that overparameterised networks could fit training data in many ways, and that generalisation did not follow automatically from low training loss. Flatness had been discussed as a candidate correlate of generalisation, but turning that intuition into a scalable training algorithm was non-trivial: true second-order methods were expensive, and simply adding noise or weight decay was a blunt instrument. SAM offered a middle path — a first-order procedure whose objective explicitly rewarded uniformly low neighbourhoods — at a cost of roughly one extra gradient evaluation per step. That made the flatness preference something practitioners could try on ordinary hardware rather than a theoretical aside.
Why it matters now
Modern training still faces the same surplus of low-loss solutions, now at larger scale. Optimisers that change which basin is found, rather than only how fast the loss drops, remain one of the few levers that improve test behaviour without changing the architecture or the data. SAM and its variants continue to appear where generalisation under distribution shift or limited data is the bottleneck. The connection to edge-of-stability dynamics also keeps the idea relevant: as learning-rate schedules and batch sizes push models toward sharpness thresholds, an update rule that penalises sharp neighbourhoods is a direct counterweight rather than an incidental regulariser.
The surprising detail
The update does not descend on the gradient at the current weights. It descends on the gradient evaluated at a deliberately worsened neighbour — the point inside the ball where loss is estimated to be highest. The parameters that move are the original ones; the perturbed ones are discarded after the second backward pass. The algorithm improves the point you keep by measuring the slope at a point you throw away.
What is disputed
Flatness is a correlate of generalisation under particular definitions of neighbourhood and scale; it is not a complete theory. Different measures of sharpness can disagree, and the SAM radius is a chosen hyperparameter rather than a quantity the data identify. The edge-of-stability results describe typical gradient-descent dynamics; they support the relevance of sharpness but do not by themselves prove that SAM's neighbourhood objective is the unique or optimal remedy.
Remember this
SAM does not just chase low loss at a point; it prefers weights whose whole local neighbourhood stays low, and that preference often improves generalisation.
Test yourself
A run of ordinary SGD and a run of SAM both finish at nearly the same training loss. On a held-out test set the SAM model is clearly better. In terms of the loss surface and the update rule, what difference between the two solutions is SAM designed to produce, and how does the two-gradient step encourage it?
SAM is designed to land in a flatter region: one where loss remains low throughout a small ball around the weights, not only at the exact point. The first gradient finds a worst-case perturbation inside that ball; the second gradient, taken at the perturbed weights, updates the original parameters so that even that bad neighbour is driven down. Ordinary SGD only ever sees the loss at the current point, so it can happily settle in a sharp ravine that matches training loss but travels poorly to new data.
Go deeper
- [2010.01412] Sharpness-Aware Minimization for Efficiently Improving Generalization · arxiv.org
- [2103.00065] Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability · arxiv.org
Image: Original diagram, The Daily Triptych. Licence: Original work. Source.