预训练中临界批大小如何缩放？
机器学习
2025-04-22 v4 人工智能
最优化与控制
机器学习
摘要
在给定资源下训练大规模模型需要仔细设计并行策略。特别是，临界批大小（CBS）的效率概念涉及时间与计算之间的折衷，它标志着一个阈值，超过该阈值后更大的数据并行度会导致收益递减。为了将其操作化，我们提出了一种 CBS 的度量方法，并在 C4 数据集上预训练了一系列参数量从 8500 万到 12 亿的自回归语言模型。通过广泛的超参数扫描以及对批大小、动量、学习率及其调度等因素的仔细控制，我们系统地研究了规模对 CBS 的影响。然后我们拟合了关于模型和数据大小的缩放律，以解耦它们的影响。总体而言，我们的结果表明，CBS 主要随数据大小而不是模型大小缩放，我们通过分析神经网络的无限宽度极限和无限维最小二乘回归在理论上证明了这一发现。作为独立关注点，我们强调了常见超参数选择和策略对于研究超越固定训练时长的大规模预训练的重要性。
引用
@article{arxiv.2410.21676,
title = {How Does Critical Batch Size Scale in Pre-training?},
author = {Hanlin Zhang and Depen Morwani and Nikhil Vyas and Jingfeng Wu and Difan Zou and Udaya Ghai and Dean Foster and Sham Kakade},
journal= {arXiv preprint arXiv:2410.21676},
year = {2025}
}
备注
ICLR 2025, Blog post: https://kempnerinstitute.harvard.edu/research/deeper-learning/how-does-critical-batch-size-scale-in-pre-training-decoupling-data-and-model-size