Core ML

Learning Rate Warmup

Ramp the learning rate up from near zero over the first few hundred steps, because the earliest updates run on the least trustworthy gradient estimates.

The learning rate ramps linearly from zero up to its peak over the warmup window, then hands off to a decay schedule that eases it back down for the rest of training

Why Does This Exist?

Train a 12-layer transformer from scratch at its intended peak learning rate from step zero, and the first few hundred steps are often where the run dies — loss spikes, then nan, on a model that would otherwise train cleanly for the next hundred thousand steps at that exact same rate. Nothing about the rate itself was wrong; it was wrong specifically for the state the weights were in at step one.

At initialization, weight initialization has set every layer's variance to a sane starting point, but the network hasn't taken a single training step yet — the first gradients it computes are a rough, noisy first impression of the loss surface, not a settled read on it. Adam's second moment $v_t$ makes this worse in a specific way: it starts at zero and needs several steps of bias correction to become a trustworthy estimate of a parameter's gradient scale, so the very first updates can be scaled by an estimate that hasn't stabilized yet. A large step, taken from an unstable estimate, on weights that have never been trained, is the combination that blows up.

Think of It Like This

Warming up a cold engine before redlining it

A car engine that's been sitting overnight has cold oil, thick and slow to circulate, and metal parts that haven't reached their normal operating clearances. Floor the accelerator immediately and you're running the engine hard before the oil has had a chance to protect the parts it's supposed to protect — that's how engines get damaged.

Idle for a minute or two first, letting the oil warm and thin out and circulate properly, and the same hard acceleration afterward is completely fine — nothing about the car changed except that it had a chance to reach a stable operating state first. Learning rate warmup is that idle period: hold off on the full-strength updates until the model's gradient estimates have had a few steps to settle, then apply full strength safely.

How It Actually Works

The ramp itself

Warmup is almost always linear: over the first $W$ steps, the learning rate rises in equal increments from (near) zero up to the schedule's intended peak rate $\eta_{\max}$ :

\eta_t = \eta_{\max} \cdot \frac{t}{W}, \qquad t \le W

$t$ is the current step, $W$ the number of warmup steps — commonly a few hundred to a few thousand, depending on model size and how unstable early training tends to be for that architecture. Once $t$ passes $W$ , control hands off to whatever decay schedule is running the rest of training: cosine annealing, step decay, or the flat plateau of warmup-stable-decay.

Why this specifically fixes the Adam instability

Trace what warmup does to Adam's update in the first few steps. The raw step size is $\eta_t \cdot \hat{m}_t / (\sqrt{\hat{v}_t} + \epsilon)$ — and at step 1, $\eta_1 = \eta_{\max}/W$ is tiny by construction, whatever $W$ is set to. The unreliable bias-corrected estimates from the first few steps get multiplied by a rate that's still near zero, so their damage is bounded exactly during the window where they're least trustworthy. By the time $\eta_t$ has ramped up to something large, $t$ steps of bias correction have already passed and the moment estimates have had a chance to become reasonable.

Where it matters most

Warmup is close to mandatory for transformer training from scratch — the original architecture's training recipe used it, and it remains standard practice across nearly every large language model trained since. It matters less for small, well-behaved convolutional networks trained with plain SGD, where the instability warmup exists to prevent is milder to begin with. The bigger the model and the more aggressive the peak learning rate, the more warmup earns its place.

Show Me the Code

The same ten steps, with and without a 10-step linear warmup, showing how much smaller the earliest updates are.

1import numpy as np2
3
4def warmup_lr(step: int, warmup_steps: int, lr_max: float) -> float:5    return lr_max * min(1.0, (step + 1) / warmup_steps)6
7
8steps = np.arange(10)9no_warmup = np.full(10, 1e-3)10warmup = np.array([warmup_lr(int(s), 10, 1e-3) for s in steps])11
12print(f"first 5 steps, no warmup: {np.round(no_warmup[:5], 6)}")13print(f"first 5 steps, warmup:    {np.round(warmup[:5], 6)}")14print(f"cumulative rate, no warmup: {no_warmup.sum():.4f}")15print(f"cumulative rate, warmup:    {warmup.sum():.4f}")16# -> first 5 steps, no warmup: [0.001 0.001 0.001 0.001 0.001]17# -> first 5 steps, warmup:    [0.0001 0.0002 0.0003 0.0004 0.0005]18# -> cumulative rate, no warmup: 0.010019# -> cumulative rate, warmup:    0.0055

Warmup spends roughly half the "budget" of those first ten steps compared to running at full rate the whole time — that reduction is concentrated exactly where the moment estimates are least reliable.

Watch Out For

Setting warmup steps as a fixed number copied from a different model size

A warmup length of 4,000 steps tuned for a large model, copied unchanged onto a much smaller one trained on less data, spends a large fraction of that smaller run's entire training budget still ramping up. Warmup length should scale with how unstable the specific model's early training actually is, not be treated as a universal constant — check the loss curve over the warmup window itself before assuming the number is right.

Adding warmup but leaving the rest of the schedule untouched

Warmup only controls the ramp-up; it says nothing about what happens after. A run with warmup added but no decay schedule afterward still trains at a flat peak rate for the rest of training, which reintroduces the exact "one rate for the whole run" problem warmup doesn't solve on its own. Warmup and decay are two separate decisions and both are usually needed together — see learning rate scheduling for the decay half.

The Quick Version

Training instability in the first few hundred steps often comes from an optimizer's moment estimates not having stabilized yet, not from the learning rate being wrong for the rest of the run.
Warmup ramps the rate linearly from near zero up to its peak over a fixed number of steps, bounding the damage while those estimates settle.
It's close to mandatory for transformer training from scratch, and less critical for smaller, plainly-optimized models.
Warmup length should scale with model size and instability, not be copied as a fixed constant.
Warmup handles the ramp-up; a separate decay schedule still has to handle the rest of the run.