Batch size, schedules and regularisation in practice

You know how the gradients get computed. This is about the handful of settings that decide whether a run works at all. Batch size and learning rate move together, so changing one alone destabilises training. Schedules exist because the step size that helps at the start is wrong by the end. Weight decay and early stopping are how you stop a model memorising its training set. These are the knobs you will actually turn, and the usual reason a run that ought to work does not.

Batch size and learning rate move together#

A larger batch gives a less noisy gradient estimate, which allows a larger learning rate. A common heuristic scales the learning rate with batch size. Change one without the other and training can become unstable or unexpectedly slow.

The reason is worth stating precisely, because the heuristic is repeated far more often than it is explained. A gradient computed on a batch is an estimate of the true gradient, and averaging more samples reduces the noise in that estimate — by the square root of the batch size, which is where the square-root version of the scaling rule comes from. A noisier estimate has to be trusted less, so the step taken along it must be shorter. Batch size therefore does not set the step size by itself; it sets how far you can afford to trust one.

Measure it rather than assuming it. The same two-layer network and the same 200 training examples, at three batch sizes and five learning rates, reporting validation loss after sixty epochs.

                lr 0.0125   0.025   0.05    0.1     0.2
batch 8           0.616      0.474   1.060   nan     nan
batch 32          0.606      0.571   0.487   1.135   nan
batch 128         0.581      0.590   0.485   1.625   nan

So use the rule as a direction rather than a formula. Doubling the batch means you should try a larger rate, and expecting the doubled rate to work is optimistic. Note also the row for batch 8, where 0.05 gives 1.060 while 0.025 gives 0.474: the failure at a slightly too large rate is not a crash but a result twice as bad, arrived at without any warning in the training loss.

Schedules#

• Warmup: start small and ramp up, avoiding large destabilising updates while weights are random.
• Cosine decay: reduce the rate smoothly towards zero, letting the model settle into a good region.
• Step decay: drop the rate at fixed points. Simpler, and the timing needs tuning.

A schedule is four lines, and writing them out removes the mystery from the names. Warmup is a straight line from near zero to the target rate over the first few hundred or few thousand steps. Cosine decay then follows half a cosine wave from the target down to zero over the remaining steps. Every large training run you read about is some version of exactly this.

import math

def lr_at(step, base=3e-4, warmup=2000, total=100_000):
    if step < warmup:
        return base * step / warmup                       # straight line up
    progress = (step - warmup) / (total - warmup)
    return base * 0.5 * (1 + math.cos(math.pi * progress))  # cosine down to zero

for step in (0, 1000, 2000, 25_000, 50_000, 100_000):
    print(step, round(lr_at(step), 8))

# 0        0.0
# 1000     0.00015
# 2000     0.0003
# 25000    0.00026104
# 50000    0.00015481
# 100000   0.0

Weight decay#

A penalty that pulls weights towards zero, discouraging reliance on any single large weight. With Adam, the correct implementation is decoupled weight decay, as in AdamW; applying it naively through the loss interacts badly with Adam's per-parameter scaling.

What it buys is measurable and it is not free. On a network with 200 training examples, trained to memorise them completely, the best validation loss reached during training was 0.611 without weight decay, 0.583 at a decay of 0.01, and 0.549 at 0.05. But at 0.05 the training loss only reached 0.406 rather than zero, and the final validation loss was 1.02 — the model had been pulled so hard towards zero that it was still improving, or getting worse, in ways that had nothing to do with the data. Weight decay is a dial with a wrong setting on both sides of the right one.

Early stopping#

Track validation loss and keep the checkpoint where it was lowest. It costs nothing and prevents the model from continuing to overfit after it has stopped improving on unseen data.

The shape of the validation curve in that same experiment is what makes the case. It reached its lowest point, 0.611, at epoch 26; it was still being trained at epoch 200, by which point the training loss was 0.0000 — every one of the 200 examples memorised exactly — and the validation loss had risen to 0.692. The last 174 epochs made the model thirteen per cent worse on data it had not seen, and consumed most of the compute of the run doing it.

Two practical points. Keep the best checkpoint rather than the last one, because those are different models and only one of them is the one you measured. And do not stop at the first epoch where validation rises: early validation curves are noisy, and the standard rule is patience, meaning stop only after some number of epochs with no improvement on the best value seen so far.

best, best_epoch, patience, waited = float("inf"), 0, 10, 0

for epoch in range(max_epochs):
    train_one_epoch(model, optimiser, train_loader)      # your own loop
    score = validate(model, val_loader)                  # lower is better
    if score < best:
        best, best_epoch, waited = score, epoch, 0
        torch.save(model.state_dict(), "best.pt")        # keep the model, not just the number
    else:
        waited += 1
        if waited >= patience:
            break

model.load_state_dict(torch.load("best.pt"))             # go back to the best one

The order to set these in#

1 batch size      largest that fits and keeps the hardware busy
2 learning rate   sweep at that batch size; this dominates everything else
3 schedule        warmup if the model is large or the batch is; then decay
4 weight decay    only once 1 to 3 are settled; try 0, 0.01, 0.1
5 early stopping  always on, patience of a few epochs, keep the best checkpoint