Learning rate schedules

A fixed learning rate is rare in modern deep learning. Almost every production training run changes the learning rate over time — small at the start, big in the middle, small at the end. Get this right and the same model trains to a better loss; get it wrong and it diverges or plateaus.

The intuition

Think of optimization as walking downhill:

• Big steps early: you’re far from the minimum, want to move fast.
• Small steps at the end: you’re close, want to settle in.
• Tiny steps at the very start (warmup): your weights are random, big steps can blow up.

Three rules, three schedule components.

The schedules, side by side

import numpy as np

total_steps = 1000
warmup_steps = 100
base_lr = 1e-3

steps = np.arange(total_steps)

def constant(s):
    return base_lr * np.ones_like(s, dtype=float)

def step_decay(s):
    # Drop by 10x every 400 steps
    return base_lr * (0.1 ** (s // 400))

def cosine_with_warmup(s):
    lr = np.empty_like(s, dtype=float)
    # Warmup: linear ramp 0 -> base_lr
    warm_mask = s < warmup_steps
    lr[warm_mask] = base_lr * (s[warm_mask] + 1) / warmup_steps
    # Cosine decay from base_lr to 0
    progress = (s[~warm_mask] - warmup_steps) / (total_steps - warmup_steps)
    lr[~warm_mask] = 0.5 * base_lr * (1 + np.cos(np.pi * progress))
    return lr

for name, fn in [("constant", constant), ("step_decay", step_decay),
                 ("cosine+warmup", cosine_with_warmup)]:
    lrs = fn(steps)
    print(f"{name:14s} start={lrs[0]:.2e}  mid={lrs[500]:.2e}  end={lrs[-1]:.2e}")

constant       start=1.00e-03  mid=1.00e-03  end=1.00e-03
step_decay     start=1.00e-03  mid=1.00e-04  end=1.00e-05
cosine+warmup  start=1.00e-05  mid=5.87e-04  end=3.05e-09

Each schedule encodes a different prior:

Why warmup is mandatory for transformers

Warmup means starting with a near-zero learning rate and ramping it up gradually over the first few hundred or thousand steps, before the main schedule kicks in.

Transformer training is notoriously brittle in the first few hundred steps. Attention weights are random, layer norms haven’t found their scale, and gradients can be enormous. Hitting them with a learning rate of 3e-4 from step 0 frequently causes the loss to NaN within the first epoch.

Warmup fixes this. Start at lr ≈ 0, ramp linearly to your target lr over the first ~1–5% of training, then start decaying. The model gets a chance to stabilize before you start moving fast.

Cosine decay in detail

Cosine decay from a max lr to 0 follows:

lr(t) = 0.5 · lr_max · (1 + cos(π · t / T))

where t is the current step (after warmup) and T is the total remaining. It’s smooth, has zero derivative at both endpoints (so the transition is graceful), and empirically beats step decay on most modern workloads. It’s the default in basically every LLM training codebase.

ReduceLROnPlateau — the lazy option

You don’t know how long training will take. You just want to:

1. Train at lr.
2. When validation loss stops improving for N epochs, drop lr by 10x.
3. Repeat until done.

scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer, mode="min", factor=0.1, patience=5,
)

for epoch in range(epochs):
    train_one_epoch(...)
    val_loss = validate(...)
    scheduler.step(val_loss)

Great for fine-tuning, kaggle-style tabular projects, and anywhere you don’t have a known total step count. Bad for large-scale pretraining where you want a deterministic, reproducible schedule.

In PyTorch

from torch.optim.lr_scheduler import LambdaLR
import math

def cosine_with_warmup(step):
    if step < warmup_steps:
        return step / warmup_steps
    progress = (step - warmup_steps) / (total_steps - warmup_steps)
    return 0.5 * (1 + math.cos(math.pi * progress))

optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
scheduler = LambdaLR(optimizer, cosine_with_warmup)

for step in range(total_steps):
    train_step(...)
    optimizer.step()
    scheduler.step()      # CALL EVERY STEP, not every epoch
    optimizer.zero_grad()

Picking one (the short version)

• Pretraining a transformer: cosine + linear warmup (1–5% of steps).
• Fine-tuning: cosine + brief warmup (~100 steps), lower lr (1e-5 to 5e-5).
• CNN on ImageNet: step decay or cosine — both work.
• Tabular / small: ReduceLROnPlateau, or don’t bother.

In one breath

• Almost every modern run changes the learning rate over time: tiny at the very start (warmup), big in the middle, small at the end.
• Cosine decay (smooth from lr_max to 0) is the modern default; step decay is the classic CNN staircase; ReduceLROnPlateau is the lazy “drop when val stalls” option.
• Warmup is mandatory for transformers — random early weights produce huge gradients, so ramp lr from ~0 over the first 1–5% of steps or the loss NaNs.
• Most schedules expect scheduler.step() after every optimizer step, not every epoch — mixing those up silently ruins the whole run.
• A higher lr speeds early progress but overshoots near the minimum; large-to-small gives fast early movement and precise final convergence.

Quick check