Learning Rate Warmup and Cosine Decay Schedule in Python

For each step in steps (0-based), return the learning rate:

if step < warmup: base_lr × (step + 1) / warmup;
otherwise, with progress = min(1, (step − warmup) / (total − warmup)): min_lr + ½(base_lr − min_lr)(1 + cos(π · progress)).

The arguments are steps, base_lr, min_lr, warmup, total, with 0 ≤ warmup < total.

Examples

Input steps = [0, 1, 2, 3], base_lr = 1, min_lr = 0, warmup = 4, total = 10

Output [0.25, 0.5, 0.75, 1]

Warm-up over four steps: 0.25, 0.5, 0.75, 1.0.

Input steps = [4, 7, 10], base_lr = 1, min_lr = 0, warmup = 4, total = 10

Output [1, 0.5, 0]

Cosine from the peak: 1 at step 4, ½ halfway, 0 at the end.

Input steps = [0, 50, 100, 150], base_lr = 0.0003, min_lr = 3.000e-5, warmup = 0, total = 100

Output [0.0003, 0.0002, 3.000e-5, 3.000e-5]

No warm-up, and the rate stays at the floor after the last step.

Constraints

0 ≤ warmup < total ≤ 10⁷
0 ≤ min_lr ≤ base_lr
Checked to a relative tolerance of 10⁻¹²

Hints

Hint 1

Handle the two phases with one branch on step < warmup.

Hint 2

During warm-up, step 0 should already have a non-zero rate: that is why it is (step + 1) / warmup.

Hint 3

Clamp progress at 1 so steps beyond total stay at min_lr instead of rising again.

Where candidates go wrong

Using step / warmup, which gives a learning rate of zero at step 0 and wastes the first update.
Forgetting to clamp progress, so the cosine comes back up after the end of the schedule.
Decaying to zero instead of to min_lr.

The idea

Warm-up grows the rate linearly so that the first steps are small, which matters for Adam, whose variance estimate is unreliable at the start. After it, the cosine term moves smoothly from 1 to −1 as progress goes from 0 to 1, so the rate falls from base_lr to min_lr slowly at first, fastest in the middle, and slowly again near the end. Clamping progress keeps the rate at the floor if training runs on.

Time O(len(steps)) · Space O(len(steps))

Reference solution

import math

def lr_schedule(steps, base_lr, min_lr, warmup, total):
    out = []
    for s in steps:
        if s < warmup:
            out.append(base_lr * (s + 1) / warmup)               # step 0 already gets a small rate
        else:
            progress = min(1.0, (s - warmup) / (total - warmup))   # clamp: stay at the floor afterwards
            out.append(min_lr + 0.5 * (base_lr - min_lr) * (1 + math.cos(math.pi * progress)))
    return out

Console⌘/Ctrl + Enter runs

Python runs in your browser with NumPy. The first run downloads the runtime, which takes a few seconds.

Why interviewers ask it

It is the schedule behind most transformer training runs, and writing it exposes off-by-one errors immediately: the first step, the last warm-up step and the final step each have a defined value, and interviewers check them.

The theory behind it is in the lesson Dropout, gradient clipping and learning-rate schedules.

Keep going

Easy8 min

Linear warm-up with cosine decay