Sharpness-Aware Minimization: Train for the Worst Loss in a Neighborhood, Not a Single Point

Ordinary training minimizes the loss L(w) at a single point in weight space, so it happily dives into the deepest hole it finds — often a sharp, narrow minimum where the loss is tiny at the bottom but rockets up the instant the weights are nudged. Here's the catch: the test loss is a slightly shifted copy of the train loss. A sharp minimum sits on a knife-edge — a small shift and the loss explodes — so it generalizes badly. A wide, flat minimum is robust: perturb the weights and the loss barely moves, so the train and test surfaces stay close.

Sharpness-Aware Minimization (Foret et al., 2021) bakes that in by changing the objective from "low loss" to "low loss in a whole ρ-neighbourhood":

min_w  max_{||δ|| ≤ ρ}  L(w + δ)

The inner max is one normalized ascent

The inner max asks "what is the worst loss within radius ρ of me?" To first order, L(w+δ) ≈ L(w) + δ·∇L, which is largest when δ points along the gradient. Normalize it to length ρ and you've solved the inner problem in closed form — no inner loop.

def ascent(w, rho):
    g   = grad(w)
    n   = np.linalg.norm(g) + 1e-12   # guard the zero-gradient case
    eps = rho * g / n                 # ||eps|| == rho exactly (unit dir x rho)
    return w + eps                    # w_adv: the worst nearby point

Ascend to the worst point, then descend the gradient there

The key move: climb to w_adv = w + ρ·ĝ, take the gradient at w_adv, then step from the original w along that perturbed gradient. Because that gradient is the gradient of the worst-case loss, descending it flattens the whole neighbourhood, not just the point.

def sam_step(w, lr, rho):
    w_adv = ascent(w, rho)            # step 1: worst point in the rho-ball
    g_adv = grad(w_adv)              # step 2: gradient THERE (2nd eval)
    return w - lr * g_adv            #         descend it from the ORIGINAL w

Contrast that with plain SGD, which uses the gradient at w (one eval) and knows nothing about the neighbourhood — so it dives into whatever minimum is nearest, sharp or not:

def sgd_step(w, lr):
    return w - lr * grad(w)          # gradient at w; sharpness invisible

So each SAM update is two gradient evaluations — about 2× the cost of a normal step — and it simply wraps any base optimizer. In the PyTorch form, first_step perturbs the weights to w_adv and stashes the offset; second_step restores w and lets the base optimizer (SGD, Adam, …) apply its normal rule using the gradient computed at w_adv.

Why it works, watchably

Run SGD and SAM from the same start on a landscape with a sharp deep dip and a wide flat basin. SGD slides straight into the sharp dip. SAM's worst-case step climbs the near wall, steps over the ridge, and settles in the flat basin — a slightly higher training loss but far lower sharpness, which is what generalizes. Shrink ρ toward 0.1 and SAM can no longer "see" over the dip's rim and gets trapped alongside SGD; push ρ up and it escapes ever more decisively. That single knob is the amount of sharpness-awareness you're buying.

Follow-ups sharpen the idea — ASAM makes ρ adaptive and scale-invariant, and efficient/Look-SAM variants amortize the extra pass. But the core is exactly this: swap "low loss here" for "low loss all around here," solve the inner max with one normalized ascent, and descend the gradient from the worst nearby point.

Drop a start point, tune ρ, and watch SGD fall in while SAM steps over the ridge, live: https://dev48v.infy.uk/dl/day57-sharpness-aware-minimization.html