Module 5 — Optimizers: stochastic gradient descent, momentum, Adam

The gradients are computed. What remains is deciding how far to move the weights, and that decision separates a training run that converges in an hour from one that never gets there.

The basic rule and its one critical hyperparameter

Gradient descent updates each parameter in the direction opposite its gradient:

$$w \leftarrow w - \eta \frac{\partial \mathcal{L}}{\partial w}$$

The learning rate $\eta$ is by far the most important hyperparameter in deep learning. Too large, the loss oscillates or diverges to NaN; too small, training crawls or stalls on a plateau. No other setting has a comparable effect, and it is where you should begin.

Batch size: the noise-versus-speed trade-off

Three regimes are distinguished by the number of observations used per update.

Full-batch gradient uses the whole dataset: an exact direction, but a single update per pass, and prohibitive memory use. Pure stochastic gradient uses one observation: very frequent updates but a very noisy direction. The mini-batch, in practice 32 to 512 observations, combines the advantages and is the universal standard.

One counter-intuitive point deserves mention: mini-batch noise is useful. It helps escape narrow local minima and acts as light regularization. Very large batches converge faster in number of passes but often generalize slightly worse, which explains why batch size is not increased indefinitely even when memory would allow it.

Momentum, or why accumulate impetus

Plain descent zigzags in narrow valleys: the gradient points at the walls rather than down the floor. Momentum fixes this by accumulating a velocity:

$$v \leftarrow \beta v + (1 - \beta)\nabla, \qquad w \leftarrow w - \eta v$$

With $\beta = 0.9$, the update retains a moving average of recent gradients. The perpendicular oscillations, which alternate in sign, cancel out; the constant component down the valley floor accumulates. The rolling-ball analogy is accurate: it rolls through small irregularities instead of stopping in them.

Adaptive rates and Adam

Not all parameters need the same step. A rare feature deserves bolder updates than an omnipresent one. Adaptive methods therefore give one rate per parameter, derived from the history of its gradients.

Adam is today's default starting point. It combines the two previous ideas: momentum on the gradient (first moment) and momentum on the squared gradient (second moment), the latter used to normalize the step.

$$m \leftarrow \beta_1 m + (1-\beta_1)\nabla, \qquad v \leftarrow \beta_2 v + (1-\beta_2)\nabla^2, \qquad w \leftarrow w - \eta \frac{\hat{m}}{\sqrt{\hat{v}} + \epsilon}$$

A parameter whose gradients are consistently large sees its step reduced by the division; one with weak gradients sees its step raised. The defaults $\beta_1 = 0.9$, $\beta_2 = 0.999$ and $\epsilon = 10^{-8}$ work in the vast majority of cases and are rarely tuned.

AdamW corrects a real flaw in Adam: the L2 penalty, when added to the gradient, ends up divided by $\sqrt{\hat{v}}$ and loses its regularizing effect. AdamW applies it separately — decoupled weight decay — and should be preferred whenever a weight penalty is used. It is the optimizer of today's large language models.

import torch

optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)

for epoch in range(epochs):
    for X_batch, y_batch in loader:
        optimizer.zero_grad()                # without this, gradients accumulate
        loss = criterion(model(X_batch), y_batch)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
    scheduler.step()                         # once per epoch, not per batch

Varying the learning rate

Keeping $\eta$ constant is rarely optimal: you want to move fast at first, then refine. Two mechanisms have become standard.

Warmup progressively raises the rate over the first few hundred iterations. At startup the weights are random and the gradients erratic; applying the nominal rate immediately can destabilize training lastingly. This is particularly true for transformers, where warmup is not optional.

Decay then reduces the rate: in steps, or along a cosine curve, which is today the most widespread choice. The intuition is annealing: large steps to find the right region, small steps to settle in it.

A tuning method that works

Look for the learning rate before anything else, and do it by orders of magnitude: $10^{-2}$, $10^{-3}$, $10^{-4}$. A quick, reliable test is to run a few hundred iterations while progressively raising the rate, then plot the loss: the right order of magnitude sits just before the point where it turns back up. For AdamW, $3 \times 10^{-4}$ is a reasonable starting point on most architectures. And tune batch size only afterwards, knowing that doubling the batch often justifies raising the rate.

Summary

• The learning rate is the most decisive hyperparameter; tune it first and by orders of magnitude.
• The mini-batch of 32 to 512 is the standard; its noise is useful and acts as light regularization.
• Momentum cancels oscillations and accumulates impetus down the valley; adaptive methods give one step per parameter.
• AdamW is the sensible default, preferable to Adam with weight decay; warmup then cosine decay is the usual schedule.

Next module: weight initialization and the vanishing gradient problem, two topics that the structure of backpropagation made predictable.