Module 6 — Optimizers and learning-rate scheduling

The Fashion-MNIST baseline from module 5 uses Adam at a fixed learning rate of $10^{-3}$. That works, and it is a fine default. But learning rate is the single hyperparameter that most influences training outcome, and choosing the right optimiser plus the right schedule for the shape of the problem is what separates a network that plateaus at 89 % from one that reaches 92 %. This module maps out the four optimisers you actually need, the three schedules that cover 95 % of cases, and the two mistakes that make either one useless.

SGD, Adam, AdamW: three tools, three habits

The zoo of PyTorch optimisers is large; the working set is small.

SGD with momentum is the reference for image classification with modern architectures. Its update rule adds a fraction of the previous update to the current gradient, which lets the optimiser build up speed in consistent directions and damp oscillations in noisy ones. It generalises exceptionally well when tuned carefully, but it demands a schedule and a warm-up.

optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)

Adam maintains a per-parameter learning rate based on running estimates of the gradient's first and second moments. That per-parameter adaptivity makes it robust to a wide range of learning rates and to unnormalised data, which is why it has become the default when time-to-first-good-result matters more than absolute peak performance.

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

AdamW is Adam with correctly decoupled weight decay. In the original Adam, weight decay is folded into the L2 gradient penalty, which the adaptive rescaling then breaks. AdamW applies decay directly to the weights, after the Adam update, and the difference is real: it usually generalises better on the same schedule.

optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-2)

The rule of thumb the field has converged on: AdamW for transformers and modern classifiers, SGD with momentum for classical ResNets, plain Adam only when you need Adam-with-broken-decay for reproducibility with an old codebase.

weight_decay: regularisation done right

Weight decay penalises large weights. Every step, weights are pulled slightly toward zero, which discourages the network from relying too heavily on any single feature. It is one of the two regularisers that consistently improve generalisation on real problems, the other being dropout.

Values worth knowing: $10^{-4}$ for SGD on classical vision, $10^{-2}$ for AdamW on transformers, $0$ when in doubt on a small problem. weight_decay on Adam is technically valid but conceptually wrong for the reason above, and switching to AdamW is the two-character fix.

Bias and norm parameters usually should not decay

A common refinement is to exclude bias terms and normalisation parameters (gamma, beta) from weight decay. Applying decay to them shrinks them toward zero, which harms training with no benefit. The pattern is one parameter group with decay and one without:

decay, no_decay = [], []
for name, p in model.named_parameters():
    if p.ndim <= 1 or name.endswith(".bias"):
        no_decay.append(p)
    else:
        decay.append(p)
optimizer = torch.optim.AdamW(
    [{"params": decay, "weight_decay": 1e-2},
     {"params": no_decay, "weight_decay": 0.0}],
    lr=3e-4,
)

Schedules: StepLR, CosineAnnealingLR, OneCycleLR

A schedule is a function that changes the learning rate over the course of training. Three cover almost every situation.

StepLR divides the learning rate by a factor every fixed number of epochs. It is the schedule of the original ResNet paper. Simple, predictable, and enough when the total budget is fixed and long.

scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)

CosineAnnealingLR decays the learning rate along a cosine from its starting value to (near) zero, over a fixed number of epochs. It requires no tuning beyond total epochs and initial rate, and reliably outperforms StepLR on most modern architectures.

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)

OneCycleLR warms up the learning rate for the first third of training, then anneals it to nearly zero. Combined with a matching momentum schedule, it enables training with much higher peak rates than a flat schedule tolerates, and often converges in fewer epochs. It is the schedule to reach for when compute is scarce.

scheduler = torch.optim.lr_scheduler.OneCycleLR(
    optimizer, max_lr=1e-2, epochs=20, steps_per_epoch=len(train_loader)
)

ReduceLROnPlateau sits apart. It watches a metric — validation loss, typically — and cuts the learning rate when it stops improving. It is reactive rather than scheduled, useful when the total number of epochs is not known in advance, but requires you to remember to feed it the metric explicitly:

scheduler.step(val_loss)      # not scheduler.step()

Where to call scheduler.step()

This is where the two frequent mistakes appear.

StepLR, CosineAnnealingLR and ReduceLROnPlateau are epoch-level schedulers: one scheduler.step() per epoch, at the end. OneCycleLR is a step-level scheduler: one scheduler.step() per iteration, after optimizer.step(). Mixing them up gives a schedule that fires 469 times faster than intended (or 469 times slower), and the visible symptom is a flat or wildly oscillating loss with no obvious cause.

for epoch in range(epochs):
    for x, y in train_loader:
        optimizer.zero_grad()
        loss = criterion(model(x), y)
        loss.backward()
        optimizer.step()
        scheduler.step()           # only for OneCycleLR
    # scheduler.step()             # for StepLR / CosineAnnealingLR

The second frequent mistake is calling scheduler.step() before optimizer.step(). PyTorch emits a warning about "detected calls to scheduler.step() before optimizer.step()" — one that many teams manage to ignore for months. The consequence is that the first update runs at the schedule's post-step rate, and every subsequent update lags by one step. Read the warning; it is telling you something real.

Finding a starting learning rate

Instead of guessing, sweep. A learning-rate range test — training briefly while increasing the rate exponentially and plotting loss against rate — reveals the largest rate the model tolerates without diverging. The elbow of the curve, one order of magnitude below the divergence point, is a strong first guess. Libraries automate this; a hand-rolled version fits in twenty lines. Once identified for one architecture on one dataset, that rate transfers well within the same family, and the sweep does not need to be repeated.

Log the actual learning rate each epoch

optimizer.param_groups[0]["lr"] reads the current rate. Printing it alongside train and val loss makes it obvious whether the schedule is behaving as you configured it. Half of scheduler bugs are visible in that column and invisible everywhere else.

Applied to Fashion-MNIST

For our MLP, the practical recipe is:

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

Fifteen epochs, cosine decay from $3 \times 10^{-3}$ to near zero, weight decay at $10^{-4}$. On a laptop CPU this reaches around 91 % validation accuracy, up from 89 % with the module 5 defaults, and does so without dramatic tuning.

In summary

• AdamW for modern classifiers and transformers, SGD with momentum for classical vision, plain Adam only for reproducing old code; the three optimisers cover the vast majority of real work.
• weight_decay is a real regulariser but must be applied correctly: AdamW decouples it from the gradient rescaling, and biases plus normalisation parameters are usually excluded from decay.
• StepLR and CosineAnnealingLR step once per epoch, OneCycleLR steps once per iteration; mixing them up is invisible in code and disastrous in results.
• Find the initial learning rate with a range test rather than a guess, and log the current rate every epoch to catch schedule bugs early.

Next module: moving training to a GPU and enabling mixed precision, the two changes that yield the largest wall-clock speed-up.