Fast Catch-Up, Late Switching: Optimal Batch Size Scheduling via Functional Scaling Laws
Not on the map. Why?
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.
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
63 of 72 references.
References · 63
- 2020Language Models are Few-Shot Learnerscited by 32,419
- 2024The Llama 3 Herd of Modelscited by 14,715
- 2019Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformercited by 14,599
- 2023LLaMA: Open and Efficient Foundation Language Modelscited by 14,132
- 2019RoBERTa: A Robustly Optimized BERT Pretraining Approachcited by 13,372
- 2020Scaling Laws for Neural Language Modelscited by 6,586
- 2021RoFormer: Enhanced Transformer with Rotary Position Embeddingcited by 4,703
- 2022Training Compute-Optimal Large Language Modelscited by 3,019
- 2019Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelismcited by 1,868
- 2017Accurate, Large Minibatch SGD: Training ImageNet in 1 Hourcited by 1,492
- 2021Scaling Vision Transformerscited by 999
- 2024MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategiescited by 723
38 more
- 2024DeepSeek LLM: Scaling Open-Source Language Models with Longtermismcited by 642
- 2021Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LMcited by 632
- 2017Deep Learning Scaling is Predictable, Empiricallycited by 582
- 2020Scaling Laws for Autoregressive Generative Modelingcited by 522
- 2025GLM-4.5: Agentic, Reasoning, and Coding (ARC) Foundation Modelscited by 467
- 2024GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projectioncited by 463
- 2023Scaling Data-Constrained Language Modelscited by 403
- 2021Explaining Neural Scaling Lawscited by 363
- 2017Don't Decay the Learning Rate, Increase the Batch Sizecited by 320
- 2018An Empirical Model of Large-Batch Trainingcited by 281
- 2024Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durationscited by 170
- 2025MiniMax-01: Scaling Foundation Models with Lightning Attentioncited by 166
- 2023Beyond Chinchilla-Optimal: Accounting for Inference in Language Model Scaling Lawscited by 159
- 2017The Power of Interpolation: Understanding the Effectiveness of SGD in Modern Over-parametrized Learningcited by 150
- 2024Nemotron-4 340B Technical Reportcited by 144
- 2023Scaling Laws for Generative Mixed-Modal Language Modelscited by 134
- 2018Stochastic Modified Equations and Dynamics of Stochastic Gradient Algorithms I: Mathematical Foundationscited by 114
- 2024A Dynamical Model of Neural Scaling Lawscited by 111
- 2024Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspectivecited by 81
- 2024How Does Critical Batch Size Scale in Pre-training?cited by 73
- 2018Statistical Optimality of Stochastic Gradient Descent on Hard Learning Problems through Multiple Passescited by 69
- 20244+3 Phases of Compute-Optimal Neural Scaling Lawscited by 68
- 2022The alignment property of SGD noise and how it helps select flat minima: A stability analysiscited by 64
- 2025LongCat-Flash Technical Reportcited by 58
- 2024Scaling Laws in Linear Regression: Compute, Parameters, and Datacited by 57
- 2024How Feature Learning Can Improve Neural Scaling Lawscited by 56
- 2024APOLLO: SGD-like Memory, AdamW-level Performancecited by 54
- 2023Optimal Linear Decay Learning Rate Schedules and Further Refinementscited by 43
- 2021Last Iterate Risk Bounds of SGD with Decaying Stepsize for Overparameterized Linear Regressioncited by 38
- 2024Nemotron-4 15B Technical Reportcited by 36
- 2024Shortcut-connected Expert Parallelism for Accelerating Mixture-of-Expertscited by 32
- 2025The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre-Trainingcited by 30
- 2021Power-law escape rate of SGDcited by 29
- 2024Scaling Law with Learning Rate Annealingcited by 29
- 2025Scaling Laws for Gradient Descent and Sign Descent for Linear Bigram Models under Zipf's Lawcited by 24
- 2025Critical Batch Size Revisited: A Simple Empirical Approach to Large-Batch Language Model Trainingcited by 23
- 2025A Multi-Power Law for Loss Curve Prediction Across Learning Rate Schedulescited by 22
- 2025Scaling Collapse Reveals Universal Dynamics in Compute-Optimally Trained Neural Networkscited by 22
50 most cited of 63 All 63 (JSON)
Built on by
Citing papers per year: 2026: 6.
6 papers.
- 2026Hyperball May Not Be a Free Lunchcited by 4
- 2026Effective Learning Rate Governs Loss Dynamics in Language Model Pretrainingcited by 3
- 2026How Should LLMs Consume High-Quality Data? Optimal Data Scheduling via Quality-Aware Functional Scaling Lawscited by 0
- 2026Dynamics of Stochastic Momentum with Sparse Updates in High Dimensionscited by 0
- 2026Curvature-Conditioned Multiscale Momentum with Sphere Constraints for LLM Pretrainingcited by 0
- 2026Learning with Synthetic Data via SGD in High-Dimensional Linear Regressioncited by 0
- 2026Learning with Synthetic Data via SGD in High-Dimensional Linear Regression9 Sep 2026
- 2026Curvature-Conditioned Multiscale Momentum with Sphere Constraints for LLM Pretraining28 Aug 2026
- 2026Effective Learning Rate Governs Loss Dynamics in Language Model Pretraining25 Aug 2026
- 2026Hyperball May Not Be a Free Lunch24 Jul 2026
- 2026Dynamics of Stochastic Momentum with Sparse Updates in High Dimensions27 May 2026
- 2026How Should LLMs Consume High-Quality Data? Optimal Data Scheduling via Quality-Aware Functional Scaling Laws25 May 2026
Similar, not cited
Close in wording, no citation either way
- 2026Test-Time Scaling Makes Overtraining Compute-Optimal
- 2026When Do Larger Batches Help Scale LLM Reinforcement Learning?
- 2025Predictable Scale: Part I, Step Law -- Optimal Hyperparameter Scaling Law in Large Language Model Pretraining
- 2026Understanding Quantization of Optimizer States in LLM Pre-training: Dynamics of State Staleness and Effectiveness of State Resets
- 2025Kinetics: Rethinking Test-Time Scaling Laws
- 2025How to Upscale Neural Networks with Scaling Law? A Survey and Practical Guidelines
- 2024Rethinking Conventional Wisdom in Machine Learning: From Generalization to Scaling
- 2025Scaling Behaviors of LLM Reinforcement Learning Post-Training: An Empirical Study in Mathematical Reasoning