Training Compute-Optimal Large Language Models

Jordan Hoffmann; Sebastian Borgeaud; Arthur Mensch; Elena Buchatskaya; Trevor Cai; Eliza Rutherford; Diego de Las Casas; Lisa Anne Hendricks; Johannes Welbl; Aidan Clark; Tom Hennigan; Eric Noland; Katie Millican; George van den Driessche; Bogdan Damoc; Aurelia Guy; Simon Osindero; Karen Simonyan; Erich Elsen; Jack W. Rae; Oriol Vinyals; Laurent Sifre

论文 · arXiv 2203.15556

训练计算最优的大型语言模型

Training Compute-Optimal Large Language Models

Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland et al.

发表于 2022-03-29cs.CLcs.LG arXiv PDF

用 XYZNOT 阅读（全文翻译与 AI 讲解）

三句话总结

🎯 确定在固定计算预算下模型参数量与训练 Token 数量之间的最优权衡。
🔧 通过最小损失包络线、IsoFLOP 切片和参数化损失建模，在超过 400 次 Transformer 训练运行中经验拟合缩放法则（Scaling Laws）。
📈 参数量和 Token 数应同等比例扩展（ $a \approx 0.5$ ， $b \approx 0.5$ ）。在 1.4T Token 上训练的 70B Chinchilla 模型在相同计算预算下匹敌 Gopher，并在 MMLU 上达到 67.6%。

摘要

我们研究了在给定计算预算下训练基于 Transformer 的语言模型的最优模型大小和 Token 数量。我们发现当前的大型语言模型存在严重的欠训练现象，这是近期过度关注扩展语言模型参数规模而保持训练数据量不变所导致的结果。通过在 50 亿至 5000 亿 Token 上训练超过 400 个参数量从 7000 万到 160 亿不等的语言模型，我们发现要实现计算最优训练，模型大小和训练 Token 数量应等比例扩展：模型大小每翻一倍，训练 Token 数量也应翻倍。我们通过训练一个预测为计算最优的模型 Chinchilla 来验证这一假设，该模型使用与 Gopher 相同的计算预算，但参数量为 70B，数据量增加了 4 倍。在广泛的下游评估任务中，Chinchilla 全面且显著地优于 Gopher (280B)、GPT-3 (175B)、Jurassic-1 (178B) 以及 Megatron-Turing NLG (530B)。这也意味着 Chinchilla 在微调和推理时占用的计算资源显著减少，极大地方便了下游应用。值得一提的是，Chinchilla 在 MMLU 基准测试上达到了 67.5% 的业界最优平均准确率，相较于 Gopher 提升了超过 7%。

英文原文摘要

We investigate the optimal model size and number of tokens for training a transformer language model under a given compute budget. We find that current large language models are significantly undertrained, a consequence of the recent focus on scaling language models whilst keeping the amount of training data constant. By training over 400 language models ranging from 70 million to over 16 billion parameters on 5 to 500 billion tokens, we find that for compute-optimal training, the model size and the number of training tokens should be scaled equally: for every doubling of model size the number of training tokens should also be doubled. We test this hypothesis by training a predicted compute-optimal model, Chinchilla, that uses the same compute budget as Gopher but with 70B parameters and 4 $\times$ more more data. Chinchilla uniformly and significantly outperforms Gopher (280B), GPT-3 (175B), Jurassic-1 (178B), and Megatron-Turing NLG (530B) on a large range of downstream evaluation tasks. This also means that Chinchilla uses substantially less compute for fine-tuning and inference, greatly facilitating downstream usage. As a highlight, Chinchilla reaches a state-of-the-art average accuracy of 67.5% on the MMLU benchmark, greater than a 7% improvement over Gopher.

论文解读

问题与动机

先前的工作（尤其是 Kaplan et al., 2020）指出，随着计算量的增加，模型参数应大幅扩展（ $N \propto C^{0.73}$ ），而 Token 数量的增长则要慢得多（ $D \propto C^{0.27}$ ）。因此，像 GPT-3 (175B)、Jurassic-1 (178B)、Gopher (280B) 和 MT-NLG (530B) 等大型模型均普遍只在约 3000 亿 Token 上进行训练。然而，相对于它们的计算投入，这些大型模型处于严重的欠训练状态。本研究重新审视了这一权衡：在固定的计算预算 $C$ 下，应如何分配参数量 $N$ 和 Token 量 $D$ 才能使测试损失最小化？

核心思想

等比例扩展：模型参数量每翻一倍，训练 Token 数量也应翻倍（ $N_{o pt} \propto C^{0.5}$ ， $D_{o pt} \propto C^{0.5}$ ）。
余弦学习率调度对齐：学习率调度必须与实际的训练步数（Horizon）相匹配。如果使用固定长周期（如 130B Token）衰减过程中的中间检查点，会高估较小 Token 数量时的损失，从而扭曲经验缩放估计。
下游效率：在预训练计算量相同的情况下，使用大幅增加的 Token 训练更小的模型，不仅能取得相同或更好的准确率，还能显著降低显存占用、微调成本和推理延迟。

方法（及核心公式）

作者使用三种互补的策略评估了 400 多个模型（参数量从 70M 到 16B，训练数据覆盖 5B 到 500B Token）：

训练曲线极小值法：针对固定模型架构设置 4 种不同的训练步数。提取每个 FLOP 的最小损失包络线并拟合幂律公式： $N_{o pt} \propto C^{a}, D_{o pt} \propto C^{b}$ 得出 $a = 0.50$ , $b = 0.50$ 。
IsoFLOP 曲线剖面：固定 9 个离散预算（ $6 \times 1 0^{18}$ 到 $3 \times 1 0^{21}$ ）的总 FLOP 数量，并将余弦衰减与训练步数对齐。在参数量维度上拟合抛物线以寻找损失极小值，得出 $a = 0.49$ , $b = 0.51$ 。
参数化损失拟合：使用风险分解建模平滑损失： $\hat{L} (N, D) ≜ E + \frac{A}{N ^{α}} + \frac{B}{D ^{β}}$ 其中 $E$ 代表自然文本熵， $A / N^{α}$ 代表表征能力的不足， $B / D^{β}$ 刻画了提前停止/有限数据导致的次优性。在对数残差上最小化 Huber loss，得到 $E = 1.69$ ， $A = 406.4$ ， $B = 410.7$ ， $α = 0.34$ ， $β = 0.28$ 。在约束条件 $FLOPs (N, D) \approx 6 N D$ 下： $a = \frac{β}{α + β} = 0.46, b = \frac{α}{α + β} = 0.54$

结果（及核心数据）

Chinchilla 70B：使用与 Gopher 相同的 FLOP 预算（ $5.76 \times 1 0^{23}$ FLOPs，按精确计算则为 $6.3 \times 1 0^{23}$ ）进行训练，采用 70B 参数和 1.4 万亿 Token。
MMLU 基准测试：Chinchilla 达到 67.6% 的平均 5-shot 准确率，超越了 Gopher (60.0%) 和 GPT-3 (43.9%)，并超过了 2023 年 6 月的人类专家预测值（63.4%）。它在 4 项不同的 MMLU 任务中取得了 $> 90%$ 的准确率。
下游基准测试：在 BIG-bench（65.1% 对 54.4%）、LAMBADA（77.4% 对 74.5%）以及闭卷 Natural Questions（5-shot 31.5% 对 24.5%）上均超越了 Gopher。
偏差与毒性：在无提示（Unprompted）毒性方面，Gopher 与 Chinchilla 相当（中位数 0.064 对 0.066），而在 Winogender 准确率上，整体表现从 71.4% 提升到了 78.3%。

局限性与开放问题

受训练成本限制，研究仅在最大规模的前沿对齐点上直接验证了两个模型（Gopher 280B 与 Chinchilla 70B）。
分析假设在网络数据集上进行单轮（Single-epoch）训练；多轮重复数据的表现尚未探索。
实验数据显示在高计算预算下 $lo g (N_{o pt})$ 存在轻微的负曲率（凹性），这表明幂律可能仍会微弱高估最优模型参数规模。

关键术语

Chinchilla: 一个拥有 700 亿参数的自回归语言模型，在与 Gopher 相同的计算预算下使用 1.4 万亿 Token 训练而成。
Gopher: 一个在 3000 亿 Token 上训练的 2800 亿参数 Transformer 语言模型，作为本研究中同等计算量的基线。
IsoFLOP profiles: 一种在保持训练总 FLOPs 不变的情况下改变模型参数量，以寻找损失最小参数规模的方法论。
Compute-optimal frontier: 将可用训练 FLOPs 与参数量及训练 Token 量的最优分配联系起来的数学曲线。
Parametric loss function: 将语言模型损失统计分解为不可约贝叶斯熵、参数容量缺陷和 Token 样本优化误差的函数。
Huber loss: 一种鲁棒的损失函数，在方法 3 中用于拟合经验语言模型训练实验，同时减弱小计算规模下残差离群值的影响。
MassiveText: 包含网页、书籍、新闻、代码和维基百科的多领域预训练语料库，用于训练 Gopher 和 Chinchilla。
FLOPs calculation: 乘加运算预算，在前向传播和反向传播中约估算为 $6 N D$ ，或通过逐层精确核算得出。

常见问题

为什么 Kaplan et al. 得出的结论是参数量的扩展速度应快于 Token？

Kaplan et al. 使用了固定的余弦学习率调度（例如衰减目标设为 130B Token），并使用较小 Token 数据量训练时的中间未收敛评估点。由于当周期长度明显超过总步数时余弦学习率调度效果会劣化，导致这些中间损失被人为抬高，从而低估了用较少 Token 训练较小模型的有效性。

根据 Chinchilla，计算最优扩展的经验法则是什么？

模型参数量和训练 Token 数量应按大致相同的比例进行扩展：计算预算每增加 $2 \times$ ，模型大小应增加约 $2^{0.5} \approx 1.41 \times$ ，训练 Token 数量也应增加约 $2^{0.5} \approx 1.41 \times$ 。

训练一个计算最优的 175B 模型建议使用多少 Token？

根据论文的预测，一个 1750 亿参数的模型大约需要 3.7 到 4.3 万亿 Token 以及 $3.85 \times 1 0^{24}$ 至 $4.54 \times 1 0^{24}$ FLOPs 才能达到计算最优，这远超 GPT-3 通常采用的 3000 亿 Token。

Gopher 与 Chinchilla 之间做了哪些架构上的改动？

除了将网络层数和隐藏层维度从 280B 缩减至 70B 之外，Chinchilla 将 Adam 优化器替换为了 AdamW，在分布式优化器状态中保留了 float32 主权重，并移除了 SentencePiece 分词器中的 NFKC 正规化。

本页解读由 AI 根据论文生成，可能有错误；细节请以原文为准。

用 XYZNOT 阅读（全文翻译与 AI 讲解）