Learning Rate Warmup › Warmup for Large Batch Training

The relationship between warmup and batch size is one of the most practically important aspects of this topic. When you increase batch size to speed up training, you must also increase the learning rate, and warmup becomes significantly more important.

As covered in more detail in the Large Batch Training chapter, the linear scaling rule states that when you multiply batch size by a factor $k$, you should multiply the learning rate by the same factor $k$ to maintain equivalent training dynamics:

$$\eta_{\text{new}} = k \cdot \eta_{\text{base}}$$

where:

• $k = B_{\text{new}} / B_{\text{base}}$ is the batch size scaling factor (new batch size divided by the baseline)
• $\eta_{\text{base}}$ is the learning rate that was tuned and validated for the reference batch size $B_{\text{base}}$
• $\eta_{\text{new}}$ is the learning rate to use at the scaled batch size

The intuition behind this rule is that a larger batch provides a lower-variance gradient estimate, because you are averaging over more samples. This makes each gradient step more informative and more representative of the true gradient direction. You can therefore take a proportionally larger step without the same risk of moving in a noisy direction. The learning rate compensates for the reduced noise by taking a larger step, maintaining roughly the same effective update per token seen.

The linear scaling rule is an approximation that holds well for moderate batch size increases but begins to break down at very large batch sizes. For most practical purposes (batch sizes up to ~8,192), it is the standard starting point.

The linear scaling rule interacts with warmup in a critical way. When you scale up the learning rate proportionally with batch size, the peak learning rate $\eta_{\max}$ becomes much larger. A 16x batch size increase means a 16x larger peak learning rate, which amplifies all the instability problems of early training:

• Adam's unreliable early statistics lead to even larger initial updates when multiplied by the higher learning rate
• Gradient noise affects updates by a larger magnitude, so a bad step in the wrong direction has worse consequences
• Attention saturation happens faster if a large step locks in degenerate patterns before meaningful representations form

This is why the scaling rule must be paired with a proportionally longer warmup. If you double the batch size (and thus the learning rate), you should also roughly double the warmup duration. The general recommendation from Goyal et al. (2017) is to scale warmup duration linearly with the learning rate scaling factor:

$$T_{\text{warmup,\,new}} = k \cdot T_{\text{warmup,\,base}}$$

where:

• $T_{\text{warmup,\,new}}$ is the warmup duration to use at the new (larger) batch size
• $T_{\text{warmup,\,base}}$ is the warmup duration tuned for the baseline batch size
• $k$ is the same batch size scaling factor from the learning rate rule above

For example, if your baseline configuration uses a batch size of 256 with a learning rate of $10^{-4}$ and 1,000 warmup steps, and you scale to a batch size of 4,096 (a 16x increase), you should use:

• Learning rate: $16 \times 10^{-4} = 1.6 \times 10^{-3}$
• Warmup steps: $16 \times 1000 = 16,000$

Without the warmup increase, the 1.6e-3 learning rate at step 1 would produce enormous updates on an optimizer with unreliable statistics, almost certainly causing training instability.

The Goyal et al. 2017 paper "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour" made warmup standard practice for large-batch training. The paper's contribution was twofold: it formulated the linear scaling rule precisely, and it demonstrated empirically that gradual warmup was necessary to make the rule work in practice.

The key experiment showed that training ResNet-50 on 256 GPUs with a batch size of 8,192 failed without warmup. The loss would spike and the model would fail to converge. With 5 epochs of gradual linear warmup, the model trained successfully and matched the accuracy of a model trained with a 32x smaller batch size. This result established warmup as a prerequisite for large-batch training, not a minor tuning convenience.

For language models, subsequent transformer training papers extended these findings. Without warmup, the initial destabilization from high learning rates can cause loss spikes that occasionally grow so large that the model never recovers, especially in deeper architectures where the gradient signal must propagate through many layers.

For very large batch sizes (beyond 8,192), the linear scaling rule begins to break down. The fundamental reason is that gradient noise reduction does not scale linearly with batch size: doubling the batch size reduces the gradient standard deviation by a factor of $\sqrt{2}$, not by a factor of 2. The noise reduction follows the law of large numbers, which is a square-root relationship.

At extreme batch sizes, the signal-to-noise ratio plateaus because the dominant source of gradient variance shifts from sampling noise (which scales as $1/\sqrt{B}$) to structural noise from the data distribution itself. Continuing to increase the learning rate linearly once this plateau is reached leads to diminishing or negative returns.

An alternative scaling rule uses square-root rather than linear scaling:

$$\eta_{\text{new}} = \sqrt{k} \cdot \eta_{\text{base}}$$

This is more conservative at large batch sizes and reduces the peak learning rate, which in turn reduces the required warmup duration. Both scaling rules are in active use, and the choice between them depends on the specific model architecture, training setup, and the batch size regime you are operating in. A practical heuristic is to use linear scaling up to $k \approx 32$ (batch sizes below ~8,000 if your baseline is 256) and to consider square-root scaling for larger scale factors.