Why Does Adaptive Batching Help LLM Pretraining? A Perspective from Unbounded Variance

2026-10-01

AI Summary

This study addresses the lack of theoretical foundations for adaptive batching in large language model (LLM) pretraining, where conventional unbounded variance assumptions frequently fail. We propose a generalized BG-a noise model that interpolates between bounded and quadratically growing variance, derive information-theoretic lower bounds, and construct a theoretical framework with matching upper and lower bounds. Guided by this analysis, we design an adaptive batch scheduling algorithm that dynamically adjusts batch sizes to control gradient variance growth, thereby optimizing non-convex training dynamics. Our approach overcomes the limitations of conservative variance assumptions. When applied to OLMo2 pretraining, it achieves lower validation loss under identical token budgets while requiring fewer than 10% of the iterations compared to small-batch training.

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.

Problem

• Adaptive Batching
• LLM Pretraining
• Stochastic Gradient Variance
• Noise Model
• Nonconvex Optimization

Innovation

• Adaptive Batching
• Generalized BG-a Noise Model
• Unbounded Variance
• Stochastic Optimization
• LLM Pretraining