Research map / Architectures and efficiency / RNNs and normalization
Critical Batch Size Revisited: A Simple Empirical Approach to Large-Batch Language Model 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.
Proposes branched training to directly measure the critical batch size (CBS) at multiple checkpoints of OLMo models, and uses the measured CBS to design a batch size warmup schedule that trains OLMo 1B with 43% fewer gradient steps without loss degradation.
Problem
Gradient-noise-scale proxies for the critical batch size rely on strong assumptions (SGD optimizer, well-conditioned Hessian) that do not hold for Adam-trained language models, so it is unclear whether they can guide large-batch LM pretraining.
Approach
Measure the CBS directly by launching short branched training runs at varying batch-size multipliers from a checkpoint for a fixed token budget (2B tokens), selecting the largest batch whose smoothed loss does not degrade, and use the resulting local CBS over training to schedule batch size warmup with square-root learning rate scaling.
Key results
- CBS starts near 0 at initialization, grows rapidly within the first ~50k tokens, then plateaus around 4096 documents for OLMo 1B and 7B.
- The gradient noise scale underestimates the CBS by several orders of magnitude, and its qualitative trend does not match the CBS especially for OLMo 7B.
- Batch size warmup on OLMo 1B reaches pretraining loss 2.5891 vs 2.6057 for the small-batch control while saving 43% of gradient steps.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use batch size warmup instead of a fixed batch size of 1024 | OLMo 1B training · OLMo 1B pretraining | PT loss: 2.6057 -> 2.5891 | Helped |
| Use batch size warmup instead of a fixed batch size of 1024 | OLMo 1B training · OLMo 1B mid-training | MT loss: 2.5486 -> 2.5433 | Helped |
| Use batch size warmup instead of a fixed batch size of 1024 | OLMo 1B training · OLMo 1B training | Grad. Steps Saved: 0% -> 43% | Helped |
| Use batch size warmup instead of a fixed batch size of 4096 | OLMo 1B training · OLMo 1B pretraining | PT loss: 2.5962 -> 2.5891 | Helped |
| Use batch size warmup instead of a fixed batch size of 4096 | OLMo 1B training · OLMo 1B training | Grad. Steps Saved: 75% -> 43% | Hurt |
| Use batch size warmup instead of a fixed batch size of 1024 | OLMo 1B training · Downstream tasks PT | Task BPB: 1.0112 -> 1.0316 | Hurt |
| Use batch size warmup instead of a fixed batch size of 1024 | OLMo 1B training · Downstream tasks MT | Task BPB: 0.9999 -> 1.0076 | Hurt |
| Use batch size warmup instead of a fixed batch size of 1024 | OLMo 1B training · C4 PT | Cross-entropy loss: 2.8196 -> 2.8049 | Helped |
| Use batch size warmup instead of a fixed batch size of 1024 | OLMo 1B training · C4 MT | Cross-entropy loss: 2.7622 -> 2.7597 | Helped |
| Use batch size warmup instead of a fixed batch size of 1024 | OLMo 1B training · The Pile PT | Cross-entropy loss: 2.2073 -> 2.1916 | Helped |
| Use batch size warmup instead of a fixed batch size of 1024 | OLMo 1B training · The Pile MT | Cross-entropy loss: 2.1471 -> 2.1521 | Hurt |
Surprising findings
- The gradient noise scale underestimates the measured CBS by several orders of magnitude and its trend can disagree with the CBS, especially at 7B.
- The large-batch control (4096) uses 75% fewer steps but has degraded loss, while warmup matches or beats the small-batch control.
Limitations
- Relies on the local recovery assumption that loss recovery within the 2B-token window predicts later behavior.
- Tolerance epsilon=0.01 and window size 2B tokens are set arbitrarily; results may depend on these choices.
- Warmup doubling thresholds were chosen manually from preliminary measurements and only applied to OLMo 1B.
Experimental setup
Datasets: OLMo pretraining data (Dolma) C4 The Pile
Benchmarks: C4 The Pile downstream task BPB (ARC-Easy, ARC-Challenge, MMLU, HellaSwag, PIQA, etc.)
Metrics: training loss bits-per-byte cross-entropy loss gradient steps saved critical batch size (documents)
Baselines: Gradient noise scale proxy (B_simple) Small-batch control (OLMo 1B, B=1024) Large-batch control (OLMo 1B, B=4096)
Builds on: gradient noise scale (McCandlish et al. 2018) OLMo
Closest papers by meaning
- How Does Critical Batch Size Scale in Pre-training? (ICLR 2025 · similarity 0.92)
- Small Batch Size Training for Language Models: When Vanilla SGD Works, and Why Gradient Accumulation is Wasteful (NeurIPS 2025 · similarity 0.92)
- Resolving Discrepancies in Compute-Optimal Scaling of Language Models (NeurIPS 2024 · similarity 0.89)
- Benchmarking Optimizers for Large Language Model Pretraining (NeurIPS 2026 · similarity 0.89)
- Language models scale reliably with over-training and on downstream tasks (ICLR 2025 · similarity 0.89)
- Through the River: Understanding the Benefit of Schedule-Free Methods for Language Model Training (NeurIPS 2025 · similarity 0.89)
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.