Fast Catch-Up, Late Switching: Optimal Batch Size Scheduling via Functional Scaling Laws
This work addresses the computational efficiency challenge in large-scale deep learning training, particularly for hard tasks like LLM pretraining, by providing a theoretical framework for batch size scheduling, though it builds incrementally on prior scaling law research.
The paper tackles the problem of optimizing batch size scheduling in deep learning training by using functional scaling laws, showing that for hard tasks, optimal schedules use small batches for most of training and switch to large batches late, which reduces data consumption without sacrificing performance, as validated in LLM pretraining experiments with up to 1.1B parameters and 1T tokens.
Batch size scheduling (BSS) plays a critical role in large-scale deep learning training, influencing both optimization dynamics and computational efficiency. Yet, its theoretical foundations remain poorly understood. In this work, we show that the functional scaling law (FSL) framework introduced in Li et al. (2025a) provides a principled lens for analyzing BSS. Specifically, we characterize the optimal BSS under a fixed data budget and show that its structure depends sharply on task difficulty. For easy tasks, optimal schedules keep increasing batch size throughout. In contrast, for hard tasks, the optimal schedule maintains small batch sizes for most of training and switches to large batches only in a late stage. To explain the emergence of late switching, we uncover a dynamical mechanism -- the fast catch-up effect -- which also manifests in large language model (LLM) pretraining. After switching from small to large batches, the loss rapidly aligns with the constant large-batch trajectory. Using FSL, we show that this effect stems from rapid forgetting of accumulated gradient noise, with the catch-up speed determined by task difficulty. Crucially, this effect implies that large batches can be safely deferred to late training without sacrificing performance, while substantially reducing data consumption. Finally, extensive LLM pretraining experiments -- covering both Dense and MoE architectures with up to 1.1B parameters and 1T tokens -- validate our theoretical predictions. Across all settings, late-switch schedules consistently outperform constant-batch and early-switch baselines.