How Does Critical Batch Size Scale in Pre-training?

在给定资源下训练大规模模型需要精心设计并行策略。特别是，关键批次大小（CBS）的概念涉及时间和计算之间的权衡，标志着超过这一阈值后，更大的数据并行性带来的收益将逐渐减少。为了将其付诸实践，我们提出了一种衡量CBS的方法，并在C4数据集上预训练了一系列自回归语言模型，参数规模从8500万到12亿不等。通过广泛的超参数搜索和对批量大小、动量和学习率及其调度等因素的严格控制，我们系统地研究了规模对CBS的影响。然后，我们根据模型和数据规模拟合了扩展定律，以分离它们的影响。总体而言，我们的结果表明，CBS主要随数据规模而非模型规模而变化，这一发现通过分析神经网络的无限宽度极限和无限维最小二乘回归得到了理论上的支持。此外，我们强调了在固定训练时长之外研究大规模预训练时，常见的超参数选择和策略的重要性。

论文探讨了在给定资源下训练大规模模型时，如何通过设计并行策略来优化训练效率，特别是关注关键批量大小（Critical Batch Size, CBS）对数据并行性和计算效率的影响。这是一个重要的问题，因为找到合适的CBS可以显著提高训练效率，避免资源浪费。

论文的关键思路是通过系统地研究不同规模的自回归语言模型（从8500万到12亿参数）在C4数据集上的预训练过程，探索CBS与数据量和模型规模之间的关系。研究发现，CBS主要随数据量而非模型规模增加而增加。这一发现通过分析无限宽度神经网络和无限维最小二乘回归得到了理论支持，为理解大规模模型训练提供了新的视角。

论文通过广泛的超参数调优和严格的控制实验，确保了实验结果的可靠性。实验设计考虑了批大小、动量、学习率及其调度等多个因素，以全面评估CBS的影响。此外，论文还提出了一个衡量CBS的方法，并通过拟合缩放定律来分离数据量和模型规模的影响。这些方法和发现对于未来的大规模预训练研究具有重要参考价值。论文使用的C4数据集是常见的大规模文本数据集，但未提及是否有开源代码。

近期在这个领域内，有几篇相关研究值得关注： 1. "Scaling Laws for Neural Language Models" - 这篇论文探讨了模型规模、数据量和训练时间之间的关系。 2. "On the Convergence of Adam and Beyond" - 研究了Adam优化器在不同设置下的收敛性，对超参数选择有重要影响。 3. "Understanding Deep Learning (Still) Requires Rethinking Generalization" - 探讨了深度学习模型的泛化能力，特别是在大规模数据集上的表现。 4. "The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks" - 提出了稀疏子网络的概念，对模型效率有重要启示。