Research map / Optimization / Optimizers
Small Batch Size Training for Language Models: When Vanilla SGD Works, and Why Gradient Accumulation is Wasteful
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.
Shows that small batch sizes (down to 1) train stably and are more robust to hyperparameters when Adam's second-moment half-life is held fixed in tokens, enabling competitive vanilla SGD and memory-efficient Adafactor fine-tuning, and argues against gradient accumulation.
Problem
Small batch sizes are commonly believed to make language model pretraining and fine-tuning unstable, motivating gradient accumulation, while other optimizer hyperparameters (beta1, beta2) are typically held fixed across batch sizes.
Approach
Scale Adam's beta2 so that its half-life measured in training tokens stays fixed across batch sizes (beta2* = beta2^(B*/B)); with this rule small batches (down to 1) train stably, are more hyperparameter-robust, and make vanilla SGD and Adafactor competitive.
Key results
- Small batch sizes (down to 1) match or beat larger batches per FLOP; with tuning, all optimizers perform best at the smallest batch size, and the gap between optimizers shrinks.
- Fixing the Adam second-moment half-life in tokens lets batch sizes 16 and 32 match 128 and 256 on the C4 setup of prior work, without extra tuning.
- At 1.3B parameters, SGD at batch size 1 performs on par with the AdamW baseline, and Adam and Adafactor at batch size 1 outperform it.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use batch size 1 instead of larger batch sizes | 30M model with tuned optimizers · FineWeb-Edu | validation loss | Helped |
| use batch size 1 instead of batch size 512 | Adam · GPT-2 | validation loss | Helped |
| scale β₂ to keep its token half-life fixed across batch sizes | Adam · C4, 19M model | performance | Helped |
| use vanilla SGD without momentum | 30M model at small batch sizes · FineWeb-Edu | validation loss | Helped |
| use SGD with batch size 1 instead of the AdamW baseline | GPT-3 (1.3B) · FineWeb | performance | No clear effect |
| use Adam or Adafactor with batch size 1 instead of the AdamW baseline | GPT-3 (1.3B) · FineWeb | performance | Helped |
| use Adafactor instead of Adam for full fine-tuning | Gemma 3 (4B) · MATH | performance-memory trade-off | Helped |
| use bfloat16 weights with stochastic rounding instead of float32 weights (ablation) | Adafactor fine-tuning · MATH | performance | No clear effect |
| keep β₂ fixed instead of keeping its token half-life fixed (ablation) | Adam at small batch sizes · FineWeb-Edu | validation loss | Hurt |
| use square-root learning-rate scaling with batch size instead of the empirically preferred slower scaling (ablation) | Adam · FineWeb-Edu | validation loss | Hurt |
| vary Adam learning rate, β₁, or β₂ away from their optimum (ablation) | Adam at larger batch sizes · FineWeb-Edu, 30M model | validation loss | Hurt |
| vary Adam learning rate, β₁, or β₂ away from their optimum (ablation) | Adam at batch size 1 · FineWeb-Edu, 30M model | validation loss | No clear effect |
Surprising findings
- Vanilla SGD without momentum or weight decay becomes competitive for LM training at small batch sizes, contrary to prior reports.
- Prior findings of poor small-batch performance with Adam appear to be an artifact of incomplete beta2 tuning.
Limitations
- Batch sizes of 1 and 2 reduce MFU by roughly 30-70%, so throughput cost is substantial.
- Large-scale validation limited to GPT-2 124M, GPT-3 1.3B pretraining and a single Gemma 3 4B fine-tuning task; the 1.3B model used minimal tuning.
- Theoretical reason why fixing the second-moment half-life in tokens generalizes across scales is not understood.
Experimental setup
Datasets: FineWeb-Edu FineWeb C4 MATH
Benchmarks: FineWeb-Edu validation loss FineWeb validation loss C4 validation loss MATH
Metrics: validation loss memory footprint loss sensitivity to hyperparameters MFU
Baselines: AdamW Adam Adafactor Muon SGD LoRA
Builds on: Adam Adafactor Muon Chinchilla scaling laws
Closest papers by meaning
- Critical Batch Size Revisited: A Simple Empirical Approach to Large-Batch Language Model Training (NeurIPS 2025 · similarity 0.92)
- Mini-batch Coresets for Memory-efficient Language Model Training on Data Mixtures (ICLR 2025 · similarity 0.92)
- Benchmarking Optimizers for Large Language Model Pretraining (NeurIPS 2026 · similarity 0.90)
- Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training (ICLR 2024 · similarity 0.90)
- Adam-mini: Use Fewer Learning Rates To Gain More (ICLR 2025 · similarity 0.89)
- The Journey Matters: Average Parameter Count over Pre-training Unifies Sparse and Dense Scaling Laws (ICLR 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.