Learning Rate: Warmup + Cosine Decay

The learning rate schedule — how the LR changes over the course of training — is critical for stable, efficient training. Modern LLMs use warmup followed by cosine decay almost universally.

Why warmup

At the start of training, weights are random and gradients are noisy. A high learning rate amplifies that noise and can permanently damage the optimization trajectory. A linear ramp from 0 to peak LR over the first ~1-5% of training steps lets the optimizer find a stable region first, then go fast.

Why cosine decay

Once warmup is done, you want to use the peak LR while you can but slowly shrink it as the model approaches a local optimum. Cosine decay — lr(t) = min_lr + 0.5 × (max_lr - min_lr) × (1 + cos(π × progress)) — provides a smooth, parameter-free way to do this. Empirically it works as well as more elaborate schedules (e.g., linear-then-flat).

The original 2017 Transformer used a different schedule (Noam scheduler: lr ∝ d_model^(-0.5) × min(step^(-0.5), step × warmup^(-1.5))). Modern open recipes (Llama, Mistral, nanoGPT) have moved to warmup + cosine.

Code

Warmup + cosine decay schedule·python

import math

def lr_schedule(step, total_steps, warmup_steps=2000,
                max_lr=3e-4, min_lr=3e-5):
    if step < warmup_steps:
        return max_lr * step / warmup_steps
    progress = (step - warmup_steps) / (total_steps - warmup_steps)
    progress = min(1.0, max(0.0, progress))
    return min_lr + 0.5 * (max_lr - min_lr) * (1 + math.cos(math.pi * progress))

# Plug into PyTorch:
# scheduler = torch.optim.lr_scheduler.LambdaLR(
#     optimizer,
#     lambda step: lr_schedule(step, total_steps) / max_lr
# )

Learning Rate: Warmup + Cosine Decay

Why warmup

Why cosine decay

Code

External links

Exercise

Progress

Comments 0