Learning-rate schedules & warmup

1

Concept

The optimizer turns gradients into parameter updates, but the learning rate decides their scale. One constant rate rarely serves an entire pretraining run well. Early training has uncalibrated activations and optimizer moments; middle training benefits from substantial progress; late training often benefits from smaller, refining updates. A schedule encodes that changing risk.

Warmup starts below the intended peak and increases over an initial interval. Linear warmup at step $t$ can use

for $t\le T_{warmup}$. The first updates are restrained while normalization scales, gradients, and Adam’s moment estimates settle. Warmup does not repair a fundamentally excessive peak rate; it only changes how that rate is reached.

After warmup, linear decay reduces the rate at a constant pace toward a chosen endpoint. It is simple and makes remaining progress easy to interpret. Cosine decay follows a half-cosine, dropping gently near the start and end of the decay region. Inverse-square-root schedules decay proportionally to $t^{-1/2}$ after a warmup construction and were used in the original Transformer recipe. Constant-with-warmup schedules retain the peak after ramp-up, which can be useful when the total horizon is uncertain but may leave late updates noisy.

The endpoint matters. Decaying exactly to zero assumes the run ends at the planned horizon; extending it provides no learning unless the schedule is changed. A nonzero minimum preserves adaptation but can keep disturbing a nearly converged model. Restarts, popular in some domains, introduce deliberate rises and are not automatically appropriate for expensive one-pass pretraining.

Define progress in the correct unit. A “1,000-step warmup” means something different when global batch changes from one million to four million tokens. Token-based schedules align with data exposure and survive some topology changes. If sequences have padding, count effective loss-bearing tokens rather than allocated slots when that distinction matters.

AdamW separates gradient-based optimization from weight decay, but scheduler and decay still interact. Parameters such as normalization scales and biases are often excluded from weight decay. Gradient clipping limits an update driven by a large gradient before the optimizer step; it is a guardrail, not a replacement for a stable schedule.

Scaling the global batch can motivate a learning-rate adjustment, but linear or square-root rules are heuristics with regime boundaries. Large batches reduce gradient noise and may require longer warmup or different optimizer tuning. Run pilot sweeps rather than extrapolating a rule across orders of magnitude.

Log the realized rate at every step alongside token count, loss, gradient norm, optimizer statistics, and skipped mixed-precision updates. When resuming, restore scheduler position. Restarting warmup accidentally after a checkpoint can change the optimization path; jumping to the wrong late-run rate can create a spike.

A schedule plotter should show warmup, peak, decay, floor, and total token horizon, with hover values in both steps and tokens. Overlaying a loss curve helps correlate instability with update scale, but correlation alone does not prove the schedule caused a spike.

The durable model is controlled momentum for learning. Warmup avoids wrenching an uncalibrated network, the central phase moves quickly through useful parameter space, and decay narrows the step size as the run approaches its planned endpoint. The schedule is part of the experiment specification and checkpoint state, not cosmetic optimizer configuration.