Sharpness-Aware Minimization: Why Flat Minima Generalise Better Than Sharp Ones

Two neural networks can reach almost exactly the same training loss and yet behave very differently on new data. One generalises well. The other collapses on the first distribution shift, the first noisy label, or the first slightly different test set. A large part of the difference is not the loss value they reached. It is the shape of the loss landscape around that value.

Sharpness-Aware Minimization, usually shortened to SAM, makes that shape an explicit part of training. Instead of asking only "how low is the loss here?", it asks "how low does the loss stay nearby?" A flat minimum, where the loss stays low in a neighbourhood, is preferred over a sharp minimum, where the loss rises steeply as soon as the weights move a little. That single change of objective explains a surprising amount of SAM's practical behaviour, and also its limits.

The intuition: two minima, same loss, different futures

Picture a one-dimensional loss curve with two valleys. Both reach the same depth. One valley is narrow and steep. The other is wide and gentle. On the training data, the two solutions look identical: same loss, same accuracy.

Now imagine that the test data are not exactly the training data. In effect, the loss landscape shifts slightly, or the weights are perturbed slightly, because the model must work on inputs it has never seen. In the sharp valley, a small shift pushes the model up a steep wall, and the loss explodes. In the flat valley, the same shift moves the model along a gentle floor, and the loss barely changes.

This is the core intuition behind flat minima. Generalisation is not only about fitting the training set. It is about robustness to the small differences between training and test conditions. A solution that is stable under perturbation is, in this sense, a safer bet than a solution that is only good at one exact point.

That intuition needs care. Flatness alone is not a complete theory of generalisation. As discussed below, flatness can be redefined by rescaling parameters, and a flat minimum in one parameterisation can look sharp in another. SAM is best understood as a useful training bias, not as a law of nature.

From intuition to an objective

Let $w$ be the model parameters and $L(w)$ be the training loss. Standard training minimises $L(w)$ directly. SAM instead minimises the worst-case loss in a neighbourhood of $w$ :

$$L_{\mathrm{SAM}}(w) = \max_{|\epsilon| \leq \rho} L(w + \epsilon)$$

Here $\epsilon$ is a perturbation of the weights, $|\epsilon|$ is its size, and $\rho$ is the neighbourhood radius. The parameter $\rho$ controls how far SAM looks. If $\rho$ is tiny, SAM is almost the same as standard training. If $\rho$ is huge, SAM may sacrifice training fit to chase an extremely flat region that is not actually useful.

The minimisation problem therefore becomes:

$$\min_{w} \max_{|\epsilon| \leq \rho} L(w + \epsilon)$$

This is a minimax problem. The inner maximisation finds the worst nearby perturbation. The outer minimisation moves the weights so that even that worst nearby point has low loss.

Solving the inner problem exactly would be expensive. SAM uses a first-order approximation that is cheap and works well in practice. Expand $L(w + \epsilon)$ to first order:

$$L(w + \epsilon) \approx L(w) + \epsilon^{\top} \nabla L(w)$$

Maximising this linear approximation over $|\epsilon| \leq \rho$ gives a closed form. With the Euclidean norm, the worst perturbation points in the direction of the gradient:

$$\hat{\epsilon} = \rho \frac{\nabla L(w)}{|\nabla L(w)|}$$

The SAM gradient is then the gradient evaluated not at $w$ , but at the perturbed point $w + \hat{\epsilon}$ :

$$\nabla L_{\mathrm{SAM}}(w) \approx \nabla L(w + \hat{\epsilon})$$

In words: take a step in the gradient direction to find the worst nearby point, compute the gradient there, and use that gradient to update the original weights. A sharp minimum will have a large gradient at the perturbed point, because the loss rises steeply nearby, so the update will push away from it. A flat minimum will have a small gradient at the perturbed point, so the update will be gentle and the solution will stay put.

The algorithm, step by step

One SAM update, in its simplest form, looks like this:

1. Compute the gradient $\nabla L(w)$ at the current weights, on the current mini-batch.
2. Form the perturbation $\hat{\epsilon} = \rho \nabla L(w) / |\nabla L(w)|$ .
3. Compute the gradient at the perturbed weights, $\nabla L(w + \hat{\epsilon})$ , on the same mini-batch.
4. Update the weights using that second gradient, with your usual optimiser step, for example stochastic gradient descent with momentum or Adam.

That is the whole idea. The cost is clear: SAM needs two gradient computations per update instead of one, so each step costs roughly twice as much as a standard step. In practice, SAM often needs fewer steps to reach a given test accuracy, or reaches a better final accuracy for the same number of epochs, but the per-step cost is real and must be counted honestly when comparing methods.

A minimal PyTorch-style sketch makes the two-pass structure explicit:

# w are the model parameters, rho is the neighbourhood radius.
# loss_fn returns the training loss on the current mini-batch.

def sam_step(model, optimiser, loss_fn, rho):
    # First pass: gradient at the current weights.
    loss = loss_fn(model)
    loss.backward()
    grads = [p.grad for p in model.parameters()]

    # Form the worst-case perturbation in the gradient direction.
    norm = sum((g ** 2).sum() for g in grads) ** 0.5
    with torch.no_grad():
        for p, g in zip(model.parameters(), grads):
            p.add_(rho * g / (norm + 1e-12))

    # Second pass: gradient at the perturbed weights.
    optimiser.zero_grad()
    loss_perturbed = loss_fn(model)
    loss_perturbed.backward()

    # Restore the original weights, then step with the perturbed gradient.
    with torch.no_grad():
        for p, g in zip(model.parameters(), grads):
            p.sub_(rho * g / (norm + 1e-12))
    optimiser.step()
    optimiser.zero_grad()

This sketch omits optimiser details, adaptive variants, and distributed-training bookkeeping, but it shows the essential point: SAM is not a new optimiser in the usual sense. It is a wrapper around an existing optimiser that changes which gradient is used for the step.

Why it helps, when it helps

SAM helps most when the gap between training and test performance is driven by sensitivity to perturbation. That includes several familiar settings.

First, noisy labels. A sharp minimum often fits label noise by contorting the decision boundary in a fragile way. A flat minimum is less able to do that, because fitting a noisy point exactly usually requires a sharp feature. By penalising sharpness, SAM makes the model less willing to memorise noise at the expense of the clean signal.

Second, small data and large models. When a model has far more capacity than the data need, there are many solutions with near-zero training loss. Standard training may land in any of them, including a sharp one. SAM biases the search toward the flatter members of that set, which tend to be more stable.

Third, distribution shift. If test inputs differ slightly from training inputs, the effective loss landscape shifts. A flat solution tolerates that shift better, for exactly the valley picture described above. SAM does not solve distribution shift, and it does not replace data augmentation or domain adaptation, but it reduces one source of fragility.

It is worth noting what SAM does not do. It does not change the model architecture, the data, or the loss function itself. It changes the training dynamics. Two runs with the same data and architecture can therefore end in different places purely because one used SAM and the other did not, even though both minimised the same nominal loss.

Limitations and common pitfalls

SAM is not a free improvement, and its failures are instructive.

The radius $\rho$ matters, and there is no universal value. Too small, and SAM does almost nothing. Too large, and it underfits, because it refuses to enter any region that is not extremely flat, even if that region fits the data well. In practice, $\rho$ is tuned on validation data, like a learning rate or weight decay, and the best value depends on the model, the batch size, and the task. Reporting a SAM result without reporting $\rho$ makes the result hard to reproduce.

Sharpness is not invariant to reparameterisation. Rescaling the parameters of a network, for example by scaling one layer up and another down in a way that leaves the function unchanged, can change the measured sharpness without changing the model at all. This means that "flat" is partly a property of the parameterisation, not only of the function. Later variants of SAM address this by using adaptive, scale-aware neighbourhoods, where the perturbation size is scaled per parameter. If you use SAM on a model with very different parameter scales across layers, an adaptive variant is usually the safer choice.

SAM costs twice the gradient computation per step. If you compare SAM for 100 epochs against standard training for 100 epochs, SAM has used roughly twice the compute. A fair comparison either matches total compute, or states the compute difference plainly. Many apparent SAM gains shrink when compute is matched, even if a real gain remains.

SAM interacts with other regularisers. Weight decay, dropout, label smoothing, data augmentation, and SAM all push the model toward simpler or more robust solutions, but they do so in different ways, and stacking all of them at full strength can underfit. When adding SAM to an already heavily regularised pipeline, it is usually better to retune the other regularisers downward rather than simply adding SAM on top.

Finally, SAM does not fix a bad loss, bad data, or a mismatched architecture. If the training labels are systematically wrong, if a key subgroup is missing from the data, or if the model cannot represent the true function, seeking a flat minimum will find a flat version of a bad solution. Flatness is a second-order concern. It improves a solution that is already roughly right. It does not create one from nothing.

Practical takeaways

Use SAM when you have a generalisation gap that looks like fragility: training accuracy is high, test accuracy lags, and small changes in data or seed move test accuracy a lot. That is the signature of a sharp-minimum problem, and SAM directly targets it.

Tune $\rho$ on validation data, and try an adaptive variant if your model has layers with very different scales. Start from the standard SAM radius range used for your architecture family, then move up if the model still overfits sharply, and down if training accuracy drops noticeably.

Count the compute. Budget SAM as two gradient passes per step, and compare it against standard training at matched compute before declaring it a win for your task. If SAM still wins at matched compute, you have a real result. If it only wins because it used twice the compute, say so.

Do not treat flatness as a certificate. A flat minimum is a good sign, not a guarantee. Validate on held-out data, check calibration and subgroup performance, and test robustness to the shifts you actually care about. SAM improves the odds. It does not remove the need to measure.

The takeaway

Standard training asks for a point with low loss. SAM asks for a neighbourhood with low loss. That change, implemented as a gradient step at a perturbed point, biases training toward flat minima that tolerate the small differences between training and test conditions. It costs an extra gradient pass, it needs its radius tuned, and its notion of sharpness depends on parameterisation, so it is not a cure-all. But when generalisation fails because the solution is fragile rather than simply wrong, seeking flatness is one of the most direct and interpretable fixes available. The loss value tells you where you are. The shape around it tells you how long you are likely to stay there.

---

Connect with me on LinkedIn: linkedin.com/in/umeshtharukaofficial