How does sharpness-aware minimization change the training objective, and what does it cost per step?

## The objective Ordinary training minimizes the empirical loss L(w) at the current weights w. Sharpness-aware minimization replaces that with a worst-case objective over a neighbourhood: minimize over w of max over ||eps|| <= rho of L(w + eps) In words: find weights such that *every* nearby weight vector within distance rho is also good. A narrow, deep crevice scores badly under this objective even though its floor is low, because a small displacement inside the ball climbs the wall. A wide basin scores well. So the objective directly targets the quantity flatness is about, instead of hoping the optimizer stumbles into a wide basin. In practice a weight-decay term is kept alongside it as usual. ## How the inner maximization is approximated The inner maximization is itself intractable — it is a constrained maximization over a high-dimensional ball at every step. It is approximated by linearizing the loss around w. To first order, L(w + eps) is about L(w) plus the inner product of eps with the gradient g at w, and the maximizer of a linear function over a ball of radius rho points along the gradient: eps_hat = rho * g / ||g|| So the recipe per update is: 1. Compute the gradient g at the current weights w (one forward and backward pass). 2. Form the perturbed weights w + eps_hat by stepping a distance rho along the normalized gradient — the *ascent* direction, deliberately uphill. 3. Compute the gradient at w + eps_hat (a second forward and backward pass). 4. Apply *that* gradient to the original weights w with the usual update rule. Step 4 is the part candidates get wrong. The perturbed point is only a probe; the update lands on the original weights. The intuition is that you descend using the slope seen from the worst nearby point, so you move in a direction that improves the whole neighbourhood rather than just where you stand. One further approximation is hidden here: differentiating the max objective exactly would also produce a term involving second derivatives, which is dropped. The method is therefore a heuristic surrogate for the worst-case objective, not an exact gradient of it — which is fine, because it demonstrably reduces measured sharpness. ## The cost Two forward-and-backward passes per update instead of one means about twice the compute and roughly twice the wall-clock per step. That is the whole cost story, and it has a sharp consequence: under a fixed compute budget you get half as many updates. Comparing sharpness-aware training against ordinary training at equal *steps* flatters it; the honest comparison is at equal *compute*, where the halved update count can eat the entire gain. Report both if you can. Several cheaper variants exist. The extra ascent pass can be run only every k steps, with ordinary updates in between. The perturbation can be computed on a subset of the batch rather than all of it. Both trade some of the effect for most of the cost. A related and slightly counter-intuitive detail: computing the perturbation separately on small chunks of the batch and averaging the resulting updates tends to work *better* than computing one perturbation from the whole batch. The worst case over a small subset is a stronger constraint than the worst case of an averaged loss, so the effective sharpness penalty is stronger. ## Choosing the radius The radius rho is the one hyperparameter that matters, and it needs a sweep on a logarithmic grid. Too small and the perturbed gradient is nearly the unperturbed one, so you have paid double for ordinary training. Too large and you are optimizing the loss at points far from any solution: training loss stops descending and the model underfits. The best value is also tied to the scale of the weights, because a fixed absolute radius means something different for a layer with large weights than for one with small weights — variants that scale the perturbation per-coordinate by the current weight magnitude exist for exactly this reason, and they make the radius easier to transfer between models. ## When to reach for it It is a reasonable tool when validation accuracy matters more than training throughput, when you already have evidence that your checkpoints are sharp, and when the schedule is long enough that halving the update count still converges. It is a bad first move when the model is under-trained, when the data pipeline is the bottleneck, or when a cheaper regularizer has not been tried. And it is not a substitute for a diagnosis: if the model is failing for a reason unrelated to the landscape, doubling per-step cost buys nothing.