41. Learning-Rate Warmup + Cosine Decay Schedule

Problem

Training a model directly at the target learning rate from a randomly-initialized state is unstable — early gradients (and, for adaptive optimizers, the second-moment estimate) are noisy before the model has seen enough data, so a too-large step can throw the parameters somewhere the optimizer never recovers from. The standard fix is a linear warmup: ramp the learning rate up from 0 over the first few hundred or thousand steps, then decay it smoothly, most commonly with a cosine curve, down to some floor for the rest of training.

Implement the schedule as a pure function of the step count. It's a small piece of math, but the boundary conditions (exactly at the warmup/decay transition, and past the end of the schedule) are easy to get subtly wrong.

Source: src/41_warmup_cosine_lr_schedule.py

def warmup_cosine_lr(step: int, warmup_steps: int, total_steps: int, base_lr: float, min_lr: float = 0.0) -> float:
    ...

Examples:
>>> warmup_cosine_lr(step=0, warmup_steps=100, total_steps=1000, base_lr=1e-3)
0.0
>>> warmup_cosine_lr(step=100, warmup_steps=100, total_steps=1000, base_lr=1e-3)
0.001

Step-by-Step Approach

1. Branch on whether step is inside the warmup window: if step < warmup_steps, the schedule is a simple linear ramp, base_lr * step / warmup_steps — 0 at the very first step, base_lr exactly at step == warmup_steps.
2. Outside warmup, remap the remaining steps onto a [0, 1] "progress" fraction: (step - warmup_steps) / (total_steps - warmup_steps).
3. Clamp that progress to [0, 1] explicitly — without the clamp, steps beyond total_steps would push progress past 1, and the cosine would start rising back up instead of staying pinned at the floor.
4. Apply the standard half-cosine curve: 0.5 * (1 + cos(pi * progress)), which is exactly 1 at progress=0 and exactly 0 at progress=1.
5. Map that curve onto [min_lr, base_lr]: min_lr + (base_lr - min_lr) * cosine.
6. Check continuity at the boundary: the warmup branch gives exactly base_lr at step == warmup_steps, and the decay branch's cosine term is exactly 1 at that same point (progress=0) — the two branches must agree exactly, with no jump.
7. Check the schedule's midpoint: at progress=0.5, cos(pi/2) = 0, so the learning rate should land exactly halfway between base_lr and min_lr — a useful independent sanity check beyond just the two endpoints.

The key insight is that both branches of the schedule have to agree exactly at the transition step, and the decay branch has to be clamped so it doesn't wrap around once training runs longer than total_steps — both are one-line fixes, but skipping either produces a schedule that looks fine on a plot until you check the exact boundary values.

Reference solution

def warmup_cosine_lr(step, warmup_steps, total_steps, base_lr, min_lr=0.0):
    if warmup_steps > 0 and step < warmup_steps:
        # Linear ramp: 0 at step=0, base_lr exactly at step=warmup_steps.
        return base_lr * step / warmup_steps

    # Past warmup: map the remaining steps onto [0, 1], clamped so anything
    # beyond total_steps stays pinned at the schedule's end (progress=1),
    # instead of the cosine wrapping back around past its own period.
    decay_span = max(1, total_steps - warmup_steps)
    progress = (step - warmup_steps) / decay_span
    progress = min(max(progress, 0.0), 1.0)

    # cos(0) = 1 -> full base_lr right at the warmup/decay boundary.
    # cos(pi) = -1 -> min_lr exactly at progress=1 (step == total_steps).
    cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
    return min_lr + (base_lr - min_lr) * cosine

Key Functions & Tricks

• math.cos(math.pi * progress) — the half-cosine curve, evaluated over [0, pi] to sweep from +1 to -1.
• min(max(progress, 0.0), 1.0) — the clamp that prevents the schedule from wrapping past its own end.
• max(1, total_steps - warmup_steps) — guards the decay-span denominator against a degenerate zero-length decay phase.
• Boundary-value testing (step=0, step=warmup_steps, step=total_steps, step>total_steps) — the right test strategy for any piecewise schedule function.
• Midpoint check via cos(pi/2) == 0 — an independent algebraic sanity check beyond just the two endpoints.

How to Recognize This Pattern