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.

Browser lab#

On a 1D nonconvex loss $L(w) = w^4 - 4w^3 + 3w^2 + \sin(4w)$ (with multiple local minima), compare GD vs SAM-style steps. Plot trajectories and final positions to show SAM avoiding the sharp minima.

import numpy as np
import matplotlib.pyplot as plt

def loss(w):
    return w**4 - 4*w**3 + 3*w**2 + 0.5 * np.sin(4*w)

def grad(w):
    return 4*w**3 - 12*w**2 + 6*w + 2 * np.cos(4*w)

# GD and SAM from different initialisation
eta = 0.01
rho = 0.15
n_steps = 300

init_vals = [0.5, 2.0, 3.5]
fig, axes = plt.subplots(1, 3, figsize=(15, 5))

for idx, w0 in enumerate(init_vals):
    # GD
    w_gd = w0
    path_gd = [w_gd]
    for _ in range(n_steps):
        w_gd = w_gd - eta * grad(w_gd)
        path_gd.append(w_gd)

    # SAM
    w_sam = w0
    path_sam = [w_sam]
    for _ in range(n_steps):
        g = grad(w_sam)
        eps = rho * g / (np.abs(g) + 1e-8)  # ascent
        g_sam = grad(w_sam + eps)           # gradient at perturbed point
        w_sam = w_sam - eta * g_sam         # descent
        path_sam.append(w_sam)

    ax = axes[idx]
    w_plot = np.linspace(-0.5, 4.5, 500)
    L_plot = loss(w_plot)
    ax.plot(w_plot, L_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_final = w0
    w_sam_final = w0
    for _ in range(n_steps):
        w_gd_final -= eta * grad(w_gd_final)
        g_s = grad(w_sam_final)
        eps = rho * g_s / (np.abs(g_s) + 1e-8)
        w_sam_final -= eta * grad(w_sam_final + eps)
    print(f"  w0={w0}: GD -> w={w_gd_final:.4f}, L={loss(w_gd_final):.4f} | SAM -> w={w_sam_final:.4f}, L={loss(w_sam_final):.4f}")