[论文解读] Small Batch Size Training for Language Models: When Vanilla SGD Works, and Why Gradient Accumulation Is Wasteful

Martin Marek, Sanae Lotfi+3 more|ArXiv.org|Jul 9, 2025

Topic Modeling被引用 3

一句话总结

论文表明小批量大小（包括1）可以稳定、鲁棒地训练语言模型，常可达到或超过大批量、以优化器为主的做法，并且不推荐梯度累积。

ABSTRACT

Conventional wisdom dictates that small batch sizes make language model pretraining and fine-tuning unstable, motivating gradient accumulation, which trades off the number of optimizer steps for a proportional increase in batch size. While it is common to decrease the learning rate for smaller batch sizes, other hyperparameters are often held fixed. In this work, we revisit small batch sizes all the way down to batch size one, and we propose a rule for scaling Adam hyperparameters to small batch sizes. In particular, rather than holding the decay rate of the second moment fixed across batch sizes, we propose to hold its half-life fixed in terms of tokens. We find that small batch sizes (1) train stably, (2) are consistently more robust to hyperparameter choices, (3) achieve equal or better per-FLOP performance than larger batch sizes, and (4) notably enable stable language model training with vanilla SGD, even without momentum, despite storing no optimizer state. Building on these results, we provide practical recommendations for selecting a batch size and setting optimizer hyperparameters. We further recommend against gradient accumulation unless training on multiple devices with multiple model replicas. Finally, we show that a small batch size combined with an optimizer with a small state size can provide the performance benefits of full fine-tuning while maintaining a similar memory footprint to LoRA.

研究动机与目标

• 研究在预训练和微调中，小批量大小是否能稳定语言模型训练。
• 考察不同优化器和超参数在不同批量大小下的鲁棒性。
• 开发基于吞吐量与内存约束的实际指南，用于缩放 Adam 超参数与选择批量大小。
• 在选择小批量大小与梯度累积时，探讨内存和硬件因素。

提出的方法

• 在批量大小从1到4096上，对SGD、Adam、Adafactor、Muon进行全面网格搜索。
• 针对每个批量大小，调整学习率与 Adam 衰减因子(beta1, beta2)，以最小化验证损失。
• 引入第二矩估计寿命 t2 的概念，并展示基于基于 token 的半寿命的 beta2 缩放规则。
• 在小批量场景下比较 Vanilla SGD（无动量）和内存高效优化器如 Adafactor。
• 在更大模型（GPT-2 124M 与 GPT-3 1.3B）以及微调场景中验证缩放启发式。
• 评估内存影响并为内存受限训练提供实际建议。

实验结果

研究问题

• RQ1 极小批量大小（降至1）是否能在没有动量或复杂优化器的情况下稳定训练语言模型？
• RQ2 为在小批量下保持性能，如何缩放 Adam 超参数？
• RQ3 小批量是否在对超参数错配的鲁棒性方面优于大批量？
• RQ4 在语言模型训练中，使用小批量与梯度累积相比，内存与硬件方面的影响？
• RQ5 这些发现是否可推广至微调和更大模型规模？

主要发现

• 在进行适当缩放超参数时，小批量在所有优化器下与大批量相匹配或超越。
• 在小批量大小下，动量需求较低，vanilla SGD 亦具竞争力。
• 通过以 token 为单位固定第二矩半寿命 t2 进行缩放，比固定 beta2 能在各批量大小下获得更好表现。
• 与使用 Adafactor 或 SGD 这样更简单的优化器相比，梯度累积通常不必要且内存效率低下。
• 对于非常大的模型，批量大小为1的 SGD 能与 AdamW 基线相匹配，且 Adafactor 提供较优的内存-性能折衷。
• 推荐的做法是使用能够最大化吞吐量的最小批量大小，除了在多设备设置下，避免使用梯度累积。