Learning Rate Warmup › A Worked Example

To make warmup concrete, let's trace through a specific training configuration and see exactly what the learning rate schedule looks like.

Consider training a medium-size transformer (300M parameters) with these settings:

• Total training steps: 100,000
• Warmup steps: 2,000
• Peak learning rate: $3 \times 10^{-4}$
• Minimum learning rate: $3 \times 10^{-5}$
• Decay schedule: cosine decay after warmup

At step $t = 1$, the warmup formula gives:

$$\eta_1 = 3 \times 10^{-4} \cdot \frac{1}{2000} = 1.5 \times 10^{-7}$$

This is nearly zero. Adam's statistics are unreliable at this point, and this tiny learning rate ensures the first update is tiny and harmless regardless of the gradient magnitude.

At step $t = 1000$, halfway through warmup:

$$\eta_{1000} = 3 \times 10^{-4} \cdot \frac{1000}{2000} = 1.5 \times 10^{-4}$$

Adam has now seen 1,000 gradient observations and has reasonable second-moment estimates. The learning rate is half the peak, which is a meaningful training rate while still being conservative enough for the optimizer's current state.

At step $t = 2000$, warmup complete:

$$\eta_{2000} = 3 \times 10^{-4} \cdot \frac{2000}{2000} = 3 \times 10^{-4}$$

The learning rate reaches its peak. Adam's second-moment estimates have stabilized (we crossed the $t > 1/(1 - 0.999) = 1000$ threshold), and the model has begun to develop basic representations. Cosine decay begins from this point.

At step $t = 100,000$, end of training, the cosine decay progress is $1.0$, giving:

$$\eta_{100000} = 3 \times 10^{-5} + \frac{1}{2}(3 \times 10^{-4} - 3 \times 10^{-5})(1 + \cos(\pi)) = 3 \times 10^{-5}$$

The learning rate ends at the minimum, one-tenth of the peak. This final low learning rate allows the model to fine-tune its parameters without overshooting the local minimum it has found during the main training phase.

Now consider what happens when you scale the same 300M parameter model to a larger compute cluster, increasing batch size from 256 to 4,096 (a 16x increase). Applying the linear scaling rule:

New peak learning rate: $16 \times 3 \times 10^{-4} = 4.8 \times 10^{-3}$

New warmup steps: $16 \times 2000 = 32,000$

The first update at step 1 with the new configuration is:

$$\eta_1^{\text{new}} = 4.8 \times 10^{-3} \cdot \frac{1}{32000} = 1.5 \times 10^{-7}$$

Notice that the first-step learning rate is identical to the original configuration. This is not a coincidence. When both the peak learning rate and the warmup duration scale by the same factor $k$, their ratio stays constant, and the per-step learning rate at any fractional position within warmup is preserved. This is the mathematical reason the two configurations are training-equivalent during warmup: the gradient step sizes at each warmup fraction are the same, and the model sees the same effective learning rate trajectory relative to the total warmup budget.

This equivalence breaks at the level of absolute step count: the large-batch run completes the same fractional training milestone in far fewer wall-clock steps (because each step processes 16x more data), but the per-token compute and gradient statistics at each checkpoint are comparable. This is why the linear scaling rule with proportional warmup maintains training quality: it preserves the statistical structure of the training trajectory across batch size changes.

To make the optimizer-statistics argument tangible, let's trace exactly what happens to Adam's second-moment estimate $v_t$ for a single parameter across the first 2,000 steps of warmup.

Suppose a particular embedding parameter has a true gradient standard deviation of $\sigma = 0.02$. The true gradient variance is $\sigma^2 = 4 \times 10^{-4}$. We want Adam's second-moment estimate to converge to this value before taking large steps.

With $\beta_2 = 0.999$, the bias-corrected second moment at step $t$ converges toward the true variance as:

$$\mathbb{E}[\hat{v}_t] = \sigma^2 \cdot \frac{1 - \beta_2^t}{1 - \beta_2^t} = \sigma^2$$

But the expected value alone does not tell the whole story. The variance of $\hat{v}_t$ around its expected value is the issue. At step $t = 10$ (very early in warmup), the estimate is based on only 10 noisy gradient observations, heavily weighted toward the initial zero. The relative error in the estimate can easily exceed 50%, meaning the optimizer might behave as if the true gradient variance is $2 \times 10^{-4}$ or $8 \times 10^{-4}$ instead of $4 \times 10^{-4}$. A 2x error in the second-moment estimate translates directly to a $\sqrt{2}$ error in the step size, since Adam's update scales as $1 / \sqrt{v_t}$.

By step $t = 1000$, the exponential moving average has accumulated enough samples that the estimate is reliable. The initialization influence has decayed to $\beta_2^{1000} = 0.999^{1000} \approx 0.37$, meaning only 37% of the second-moment estimate's weight is still on the zero initialization, while 63% comes from actual gradient observations. At this point, the step size variation due to estimation error is under 10%, which is acceptable for taking large updates safely.

This is precisely why warmup duration and the Adam $\beta_2$ parameter are tightly linked. If you reduce $\beta_2$ to 0.99, the initialization influence at step 100 is already $0.99^{100} \approx 0.37$, matching what takes 1,000 steps with $\beta_2 = 0.999$. A lower $\beta_2$ effectively speeds up the statistics stabilization, allowing shorter warmup. This relationship is exploited in some modern training setups that use $\beta_2 \in [0.95, 0.99]$ to enable faster adaptation during early training and continued training stages.