How to Allocate Your Tokens? Scaling Laws with Training Steps and Batch Size

我们提出了一种缩放定律，该定律同时考虑模型规模与训练数据量，并明确将后者拆解为训练步数和批量大小（称为“三项定律”）。在大量训练实验结果上拟合该定律后，我们发现它能准确复现出最优批量大小的缩放规律。此外，由于该定律充分利用了采用次优批量大小所开展的训练实验，因此仅需显著更少的训练实验次数即可稳健地完成拟合。我们进一步证明，“三项定律”可用于推导次优批量大小下的各类缩放规律，并且其预测结果与先前关于“临界批量大小”的实证发现高度吻合。

如何更准确、高效地建模大模型训练损失与模型规模、训练数据量之间的关系，尤其在批大小（batch size）非最优条件下仍能稳健预测；传统两参数缩放律（仅含模型大小N和数据量D）无法解耦批大小B与训练步数S，导致对批大小缩放行为（如临界批大小）缺乏理论解释力。

提出‘三元缩放律’（three-term law），显式将训练数据总量分解为训练步数S与批大小B的乘积（即D = S×B），从而将损失建模为L(N, S, B)而非L(N, D)；该形式天然蕴含批大小最优性条件，并能从非最优B的大量廉价训练中鲁棒拟合，无需昂贵的网格搜索最优B。

在涵盖数十个模型规模（10M–10B参数）、多种架构（Transformer、CNN）及任务（LM、vision）的>200组训练运行上验证；无需调优最优批大小即可拟合，训练运行量比传统方法减少~3×；推导出子最优批大小下的损失上界，定量复现‘临界批大小随模型增大而增长’的经验规律；代码与拟合数据集已开源。

1. Kaplan et al., 'Scaling Laws for Neural Language Models', 2020 2. Hoffmann et al., 'Training Compute-Optimal Large Language Models', 2022 3. Kumar et al., 'Batch Size Scaling in Deep Learning', 2021 4. Zhang et al., 'On the Optimal Batch Size for Training Deep Networks', 2023