0. TL; DR

本文研究了在给定计算预算下，如何优化 Transformer 语言模型的大小和训练 token 数量。研究发现，当前的大型语言模型普遍训练不足，主要原因是近期研究更多关注模型规模的扩大，而训练数据量保持不变。通过对 400 多个不同大小（7000 万到 160 亿参数）和不同训练 token 数量（50 亿到 5000 亿）的语言模型进行训练，研究发现为了实现计算最优训练，模型大小和训练 token 数量应该等比例增长：每增加一倍的模型大小，训练 token 数量也应该增加一倍。

基于这一发现，研究者训练了一个名为 Chinchilla 的模型，其参数为 700 亿，训练 token 数量为 1.4 万亿，计算预算与 Gopher 相同，但 Chinchilla 在多个下游任务上的表现显著优于 Gopher（2800 亿参数）、GPT-3（1750 亿参数）、Jurassic-1（1780 亿参数）和 Megatron-Turing NLG（5300 亿参数），并且在微调和推理时计算成本更低，极大地促进了下游应用。

1. 背景介绍

近年来，大型语言模型（LLMs）取得了显著进展，模型规模不断扩大，参数数量从数十亿到数千亿不等。这些基于 Transformer 架构的自回归语言模型在多种任务上表现出色，包括零样本、少样本和微调等评估协议。然而，训练这些大型语言模型的计算和能源成本巨大，且随着模型规模的增加而增加。在实践中，训练这些模型的计算预算通常是预先确定的，因此准确估计给定计算预算下的最佳模型超参数至关重要。

Kaplan 等人（2020）首次展示了自回归语言模型的参数数量与性能之间存在幂律关系。因此，研究领域开始训练越来越大的模型，期望性能得到提升。然而，Kaplan 等人也指出，为了实现计算最优，大型模型不应训练到最低可能的损失。尽管作者得出了相同的结论，但作者估计大型模型应该在比作者建议的训练 token 数量更多的数据上进行训练。具体来说，对于计算预算增加 10 倍，Kaplan 等人建议模型大小增加 5.5 倍，而训练 token 数量仅增加 1.8 倍。相反，作者发现模型大小和训练 token 数量应该等比例增长。

2. 方法介绍

给定一个固定的计算预算（以 FLOPs 表示），作者需要确定最优的模型大小 $N$ 和训练 token 数量 $D$。作者假设最终的预训练损失 $L(N, D)$ 是模型参数数量 $N$ 和训练 token 数量 $D$ 的函数。由于计算预算 $C$ 是训练 token 数量和模型参数的确定性函数 $FLOPs(N, D)$，作者的目标是在满足 $FLOPs(N, D) = C$ 的约束下最小化 $L(N, D)$。即：

\[(N^*, D^*) = \mathop{\arg\min}_{N, D} L(N, D) \quad \\ \text{s.t.} \quad FLOPs(N, D) = C\]

作者提出了三种方法来估计最优的模型大小和训练 token 数量。

方法一：固定模型大小，变化训练步数

作者固定了一系列模型大小（从 7000 万到 100 亿参数），并为每个模型训练了 4 种不同的训练步数。通过这些训练曲线，作者直接提取了在给定计算量下达到的最小损失。具体步骤如下：

对于每个参数数量 $N$，训练 4 个不同的模型，学习率在训练过程中按余弦衰减 10 倍，衰减周期为训练 token 数量的 16 倍。
对每个训练运行，平滑并插值训练损失曲线，得到从 FLOP 数到训练损失的连续映射。
对于每个 FLOP 数，确定哪次运行的损失最低。
在 1500 个对数间隔的 FLOP 值上，找到实现最低损失的模型大小和所需的训练 token 数量。
最后，拟合幂律关系，估计任何给定计算量下的最优模型大小和训练 token 数量。

通过这种方法，作者得到的最优模型大小和训练 token 数量的幂律关系为：

\[N \propto C^{0.50}, \quad D \propto C^{0.50}\]

方法二：IsoFLOP 分析

作者固定了 9 个不同的计算量（从 $6 \times 10^{18}$ 到 $3 \times 10^{21}$ FLOPs），并为每个计算量训练了不同大小的模型。作者探索了对于给定的 FLOP 预算，最优的参数数量是多少？

对于每个 FLOP 预算，绘制最终损失（平滑后）与参数数量的关系图。
通过拟合抛物线，直接估计在每个 FLOP 预算下达到最低损失的模型大小。
与方法一类似，拟合幂律关系，估计任何给定计算量下的最优模型大小和训练 token 数量。

通过这种方法，作者得到的最优模型大小和训练 token 数量的幂律关系为：

\[N \propto C^{0.49}, \quad D \propto C^{0.51}\]

方法三：参数化损失函数拟合

作者将所有实验的最终损失建模为模型参数数量和训练 token 数量的参数化函数。作者假设损失函数的形式为：

\[\hat{L}(N, D) = L_0 + \frac{A}{N^{\alpha}} + \frac{B}{D^{\beta}}\]

其中，$L_0$ 表示理想生成过程的损失（即自然文本的熵），$\frac{A}{N^{\alpha}}$ 表示模型大小对损失的影响，$\frac{B}{D^{\beta}}$ 表示训练 token 数量对损失的影响。作者通过最小化 Huber 损失来估计参数 $(L_0, A, B, \alpha, \beta)$。

\[\min_{L_0, A, B, \alpha, \beta} = \text{Huber}(\log \hat{L}(N, D) - \log L_0)\]

通过这种方法，作者得到的最优模型大小和训练 token 数量的幂律关系为：

\[N \propto C^{0.46}, \quad D \propto C^{0.54}\]

方法比较

三种方法虽然采用了不同的拟合方法和训练模型，但都得出了类似的结论：随着计算预算的增加，模型大小和训练 token 数量应该等比例增长。这些结果表明，当前的大型语言模型普遍过大，而训练 token 数量不足。

3. 实验分析

基于上述分析，作者训练了一个名为 Chinchilla 的模型，其参数为 700 亿，训练 token 数量为 1.4 万亿，计算预算与 Gopher 280B 相同。

Chinchilla 的训练细节如下：

使用与 Gopher 相同的模型架构和训练设置，但有一些差异：
训练数据集：使用 MassiveText 数据集，但稍微调整了子集分布，以适应更多的训练 token。
优化器：使用 AdamW 而不是 Adam，这在语言建模损失和微调后的下游任务性能上有所提升。
分词器：使用稍微修改过的 SentencePiece 分词器，不进行 NFKC 规范化，这在数学和化学表示上表现更好。
权重存储：在分布式优化器状态中存储一个 float32 的权重副本，以提高数值稳定性。

作者对 Chinchilla 进行了广泛的评估，与 Gopher 和其他大型语言模型进行了比较。评估任务包括语言建模、阅读理解、问答、常识理解等。在这些任务上，Chinchilla 的表现都超越了同等计算量下的 Gopher。

此外，Chinchilla 在生成有毒语言和性别偏见方面的表现也值得关注。尽管 Chinchilla 在生成有毒语言方面的表现略好于 Gopher，但其仍然能够生成有毒内容，尤其是在接收到有毒提示时。在性别偏见方面，Chinchilla 仍然会模仿数据集中的刻板印象，例如将某些职业与特定性别联系起来。