Learning Rate Warmup › Why Warmup Is Necessary

To understand warmup, you first need to understand what goes wrong without it. The instability of early training is not a single problem but a combination of three overlapping issues that each make the others worse.

At the start of training, all model parameters are drawn from a random distribution. For a transformer with billions of parameters, this means the function the model computes is essentially random noise. The loss is high, the gradients are large, and the curvature of the loss landscape is steep and jagged. There is no structure yet, no learned representations, and no reliable gradient signal.

The problem is that adaptive optimizers like Adam, which you encountered in the Adam Optimizer chapter, estimate gradient statistics from past updates. Adam maintains two moving averages:

• $m_t$: the first moment, tracking the exponentially weighted mean of recent gradients
• $v_t$: the second moment, tracking the exponentially weighted mean of squared gradients (a proxy for variance)

At step $t = 1$, both $m_t$ and $v_t$ are zero, having just been initialized. Adam applies bias correction to account for this, but the estimates are still unreliable because they are based on a single gradient observation. No single sample tells you much about a distribution.

To understand the magnitude of this problem, consider how Adam computes the effective parameter update. Given the corrected moment estimates, the update rule is:

$$\Delta \theta = -\eta \cdot \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}$$

where:

• $\Delta \theta$ is the parameter update applied at step $t$
• $\eta$ is the learning rate you specify
• $\hat{m}_t = m_t / (1 - \beta_1^t)$ is the bias-corrected first moment estimate (mean of recent gradients)
• $\hat{v}_t = v_t / (1 - \beta_2^t)$ is the bias-corrected second moment estimate (mean of squared gradients)
• $\epsilon$ is a small constant (typically $10^{-8}$) added for numerical stability

The bias correction factor $1 / (1 - \beta_2^t)$ deserves careful attention. With typical $\beta_2 = 0.999$, this correction factor at step $t = 1$ evaluates to:

$$\frac{1}{1 - \beta_2^1} = \frac{1}{1 - 0.999} = 1000$$

This means the corrected second moment $\hat{v}_t$ is 1000 times the raw second moment $v_t$ at step 1. Since the second moment estimate was initialized at zero and has only received a single noisy gradient sample, this correction amplifies even a tiny initial gradient variance estimate to a very large value. The result is that the effective step Adam takes at step 1 is not the step you intended when you chose your learning rate. It behaves unpredictably.

This is a fundamental property of bias correction in exponential moving averages: the correction is necessary to get unbiased estimates, but it comes at the cost of high-variance estimates in the earliest steps. The statistics simply have not had enough data yet to be trusted.

Early in training, the gradient direction is particularly unreliable. The model has not yet learned any meaningful representations, so gradients are computed over a random loss landscape. Two batches drawn at step 1 might produce gradient directions that point nearly opposite to each other, whereas by step 10,000, gradients have accumulated structure and roughly agree on which direction to move.

This unreliability of gradient direction is the fundamental bottleneck that warmup is designed to address. The key insight is that gradient variance decreases as the model begins to learn. Once the model has identified some structure in the data (even rudimentary patterns like common word frequencies or basic syntactic relationships), the gradients computed on different batches start to point in more consistent directions. In the first few hundred steps, however, the optimizer has little evidence about which direction will reduce the loss consistently.

Taking large steps when the gradient direction is uncertain is doubly harmful. First, the large step itself may land in a bad region. Second, the Adam second-moment estimate records this large, noisy gradient as a signal that the true gradient magnitude is high, which will affect subsequent updates. A bad early step can poison the optimizer's statistics for thousands of subsequent steps, slowing down training long after the immediate instability has passed.

Transformers are particularly vulnerable to early instability because of the scaled dot-product attention mechanism. As covered in the Scaled Dot-Product Attention chapter, attention computes softmax over scaled dot products of queries and keys. When weights are randomly initialized, the dot products can be arbitrarily large, pushing the softmax into saturation, where all attention mass concentrates on a single token. This creates near-zero gradients for all other tokens, essentially making the attention layer blind to most of its input.

If the learning rate is large in those first steps, the optimizer reinforces these degenerate attention patterns before the model has had a chance to develop meaningful attention distributions. Recovery from softmax saturation is slow and painful: the gradients for non-attended tokens are near-zero, so the optimizer receives no signal pushing it to attend to those tokens, and the model gets stuck in a degenerate solution.

This sensitivity is one reason why transformer training is more warmup-dependent than simpler architectures like MLPs. The softmax function introduces a hard nonlinearity that can create these attention collapse failures, and they are much easier to prevent (via warmup) than to recover from after the fact.

Warmup addresses all three problems simultaneously. By starting with a very small learning rate:

• Adam's accumulated statistics ($m_t$ and $v_t$) have time to build up reliable estimates before large steps are taken
• Gradient directions have time to stabilize as the model begins learning basic representations
• Attention weights can develop reasonable distributions before the optimizer locks them in

The core intuition is that you want the optimizer to be in a reliable, well-characterized state before making large, consequential weight updates. The first few hundred warmup steps are almost like a reconnaissance phase: the model learns the rough terrain of the loss landscape at a safe, slow pace before committing to the large steps that characterize the main training run.

There is also a subtler benefit that is often overlooked. During warmup, the model learns from every batch at a learning rate that is small enough to prevent overcommitting to any single batch's gradient. This gives the optimizer a better picture of the average gradient direction across many batches before it starts making large updates based on that direction. It is analogous to averaging multiple noisy measurements before acting on them.