Fast Catch-Up, Late Switching: Optimal Batch Size Scheduling via Functional Scaling Laws

TL;DR

This paper introduces a theoretical framework based on functional scaling laws to optimize batch size schedules in large-scale deep learning, revealing task-dependent strategies and a late-switching phenomenon that improves efficiency.

Contribution

It provides a principled analysis of batch size scheduling using FSL, characterizes optimal schedules based on task difficulty, and uncovers the late-switching mechanism validated by extensive experiments.

Findings

Late switching improves training efficiency across models.

Large batches can be deferred to late training without performance loss.

The fast catch-up effect explains the late-switching advantage.

Abstract

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.