Question

《Training Compute-Optimal Large Language Models》

1. 最近，一系列大型语言模型（Large Language Model: LLM ）被引入，目前最大的稠密语言模型拥有超过 500B 参数。这些大型自回归 transformer 模型在许多任务上使用各种评估协议（如 zero-shot 、few-shot 和 fine-tuning ），都表现出令人印象深刻的性能。

   训练大型语言模型的计算成本和能源成本很高，并随着模型规模的增加而上升。在实践中，分配的 training compute budget 通常是事先知道的：有多少个 accelerators 可用，以及用多久。由于这些大型模型通常只训练一次，因此准确估计给定 computational budget 条件下的最佳模型超参数至关重要（《Scale efficiently: Insights from pre-training and fine-tuning transformers》）。

   《Scaling laws for neural language models》 表明：自回归语言模型的参数数量与它的性能之间存在幂律（power law）关系。因此，该领域一直在训练越来越大的模型，并预期带来性能的提高。《Scaling laws for neural language models》 的一个值得注意的结论是：大型模型不应该被训练到其最低的损失（test loss ），以达到 compute optimal 。虽然我们得出了同样的结论，但我们估计大型模型的训练应该比作者建议的多得多的 training tokens 。具体来说，在 computational budget 增加十倍的情况下，他们建议模型的大小应该增加 5.5 倍，而 training tokens 的数量只应该增加 1.8 倍。相反，我们发现模型的大小和 training tokens 的数量应该以相等的比例进行 scale 。

   《Scaling laws for neural language models》 的结论：

   $$\begin{array}{}\text{(34)}& N\propto C_{min}^{0.73},B\propto C_{min}^{0.24},S\propto C_{min}^{0.03}\end{array}$$

   当计算预算增加十倍时，模型大小需要增加${10}^{0.73}=5.37$$ 10^{0.73} = 5.37 $ 倍；training tokens 数量为$E=B\times S$$ E=B\times S $ ，需要增加${10}^{0.27}=1.86$$ 10^{0.27} = 1.86 $ 倍。

   遵循 《Scaling laws for neural language models》 和 GPT-3 的训练设置，最近训练的许多大型模型已经训练了大约 300B token （如下表所示），这与增加计算量时主要增加模型规模的做法是一致的。

   

   在这项工作中，我们重新审视了这个问题： 在固定的 FLOPs 预算下，应该如何权衡模型大小和 training tokens 数量？为了回答这个问题，我们将最终的预训练损失$L(N,D)$$ L(N,D) $ 作为模型参数数量$N$$ N $ 和 training tokens 数量$D$$ D $ 的函数。由于 computational budget$C$$ C $ 是模型参数数量$N$$ N $ 和 seen training tokens 数量$D$$ D $ 的确定性函数$\text{FLOPs}(N,D)$$ \text{FLOPs}(N, D) $ ，我们感兴趣的是在$\text{FLOPs}(N,D)=C$$ \text{FLOPs}(N, D) = C $ 的约束条件下最小化$L$$ L $ ：

   $$\begin{array}{}\text{(35)}& N_{\text{opt}}(C),D_{\text{opt}}(C)=\arg \underset{N,Ds.t.\text{FLOPs}(N,D)=C}{min}L(N,D)\end{array}$$

   其中：$N_{\text{opt}}(C),D_{\text{opt}}(C)$$ N_\text{opt}(C), D_\text{opt}(C) $ 描述了 computational budget$C$$ C $ 的最佳分配。

   注意：这里的$D$$ D $ 指的是 training tokens 数量，这跟 《Scaling Laws for Neural Language Models》 中的符号不同。在 《Scaling Laws for Neural Language Models》 中，$D$$ D $ 表示训练集的 tokens 数量，$E=B\times S$$ E=B\times S $ 才是 training tokens 数量。

   这里的$L$$ L $ 指的是 test loss 。

   我们根据 400多个模型的 loss 来经验性地估计这些函数，这些模型的参数规模从低于 70M 到超过 16B ，并在从 5B 到超过 400B 个 token 上进行训练，每个模型配置都在几个不同的 training horizons 进行训练。我们的方法导致了与 《Scaling laws for neural language models》 的结果大不相同。我们在下图中强调了我们的结果，并在后面章节中强调了我们的方法有什么不同。

   

   基于我们估计的 compute-optimal 边界，我们预测，对于用于训练 Gopher 的 compute budget ，最优模型应该小 4 倍，同时在 4 倍的 tokens 上进行训练。我们通过在 1.4T tokens 上训练一个更加 compute-optimal 的 70B 模型，称为 Chinchilla ，来验证这一点。Chinchilla 不仅比其更大的同类模型 Gopher 表现出色，而且其减少的模型尺寸大大降低了推理成本，并大大促进了下游在更小硬件上的使用。大型语言模型的能源成本通过 inference和微调的使用而摊销。因此，一个经过优化训练的小型模型，其好处要超过改善模型性能的直接好处。

2. 相关工作：

   • 大型语言模型：在过去的几年里，各种大型语言模型已经被引入。这些模型包括 dense transformer models 、以及 mixture-of-expert (MoE) models 。最大的 dense transformers 已经超过 500B 参数（《Using Deepspeed and Megatron to Train Megatron-turing NLG 530b, A Large-Scale Generative Language Model》）。

     训练越来越大的模型的动力是很明显的：到目前为止，增加语言模型的规模已经在许多语言建模任务中提高了 SOTA 。然而，大型语言模型面临着一些挑战，包括它们巨大的计算需求（training 和 inference 的成本随着模型的大小而增加），以及需要获得更多的高质量训练数据。事实上，在这项工作中，我们发现更大的、高质量的数据集将在语言模型的任何 further scaling 中发挥关键作用。

   • 对 scaling 行为进行建模：了解语言模型的 scaling 行为及其迁移特性（transfer property）在最近的大型模型的发展中一直很重要。

     《Scaling laws for neural language models》 首次展示了模型大小与 loss 之间在许多数量级上的 predictable relationship 。作者研究了在给定的计算预算（compute budget ）下选择最佳的模型大小来训练的问题。与我们类似，他们通过训练各种模型来解决这个问题。我们的工作在几个重要方面与《Scaling laws for neural language models》 不同：

     • 首先，作者对所有的模型都使用了固定的 training tokens 数量和学习率调度；这使得他们无法建模这些超参数对 loss 的影响。相比之下，我们发现，无论模型大小如何，将学习率调度设置为与 training tokens 数量大致匹配的结果是最好的 final loss ，如下图所示。

       《Scaling Laws for Neural Language Models》 论文中并未固定 training tokens 数量，所以这里的结论是从何而来？

       上面三个子图：training steps = 8M ；下面三个子图：training steps = 12.5M 。每条曲线分别代表余弦学习率的周期长度为 training steps 数量$k$$ k $ 倍，$k=1/1.1/1.25/1.5/2.0/5.0$$ k=1/1.1/1.25/1.5/2.0/5.0 $ 。

       对于一个固定的 learning rate cosine schedule，从训练刚开始到 130B tokens ，中间的 loss estimates （对于$D^{\prime }\ll 130B$$ D^\prime \ll 130 B $ ）高估了用与$D^{\prime }$$ D^\prime $ 相匹配的 schedule length 所训练的模型的 loss 。使用这些中间的 loss 的结果是，低估了在少于130B tokens 的数据上训练模型的有效性，并最终促成了这样的结论：随着计算预算的增加，模型规模的增加应该快于训练数据规模。相反，我们的分析的预测表明，这两个数量应该以大致相同的速度 scale 。

       给定一个 training tokens 数量，由于学习率调度的 non-optimal 选择，导致了训练得到的 training loss 要比最优的 training loss 更高。

       注意，这里是 training loss，而 《Scaling Laws for Neural Language Models》 中的指标是 test loss 。

       

     • 其次，我们包括参数高达 16B 的模型，因为我们观察到 FLOP-loss 边界有轻微的曲率。这意味着来自非常小的模型的预测与来自较大模型的预测不同。事实上，我们分析中使用的大多数模型都有超过 500M 个参数。相比之下，《Scaling laws for neural language models》 的大多数 runs 都要小得多，许多模型的参数不到 100M 。

       不同颜色的 training loss 曲线代表了不同大小的模型。

       在下图中，我们显示了使用第一批三个frontier-points 、中间三个frontier-points 、以及最后三个frontier-points 的线性拟合（橙色虚线、绿色虚线、蓝色虚线）。在这项工作中，我们没有考虑到这一点，我们把它作为有趣的未来工作，因为它表明，对于大的 FLOPs 预算，甚至更小的模型也可能是最佳的。

       

     最近，《Unified scaling laws for routed language models》 专门研究了 Mixture of Expert 语言模型的 scaling 特性，表明随着模型规模的增加，expert 数量的 scaling 会减弱。他们的方法将 loss 作为两个变量的函数：模型规模、专家数量。然而，与《Scaling laws for neural language models》 一样，该分析是在固定数量的 training tokens 下进行的，有可能低估了 branching 的改进。

   • 为大型模型而估计超参数：在选择语言模型，并且选择训练模型的程序时，模型大小和 training tokens 数量并不是唯一需要选择的两个超参数。其他重要的超参数包括学习率、学习率调度、batch size 、优化器、宽深比（width-to-depth ratio ）。在这项工作中，我们专注于模型大小和 training steps 数量，我们依靠现有的工作、以及所提供的实验启发式方法来确定其他必要的超参数。

     • 《Tuning large neural networks via zero-shot hyperparameter transfer》研究了如何选择这些超参数来训练自回归变 transformer ，包括学习率和 batch size。

     • 《An empirical model of large-batch training》 发现最佳 batch size 和模型大小之间只有微弱的依赖性。

     • 《Measuring the effects of data parallelism on neural network training》 提出，使用比我们使用的更大的 batch size 是可能的。

     • 《The depth-to-width interplay in self-attention》研究了各种 standard model size 的最佳 depth-to-width ratio 。我们使用比所提议的深度稍小的模型，因为这可以在我们的硬件上转化为更好的 wall-clock performance 。

   • 改进的模型架构：最近，人们提出了各种有前途的替代传统 dense transformer 的方法。例如，通过使用条件计算， 大型 MoE 模型 （如 1.7T 参数的 Switch transformer 、1.2T 参数的 GLaM 模型，以及其他模型）能够提供较大的有效 model size ，但是使用相对较少的 training FLOPs 和 inference FLOPs 。然而，对于非常大的模型，routed model 的计算优势似乎在减少（《Unified scaling laws for routed language models》）。

     改进语言模型的一个正交方法是用明确的检索机制来 augment transformer 。这种方法有效地增加了训练期间看到的 tokens 数量（在 《Improving language models by retrieving from trillions of tokens》 中增加了 10 倍）。这表明，语言模型的性能可能比以前认为的更加依赖于训练数据的大小。

44.1 方法

1. 我们提出三种不同的方法来回答问题： 在固定的 FLOPs 预算下，应该如何权衡模型大小和 training tokens 数量？在所有三种情况下，我们首先训练一系列模型，改变模型大小和 training tokens 数量，并使用所得的训练曲线来拟合一个empirical estimator ，该 empirical estimator 关于应该如何 scale 。我们假设计算量和模型大小之间存在幂律关系，正如《Unified scaling laws for routed language models》、《Scaling laws for neural language models》 所做的那样，尽管未来的工作可能希望针对大的模型规模在这种关系中包括潜在曲率（potential curvature）。

   所有三种方法的预测结果都是相似的，并且表明参数数量和 training tokens 数量应该随着计算量的增加而同样增加，这个比例如下表所示（括号里的值分别是 10% 分位、90% 分位的统计结果）。这与以前关于这个主题的工作形成了显著的对比，值得进一步研究。

   注意，《Scaling laws for neural language models》 研究的是 test loss，而这里研究的是 training loss 。

   

44.1.1 方法一：固定模型规模并改变 training tokens 数量

1. 在我们的第一种方法中，我们针对一系列的固定规模的模型（模型参数从 70M 到 10B ）改变训练步数，为每个模型训练 4 个不同数量的 training sequences （即，training seen examples 数量）。从这些 runs 中，我们能够直接提取在定数量的 training FLOPs 下实现的最小 loss 的估计。

   四个不同数量的 training sequences 是如何实现的？如果是非常大的数据集，那么很可能最长的 training steps 都没能过遍历一个 epoch，此时，每个样本仅被训练一次；如果是较小的数据，那么很可能最短的 training steps 都遍历了多个 epoch，此时每个样本被训练多次。

   论文并未说明预训练的 datasets 是如何构建的。

2. 我们对最小的模型使用最大的学习率为$2\times {10}^{-4}$$ 2\times 10^{-4} $ ，对最大的模型使用最大的学习率为$1.25\times {10}^{-4}$$ 1.25\times 10^{-4} $ 。在所有情况下，使用余弦学习率调度，学习率在训练结束时相比初始学习率下降了 10 倍。我们假设余弦周期的长度应该与训练步数大致匹配。我们发现，当余弦周期长度超出训练步数 25% 以上时，性能会明显下降，如下图所示。 我们使用窗长为 10 步的高斯平滑来平滑训练曲线。

   

3. 对于每个参数数量$$