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.

Choosing a length, and choosing a shape

Two floors set the length. The optimizer floor wants t comparable to 1/(1 − β_2), so β_2 = 0.999 argues for roughly 500–2000 steps, and a smaller β_2 genuinely permits a shorter ramp. The architectural floor scales with depth and batch size. Because these are floors, the usual heuristics — 1–4% of total steps, or a fixed budget of warmup samples — are compatible with the theory rather than rivals to it; the sample-based form is better, lengthening the step count automatically when batches shrink.

On shape, linear from near-zero is the default and usefully conservative: √v̂ improves in reliability like √t while a linear ramp trusts the optimizer only like t. Exponential ramps rise too slowly at first and too fast at the end, inverting that match; a constant low rate followed by a jump reintroduces the shock at the transition. Start at a small positive rate, not exactly zero.

Too-short warmup: a worked example

Take a 124M-parameter model, initialization standard deviation 0.02, peak α = 3×10^-4, β_2 = 0.999. Because the early update is a sign step, every coordinate moves by α regardless of its gradient — 1.5% of a typical weight’s initial magnitude, per step. With no warmup, 50 steps of partially consistent signs displace a coordinate by up to 50 × 3e-4 = 0.015, about 0.75 initial standard deviations: the network is substantially rewritten before its variance estimate is worth anything.

A 2000-step linear ramp changes the arithmetic. Cumulative rate over the first 50 steps is α · (50·51/2)/2000 = 0.64 α versus 50α — a 78× reduction exactly where the estimate is worst. The point is not the total: averaged over the whole ramp, warmup only halves the distance travelled. The point is when it is travelled — after √v̂ is accurate to a few percent instead of tens.

Loud failure, silent failure, and the RAdam question

Too-short warmup fails in three ways of decreasing visibility. A NaN or outright divergence in the first few hundred steps is obvious. A loss spike — loss and gradient norm jump, then the run recovers over a few thousand steps — is visible if you are watching. Worst is the silent case: no spike, no alarm, just a final loss a couple of percent worse than the same run with a longer ramp, because the model was kicked into a poorer basin at step 30 and optimized competently within it ever after. Only an ablation reveals that one, which is much of why warmup gets cargo-culted: its absence is often not visibly fatal.

RAdam acts on exactly this analysis: rather than damping the rate from outside, it estimates how many effective samples the second-moment average has accumulated and withholds the adaptive term until that count is large enough, falling back to a momentum-SGD-like update meanwhile. The mechanism targets the right quantity but does not fully deliver — independent evaluations put it at roughly matching, not beating, a tuned linear warmup, which reaches the same place with one fewer moving part — and it addresses only the optimizer cause. Production recipes have kept warmup.