Warmup then cosine: the learning-rate shape of big training runs

Build a Small GPT From Scratch · Part 3 of 8

Almost every large training run starts the learning rate near zero, climbs to a peak, then slowly brings it back down. Here is what that shape looks like with real numbers, and why it exists.

Warmup

At the very start the weights are random and the optimizer has no reliable estimate of the typical size of a gradient. Adam, the usual optimizer, keeps running averages of gradients and of their squares; early on those averages are based on very little data. A full-size step taken on such estimates can throw a deep transformer into a bad region from which it never recovers. So the learning rate warms up: it rises linearly from zero to the peak over a short stretch, often one to two per cent of the run.

Cosine decay

After the peak, the rate follows half a cosine wave down to a small floor. Large steps early move quickly through the loss landscape; small steps late let the weights settle instead of bouncing around the bottom. The formula is

lr = final + (peak − final) × (1 + cos(π × progress)) / 2

where progress goes from 0 to 1 between the end of warmup and the end of the schedule.

The schedule as code

import math

def lr_at(step, warmup=5, total=20, peak=0.3, floor=0.01):
    if step < warmup:                                   # linear warmup from 0 to the peak
        return peak * step / warmup
    if step < total:                                    # cosine decay from the peak to the floor
        progress = (step - warmup) / (total - warmup)
        return floor + (peak - floor) * (1 + math.cos(math.pi * progress)) / 2
    return floor                                        # hold at the floor afterwards

print([round(lr_at(s), 3) for s in range(0, 22)])
# [0.0, 0.06, 0.12, 0.18, 0.24, 0.3, 0.297, 0.287, 0.272, 0.252, 0.227, 0.2, 0.17,
#  0.14, 0.11, 0.083, 0.058, 0.038, 0.023, 0.013, 0.01, 0.01]

These are the toy numbers of a 20-step example. A real run has thousands of steps, and the peak is far smaller than 0.3 for a transformer, but the shape is identical.

Takeaways

• Warmup raises the learning rate gradually while the optimizer’s statistics are still unreliable.
• Cosine decay brings it down smoothly to a small floor.
• The whole schedule is a few lines of code and the shape is the same at any scale.
• Instability early in a run often points to the peak rate or the warmup length.

Next in the series (Part 4): Gradient accumulation: a big batch on a small GPU