mlsanityData to 14 Oct 2026

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

arXiv 2602.14208 ↗First posted 15 Feb 2026cs.LG math.OC stat.MLcited by 6

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

  1. 2020Language Models are Few-Shot Learnerscited by 32,419
  2. 2024The Llama 3 Herd of Modelscited by 14,715
  3. 2019Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformercited by 14,599
  4. 2023LLaMA: Open and Efficient Foundation Language Modelscited by 14,132
  5. 2019RoBERTa: A Robustly Optimized BERT Pretraining Approachcited by 13,372
  6. 2020Scaling Laws for Neural Language Modelscited by 6,586
  7. 2021RoFormer: Enhanced Transformer with Rotary Position Embeddingcited by 4,703
  8. 2022Training Compute-Optimal Large Language Modelscited by 3,019
  9. 2019Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelismcited by 1,868
  10. 2017Accurate, Large Minibatch SGD: Training ImageNet in 1 Hourcited by 1,492
  11. 2021Scaling Vision Transformerscited by 999
  12. 2024MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategiescited by 723
38 more
  1. 2024DeepSeek LLM: Scaling Open-Source Language Models with Longtermismcited by 642
  2. 2021Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LMcited by 632
  3. 2017Deep Learning Scaling is Predictable, Empiricallycited by 582
  4. 2020Scaling Laws for Autoregressive Generative Modelingcited by 522
  5. 2025GLM-4.5: Agentic, Reasoning, and Coding (ARC) Foundation Modelscited by 467
  6. 2024GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projectioncited by 463
  7. 2023Scaling Data-Constrained Language Modelscited by 403
  8. 2021Explaining Neural Scaling Lawscited by 363
  9. 2017Don't Decay the Learning Rate, Increase the Batch Sizecited by 320
  10. 2018An Empirical Model of Large-Batch Trainingcited by 281
  11. 2024Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durationscited by 170
  12. 2025MiniMax-01: Scaling Foundation Models with Lightning Attentioncited by 166
  13. 2023Beyond Chinchilla-Optimal: Accounting for Inference in Language Model Scaling Lawscited by 159
  14. 2017The Power of Interpolation: Understanding the Effectiveness of SGD in Modern Over-parametrized Learningcited by 150
  15. 2024Nemotron-4 340B Technical Reportcited by 144
  16. 2023Scaling Laws for Generative Mixed-Modal Language Modelscited by 134
  17. 2018Stochastic Modified Equations and Dynamics of Stochastic Gradient Algorithms I: Mathematical Foundationscited by 114
  18. 2024A Dynamical Model of Neural Scaling Lawscited by 111
  19. 2024Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspectivecited by 81
  20. 2024How Does Critical Batch Size Scale in Pre-training?cited by 73
  21. 2018Statistical Optimality of Stochastic Gradient Descent on Hard Learning Problems through Multiple Passescited by 69
  22. 20244+3 Phases of Compute-Optimal Neural Scaling Lawscited by 68
  23. 2022The alignment property of SGD noise and how it helps select flat minima: A stability analysiscited by 64
  24. 2025LongCat-Flash Technical Reportcited by 58
  25. 2024Scaling Laws in Linear Regression: Compute, Parameters, and Datacited by 57
  26. 2024How Feature Learning Can Improve Neural Scaling Lawscited by 56
  27. 2024APOLLO: SGD-like Memory, AdamW-level Performancecited by 54
  28. 2023Optimal Linear Decay Learning Rate Schedules and Further Refinementscited by 43
  29. 2021Last Iterate Risk Bounds of SGD with Decaying Stepsize for Overparameterized Linear Regressioncited by 38
  30. 2024Nemotron-4 15B Technical Reportcited by 36
  31. 2024Shortcut-connected Expert Parallelism for Accelerating Mixture-of-Expertscited by 32
  32. 2025The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre-Trainingcited by 30
  33. 2021Power-law escape rate of SGDcited by 29
  34. 2024Scaling Law with Learning Rate Annealingcited by 29
  35. 2025Scaling Laws for Gradient Descent and Sign Descent for Linear Bigram Models under Zipf's Lawcited by 24
  36. 2025Critical Batch Size Revisited: A Simple Empirical Approach to Large-Batch Language Model Trainingcited by 23
  37. 2025A Multi-Power Law for Loss Curve Prediction Across Learning Rate Schedulescited by 22
  38. 2025Scaling Collapse Reveals Universal Dynamics in Compute-Optimally Trained Neural Networkscited by 22

50 most cited of 63 All 63 (JSON)

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2026: 6.

6 papers.

Similar, not cited

Close in wording, no citation either way

  1. 2026Test-Time Scaling Makes Overtraining Compute-Optimal
  2. 2026When Do Larger Batches Help Scale LLM Reinforcement Learning?
  3. 2025Predictable Scale: Part I, Step Law -- Optimal Hyperparameter Scaling Law in Large Language Model Pretraining
  4. 2026Understanding Quantization of Optimizer States in LLM Pre-training: Dynamics of State Staleness and Effectiveness of State Resets
  5. 2025Kinetics: Rethinking Test-Time Scaling Laws
  6. 2025How to Upscale Neural Networks with Scaling Law? A Survey and Practical Guidelines
  7. 2024Rethinking Conventional Wisdom in Machine Learning: From Generalization to Scaling
  8. 2025Scaling Behaviors of LLM Reinforcement Learning Post-Training: An Empirical Study in Mathematical Reasoning