Learning rate schedules and warmup

The big idea

The learning rateThe multiplier that sets how big a step the optimizer takes along the gradient. Usually the single most important hyperparameter. is the most important knob in training, and the best value changes as training goes on. Early, you are far from a good solution and big steps make fast progress. Late, you are close, and big steps keep bouncing around the bottom of the valley, because mini-batch gradients are noisy.

Think of parking a car. You drive fast down the street, slow down as you approach the space, and creep the last few centimetres. A learning rate scheduleA rule that changes the learning rate over the course of training, for example decaying it step by step or along a cosine curve. does the same: large steps first, small steps at the end.

There is one exception at the very start. A freshly initialized network, especially a transformer trained with Adam, can be thrown off by a few huge, poorly estimated early updates. So many recipes start with a short warmupA short phase at the start of training in which the learning rate rises from near zero to its peak value, avoiding unstable early updates.: the learning rate ramps up from near zero over the first few hundred or thousand steps, then the decay begins.

How it works

Step 1: Step decay

Keep the learning rate constant, then divide it by 10 at fixed points, for example at 40 and 80 percent of training. Simple and effective; it was the standard for image classifiers for years. The drawback is that you must choose when to drop, and the loss often plateaus until you do.

Step 2: Cosine decay

Follow half a cosine wave from the peak down to (nearly) zero:

It stays high for a long time, then decays smoothly, with only two settings: the peak and the total number of steps T. At t = T/2 it is exactly half the peak. Cosine decay is the most common schedule in current practice.

Step 3: Warmup

Ramp linearly from 0 to the peak over the first W steps (often 1 to 5 percent of training), then decay. Two reasons it helps. Adam's estimates of gradient scale (v) are unreliable in the first steps, so the first updates can be far too large. And early in training, before normalization statistics and weights settle, a large learning rate can push the network into a bad region it never recovers from. Warmup costs little and removes a common source of early divergence, especially for transformers and large batch sizes.

Step 4: One-cycle

One-cycleA schedule that raises the learning rate from a low value to a high peak over part of training, then lowers it to a very small value by the end. rises from a low value to the peak over roughly the first 30 percent of training, then anneals far below the starting value. The high middle phase explores and acts as a regularizer; the long low tail settles into a minimum. It often reaches good results in fewer epochs and is popular for fine-tuning and fast training runs.

Step 5: Finding the peak with a range test

A learning rate range testA short run in which the learning rate is increased exponentially every step while the loss is recorded, to find the range of rates where training makes fast progress without diverging. (or LR finder) runs a few hundred steps while multiplying the learning rate by a constant each step, from tiny (1e-6) to huge (1 or 10), and records the loss. The loss first barely moves, then falls quickly, then explodes. Choose a peak somewhat below the point where it is lowest, a common rule of thumb being about a tenth of that value, or the point where it falls fastest.

Going deeper

Code: schedules, decay versus constant, and a range test

import numpy as np

T, peak = 1000, 0.1  # total steps and peak learning rate
def step_sched(t):   return peak * 0.1 ** (t // 400)                 # drop 10x at steps 400 and 800
def cosine(t):       return peak * 0.5 * (1 + np.cos(np.pi * t / T))
def warm_cos(t, w=100):                                              # linear warmup, then cosine
    return peak * t / w if t < w else peak * 0.5 * (1 + np.cos(np.pi * (t - w) / (T - w)))
def one_cycle(t):    # up to the peak at 30%, down to almost zero
    return peak * (0.04 + 0.96 * t / 300) if t < 300 else peak * (0.5 * (1 + np.cos(np.pi * (t - 300) / 700)) * 0.999 + 0.001)
print("step     " + "".join(f"{t:>8}" for t in (0, 50, 100, 400, 700, 999)))
for name, s in [("step", step_sched), ("cosine", cosine), ("warm+cos", warm_cos), ("1cycle", one_cycle)]:
    print(f"{name:<9}" + "".join(f"{s(t):8.4f}" for t in (0, 50, 100, 400, 700, 999)))

# Linear regression trained with noisy mini-batch SGD: constant vs decaying learning rate
rng = np.random.default_rng(0)
X = rng.normal(size=(2000, 5)); w_true = rng.normal(size=5)
y = X @ w_true + rng.normal(0, 0.5, 2000)
def train(sched, steps=T, bs=8):
    w = np.zeros(5)
    for t in range(steps):
        i = rng.integers(0, 2000, bs)
        w -= sched(t) * 2 * X[i].T @ (X[i] @ w - y[i]) / bs
    return np.mean((X @ w - y) ** 2)
print(f"\nfinal MSE, constant lr 0.1:   {train(lambda t: 0.1):.4f}")
print(f"final MSE, cosine from 0.1:   {train(cosine):.4f}")
print(f"best possible (noise var):    {np.mean((X @ w_true - y) ** 2):.4f}")

# Learning rate range test: raise lr exponentially each step, watch the loss
w, lrs, losses = np.zeros(5), np.geomspace(1e-4, 2, 100), []
for lr in lrs:
    i = rng.integers(0, 2000, 64)
    w -= lr * 2 * X[i].T @ (X[i] @ w - y[i]) / 64
    losses.append(np.mean((X @ w - y) ** 2))
k = int(np.argmin(losses)); best = lrs[k]
blow = k + int(np.argmax(np.array(losses[k:]) > 10 * losses[k]))
print(f"\nrange test: loss lowest at lr {best:.3f}; 10x worse again by lr {lrs[blow]:.3f}")
print(f"rule of thumb: pick about a tenth of the lowest-loss lr, here {best / 10:.3f}")

step            0      50     100     400     700     999
step       0.1000  0.1000  0.1000  0.0100  0.0100  0.0010
cosine     0.1000  0.0994  0.0976  0.0655  0.0206  0.0000
warm+cos   0.0000  0.0500  0.1000  0.0750  0.0250  0.0000
1cycle     0.0040  0.0200  0.0360  0.0951  0.0389  0.0001

final MSE, constant lr 0.1:   0.2803
final MSE, cosine from 0.1:   0.2523
best possible (noise var):    0.2521

range test: loss lowest at lr 0.245; 10x worse again by lr 1.213
rule of thumb: pick about a tenth of the lowest-loss lr, here 0.024

The middle block is the core argument for decaying. Linear regression trained with noisy batches of 8 and a constant learning rate of 0.1 ends at MSE 0.2803, because it keeps jumping around the minimum. The same run with cosine decay from 0.1 ends at 0.2523, within 0.0002 of the best possible 0.2521 (the noise variance). Same steps, same peak: the only difference is small steps at the end.

In the range test, the loss is lowest around a learning rate of 0.245 and is 10 times worse again by 1.21, so a peak around 0.02 to 0.1 is a safe choice for this problem.

Practical notes

• Schedules are defined in steps, not epochs. If you change the batch size, the number of steps changes, so recompute T and the warmup length.
• When you increase the batch size, a larger peak learning rate is often appropriate, but warmup becomes more important.
• When resuming from a checkpoint, restore the scheduler state too, or the learning rate jumps back to the start.
• Frameworks provide these schedules ready-made (cosine with warmup, one-cycle, step decay); you call scheduler.step() after optimizer.step() in the loop from Autograd and your first PyTorch training loop .

Trade-offs

Common mistakes

• Training a transformer with Adam at full learning rate from step 0, and seeing the loss spike or diverge early.
• Defining the schedule in epochs, then changing the batch size and getting a schedule that ends too early or too late.
• Forgetting to restore the scheduler when resuming, so the learning rate restarts at the warmup.
• Picking the peak from the lowest point of a range test instead of somewhat below it, leaving no safety margin.

In the wild

Large language model training almost always uses AdamW with a linear warmup of a few thousand steps followed by cosine decay to a small fraction of the peak. Image classifiers have been trained with step decay and, more recently, cosine. Fine-tuning recipes commonly use a short warmup and linear or cosine decay over a few epochs, and fast-training competitions popularized one-cycle schedules found with a range test.

Check yourself

Recap