Week 12: Sharpness and SAM

Loss Geometry Beyond Value#

Two weight vectors $w_1$ and $w_2$ can achieve the same training loss $L(w_1) \approx L(w_2)$ but differ dramatically in test performance. Why?

Consider a 1D loss $L(w)$ with a broad, shallow basin at $w_1$ and a narrow, steep basin at $w_2$. Both are local minima with similar loss values. But:

• At $w_1$ (flat): a small perturbation $\varepsilon$ raises the loss by $O(\|\varepsilon\|^2)$ with a small coefficient — the Hessian has small eigenvalues.
• At $w_2$ (sharp): the same perturbation raises the loss by a much larger amount — the Hessian has large eigenvalues.

A model at a flat minimum is robust to parameter perturbations. Since inference involves numerical noise, quantisation, and distribution shift, flatness is a proxy for stability. The empirical observation: flat minima tend to generalise better.

Measuring Sharpness#

The sharpness of the loss at $w$ within a neighbourhood of radius $\rho$ is:

$S_\rho(w) = \max_{\|\varepsilon\| \leq \rho} L(w + \varepsilon) - L(w)$

$S_\rho(w)$ measures the worst-case loss increase from a bounded perturbation. A small $S_\rho$ means the basin is flat; a large $S_\rho$ means steep walls surround the minimum.

For small $\rho$, a second-order approximation gives:

$S_\rho(w) \approx \frac{\rho^2}{2} \lambda_{\max}(\nabla^2 L(w))$

where $\lambda_{\max}$ is the largest Hessian eigenvalue. Sharpness is dominated by the most sensitive direction — a single sharply curving direction makes the entire basin sharp, even if other directions are flat.

SAM#

Sharpness-Aware Minimisation (SAMSharpness-Aware Minimization) explicitly optimises for flat minima by solving a min-max problem:

$\min_w \; \max_{\|\varepsilon\| \leq \rho} L(w + \varepsilon)$

The inner maximisation finds the worst perturbation $\varepsilon^*$ within radius $\rho$. The outer minimisation drives down the loss at the perturbed point, which forces the optimiser toward regions where the loss is uniformly low in a neighbourhood — i.e., flat basins.

Practical SAM step (first-order approximation):

1. Ascent: Compute the gradient at $w$, and take a step in that direction: $\varepsilon^* = \rho \frac{\nabla L(w)}{\|\nabla L(w)\|}$ This is the worst-case perturbation to first order (the gradient direction is steepest ascent).

2. Descent: Compute the gradient at the perturbed point $w + \varepsilon^*$, and apply the optimiser update: $w \leftarrow w - \eta \,\nabla L(w + \varepsilon^*)$

The key difference from standard SGD/AdamAdaptive Moment Estimation (optimizer): the gradient is evaluated at $w + \varepsilon^*$, not at $w$. This costs roughly twice as much per step (two forward/backward passes), but the improvement in generalisation often justifies the overhead. In practice, SAM is typically applied on top of an existing optimiser (SGDStochastic Gradient Descent or AdamAdaptive Moment Estimation (optimizer)), replacing the gradient in step 2.

Role of $\rho$. The radius $\rho$ controls how wide a neighbourhood is considered:

• Too small: SAM reduces to standard training (no sharpness penalty).
• Too large: the perturbation overshoots the basin entirely, and the ascent step may point toward a different basin — unhelpful noise.
• A typical range for $\rho$ in image classification is $0.05$ to $0.2$ (for normalised inputs); it requires tuning.

Large-Batch Connection#

Week 10 noted that large-batch training often converges to sharper minima. The mechanism: large batches produce lower-variance gradient estimates, which behave more like full-batch gradient descent. Full-batch GD gravitates toward the nearest local minimum — often a sharp one — while noisy SGDStochastic Gradient Descent can "bounce out" of sharp basins and into flatter ones.

SAM can be seen as explicitly injecting the flatness-seeking behaviour that small-batch noise provides implicitly. Several studies have shown that SAM + large-batch training can recover (or exceed) the generalisation of small-batch training while retaining the computational benefits of large batches.

Limits of the Story#

The flatness–generalisation connection, while empirically useful, has well-documented caveats:

1. Flatness definitions disagree. Hessian-based sharpness, adversarial sharpness ($\max_{\|\varepsilon\| \leq \rho}$), and PAC-Bayes sharpness can rank minima differently. There is no single agreed-upon metric.
2. Sharp minima can generalise. By rescaling weights (e.g., multiplying one layer by $\alpha$ and the next by $1/\alpha$), one can make a network arbitrarily sharp or flat without changing its function — a reminder that parameter-space geometry does not perfectly capture function-space behaviour.
3. Correlation, not causation. Flatness and generalisation are correlated in standard training regimes, but correlation weakens under distribution shift, architecture changes, and non-standard data augmentation.
4. SAM is not a silver bullet. It improves generalisation in many settings but adds computational cost and an extra hyperparameter ($\rho$). For some tasks, standard training with good regularisation (weight decay, data augmentation, dropout) matches or exceeds SAM.

Despite these caveats, the flatness lens is a productive diagnostic tool — when your model is not generalising as expected, checking the sharpness of the converged solution can reveal whether the optimiser is getting stuck in overly narrow basins.

Knowledge check#

Browser lab#

On a 1D loss with both a deep sharp basin and a wide flat basin, compare GD vs SAM-style steps from the same three starts. Plot the trajectories and final positions. The interesting panel is the one started on the ridge between the two basins: GD rolls down into the sharp minimum, while SAM's ascent step lands on the far side of the ridge, so its descent points into the flat basin instead. SAM pays for this — it ends with a higher loss than GD, which is the honest trade: it buys a flatter neighbourhood, not a better training loss.

import numpy as np
import matplotlib.pyplot as plt

# A landscape with BOTH kinds of basin: a deep narrow one (curvature ~24) and a
# wide shallow one (curvature ~2.4). Both are genuine local minima with a ridge
# between them, so "which minimum an optimizer finds" is a real question.
def loss(w):
    return 0.02 * w**2 - 3.0 * np.exp(-4.0 * (w + 1.0)**2) - 2.4 * np.exp(-0.5 * (w - 0.2)**2)

def grad(w):
    return (0.04 * w
            + 24.0 * (w + 1.0) * np.exp(-4.0 * (w + 1.0)**2)
            + 2.4 * (w - 0.2) * np.exp(-0.5 * (w - 0.2)**2))

eta, rho, n_steps = 0.05, 0.2, 400
init_vals = [-1.4, -0.1, 0.7]

def sam_escape(w0):
    """GD and SAM trajectories from the same start."""
    w_gd, w_sam = w0, w0
    path_gd, path_sam = [w_gd], [w_sam]
    for _ in range(n_steps):
        w_gd = w_gd - eta * grad(w_gd)
        g = grad(w_sam)
        eps = rho * g / (np.abs(g) + 1e-8)   # ascent step of size rho
        w_sam = w_sam - eta * grad(w_sam + eps)
        path_gd.append(w_gd)
        path_sam.append(w_sam)
    return path_gd, path_sam

fig, axes = plt.subplots(1, 3, figsize=(15, 5))
final = {}
for idx, w0 in enumerate(init_vals):
    path_gd, path_sam = sam_escape(w0)
    final[w0] = (path_gd[-1], path_sam[-1])

    ax = axes[idx]
    w_plot = np.linspace(-2.0, 1.6, 500)
    ax.plot(w_plot, loss(w_plot), "k-", linewidth=1, alpha=0.4, label="Loss $L(w)$")
    ax.scatter(path_gd, loss(np.array(path_gd)), c="blue", s=4, alpha=0.5, label="GD")
    ax.scatter(path_sam, loss(np.array(path_sam)), c="red", s=4, alpha=0.5, label="SAM")
    ax.plot(path_gd[-1], loss(path_gd[-1]), "bo", markersize=8, label=f"GD final ({path_gd[-1]:.2f})")
    ax.plot(path_sam[-1], loss(path_sam[-1]), "ro", markersize=8, label=f"SAM final ({path_sam[-1]:.2f})")
    ax.set_xlabel("$w$")
    ax.set_ylabel("$L(w)$")
    ax.set_title(f"Start $w_0 = {w0}$")
    ax.legend(fontsize=8)
    ax.grid(True, alpha=0.2)

plt.tight_layout()
plt.show()

print("Final positions and losses:")
for w0 in init_vals:
    w_gd, w_sam = final[w0]
    print(f"  w0={w0:4}: GD -> w={w_gd:.3f}, L={loss(w_gd):.3f}"
          f" | SAM -> w={w_sam:.3f}, L={loss(w_sam):.3f}")

# How sharp is each basin? Max loss rise inside a radius-rho ball around its floor.
print("\nBasin sharpness (max loss rise within rho):")
for name, w0 in [("sharp basin", -1.4), ("flat basin", 0.7)]:
    wb = w0
    for _ in range(n_steps):
        wb = wb - eta * grad(wb)
    print(f"  {name}: w={wb:.3f}, L={loss(wb):.3f}, "
          f"sharpness={max(loss(wb + rho), loss(wb - rho)) - loss(wb):.3f}")

print("Notice: from the ridge at w0=-0.1 GD rolls into the deep sharp basin, while SAM's "
      "ascent step lands on the flat side and it settles in the wide flat basin -- a higher "
      "loss (-2.41 vs -4.19) bought for a much flatter neighbourhood (measured sharpness "
      "0.031 vs 0.458).")

Further Reading#

• Pierre Foret, Ariel Kleiner, Hossein Mobahi & Behnam Neyshabur, "Sharpness-Aware Minimization for Efficiently Improving Generalization" (ICLR 2021) — the SAM paper. Freely available at arXiv:2010.01412.
• Sepp Hochreiter & Jürgen Schmidhuber, "Flat Minima" (Neural Computation, 1997) — the original flatness argument.
• Nitish Shirish Keskar et al., "On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima" (ICLR 2017) — sharpness and generalisation.