LR Warmup Deep Dive

Two independent reasons, not one

Warmup patches two problems that co-occur in the first few hundred steps, and conflating them leads to bad decisions. The optimizer reason is that an adaptive method needs history before its per-coordinate scaling means anything: at step 1 Adam has seen one minibatch, so the denominator it divides by is an estimate with a sample size of one. The architecture reason is that the initialized network sits at a sharp, delicately balanced point, and in Post-LN transformers gradient magnitude grows with depth near the output.

They are separable. Switching Post-LN to Pre-LN removes the architectural cause outright — the companion signal-propagation article derives why — yet Pre-LN models are still trained with several hundred warmup steps. That residue is the optimizer reason, which no architecture change touches. Claiming Pre-LN removed the need for warmup collapses the two.

The first Adam step is a sign step

Write out Adam at t = 1 with bias correction and the pathology falls out in three lines. With m_0 = v_0 = 0:

m_1 = (1 − β_1) g_1        →  m̂_1 = m_1 / (1 − β_1) = g_1
v_1 = (1 − β_2) g_1^2      →  v̂_1 = v_1 / (1 − β_2) = g_1^2

Δθ_1 = −α · m̂_1 / (√v̂_1 + ε)
      = −α · g_1 / (|g_1| + ε)
      ≈ −α · sign(g_1)

Every coordinate moves by the full peak learning rate, in the sign direction of a single minibatch, and the true gradient magnitude has cancelled out: a pure-noise coordinate takes exactly the same size step as one carrying real signal. That is the correct behaviour of a scale-invariant update given one sample; it is simply not a step you want at full size.

What bias correction fixes, and what it cannot

Bias correction does not already handle this; it solves a different problem. Without it, v_1 = (1 − β_2) g^2 ≈ 0.001 g^2, so √v_1 ≈ 0.032 |g| and the update would be roughly 30× too large. Dividing by 1 − β_2^t removes exactly that, making v̂_t an unbiased estimate of E[g^2].

Unbiased is a statement about the mean over hypothetical repeats, not about the reliability of the one draw you have. The update is a ratio of two few-sample estimates, and a ratio’s variance does not vanish because numerator and denominator are individually unbiased. Bias correction makes the early step right on average while leaving it wildly variable step to step. Warmup attacks the variance term it cannot reach — the gap RAdam later tried to close analytically.

How thin is the estimate, quantitatively

Treat v̂_t as a weighted average of past g^2 with normalized weights w_k. Its effective sample size n_eff = 1 / Σ_k w_k^2 has a closed form for an exponential average:

n_eff(t) = (1 − β_2^t)^2 (1 + β_2) / [ (1 − β_2)(1 − β_2^(2t)) ]

t → ∞:  n_eff → (1 + β_2)/(1 − β_2) ≈ 2000   (β_2 = 0.999)
t small:  n_eff ≈ t

For a Gaussian gradient coordinate, Var(g^2)/E[g^2]^2 = 2, so the relative standard deviation of v̂ is √2 / √n_eff, and of √v̂ about half that:

The denominator is trustworthy only after roughly 1/(1 − β_2) steps — a thousand at the default — and that is optimistic: early gradients are non-stationary, so the estimate also chases a moving target.

The curvature argument

The second cause is geometric. Gradient descent on a quadratic with top Hessian eigenvalue λ_max is stable only while η < 2/λ_max; above that the component along the sharpest direction grows every step and the loss diverges geometrically. Adam replaces λ_max with the top eigenvalue of the preconditioned Hessian, but the threshold structure is identical: curvature puts a ceiling on step size.

A randomly initialized transformer is not near a nice basin. Its surface has directions far sharper than the ones it settles into, and sharpness is not static — networks reliably increase curvature early before levelling off near the stability boundary. The peak rate you tuned for the bulk of training is chosen for a region the model has not reached. Warmup is the admission that η_peak describes a well-conditioned regime you must travel to first.

The Post-LN gradient argument

In a Post-LN block the normalization sits on the residual path, so the backward pass inherits a multiplicative depth recursion: gradient magnitudes near the output layers at initialization scale with depth. A full-size first step there moves the weights far enough to destroy the near-critical configuration the initialization set up, and the recursion does not recover on its own. In Pre-LN the residual stream is unnormalized and grows like √L, so each block’s LayerNorm divides by a factor growing with depth — a built-in 1/√L gradient damper.

The two arguments therefore predict different things. Post-LN warmup requirements grow with depth; Adam’s do not, since n_eff(t) depends only on β_2 and the step index. If lengthening warmup helps far more on a deep model than a shallow one, the architectural cause dominates — switch to Pre-LN rather than warm up longer.

Large batches: warmup stops being optional

The linear scaling rule says that multiplying batch size by k calls for multiplying the learning rate by k, keeping expected per-sample progress constant. Its justification is that k small steps on nearly identical weights approximate one step on the averaged gradient — which holds only while weights barely change across those k steps.

That assumption fails hardest at the start, when weights move fastest, so scaling the rate is precisely wrong on the first steps and roughly right later. Gradual warmup bridges that gap, which is why large-batch recipes treat it as mandatory. Two effects compound: the scaled rate is larger in absolute terms, and a large-batch run reaches a given token count in far fewer optimizer steps, so n_eff(t) — which counts steps, not tokens — is still tiny where a small-batch run would have settled. Bigger batch, larger rate, fewer steps of history: three reasons pointing the same way.