Learning Rate Warmup › Warmup and Optimizer Choice

The case for warmup changes depending on which optimizer you use. Warmup was originally motivated largely by the properties of Adam, but the same principles apply, with different degrees of urgency, across the optimizer landscape.

For Adam and AdamW, which are the overwhelmingly dominant choice for language model training, warmup is essentially mandatory. The bias correction problem described earlier applies directly, and the second-moment stabilization argument gives you a concrete lower bound for warmup duration. The standard $\beta_2 = 0.999$ configuration requires at least 1,000 steps of warmup, and most runs use 2,000-5,000 steps.

The AdamW variant, which corrects the weight decay implementation (decoupling it from the gradient update), does not change the warmup requirements relative to Adam. The moment estimation dynamics are identical. AdamW is uniformly preferred over Adam for large-scale pretraining because of better regularization behavior, but your warmup settings should be the same.

For SGD with momentum, warmup is less theoretically necessary but still empirically beneficial. SGD does not have second-moment estimation, so the bias correction issue does not apply. However, SGD still accumulates a momentum buffer that can be problematic at high learning rates in early training. The momentum buffer starts at zero and can oscillate when the gradients are noisy and the learning rate is large, producing oscillatory parameter trajectories that slow convergence. Linear warmup reduces this risk.

The Goyal et al. 2017 paper used SGD for its ImageNet training experiments, establishing that warmup is beneficial for SGD as well, particularly at large batch sizes. For language models, SGD is rarely used (Adam dominates), but the principle transfers: warmup is a safe default for any optimizer.

Newer optimizers like Adafactor and Sophia are designed with large-scale training stability in mind, and they often incorporate warmup-like behavior internally. Adafactor, which factorizes the second-moment estimate to reduce memory usage, has built-in scale-invariant updates that partially compensate for early estimation noise. Some Adafactor implementations use a relative learning rate that grows as $\min(10^{-2}, 1/\sqrt{t})$, which is effectively an automatic warmup. If you use Adafactor with its default schedule, you may not need to add explicit warmup. But when using a fixed learning rate with Adafactor, the same warmup logic applies.