重新审视临界批量大小：一种面向大批量语言模型训练的简单经验方法
机器学习
2025-11-07 v3
摘要
在规模化训练语言模型时，合适的批量大小很重要：为了快速训练，需要较大的批量大小，但过大的批量大小会损害 token 效率。为了权衡这一矛盾，McCandlish 等人（2018）提出，可以根据训练期间的梯度噪声尺度来估计一个临界批量大小（CBS），在低于该值时训练不会显著降低损失。虽然他们的方法已在实践中被采用（例如在训练 GPT-3 时），但证明将梯度噪声作为 CBS 的代理需要强假设，这使得其方法在实践中是否可信变得不明确，从而限制了其适用性。在本文中，我们引入了一种简单的经验方法来直接测量 CBS，并展示了 CBS 在训练过程中的演变。将我们的方法应用于 OLMo 模型，我们发现 CBS 在初始化时接近 0，起初迅速增加，然后随着训练进行趋于平稳。此外，我们发现这一趋势在不同模型规模（1B 和 7B）下均成立，这表明来自小规模训练运行的 CBS 可以为更大规模的训练运行提供参考。我们关于 CBS 如何随训练变化的发现，促使将批量大小预热作为一种可靠地以大批量训练语言模型的自然方法：从较小的批量大小开始，随着 CBS 的增长而增加它。为了验证这一主张，我们使用批量大小预热训练 OLMo 1B，以比原始训练运行少 43% 的梯度步数达到了略好的损失。这表明我们的框架如何能够应用于可靠地以更大批量大小训练语言模型，在不损害性能的情况下增加数据并行度。
引用
@article{arxiv.2505.23971,
title = {Critical Batch Size Revisited: A Simple Empirical Approach to Large-Batch Language Model Training},
author = {William Merrill and Shane Arora and Dirk Groeneveld and Hannaneh Hajishirzi},
journal= {arXiv preprint arXiv:2505.23971},
year = {2025}
}
备注
Neurips 2025