Schedules, Weight Averaging, and Flat Minima

The schedule is part of the algorithm

A beginner treats the learning-rate schedule as a side setting. At graduate level it is part of the optimizer's identity: the same update rule with a constant rate, a step decay, or a cosine curve produces qualitatively different trajectories and different final solutions. Early training wants a large rate to traverse the landscape quickly; late training wants a small rate to settle into a basin without rattling around its walls. The schedule is how you interpolate between those regimes.

There is also a subtle dynamical story. Modern analyses show large-batch full-network training often hovers at the edge of stability — the learning rate sits right where the sharpest curvature direction is barely stable, so the loss oscillates while still descending on average. See edge of stability for the full picture; the practical upshot is that a schedule which lowers the rate near the end is also lowering the curvature threshold the iterate is allowed to occupy, gently steering it toward flatter regions.

\lambda_{\max}\!\left(\nabla^2 L(\theta)\right) \;\approx\; \frac{2}{\eta}

At the edge of stability the largest Hessian eigenvalue — the sharpness — hovers near two over the learning rate.

Cosine annealing and warm restarts

Cosine annealing decays the learning rate from its peak to near zero along a half-cosine curve. Compared with step decay it spends more time at intermediate rates and lands softly, which empirically yields cleaner convergence for deep networks. It is almost always paired with a short linear warmup at the very start — a few hundred to a few thousand steps — that ramps the rate up from near zero so the early, high-variance updates do not blow up freshly initialized weights.

\eta_t = \eta_{\min} + \tfrac{1}{2}\left(\eta_{\max}-\eta_{\min}\right)\!\left(1+\cos\!\frac{T_{\mathrm{cur}}}{T_{\max}}\pi\right)

Cosine annealing decays the rate from its peak toward zero along a half-cosine curve.

def lr(t, base, warmup, total):
    if t < warmup:
        return base * t / warmup            # linear warmup
    p = (t - warmup) / (total - warmup)     # 0 -> 1
    return 0.5 * base * (1 + cos(pi * p))    # cosine decay to ~0

Warmup then cosine decay — the default schedule for most transformer pretraining.

The warm restart variant periodically resets the rate back to its peak. Each restart kicks the iterate out of its current basin, and the subsequent decay lets it settle into a possibly different, possibly better minimum. A bonus: the snapshots taken just before each restart form a free ensemble of diverse models — a cheap way to get ensemble-like robustness from a single training run.

Averaging weights, not just predictions

Stochastic Weight Averaging (SWA) is almost too simple to believe. Run training as usual into a region of good solutions, then switch to a constant or cyclic learning rate and average the weights visited along the way. Because the SGD iterate bounces around the rim of a wide basin rather than sitting at its bottom, the average of those points lands closer to the center — a flatter, wider solution that tends to generalize better than any single snapshot.

The same instinct lives at the inner loop in Lookahead, which maintains two sets of weights: a fast set updated by any base optimizer for `k` steps, and a slow set that, every `k` steps, takes a small step toward where the fast weights ended up. The slow weights are effectively a running average of the fast trajectory, which damps the noise of the inner optimizer and frequently improves stability with almost no tuning.

Sharpness-aware minimization: optimizing the neighborhood

All three tools above aim at flat minima indirectly. Sharpness-Aware Minimization (SAM) optimizes for flatness directly. Instead of minimizing the loss at a point, it minimizes the worst-case loss within a small ball around the current weights. Each step is two passes: first find the nearby point where the loss is highest by taking a small ascent step in the gradient direction, then take the descent step using the gradient at that perturbed point.

\min_{\theta}\; \max_{\lVert\epsilon\rVert_2 \le \rho}\; L(\theta + \epsilon)

SAM minimizes the worst-case loss within a radius-ρ ball, directly favoring flat minima.

g  = grad(loss, w)
e  = rho * g / (g.norm() + eps)   # ascent to the worst nearby point
g2 = grad(loss, w + e)            # gradient at the perturbed point
w  = w - lr * g2                  # descend using the sharpness-aware gradient

SAM costs two forward-backward passes per step but penalizes sharp directions explicitly.

SAM connects directly to the literature on flat minima and generalization: solutions in wide, flat basins are robust to weight perturbation and tend to transfer better to held-out data. SAM essentially folds that prior into the loss. The cost is the doubled gradient computation, so variants spread the perturbation over micro-batches or apply it only periodically to recover most of the benefit at a fraction of the overhead.

One frame: optimize for the basin, not the point

Tie the section together. Plain training minimizes loss at the final iterate. Schedules, SWA, Lookahead, and SAM all push that objective sideways toward the geometry of the surrounding region. They are partly cashing out the same phenomenon that gradient descent already exhibits on its own — see implicit regularization — but they make the bias toward flat, robust solutions explicit and controllable instead of accidental.

Always warm up, then anneal — cosine is a safe default; add warm restarts if you want free snapshot ensembles.
Add SWA in the final phase for a near-free generalization bump and no inference cost.
Reach for SAM when generalization is the bottleneck and you can afford the doubled gradient.