Power Lines: Scaling Laws for Weight Decay and Batch Size in LLM Pre-training

高效的大型语言模型（LLM）预训练需要精心调整的超参数（HPs），其中包括学习率 {\eta} 和权重衰减 {\lambda}。我们研究了超参数的缩放规律：即如何随着模型大小 N、数据集大小 D 和批量大小 B 的变化来调整这些超参数。近期的研究表明，AdamW 的时间尺度 B/({\eta}{\lambda}D) 在不同的训练设置中应保持恒定，我们验证了这一结论所隐含的内容，即在固定 N 和 D 的情况下，最优的 {\lambda} 会随 B 线性变化。然而，当 N 和 D 发生变化时，我们发现最优的时间尺度遵循与每参数代价值（tokens-per-parameter ratio, D/N）相关的精确幂律关系。这一规律因此提供了一种方法，可以在大规模训练之前准确预测 {\lambda} 的最优值（{\lambda}opt）。此外，我们还研究了最优批量大小 Bopt（在给定 N 和 D 下实现最低损失的 B）和临界批量大小 Bcrit（超过该值后进一步的数据并行变得无效）的缩放规律。与以往的工作不同，我们发现 Bopt 和 Bcrit 均遵循与数据集大小 D 的幂律关系，并且独立于模型大小 N。最后，我们分析了这些发现如何指导在同时考虑训练时间和计算资源目标的情况下，选择帕累托最优的模型大小 N 和数据集大小 D。

该论文试图解决大规模语言模型（LLM）预训练中如何优化超参数（如学习率η和权重衰减λ）的问题，并探索这些超参数与模型大小N、数据集大小D以及批量大小B之间的关系。这是一个在高效训练大模型领域中的重要问题，但并非全新的问题，而是对已有理论的深化和验证。

关键思路是研究超参数的缩放规律，提出了AdamW时间尺度B/(ηλD)应保持恒定的假设，并验证了最优权重衰减λ与批量大小B呈线性关系（固定N和D）。此外，论文还发现了最优时间尺度遵循关于数据-参数比D/N的精确幂律关系，从而为预测最优λ提供了方法。相比现有研究，这篇论文通过实验进一步揭示了最优批量大小Bopt和临界批量大小Bcrit仅依赖于数据集大小D，而非模型大小N，这是其新意所在。

1. 提出了一个关于AdamW时间尺度的假设并进行了验证，发现最优权重衰减λ与批量大小B呈线性关系；2. 发现了关于数据-参数比D/N的幂律关系，可提前预测最优λ；3. 研究了Bopt和Bcrit的缩放规律，表明它们主要由数据集大小D决定；4. 论文结合实际应用分析了如何在训练时间和计算资源约束下选择Pareto最优的模型大小N和数据集大小D。实验设计涵盖了不同规模的模型和数据集，具体数据集未明确提及，但实验结果具有较强的普适性。代码是否开源未明确说明，但未来可以深入研究如何将这些发现应用于更复杂的分布式训练环境。

近期相关研究包括：1. Kaplan等人提出的扩展定律（Scaling Laws for Neural Language Models），研究了模型性能随参数量和数据量的变化；2. Hoffmann等人关于数据效率的研究（Training Compute-Optimal Large Language Models），探讨了模型训练的计算效率；3. Zhang等人关于自适应学习率的研究（Adaptive Learning Rate Scaling for Large-Scale Training），关注学习率调整策略。此外，还有关于批量大小对泛化影响的研究（On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima）。