Research map / Architectures and efficiency / RNNs and normalization
How Does Critical Batch Size Scale in Pre-training?
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Measures critical batch size (CBS) for autoregressive LMs (85M–1.2B) on C4 and finds CBS grows mainly with data size rather than model size, supported by infinite-width and least-squares theory.
Problem
How the critical batch size, the point where larger data parallelism yields diminishing returns, scales with model size and data size in language model pre-training, which prior work has not decoupled.
Approach
Defines CBS as the batch size where steps to reach a target loss exceed linear-scaling prediction by 20%, uses constant+EWA schedule to train beyond fixed durations, and fits scaling laws separately in model size and data size, with theory from muP infinite-width limits and mini-batch SGD on linear regression.
Key results
- Fitted CBS scaling with model size under Chinchilla settings: B* = 93.20 * N^0.47 (N in millions).
- With fixed data size, CBS barely depends on model size: B* = 621.341 * N^0.087.
- CBS scales with data size: B* = 22.91 * D^0.47 (D in billions tokens), and constant+EWA matches cosine schedule efficiency especially at large batch sizes.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Scale model size and data size together in Chinchilla-optimal training | Autoregressive LMs on C4 · C4, Chinchilla-optimal training | CBS: B^* = 93.20 * N^{0.47} | Helped |
| Scale model size while holding data size fixed | Autoregressive LMs on C4 · C4, fixed data size | CBS: B^*=621.341 * N^{0.087} | No clear effect |
| Scale data size while holding model size fixed | 302M autoregressive LM on C4 · C4, varying training duration | CBS: B^* = 22.91 * D^{0.47} | Helped |
| Add EWA to a constant learning rate (ablation) | Autoregressive LM pre-training · C4, varying batch sizes | training efficiency | Helped |
| Use Constant+EWA instead of WSD scheduling (ablation) | 151M autoregressive LM · C4, especially large batch sizes | training efficiency | Helped |
| Vary context length from 2^9 to 2^12 (ablation) | 151M autoregressive LM · C4, multiple batch sizes | relative optimization efficiency | No clear effect |
| Scale model size by increasing depth rather than width (ablation) | 151M to 604M autoregressive LMs · C4, compute-optimal training | CBS | No clear effect |
| Use Adam momentum β1=0.95 (ablation) | 151M autoregressive LM · C4, varied batch sizes | evaluation loss: <0.02 gain on eval loss | Helped |
| Remove Adam momentum (ablation) | 151M autoregressive LM · C4, varied batch sizes | optimization efficiency | Hurt |
| Use a larger β2 for small-batch, longer-duration training (ablation) | Autoregressive LM pre-training · C4, small batch sizes and long durations | training efficiency | Helped |
| Add weight decay (ablation) | Constant learning rate with EWA · C4, 151M proxy models | critical batch size | No clear effect |
Surprising findings
- CBS remains nearly invariant when scaling model size with fixed data, contrary to the expectation that larger models need larger batches.
- SGD without momentum was significantly worse than Adam, so Adam was used throughout.
Limitations
- Experiments limited to models up to 1.2B parameters and fully synchronous data parallelism with frequent communication.
- Theory relies on a well-specified Gaussian linear regression and infinite-width limits, not full transformer training.
- CBS depends on the chosen 20% overhead threshold and B_opt=256 choice.
Experimental setup
Datasets: C4
Benchmarks: C4 holdout validation loss
Metrics: steps to reach target validation loss critical batch size fitted scaling-law exponents
Baselines: cosine schedule warmup-stable-decay schedule schedule-free optimizer trapezoidal schedule
Builds on: Chinchilla scaling Kaplan scaling laws maximal update parameterization gradient noise scale schedule-free optimizer
Closest papers by meaning
- Critical Batch Size Revisited: A Simple Empirical Approach to Large-Batch Language Model Training (NeurIPS 2025 · similarity 0.92)
- Power Lines: Scaling laws for weight decay and batch size in LLM pre-training (NeurIPS 2025 · similarity 0.90)
- Resolving Discrepancies in Compute-Optimal Scaling of Language Models (NeurIPS 2024 · similarity 0.89)
- Small Batch Size Training for Language Models: When Vanilla SGD Works, and Why Gradient Accumulation is Wasteful (NeurIPS 2025 · similarity 0.89)
- Language models scale reliably with over-training and on downstream tasks (ICLR 2025 · similarity 0.88)
- The Journey Matters: Average Parameter Count over Pre-training Unifies Sparse and Dense Scaling Laws (ICLR 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.