Learning Rate, Decay

The learning rate is the single most consequential hyperparameter in a training run, and a constant value is a compromise. Early on you want large steps to cover ground quickly; near the end you want small steps so the loss can settle instead of bouncing. Decaying the rate over time gets both.

The idea

A schedule is a function η(t) that replaces the fixed learning rate. The common shapes:

• Step decay — divide by 10 every $N$ epochs. Simple, common in classic vision training.
• Cosine decay — follow a cosine from η_max to η_min over the whole run. Smooth, no cliff, and the dominant choice for language models.
• Inverse square root — η_t = η_0 / √t. Used by the original Transformer, aggressive at the start and flat later.

The cosine form is

Warmup is usually prepended: for the first few hundred or thousand steps, ramp linearly from roughly zero up to η_max. Gradients early in training are large and poorly aligned with the loss landscape, and a full-size step can wreck the initialization. Warmup plus cosine decay is the standard recipe for pretraining.

The schedule must be applied every step, not every epoch: loop over opt.param_groups and set g["lr"] = lr_at(step) before calling opt.step().

Worked example

Run cosine decay with η_max = 1e-3, η_min = 1e-4, and T = 100 steps:

Half the run happens above 8.68e-4, and the last quarter drops from 5.5e-4 to 1e-4, which is where the model gets to polish its weights.

For comparison, step decay with a factor of 0.1 every 10 epochs gives 1e-3 → 1e-4 → 1e-5 at epochs 10 and 20, and inverse square root with η_0 = 1e-3 gives 1e-3, 5e-4, 1e-4, 3.16e-5 at steps 0, 3, 99, and 999.

In code

import math

def cosine_lr(t, t_max=100, lr_max=1e-3, lr_min=1e-4):
    return lr_min + 0.5 * (lr_max - lr_min) * (1 + math.cos(math.pi * t / t_max))

def warmup_cosine(t, warmup=10, t_max=100, lr_max=1e-3, lr_min=1e-4):
    if t < warmup:
        return lr_max * (t + 1) / warmup
    return cosine_lr(t - warmup, t_max - warmup, lr_max, lr_min)

for t in (0, 25, 50, 75, 100):
    print(t, f"{cosine_lr(t):.3e}")
print([f"{warmup_cosine(t):.3e}" for t in (0, 4, 9, 10)])
print(cosine_lr(50) == 5.5e-4)   # True

Check yourself

1. Why does a learning rate that is too high produce a loss that bounces instead of settling?
2. What problem does warmup solve, and why is it more important for large models trained on large batches?
3. At t = T the cosine schedule returns η_min rather than 0. Why keep a nonzero floor?

Key takeaways

• A constant learning rate is a compromise; schedules let a run be fast first and precise later.
• Cosine decay dominates modern training, usually with a short linear warmup in front.
• The schedule belongs in the step loop — set param_groups[i]["lr"] before every optimizer step.