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

arXiv: https://arxiv.org/abs/2602.14208 · HTML: https://mlsanity.com/paper/2602.14208/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 63 of its 72 references that are in the map.
- First posted on arXiv 15 Feb 2026.

## Builds on (most cited first, 30 of 63)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://mlsanity.com/paper/1706.02677/) (2017)
- [Scaling Vision Transformers](https://mlsanity.com/paper/2106.04560/) (2021)
- [MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies](https://mlsanity.com/paper/2404.06395/) (2024)
- [DeepSeek LLM: Scaling Open-Source Language Models with Longtermism](https://mlsanity.com/paper/2401.02954/) (2024)
- [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://mlsanity.com/paper/2104.04473/) (2021)
- [Deep Learning Scaling is Predictable, Empirically](https://mlsanity.com/paper/1712.00409/) (2017)
- [Scaling Laws for Autoregressive Generative Modeling](https://mlsanity.com/paper/2010.14701/) (2020)
- [GLM-4.5: Agentic, Reasoning, and Coding (ARC) Foundation Models](https://mlsanity.com/paper/2508.06471/) (2025)
- [GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection](https://mlsanity.com/paper/2403.03507/) (2024)
- [Scaling Data-Constrained Language Models](https://mlsanity.com/paper/2305.16264/) (2023)
- [Explaining Neural Scaling Laws](https://mlsanity.com/paper/2102.06701/) (2021)
- [Don't Decay the Learning Rate, Increase the Batch Size](https://mlsanity.com/paper/1711.00489/) (2017)
- [An Empirical Model of Large-Batch Training](https://mlsanity.com/paper/1812.06162/) (2018)
- [Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations](https://mlsanity.com/paper/2405.18392/) (2024)
- [MiniMax-01: Scaling Foundation Models with Lightning Attention](https://mlsanity.com/paper/2501.08313/) (2025)
- [Beyond Chinchilla-Optimal: Accounting for Inference in Language Model Scaling Laws](https://mlsanity.com/paper/2401.00448/) (2023)
- [The Power of Interpolation: Understanding the Effectiveness of SGD in Modern Over-parametrized Learning](https://mlsanity.com/paper/1712.06559/) (2017)
- [Nemotron-4 340B Technical Report](https://mlsanity.com/paper/2406.11704/) (2024)
- [Scaling Laws for Generative Mixed-Modal Language Models](https://mlsanity.com/paper/2301.03728/) (2023)
- [Stochastic Modified Equations and Dynamics of Stochastic Gradient Algorithms I: Mathematical Foundations](https://mlsanity.com/paper/1811.01558/) (2018)
- [A Dynamical Model of Neural Scaling Laws](https://mlsanity.com/paper/2402.01092/) (2024)

## Built on by (most cited first, 6 of 6)

- [Hyperball May Not Be a Free Lunch](https://mlsanity.com/paper/2607.22444/) (2026)
- [Effective Learning Rate Governs Loss Dynamics in Language Model Pretraining](https://mlsanity.com/paper/2608.24814/) (2026)
- [How Should LLMs Consume High-Quality Data? Optimal Data Scheduling via Quality-Aware Functional Scaling Laws](https://mlsanity.com/paper/2605.25698/) (2026)
- [Dynamics of Stochastic Momentum with Sparse Updates in High Dimensions](https://mlsanity.com/paper/2605.28961/) (2026)
- [Curvature-Conditioned Multiscale Momentum with Sphere Constraints for LLM Pretraining](https://mlsanity.com/paper/2608.28442/) (2026)
- [Learning with Synthetic Data via SGD in High-Dimensional Linear Regression](https://mlsanity.com/paper/2609.09572/) (2026)

## Abstract (arXiv)

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.
