Learning Rate Warmup and the Early Instability Window

Almost every published training recipe for a deep network begins with a learning rate that ramps from near zero to its target value over some number of steps. The ramp is not cosmetic. It exists because a randomly initialised network at step zero has no reason to produce well-scaled gradients, and the adaptive optimizers used in practice respond to badly scaled gradients by producing steps whose size is unrelated to the loss landscape.

The clearest statement of the problem comes from the Adam optimizer’s own arithmetic. Its update at step $t$ divides the momentum estimate by the square root of the second-moment estimate, and that denominator is a biased estimator whose bias correction divides by $1-\beta_2^t$. At $t=1$ with $\beta_2 = 0.999$, the correction multiplies the estimate by roughly a thousand. The ratio that determines the size of the update is therefore computed from one observation, and there is nothing to average.

The arithmetic of a first step

Adam maintains $m_t = \beta_1 m_{t-1} + (1-\beta_1) g_t$ and $v_t = \beta_2 v_{t-1} + (1-\beta_2) g_t^2$, and updates $\theta \leftarrow \theta - \eta , \hat{m}_t / (\sqrt{\hat{v}_t} + \epsilon)$ where hatted quantities are bias corrected. At the first step, $\hat{m}_1 = g_1$ and $\hat{v}_1 = g_1^2$ exactly. The update is therefore $\eta , g_1 / (|g_1| + \epsilon)$, which is $\pm \eta$ for every coordinate whose gradient exceeds $\epsilon$ in magnitude.

This is the central fact: the first Adam step has magnitude $\eta$ per coordinate regardless of the gradient’s scale. A parameter whose gradient is a thousand times larger than another’s receives the same step. If $\eta$ is set to a value appropriate for a converged model, say $10^{-4}$, the first step moves every coordinate by $10^{-4}$. If $\eta$ is set to the large value a practitioner might choose for fast initial progress, every coordinate moves by that amount, including coordinates whose correct direction is nearly zero and whose gradient is dominated by initialization noise.

The bias correction is what makes the $t=1$ step scale-free, and it fades quickly: by the time $t$ reaches a few hundred with $\beta_2$ near $0.999$ the correction is mild. Warmup covers exactly the window in which the second-moment estimate is built from too few samples to be meaningful. After roughly $1/(1-\beta_2)$ steps the estimate has seen a full effective window; before that, its relative error is large.

A bound rather than a recipe

Adam’s update satisfies $|\Delta\theta_i| \le \eta$ per coordinate, since $\hat{m}_i / \sqrt{\hat{v}_i}$ has magnitude at most about one whenever $|\hat{m}_i| \le \sqrt{\hat{v}_i}$, which holds when the momentum estimate does not exceed the root-mean-square estimate. The bound is violated when $\beta_1$ is large and the recent gradients are much smaller than the moments stored in $\hat{v}$: the numerator reflects a trend while the denominator reflects a history. Right after a large gradient, the second moment is inflated and later, smaller gradients produce steps proportionally larger relative to their own scale. Warmup does not fix this asymmetry; it limits the absolute damage during the period when the estimate is least reliable.

What warmup bounds

Three failure modes are commonly attributed to the absence of warmup, and each has a different mechanism.

The first is divergence or loss spike within the first few dozen steps: the parameters move far enough that the network’s representation degrades faster than the loss can follow, and the next gradients are computed from a worse point. With a step bounded by $\eta$ per coordinate and a parameter initialized at scale $1/\sqrt{n}$, a learning rate of order $1$ can move each weight by an amount comparable to the weight itself in one step. The linear ramp keeps the product of step size and step count small while the estimate settles.

The second is the destruction of the momentum and second-moment state: a large early step produces gradients that are wildly unrepresentative, and those gradients enter $v_t$ with weight $1-\beta_2$, which for the first few steps is a substantial fraction of the estimate. The estimate then takes $O(1/(1-\beta_2))$ steps to wash out. Warmup keeps the unrepresentative gradients from moving the parameters, though not from entering the estimate.

The third is specific to residual and normalised architectures: the scale of the initial forward pass is set by the initialization, and a large step at the start perturbs a carefully balanced signal before any normalization buffer has converged, interacting with the mismatch described in batch normalization and its train-eval gap.

def lr_at(step, base_lr, warmup_steps):
    # linear ramp to base_lr, then constant; the caller appends decay
    if step >= warmup_steps:
        return base_lr
    return base_lr * (step + 1) / max(1, warmup_steps)

The +1 matters: at step zero the rate is $1/\text{warmup_steps}$ of the target rather than zero, since a zero step wastes an iteration and some implementations index steps from one.

Coupling warmup to model width

The length of the ramp is not a free parameter in the transformers that made warmup standard. The original transformer schedule was $\eta_t = d_{\text{model}}^{-0.5} \min(t^{-0.5}, t \cdot t_{\text{warmup}}^{-1.5})$, in which the ramp length is not proportional to total training length but is roughly a fixed few thousand steps. That form implies that at the peak, $\eta \approx d_{\text{model}}^{-0.5} t_{\text{warmup}}^{-0.5}$; increasing width lowers the peak learning rate, and the ratio of warmup length to peak rate is constrained by the schedule’s own shape rather than by an external choice.

The variance argument for warmup in adaptive methods gives a related scaling. The relative error of the second-moment estimate depends on the number of effective samples in the moving average, and the gradients’ variance does not shrink with width in the same way their mean does; a wider layer’s mean gradient direction is better defined while the per-coordinate noise can remain large. The analysis behind the rectified form of Adam attributes early-training instability to the variance of the adaptive denominator and prescribes a longer warmup for settings where that variance is large. The practical implication is that a warmup length tuned for one width is not transferable to another without rechecking.

Where the ramp must not be confused with the schedule

Warmup and decay are separate transformations of the same rate, and conflating them produces a schedule that no single mechanism explains. The ramp is a protection for the optimizer’s state estimation and for the initialization’s balance. The decay is a statement about the loss surface: further progress requires smaller steps because the gradient noise around a minimum has a magnitude that does not vanish. Warmup lengths are typically 1–5% of total steps in vision training and a few thousand steps in language model pretraining, where total length is much larger; the ratio is not a constant, and a long warmup on a short run is equivalent to never reaching the intended rate.

Two degenerate cases are worth naming. A warmup that starts at exactly zero and reconstructs the rate from step index means the first step performs no update; some implementations therefore define the schedule at step $1$ rather than $0$, and checkpoints resumed from a saved step index can produce a discontinuity if the convention differs between save and load. A warmup that is applied only to the global rate while a per-group rate multiplier is not, as when the bias and normalization parameters use a different rate, scales the ramp differently for each group — which is usually intended, since the higher-rate group is the one the protection targets, but is a frequent source of misconfiguration.

The alternative to warmup is to remove its cause: use a fresh optimizer state with a second-moment estimate that does not need correcting, for instance by initialising the estimate from a short calibration pass, or by using an optimizer whose step size does not depend on a per-coordinate ratio. Those are different mechanisms with different costs, and the decoupled weight decay discussion shows the same pattern of a term that interacts with adaptive scaling in ways that the naive analysis misses.