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

技术摘要：基于训练步数与批次大小的缩放法则

问题陈述
现有的大型语言模型（LLM）缩放法则（如 Chinchilla 法则）将测试损失建模为模型规模（$N$）和总训练数据量（$D$）的函数。虽然这些法则在确定模型规模 $N$ 与数据量 $D$ 之间的最优计算分配方面非常有效，但它们并未明确考虑训练超参数，特别是批次大小（$b$）和训练步数（$K$）。以往对超参数缩放的建模尝试通常无法与基于损失的缩放法则集成，或者依赖于“临界批次大小”（critical batch size）的概念，而在某些理论表述中，该概念暗示最优批次大小为 1——这与观察到最优批次大小随 Token 总预算缩放的经验结果相矛盾。此外，为拟合最优超参数的缩放法则通常需要针对各种 Token 预算进行详尽的批次大小搜索，这会产生巨大的计算成本。

方法论
作者提出了一种三项缩放法则（3TL），该法则明确地将训练数据预算 $D$ 分解为以 Token 为单位的批次大小（$M = b \times s$，其中 $s$ 为序列长度）和训练步数（$K$），即 $D = MK$。所提出的损失函数$L$ 的形式为：

$$L(N, M, K) = E + \frac{A}{N^\alpha} + \frac{B}{M^\beta} + \frac{C}{K^\gamma}$$

其中 $(E, A, B, C, \alpha, \beta, \gamma)$ 是通过训练运行拟合得到的参数。

作者将此 3TL 与**两项法则（2TL）**方法进行了对比。在 2TL 中，模型规模 $N$ 是固定的，并针对每个 $N$ 使用形式 $L(M, K) = E + B/M^\beta + C/K^\gamma$ 独立进行拟合。

关键方法步骤包括：

1. 最优批次大小的推导： 通过在约束条件 $D=MK$ 下对 3TL 关于 $M$ 求极小值，作者推导出最优批次大小 $M^\star$ 随 $D$ 按幂律缩放：$M^\star \propto D^{\frac{\gamma}{\beta+\gamma}}$。
2. 数据利用： 不同于以往需要通过全量搜索来确定每个配置下的单一最优批次大小的方法，3TL 可以使用次优批次大小的数据进行拟合。这使得该法则能够以显著更少的训练运行次数实现稳健拟合。
3. 次优区间建模： 为了解决理论最优批次大小在硬件约束下不可行的问题，作者提出了一个两阶段拟合程序。首先，他们使用 3TL 来指导一个简化的函数形式 $L(b) = \tilde{E} + \tilde{A}b^{-\tilde{\alpha}} + \tilde{B}b^{\tilde{\alpha}}$（假设 $\gamma \approx \beta$）。其次，他们针对 $\epsilon$-次优批次大小区间（其中 $\epsilon$ 代表 5% 的计算浪费阈值）拟合一个幂律。
4. 临界批次大小分析： 作者分析了达到目标损失 $\bar{L}$ 所需的步数 $K_{\bar{L}}(b)$ 随批次大小的变化行为，并将 3TL 的预测与关于临界批次大小的经验发现进行了对比。

关键结果
作者在两个数据集上评估了其方法：Li (Li et al., 2025a) 和 OpenEuroLLM。

• 最优批次大小缩放： 3TL 成功恢复了最优批次大小 $M^\star$ 的缩放特性。在 Li 数据集上，推导出的缩放关系 $M^\star \approx 0.667 \cdot D^{0.566}$ 与 Li 等人 (2025a) 的 “Step-Law” ($M^\star \approx 0.58 \cdot D^{0.571}$) 高度一致，后者是通过全量批次大小搜索得出的。
• 计算效率： 一个主要发现是，3TL 仅需传统方法 28% 的训练运行次数即可完成拟合。通过将数据集掩码化，仅保留每个 $(N, D)$ 组合中的 2 个批次大小（而非进行全量搜索），3TL 仍能恢复正确的 $M^\star$ 缩放法则。相比之下，如果不进行损失曲面建模，直接对少数可用的“最佳”批次大小进行幂律拟合，会导致缩放法则极其不稳定且不准确。
• 次优批次大小： 两阶段拟合程序成功模拟了计算浪费不超过 5% 的次优批次大小区间。作者发现，该区间的宽度在批次大小上约为 $2^2$（即 4 倍因子），且其缩放行为在不同模型规模和数据集之间保持一致。
• 临界批次大小： 3TL 正确模拟了这样一个现象：即在超过临界点后，达到目标损失所需的步数对批次大小增加的敏感度降低。至关重要的是，该模型表明这种临界行为随 Token 预算 $D$ 缩放，但在很大程度上与模型规模 $N$ 无关，这与 Zhang 等人 (2025) 的经验观察相符。
• 与 Chinchilla 的比较： 当将 3TL 简化为标准的 Chinchilla 形式（通过将最优 $M^\star$ 代回损失方程）时，隐含的对于模型规模 $N$ 的计算最优缩放与直接拟合 Chinchilla 法则的结果有所不同。3TL 建议了一个较小的 $D$ 缩放指数，这表明缩放法则形式的选择会影响推导出的最优配置。

意义与主张
本文声称，三项法则提供了一个统一的框架，弥合了 Chinchilla 式缩放法则（模型规模 vs 数据）与超参数缩放法则（批次大小）之间的鸿沟。其意义主要体现在三个方面：

1. 稀疏数据的稳健性： 利用次优批次大小来拟合该法则的能力大幅降低了推导缩放法则的计算成本，使得在无需穷举搜索的情况下确定最优批次大小成为可能。
2. 理论一致性： 该公式在结构上与随机优化中的收敛界限相似，为经验性的缩放分析与优化理论之间提供了理论桥梁。
3. 实际应用性： 不同于以往暗示最优批次大小为 1 或无法描述次优区间的模型，3TL 提供了一个非平凡的最优批次大小，并明确模拟了次优区间内的性能，这对于硬件约束下的实际训练至关重要。

作者也指出了局限性，包括不同数据集之间可能存在参数值的不一致性（例如不可约损失项 $E$），以及该法则仍需要为每个配置确定最优学习率。他们还承认 3TL 隐含的计算最优缩放与直接的 Chinchilla 拟合有所不同，这表明它可能不是确定最优模型规模的唯一工具，但在理解步数、批次大小与数据之间的相互作用方面具有优越性。