Adaptive batching improves LLM training by managing noise growth
Why Does Adaptive Batching Help LLM Pretraining? A Perspective from Unbounded Variance
Machine Learning
Summary
Training large language models often involves picking the right batch size, but why bigger batches help isn’t fully known. The authors found that the noise in training doesn’t stay constant and can increase as the model learns. They introduced a new mathematical model to describe how this noise grows and showed that changing batch sizes dynamically can make training more efficient. Their new method reaches better performance using fewer training steps than fixed batch sizes.
What this means in practice
- •For machine learning engineers: Optimize LLM pretraining schedules by adjusting batch sizes to control noise growth and improve training efficiency.
- •For data center operations: Reduce computational iterations needed for LLM training by using adaptive batch size strategies that maintain model quality with fewer steps.
Authors
Arda Fazla, Antesh Upadhyay, Ege C. Kaya, M. Berk Sahin, Abolfazl Hashemi
Abstract
Increasing the batch size during training is a common practice in large language model (LLM) pretraining, yet the theoretical justification behind its success is not well understood. Analyses of stochastic optimization often assume uniformly bounded stochastic gradient variance, yet recent evidence suggests that this assumption fails in many practical nonconvex problems. The Blum--Gladyshev (BG-$0$) noise model relaxes this assumption by allowing the variance to grow quadratically with the distance from initialization, suggesting that batch size schedulers can help by controlling the variance growth during training. However, this growth can be overly conservative in practice. We empirically investigate variance growth in LLM pretraining and observe that a generalized BG model with a tunable growth exponent provides a tighter description of practical noise behavior. Motivated by this observation, we introduce the generalized BG-$a$ noise model, which interpolates between bounded variance ($a=0$) and BG-$0$ noise ($a=2$). Under $L$-smoothness, we derive an information-theoretic lower bound with growth-dependent oracle complexity $Ω(ε^{-(4+a)})$ and establish a matching upper bound in $ε$-dependence by increasing the batch size as the iterates move away from initialization. Finally, we propose an adaptive batch scheduler that controls variance growth through dynamic batch size adjustments during training. In pretraining OLMo2 models of up to 1B parameters on C4, our scheduler achieves a lower validation loss than both small and large batch training under matched token budgets, while using less than 10\% of the iterations of small batch training.