Train Neural Networks with Sharpness-Aware Minimization

A neural network can reach low training loss at parameter values where a small change to the weights makes the loss rise sharply. Standard optimization does not directly discourage this behavior: it mainly asks whether the loss is low at the current parameters.

Sharpness-Aware Minimization (SAM) changes the training objective. Instead of optimizing only the loss at the current weights, it approximately optimizes the worst loss in a small neighborhood around them. The practical idea is simple: find a nearby parameter perturbation that makes the current mini-batch harder, then update the original model using the gradient measured at those perturbed parameters.

This article builds SAM from that mental model. You will see why one training step needs two gradient evaluations, how the perturbation radius changes the objective, what SAM does and does not guarantee, and when its extra training cost is worth considering.

Standard training optimizes one point in parameter space

Let a model have parameters w, and let L(w) be the loss on the current mini-batch. A conventional optimizer uses the gradient

g = gradient of L(w)

to choose an update that should reduce the loss.

This treats the current parameter vector as the point of interest. Imagine two candidate solutions with the same loss:

solution A: nearby weight changes -> loss stays similar
solution B: nearby weight changes -> loss rises quickly

Both can look equally good if you inspect only L(w). SAM asks a different question: how bad can the loss become after a small allowed change to the parameters?

That neighborhood matters because a low value at one exact point does not describe the local shape of the loss surface. SAM explicitly includes that local sensitivity in the optimization problem.

Replace point loss with neighborhood loss

A common form of the SAM objective is

minimize over w:  max over ||epsilon|| <= rho:  L(w + epsilon)

Here:

• w is the parameter vector being trained;
• epsilon is a temporary parameter perturbation;
• rho is the maximum allowed perturbation norm.

The inner problem searches the neighborhood for parameters with high loss. The outer problem then tries to make that nearby worst case smaller.

This is a local robustness objective in parameter space. It should not be confused with adversarial training on input examples, which perturbs inputs rather than model parameters.

The radius rho determines what “nearby” means. If rho is very small, the objective approaches ordinary local optimization. A larger radius asks the model to tolerate a wider parameter neighborhood, but it can also make the perturbation dominate the useful training signal. rho is therefore a hyperparameter, not a universally optimal constant.

Approximate the hardest nearby parameters

Solving the inner maximization exactly at every training step would generally be impractical. SAM uses a first-order approximation.

Suppose the gradient at the current parameters is g. For an L2-bounded perturbation, the first-order direction that increases the loss most rapidly is the normalized gradient direction. The perturbation is approximately

epsilon = rho * g / ||g||_2

when the gradient norm is nonzero.

Consider a two-parameter teaching example:

g = [3, 4]
rho = 0.1

The L2 norm of g is 5, so

epsilon = 0.1 * [3, 4] / 5
        = [0.06, 0.08]

The temporary parameters are therefore

w_perturbed = w + [0.06, 0.08]

The numbers are deliberately simple, but the important point is general: the first gradient is used to construct a local uphill perturbation. It is not yet the gradient used for the final optimizer update.

A SAM step needs two gradient evaluations

The simplest useful SAM step can be described as follows:

1. Compute loss L(w) on a mini-batch.
2. Compute g = grad L(w).
3. Build epsilon from g and rho.
4. Temporarily evaluate the model at w + epsilon.
5. Compute g_sam = grad L(w + epsilon).
6. Restore the original parameters w.
7. Apply the base optimizer update to w using g_sam.

The distinction in steps 4 through 7 is important. The perturbation is temporary. SAM does not permanently add epsilon and then optimize from that perturbed model. The base optimizer updates the original parameters using information obtained from the nearby higher-loss point.

Conceptually, the first backward pass asks:

Which nearby direction makes this batch worse?

The second asks:

From that difficult nearby point, which gradient should guide the real update?

That two-stage process is the core of SAM.

Why the extra gradient changes the update

If ordinary training and SAM both used the gradient at w, they would make essentially the same local decision. SAM instead measures the update direction after moving to w + epsilon.

Suppose the loss is highly sensitive in one region. The gradient at the perturbed point can differ substantially from the gradient at the original point. Optimizing with that second gradient encourages an update that reduces not just the original loss, but the elevated loss found nearby.

This is why SAM is often described as seeking parameter regions with lower neighborhood loss or lower sharpness. That description is more precise than saying it simply “finds flat minima.” Flatness has multiple mathematical definitions, can depend on parameterization, and does not by itself provide a universal guarantee about generalization.

SAM directly optimizes its chosen neighborhood-based objective. Any improvement in validation or test performance is an empirical outcome that should be measured for the model, data, and training setup at hand.

Keep the same mini-batch for both passes

A useful implementation detail follows from the objective. The perturbation is supposed to expose a difficult nearby parameter point for the loss being considered. If the first gradient comes from one mini-batch and the second gradient from an unrelated mini-batch, two sources of change become mixed together:

parameter perturbation
mini-batch sampling difference

Using the same mini-batch for both gradient evaluations keeps the second loss tied to the neighborhood constructed from the first loss.

Stochastic layers and data augmentation deserve similar attention. If two forward passes use substantially different random realizations, the measured loss difference contains both parameter sensitivity and randomness. Exact handling depends on the training framework and experiment, but the general principle is to know which randomness is intentional rather than assuming both passes evaluate identical functions.

The main cost is additional training compute

Ordinary mini-batch training typically performs one forward/backward gradient evaluation per optimizer step. Basic SAM performs two. That makes its training-time cost materially higher than the base optimizer.

The wall-clock increase is not guaranteed to be exactly 2x. Data loading, communication, compiler behavior, hardware utilization, and optimizer overhead also contribute to runtime. Still, the second forward/backward pass is real work and should be included in training budgets.

SAM does not inherently require two model copies at inference time. Once training is complete, the resulting model can be served like an ordinarily trained model with the same architecture. The extra cost is primarily a training concern unless the implementation introduces other changes.

This distinction matters when comparing alternatives. If training compute is scarce but inference cost is the main constraint, SAM may still be acceptable. If training is already the dominant cost, the additional pass can be a decisive disadvantage.

Tune the perturbation radius with the optimizer

rho changes the objective, so it should be validated rather than treated as a cosmetic setting.

With a very small rho, w + epsilon remains close to w, and the SAM gradient may resemble the ordinary gradient. With an excessively large value, the temporary point may represent a neighborhood that is not useful for the intended optimization problem, making training harder or hurting validation performance.

The appropriate value can interact with other choices, including:

• the base optimizer;
• learning rate and schedule;
• weight decay;
• batch size;
• architecture and normalization;
• the scale and parameterization of model weights.

For this reason, a fair experiment should not assume that hyperparameters tuned for the baseline transfer unchanged to SAM. At minimum, compare validation performance under a small, controlled search over rho and any learning-rate adjustments that the training setup requires.

Do not confuse SAM with gradient clipping

SAM and gradient clipping can both involve gradient norms, but they solve different problems.

Gradient clipping modifies a gradient when its norm is too large, commonly to control unstable updates. A simplified global-norm form is

g_clipped = g * min(1, threshold / ||g||)

SAM uses a normalized gradient to construct a parameter perturbation and then computes another gradient at the perturbed parameters.

So these operations answer different questions:

gradient clipping: should this update magnitude be limited?
SAM: what update looks useful after testing a difficult nearby parameter point?

One does not automatically replace the other. A training system may have reasons to use both, but their interactions should be tested rather than assumed harmless.

Common implementation mistakes

A SAM implementation can look plausible while optimizing something different from the intended procedure.

Updating from the perturbed weights

If the base optimizer step is applied while w + epsilon is still treated as the persistent model state, the temporary adversarial perturbation becomes part of the real update. Restore the original parameters before applying the base optimizer step.

Reusing the first gradient for the final update

The second gradient evaluation is the point of SAM. Constructing epsilon and then stepping with the original g does not measure how the gradient changes at the perturbed point.

Forgetting the zero-gradient case

The expression g / ||g|| is undefined when the norm is zero. Implementations normally need numerical protection, such as checking the norm or using an appropriate small stabilizer consistent with the chosen implementation.

Comparing training loss without comparing validation quality

SAM changes the training objective. Lower ordinary training loss is therefore not the only meaningful comparison. Track the task metric and validation loss that represent the actual deployment goal, together with training time and resource use.

Treating lower sharpness as a guarantee

The geometry of neural-network parameter spaces is subtle. A neighborhood metric can change under reparameterization, and a lower measured sharpness under one definition does not prove better behavior on unseen data. Treat SAM as an optimization method whose value must be established experimentally, not as a theorem that a trained model will generalize better.

When SAM is worth trying

SAM is a reasonable experiment when a model already trains successfully but validation quality is the limiting concern, and additional training compute is available. It is especially useful as a controlled comparison: keep the architecture and data pipeline fixed, replace the training step with SAM, tune the relevant hyperparameters, and measure whether the validation gain justifies the extra cost.

A simpler optimizer is usually the better starting point when the baseline is not yet healthy. If training diverges, labels are poor, the data split leaks information, the model is badly underfit, or the evaluation metric does not match the product goal, SAM addresses none of those root causes.

It is also less attractive when training throughput is the dominant constraint and a small possible quality gain would not repay an additional gradient evaluation per step.

A practical evaluation checklist

When testing SAM, compare it with a well-tuned baseline rather than an intentionally weak one. Record at least:

• the same train and validation splits;
• the same model architecture;
• base optimizer and learning-rate settings;
• rho and any SAM-specific choices;
• validation metric at the selected checkpoint;
• total training steps or examples processed;
• wall-clock training time and compute use.

If SAM improves the target metric, the result is easier to interpret when you can also state what extra training cost produced that improvement. If it does not improve the target metric, the neighborhood objective may simply not be a useful trade-off for that workload.

Conclusion

Sharpness-Aware Minimization changes neural-network training from optimizing loss at one parameter vector to approximately optimizing the worst loss in a small neighborhood around that vector. Its practical mechanism is a two-gradient step: use the first gradient to construct an uphill parameter perturbation, then use the gradient at the perturbed point to update the original weights.

That mechanism makes SAM easy to reason about and expensive enough to evaluate carefully. The perturbation radius changes what neighborhood the optimizer cares about, the second gradient adds substantial training work, and improved generalization is something to measure rather than assume. When a strong baseline is already in place, SAM offers a focused way to test whether local parameter robustness is worth additional training compute.