Large Batch Training › Limitations and Practical Considerations

The linear scaling rule is an approximation, and its applicability depends on assumptions that are often only partially satisfied. Understanding where it breaks down guides practical decisions.

The most significant limitation is the relationship between batch size and generalization. While careful engineering can close most of the accuracy gap between large and small batch training, the gap rarely disappears entirely. Training with very large batches tends to converge to sharper minima. Sharp minima are more sensitive to distribution shift between training and test sets, leading to worse generalization at deployment. In practice, large batch pre-trained models often require more aggressive fine-tuning data augmentation to match the downstream performance of models trained with smaller batches.

A second practical constraint is compute efficiency. The critical batch size is typically much smaller than what modern distributed training hardware can support at full utilization. For example, a 1,024-GPU cluster might require a batch size of 32,768 per step just to keep all GPUs busy, while the critical batch size for the task might be 2,048. The extra 16x scaling is necessary for hardware efficiency but does not improve convergence speed per example. In these regimes, engineers often accept some efficiency loss in exchange for better model quality, choosing a batch size that balances hardware utilization with convergence efficiency.

LAMB's layer-wise trust ratio addresses the adaptive optimizer scaling problem, but it introduces its own sensitivity. The trust ratio computation depends on the L2 norms of weights and Adam update directions. Early in training, when weights are small and gradients are large, the trust ratio can produce very small effective learning rates for some layers. This is usually mitigated by trust ratio clipping and careful initialization, but it requires tuning. LAMB's original validation was primarily on BERT and ResNets. For other architectures with unusual weight scale distributions across layers (especially models with heavily normalized activations), the trust ratio may behave differently and require additional tuning.

Batch normalization also interacts with batch size in complex ways. Batch normalization statistics (mean and variance computed over the batch) change significantly as batch size changes. With very large batches, the statistics become very accurate, but they also incorporate more diversity in the batch, potentially changing the effective normalization behavior. Ghost batch normalization sidesteps this by computing BN statistics on smaller sub-batches, but it introduces the hyperparameter of ghost batch size. For language models using layer normalization rather than batch normalization, this issue is less acute.

The field has also moved toward gradient accumulation as an alternative to true large batch training. Rather than physically processing a large batch on one step, you accumulate gradients from multiple smaller batches before applying the update. This achieves the same effective batch size without requiring hardware that can fit the large batch in memory. Gradient accumulation is straightforward for plain SGD and Adam, but interacts subtly with batch normalization statistics and some regularization techniques. The next chapter on gradient accumulation explores this approach in depth.

Despite these limitations, large batch training has become indispensable for modern AI development. Training GPT-3 with batch sizes of millions of tokens, or BERT with batch sizes of hundreds of thousands, was only possible because of the combination of linear scaling, LAMB, warmup, and careful monitoring of training stability. The techniques in this chapter represent the engineering foundation that made the scaling era of language models possible.