Learning Rate Schedules — Why Warmup and Why Cosine

The learning rate is not a fixed number — it is a curve you design across the whole run. Why we deliberately start slow (warmup), why we come down along a cosine, and what else has to move when batch size changes. Equations, live figures, PyTorch code, and the mistakes that actually break runs.

Attention Is All You Need

An analogy: pulling out of an icy lot and backing into a space

You are backing a car into a space across a snow-covered parking lot. Right after you pull out you have no idea whether the surface is ice, so you ease onto the throttle and feel for grip. Across the open stretch you carry real speed. As the lines come up you slow down — carry that speed in and you overshoot, correct, overshoot again, and never settle.

Training a neural network moves through those same three phases. Easing onto the throttle is warmup, the open stretch is the peak learning rate, and slowing for the fine adjustment is decay (cosine and friends). A learning rate schedule is that throttle work, written down as a function of the training step.

What the learning rate actually decides

Training is one thing repeated: nudge the weights. A single nudge looks like this.

Here $\theta_t$ are the parameters at step $t$, $\nabla L(\theta_t)$ is the gradient — the arrow giving the direction the loss $L$ decreases and how steeply — and $\eta_t$ is the learning rate. Equation (1) says only this: from where you stand, take one step downhill of size $\eta_t$.

Written out in words: new weights = current weights − (stride) × (downhill direction). The minus sign is what makes this a descent rather than a climb, and the multiplication is a division of labour — the gradient chooses the direction, the learning rate chooses how far you go in it. The gradient itself is unpacked in Loss Functions and Optimization.

The crucial part is that the gradient tells you only a direction and a steepness — nothing about how far you must travel to reach the bottom. That distance is what the learning rate supplies, and the right value changes with where you are. Hence the subscript $t$ on $\eta$: not a constant, but a value chosen fresh each step.

What goes wrong with one fixed stride

Too small, and you head the right way but never arrive. GPU hours are finite, so that is failure in practice.

Too large, and you clear the valley in one stride and land higher up the opposite slope. The gradient there still points into the valley, so you leap across again. Each crossing widens the swing until the loss diverges into NaN.

The awkward part is that you cannot see the boundary in advance. Feel for it in the figure below: raise the slider and convergence speeds up — until one notch past a threshold the ball flies out of the valley entirely.

Two things fall out. There is a ceiling on a good learning rate, and the closer you get to the bottom, the smaller the right stride becomes. So: big early, small late. That is where decay comes from.

Why not simply run a small learning rate throughout and stay safe? Training does still work that way, but the loss you reach on the same compute budget is clearly worse. A schedule is not a device for avoiding crashes; it is a device for reaching the lowest loss within a fixed number of GPU hours.

And there is one more twist. Almost no modern run uses a plain downward slope — it deliberately starts small, climbs, and only then comes down. That odd hump is warmup plus cosine.

warmup — deliberately crawling for the first few hundred steps

Warmup is simple. Start near zero and raise the rate linearly to the peak $\eta_{\max}$ over a chosen number of steps $T_{\text{warm}}$.

Equation (2) says the rate rises in proportion to how far through the warmup you are. With a 1,000-step warmup, step 100 sits at 10% of the peak. That is the whole rule.

Put in words, $t / T_{\text{warm}}$ is a progress meter that runs from 0 to 1 across the run-up, and the learning rate at any moment is simply that meter's reading times the peak value.

The real question is why anyone would do this. Three reasons, and they stack rather than compete.

1. The starting weights are random noise. The loss is high and the gradients are large, but nothing guarantees they point somewhere useful. Take a big step from there and you can land worse off than random initialization left you. Treat the first few steps as gathering information rather than making progress.

2. Adam's internal statistics are still estimated from a handful of gradients. Adam tracks a running estimate of the squared gradient and divides by its square root to set a per-parameter effective stride. Right after the start that estimate is small and noisy, and dividing by a small number sends the effective learning rate spiking. The RAdam paper (Liu et al., 2020) singled out this early variance as the main reason warmup helps. No single explanation has won out, but the phenomenon — Adam-family optimizers tend to break early without warmup — reproduces widely.

3. The bigger the batch, the more the first step counts. A larger batch means less gradient noise, so every parameter moves the same way at once. Marching everything together in a direction not yet shown to be correct is exactly the risky case.

A worked example: the hump from the original Transformer paper

The classic case of folding the climb and the descent into one expression is the schedule from the original Transformer paper (Vaswani et al., 2017).

$d_{\text{model}}$ is the model's hidden dimension. Read equation (3) as: rise in proportion to $t$ until $T_{\text{warm}}$, then fall off as one over $\sqrt{t}$. Taking the min of the two terms is what turns the hump into a single formula. The paper used $T_{\text{warm}} = 4000$.

The leading $d_{\text{model}}^{-0.5}$ is a size discount, which says the wider the model, the shorter every stride across the whole run. Inside the min sit two curves — the rising line of the climb and the falling curve of the descent — and at each step whichever is currently lower is the one you get.

cosine — why a cosine, of all things

The standard way to come down after the peak is cosine decay, in the form proposed in Loshchilov and Hutter's SGDR paper (2017).

$\eta_{\min}$ is the floor you keep at the end and $T_{\text{total}}$ is the total number of training steps. Equation (4) says: map progress from the end of warmup to the finish line onto an angle from $0$ to $\pi$, and let the cosine carry the rate smoothly from $\eta_{\max}$ down to $\eta_{\min}$. Halfway through, you are at exactly the midpoint value.

Stated in words, the $1 + \cos(\cdot)$ term is a dial that slides gently from 2 to 0 as the run proceeds. Halving it, multiplying by the gap between the peak and the floor, and adding the floor back on is the three-step recipe that converts that dial reading into an actual learning rate.

So why not a straight line or a staircase? Because a cosine is flat at both ends and steepest in the middle. Just past the peak you keep speed and cover ground; near the end you approach zero slowly, which is what fine adjustment needs. A step decay approximates the same shape with discrete drops, but it forces a human to pick when to drop and by what factor. A cosine is fully determined once you fix the peak and the total step count — fewer knobs to tune is the practical reason it wins.

That same property is also the trap. The Chinchilla paper (Hoffmann et al., 2022) states explicitly that the cosine cycle length should be matched to the actual number of training steps. Set the cycle longer and stop early, and the run ends before the rate has come down, which costs you final loss. Same model, same data, different conclusion — purely from a schedule setting. How much that mattered in the scaling-law experiments is covered in Scaling Laws from Scratch.

More recently, WSD (Warmup–Stable–Decay) schedules — a long constant plateau followed by a short decay — have also come into use. When you want to change how many more steps you will run mid-flight, as in continued pretraining, not being locked to a total step count is the point.

The batch size connection — two knobs that move together

The $\nabla L$ in equation (1) is really the average of the per-sample gradients in the batch. Enlarge the batch and that average gets less noisy. A step taken with a large batch is "the same stride, but in a more trustworthy direction" — and if it is more trustworthy, you can afford to stride further.

The simplest formulation of this is Goyal et al.'s linear scaling rule (2017): multiply the batch size by $k$ and multiply the learning rate by $k$ too. Combining that rule with a gradual warmup over several epochs, they trained ResNet-50 on ImageNet at very large batch sizes. Warmup shows up here as the necessary partner, again.

The proportionality does not extend forever. McCandlish et al.'s An Empirical Model of Large-Batch Training (2018) introduced the critical batch size, set by how large the gradient noise is. Below it, growing the batch buys you a proportional cut in step count; above it, the step count barely moves and you are simply burning compute. The noise is already small enough that a more precise average buys nothing.

The practical takeaway is clean: when the effective batch grows, the peak learning rate should go up with it. Under distributed data parallelism the effective batch is silently multiplied by world_size, so adding nodes moves this knob whether you meant to or not (see Distributed Training from Scratch). You will also hear that SGD follows the linear rule while Adam follows a square-root rule; that is an empirical observation that depends on the task and the implementation, not a settled law. Confirm it with a small-scale run.

Writing it: warmup plus cosine

In PyTorch you hand LambdaLR a function returning a multiplier on the peak learning rate, and you are done.

import math
from torch.optim.lr_scheduler import LambdaLR

def warmup_cosine(step, warmup, total, min_ratio=0.1):
    if step < warmup:
        return step / max(1, warmup)                     # linear 0 → 1
    p = (step - warmup) / max(1, total - warmup)         # progress 0 → 1
    return min_ratio + (1 - min_ratio) * 0.5 * (1 + math.cos(math.pi * p))

sched = LambdaLR(opt, lambda s: warmup_cosine(s, warmup=2000, total=100_000))
# call sched.step() right AFTER opt.step() — never on an accumulating micro-batch

That last comment marks the single most common way this breaks in practice. It is also worth seeing how the learning rate and momentum interact before you tune either.

How this shows up on the job

Who touches it, and when. The ML engineer launching a pretraining, continued-pretraining, or fine-tuning job sets this in the config before submitting. Changing the learning rate mid-run makes everything after it a different experiment, so the ten minutes before you hit submit are essentially the whole decision.

The names you actually type. In Hugging Face TrainingArguments: learning_rate (which is the peak), lr_scheduler_type="cosine", warmup_ratio or warmup_steps, max_steps, gradient_accumulation_steps. Calling the helper directly, it is get_cosine_schedule_with_warmup. In plain PyTorch: LambdaLR, CosineAnnealingLR (where T_max counts calls to step(), not epochs), OneCycleLR. In a DeepSpeed config JSON: scheduler.params.warmup_num_steps and total_num_steps.

What must be logged. Not just the loss — log the learning rate itself and the gradient norm to W&B or TensorBoard. Schedule failures show up on the learning rate curve before they show up on the loss curve.

The mistakes that cost runs.

• Calling sched.step() on every micro-batch under gradient accumulation. With gradient_accumulation_steps=8 the optimizer updates once every eight batches. Stepping the scheduler every time runs the schedule 8× too fast, so most of training happens at the floor learning rate. The loss still goes down, which is why this hides so well.
• Forgetting to restore the scheduler's state_dict on resume. The rate jumps back to the peak and the loss spikes with it. On preemptible GPUs this is fatal.
• Changing the total step count after the fact. Editing max_steps changes the shape of the cosine itself, and a run whose cycle length disagrees with its real length pays for it in final loss.
• Judging a config from a checkpoint taken before the decay finishes. Loss drops visibly in the last stretch of the decay, so cutting a verdict at the 70% mark misreads it.
• Changing batch size and leaving the learning rate alone. Adding GPUs, or halving per_device_train_batch_size to dodge an OOM — both changed the effective batch.
• Making warmup too short. Early loss spikes and NaNs point first at the warmup length and the peak value. Fix those before reaching for symptomatic workarounds.

What a design review will ask. For "why is warmup needed," three points suffice: the unreliability of random initialization, the estimation variance in Adam's second moment, and how little noise a large batch carries.

Summary

• The learning rate is not a constant but a curve designed as a function of the step
• warmup insures against three instabilities: random initialization, Adam's early internal statistics, and low noise at large batch
• cosine reduces the knobs you must pick to two: the peak and the total step count
• Change the batch size and the learning rate moves too — but past the critical batch size, growing it buys little
• The implementation traps are gradient accumulation and resume. Always log the learning rate curve