Learning Rate Warmup and Decay: The Schedule That Decides Whether a Model Trains at All

The core idea and the invariant it protects

Gradient descent updates parameters by θ ← θ − α · ∇L(θ). The learning rate α is the one scalar that scales the entire step. It sits at the intersection of two failure modes: too large and the update overshoots, oscillates, or diverges to NaN; too small and training crawls, wasting compute and settling in a poor basin. A learning rate schedule is a deliberate, time-varying policy α(t) that resolves this tension by using a different rate at different phases of training.

The schedule has two conceptual halves. Warmup starts α near zero and ramps it up over the first w steps. The invariant it protects is stability: while parameters are far from any good region and gradient estimates are noisy, no single step should move θ more than a small trust-region radius. Decay then shrinks α as t grows. Its invariant mirrors stochastic-approximation theory: to converge to a minimum rather than bounce around it, the step size must shrink so the noise ball around the optimum tightens over time.

Formally, warmup + decay is a piecewise function evaluated per optimizer step:

alpha(t):
  if t < w:                 # warmup
    return alpha_max * (t / w)
  else:                      # decay (cosine example)
    p = (t - w) / (T - w)
    return 0.5 * alpha_max * (1 + cos(pi * p))

Every entry point — SGD, Adam, AdamW, LAMB — reads α(t) once per step and applies it uniformly (or per-parameter, in adaptive methods) to the update.

Why warmup stops divergence: the large-batch and adaptive story

Two independent lines of evidence made warmup standard practice. The first is large-batch training (Goyal et al., 2017, "Accurate, Large Minibatch SGD"). The linear scaling rule says: multiply the batch size by k, multiply α by k. But at batch 8,192 that means α is ~32× the batch-256 baseline value, and at step 0 — random init — the gradient direction is essentially meaningless. A full-size step in a meaningless direction destroys the network. Warmup lets the model take tiny, safe steps while the weights organize into a sane configuration, then unlocks the aggressive rate the large batch can tolerate because its gradient variance is 1/k smaller.

The second is adaptive optimizers on Transformers. Adam divides each gradient by a running estimate of its second moment, √(v̂ + ε). Early in training v̂ is computed from very few samples, so its variance is enormous — the effective step size Adam takes has a heavy tail and occasionally becomes gigantic. Liu et al. (2020, RAdam) showed this variance is the mechanistic reason Transformers need warmup: it is a crude but effective variance-reduction bandage for the first w steps. The original Transformer paper (Vaswani et al., 2017) baked warmup directly into its schedule:

alpha(t) = d_model^(-0.5) * min(t^(-0.5), t * w^(-1.5))

This is linear warmup for t ≤ w, then inverse-square-root decay for t > w, with the crossover exactly at the peak. It requires no knowledge of the total step count T — useful when you do not know in advance how long you will train.

The decay shapes and what they optimize

• Step decay: α ← α·γ every s steps (e.g. ÷10 at epochs 30, 60, 90 for ImageNet ResNet). Discontinuous, dead simple, and historically the strongest CNN result. The visible loss "cliff" at each drop is the model settling into a tighter basin.
• Cosine annealing (Loshchilov & Hutter, 2017): α(t) = ½α_max(1 + cos(π·p)), p ∈ [0,1] the fraction of remaining budget. Smooth, spends most of its time at a high rate and only anneals hard near the end, and empirically edges out step decay on many benchmarks. Requires a known total T. Cosine restarts (SGDR) reset p to 0 periodically to escape sharp minima.
• Inverse-√t: α ∝ 1/√t. The Transformer default; assumes no fixed budget.
• Polynomial decay: α_max·(1 − t/T)^k, k=1 being linear-to-zero (BERT's schedule).
• One-cycle (Smith, 2018): a single warmup-then-anneal triangle plus an inverse momentum schedule, enabling "super-convergence" at very high peak rates.

The theoretical backbone is Robbins-Monro (1951) stochastic approximation: a step-size sequence converges if Σα_t = ∞ (you can still travel arbitrarily far) and Σα_t² < ∞ (accumulated noise is finite). 1/t satisfies both; a constant rate satisfies neither, which is exactly why constant-α SGD converges only to a noise ball, not a point.

Complexity: why the schedule is free

Evaluating α(t) is O(1) time and O(1) space per step — a handful of floating-point ops (a division, a cosine, a min/max). Against the per-step cost of the forward/backward pass — O(P) for P parameters, and for a Transformer O(n²·d) in sequence length n and width d — the scheduler is unmeasurably cheap. Over a full run of T steps the schedule adds Θ(T) scalar operations total, dwarfed by the Θ(T·P) work of the optimizer.

State-based schedules cost slightly more. ReduceLROnPlateau keeps the best-so-far metric and a patience counter — O(1) extra space, O(1) update, but it must be fed a validation metric, adding an evaluation pass every k steps. Cosine and inverse-√t need only the integers t, w, T. Step decay needs s and γ. None of these change the asymptotic cost of training.

A subtle correctness point: t must be the global optimizer-step index, not the epoch or the micro-batch. With gradient accumulation over a micro-batches, α(t) should advance once per effective step, not once per micro-batch — a common bug that makes warmup finish a× too fast and silently destabilizes large-batch runs. Similarly, when resuming from a checkpoint you must restore t, or the schedule restarts and the rate jumps.

Tuning: choosing α_max, w, and T

The three knobs interact, but there is a reliable recipe.

• α_max: find it with an LR range test (Smith, 2015): train for a few hundred steps while exponentially increasing α, plot loss vs α, and pick the value roughly one order of magnitude below where the loss starts to diverge (the steepest-descent region). For AdamW on Transformers, α_max is typically 1e-4 to 3e-4 for large models, 1e-3 for small ones.
• Warmup steps w: a small fraction of total steps — commonly 1–10%. BERT used 10k warmup steps out of ~1M; GPT-3 used 375M tokens of warmup. As a rule of thumb, larger batch ⇒ longer warmup. Too-short warmup reintroduces the divergence it was meant to prevent; too-long warmup wastes budget at low rates.
• Total steps T: for cosine/poly you must commit to T up front. Undershooting T means you decay to near-zero before you have used your data; overshooting means the run ends while α is still high, leaving performance on the table. If T is genuinely unknown, use inverse-√t or plateau-based decay instead.

The linear scaling rule ties α_max to batch size B: α_max ≈ α_base · (B / B_base). It holds well up to a batch-size ceiling beyond which returns vanish — LARS (You et al., 2017) and LAMB (You et al., 2019) extend that ceiling with per-layer trust ratios, which is how BERT was trained in 76 minutes at batch 32k.

Where it runs in production

Warmup + decay is not a niche trick; it is the default in every serious training stack.

• Transformers / LLMs: GPT-2/3, BERT, T5, LLaMA all use linear warmup then cosine (or linear) decay to a small fraction of α_max. The scheduler lives in Hugging Face transformers as get_cosine_schedule_with_warmup and friends.
• Vision: the 1-hour ImageNet result, ViT, and DeiT all use warmup; ViT specifically pairs it with cosine decay and is famously warmup-sensitive.
• Frameworks: PyTorch ships torch.optim.lr_scheduler (CosineAnnealingLR, OneCycleLR, LambdaLR, ReduceLROnPlateau); Keras has LearningRateSchedule; JAX/Optax composes schedules with optax.warmup_cosine_decay_schedule. Every one applies the schedule by multiplying the base step, so it composes cleanly with weight decay, gradient clipping, and mixed precision.

The reason it is universal: the schedule is decoupled from the model. It observes only t and emits a scalar, so the exact same 20 lines work for a 5M-parameter classifier and a 500B-parameter LLM. That decoupling is why it survives every architecture change.

Pitfalls, edge cases, and variants

• The gradient-accumulation off-by-a bug (above): step the scheduler per effective step. The symptom is a warmup that visibly finishes too early plus early instability.
• Checkpoint resume resets t: if you rebuild the scheduler without restoring its step count, α jumps back to the warmup value and can blow up a converged model. Always checkpoint scheduler state alongside optimizer state.
• Decaying to exactly zero too early: cosine with a mis-set T can hit α≈0 while data remains, freezing learning. Many recipes floor α at a small α_min (e.g. 10% of α_max) rather than 0.
• Warmup and Adam bias correction interact: Adam already applies 1/(1−β₁ᵗ) and 1/(1−β₂ᵗ) bias correction, which is itself a mild implicit warmup on the second moment. RAdam (2020) makes this explicit and can remove the need for a hand-tuned warmup length — a clean alternative when you cannot afford to tune w.
• Not scaling weight decay with the schedule: in AdamW, decoupled weight decay is often scaled by α too; forgetting this changes the effective regularization as the rate anneals.
• Batch-size / rate mismatch under the linear rule: past the critical batch size the rule over-scales α and diverges even with warmup; switch to LARS/LAMB or cap the batch.

The through-line: α(t) is trivially cheap to compute but globally coupled to init, optimizer internals, batch size, and total budget. Almost every training instability report in the wild traces back to one of these coupling mistakes, not to the loss function or the architecture.

Common learning-rate schedules and where each wins

Schedule	Formula (after warmup)	Extra state	Best regime
Constant	α(t) = α_max	none	Debugging, small convex problems
Step decay	α·γ^⌊t/s⌋	step size s, factor γ	Classic CNNs (ResNet ImageNet)
Cosine annealing	½α(1+cos(πt/T))	total steps T	Fixed-budget vision & LLM pretraining
Inverse √t	α_max·√(w/t)	warmup w	Transformers (Adam), unknown T
ReduceLROnPlateau	α ← α·γ on stall	patience, best-metric	When T is unknown, metric-driven