Learning rate schedules and a tuning workflow

The learning rate matters more than the rest of the hyperparameters put together, and a single run can find a good one. This section covers that range test, what warmup and cosine decay are for, why changing the batch size means retuning the rate, and then how to spend a fixed budget of runs: a search on log scales with unpromising trials stopped early, in a set order, ending with several seeds on the top two settings. That last step is what stops you shipping a winner that was noise.

Find the learning rate first#

A learning rate range test trains for a few hundred steps while raising the learning rate exponentially from tiny to huge, and records the loss. Loss falls, then explodes. A good starting rate sits a little below where the loss falls fastest, well before it explodes.

import torch, torch.nn as nn

def range_test(model_fn, batches, lo=1e-6, hi=3.0, steps=400):
    """One run over exactly steps batches. Returns (lr, smoothed loss) per step."""
    m = model_fn()
    opt = torch.optim.AdamW(m.parameters(), lr=lo)
    mult, lr, smooth, out = (hi / lo) ** (1 / steps), lo, None, []
    for x, y in batches:
        for pg in opt.param_groups:
            pg["lr"] = lr
        opt.zero_grad()
        loss = nn.functional.cross_entropy(m(x), y)
        loss.backward()
        opt.step()
        v = loss.item()
        smooth = v if smooth is None else 0.9 * smooth + 0.1 * v
        out.append((lr, smooth))
        lr *= mult
    return out

Run on the tabular network from the tuning section, 400 steps from 0.000001 to 3, this is what comes back.

learning rate    smoothed training loss
  0.00001                 0.697     flat: nothing is happening
  0.001                   0.564     falling
  0.0028                  0.362     steepest descent
  0.026                   0.246     minimum
  0.1                     0.598     climbing back out
  1.0                     0.725     diverged

The test found in one run what the sweep found in eight, which is the point of it. Read the advice about it carefully, though. The common rule is to take an order of magnitude below the minimum, which here would give 0.0026, and the sweep says that setting is worth 0.9235 against 0.9345 at the true optimum. That is a point of accuracy given away. The range test brackets the answer within an order of magnitude; it does not hand you the answer. Use it to decide where to centre a short sweep, not to skip the sweep.

Warmup, then decay#

Large models are often unstable in the first steps, when weights and optimiser statistics are poorly set. Warmup raises the learning rate from near zero over the first steps. After that, decaying it, commonly along a cosine curve, lets training settle into a good minimum.

import math

def lr_at(step, total, peak=3e-4, warmup=500):
    if step < warmup:
        return peak * step / warmup
    progress = (step - warmup) / (total - warmup)
    return peak * 0.5 * (1 + math.cos(math.pi * progress))

lr_at(step, total=20000) with peak 3e-4 and 500 warmup steps

  step      0    0.00e+00 the ramp starts at zero
  step    250    1.50e-04 halfway up the warmup ramp
  step    500    3.00e-04 peak
  step   5000    2.62e-04
  step  10000    1.56e-04 halfway through, at 52% of peak
  step  15000    4.61e-05
  step  20000    0.00e+00 the cosine ends at zero

Whether the decay is worth anything is measurable. On the tabular network from the tuning section, cosine decay against a constant rate at the same peak, twenty epochs each.

peak lr   constant: best / final     cosine: best / final
  0.01      0.9340 / 0.9340            0.9370 / 0.9365
  0.03      0.9345 / 0.9230            0.9350 / 0.9320

Read the two columns separately. At a peak of 0.03 the constant-rate run scored 0.9345 at its best epoch and 0.9230 at its last, because a large constant rate keeps bouncing the weights around a minimum it cannot settle into. Decay closed most of that gap. On a small problem you can paper over this with early stopping on the best epoch. On a large training run there is no best epoch to go back to, because you only keep a handful of checkpoints, so the value of the schedule is that the final model is the good one.

Stop when the thing you care about stops improving#

Early stopping means keeping the weights from the epoch with the best validation score rather than the last epoch. What it saves depends on whether the validation curve turns, and the trap is that different validation metrics turn at different times. Here is a deliberately overfitting run: 600 training examples, a 512-wide network, no dropout, 60 epochs.

epoch   train loss   val loss   val accuracy
    1      0.4853      0.5655        0.7035
    3      0.2535      0.4571        0.7820
    5      0.1252      0.4017        0.8235   <- val loss is at its minimum
   10      0.0183      0.4839        0.8400   <- val accuracy is at its maximum
   20      0.0021      0.6276        0.8360
   40      0.0005      0.7440        0.8360
   60      0.0002      0.8113        0.8365

Stopping on validation loss would have stopped at epoch 5 and handed you a model 1.65 accuracy points worse than the one at epoch 10. Both curves are correct: the model is becoming more confident and more often wrong about the cases it gets wrong, which hurts cross-entropy, while its top-1 decisions keep improving slightly. Stop on the metric the product is judged by. If that metric is noisy, use patience, meaning stop after N epochs with no improvement rather than at the first downturn, and pick N from how bumpy your validation curve is rather than from a blog post.

Batch size and learning rate move together#

A larger batch gives a less noisy gradient, which usually allows a larger learning rate. If you change the batch size, retune the learning rate. Comparing two batch sizes at the same learning rate compares two badly tuned runs.

validation accuracy, 20 epochs, learning rate across the top

batch     1e-4    3e-4    1e-3    3e-3    1e-2    3e-2    1e-1
   32    0.914   0.918   0.925   0.941   0.932   0.918   0.506
  128    0.886   0.916   0.918   0.924   0.934   0.934   0.899
  512    0.769   0.896   0.916   0.922   0.928   0.927   0.919
 2048    0.663   0.735   0.882   0.913   0.925   0.926   0.886

Two things in that grid are worth more than the headline. First, 0.1 destroys the batch-32 run, landing it on 0.506, and the same rate is perfectly usable at batch 512. Anyone who fixes the learning rate and then sweeps batch size will conclude that small batches are unstable, which is a statement about their sweep. Second, 0.0001 costs the batch-2048 run 26 points against its own best and the batch-32 run under three points, because at a fixed number of epochs a larger batch takes proportionally fewer steps. The rule of thumb that survives is to move the learning rate in the same direction as the batch size, roughly linearly for moderate changes, and then check.

A search that stops bad runs early#

import optuna

def objective(trial):
    lr = trial.suggest_float("lr", 1e-5, 1e-2, log=True)
    wd = trial.suggest_float("weight_decay", 1e-6, 1e-1, log=True)
    model = build(lr=lr, weight_decay=wd)
    for epoch in range(20):
        score = train_and_validate(model)
        trial.report(score, epoch)
        if trial.should_prune():          # clearly worse than earlier trials at this point
            raise optuna.TrialPruned()
    return score

study = optuna.create_study(direction="maximize")
study.optimize(objective, n_trials=50)
print(study.best_params)

Search learning rates and weight decay on a log scale; the difference between 0.001 and 0.01 matters as much as between 0.0001 and 0.001. Pruning stops unpromising runs after a few epochs, so the budget goes to the good ones.

The saving from pruning is worth estimating before you set it up. Fifty trials at twenty epochs each is 1,000 epoch-units. Suppose pruning kills 25 of them at epoch two, and 12 of the survivors at epoch five, leaving 13 to run in full. The cost is 25 times 2, plus 12 times 5, plus 13 times 20, which is 370 epoch-units: 63% of the budget came back. The saving grows with the number of trials, and shrinks when good and bad configurations only separate late in training, which is exactly the case where pruning also throws away the winner.

Spend the budget in this order#

1. 1.Get a baseline training cleanly with default settings.
2. 2.Tune the learning rate.
3. 3.Tune regularisation: weight decay, dropout, augmentation strength.
4. 4.Only then touch architecture choices.
5. 5.Rerun the best setting with several seeds before believing it beat the runner-up.

What to skip: warmup on a small model trained with a modest learning rate, where it does nothing measurable and costs you steps; exotic schedules, since cosine with a short warmup is within noise of nearly all of them on nearly all problems; and any schedule tuning at all before the peak learning rate is right, because the schedule is a refinement of the peak and moving both at once tells you which did nothing. Worth memorising: the range test is one run, the cosine ends at zero, and the learning rate moves with the batch size.