Learning-rate schedules: warm-up, decay

No single learning rate suits the whole run

In gradient descent, the learning rate η scales every step. A large rate makes fast progress but overshoots and bounces around a low point. A small rate is precise but slow. A learning-rate schedule uses different rates at different times: a function that gives η for each step number t.

• Warm-up starts small. At the start the weights are random, so gradients are large and point in changing directions. Optimizers such as AdamW also have only a few gradients in their running averages. Large steps now can throw the weights into a region where training fails.
• Decay ends small. Late in training, each step's gradient comes from a small random batch of examples, so it is noisy. With a large rate the weights keep jumping around near a low point. Shrinking the rate lets them settle into it.

Warm-up is a straight rise; cosine decay is half a wave down

Write W for the number of warm-up steps, T for the total number of steps, and η_peak for the highest rate. During warm-up (t < W):

η_t = η_peak × t / W

After warm-up, let p = (t − W) / (T − W) be the fraction of the decay that is done, from 0 to 1. Cosine decay to zero is:

η_t = η_peak × ½(1 + cos(πp))

The factor ½(1 + cos(πp)) falls smoothly from 1 at p = 0 to 0 at p = 1. It falls slowly at first, fastest in the middle, and slowly again at the end. (A cosine is a sine wave shifted by a quarter turn.) Some schedules stop at a small floor η_min instead of 0.

The warm-up length and the decay shape are settings to choose

Try it. The figure draws the schedule over 10,000 steps with a peak of 0.001. Predict first: with no warm-up at all, what rate does the very first step use? Then find the warm-up length at which step 500 uses half the peak rate.

Other schedules exist. Step decay multiplies the rate by a fixed factor, such as 0.1, at chosen steps. The original Transformer paper used warm-up followed by a fall in proportion to 1 / √t. All of them share the same two aims: careful at the start, settled at the end.

In PyTorch, a scheduler sets the rate before each step

LambdaLR multiplies the optimizer's rate by a function of the step count. Call sched.step() once after each opt.step(). If you call it once per pass over the data instead, the function's t counts passes, not steps.

schedule_torch.py

import math
import torch

peak, warmup, total = 0.001, 1000, 10_000

def factor(step):                        # multiplies the peak rate
    if step < warmup:
        return step / warmup             # straight rise from 0
    p = (step - warmup) / (total - warmup)
    return 0.5 * (1 + math.cos(math.pi * p))  # cosine fall to 0

w = torch.zeros(1, requires_grad=True)
opt = torch.optim.AdamW([w], lr=peak)    # the optimizer that updates w, with rate peak
sched = torch.optim.lr_scheduler.LambdaLR(opt, factor)

for step in range(total):
    if step in (0, 500, 1000, 5500, 7750):    # print only at these steps
        print(step, f"{opt.param_groups[0]['lr']:.6f}")   # the optimizer's current rate
    w.grad = torch.zeros(1)              # stands in for loss.backward()
    opt.step()
    sched.step()                         # after each optimizer step

0 0.000000
500 0.000500
1000 0.001000
5500 0.000500
7750 0.000146

With this function, step 0 has a rate of exactly 0, so the first step changes nothing. Some code uses (t + 1) / W instead to avoid that wasted step.