神经语言模型的缩放定律

Jared Kaplan · 约翰霍普金斯大学 & OpenAI · jaredk@jhu.edu
Sam McCandlish · OpenAI · sam@openai.com
Tom Henighan · OpenAI · henighan@openai.com
Tom B. Brown · OpenAI · tom@openai.com
Benjamin Chess · OpenAI · bchess@openai.com
Rewon Child · OpenAI · rewon@openai.com
Scott Gray · OpenAI · scott@openai.com
Alec Radford · OpenAI · alec@openai.com
Jeffrey Wu · OpenAI · jeffwu@openai.com
Dario Amodei · OpenAI · damodei@openai.com

致谢:同等贡献。特别鸣谢:Jared Kaplan 和 Sam McCandlish 主导了本研究;Tom Henighan 贡献了 LSTM 实验;Tom Brown、Rewon Child、Scott Gray 和 Alec Radford 开发了优化的 Transformer 实现;Jeff Wu、Benjamin Chess 和 Alec Radford 开发了文本数据集;Dario Amodei 全程提供指导。

摘要(Abstract)

我们研究了语言模型性能(以交叉熵损失衡量)的经验缩放定律。损失与模型大小、数据集大小以及训练所用的计算量之间呈幂律(power-law)关系,某些趋势跨越了超过七个数量级。网络宽度或深度等其他架构细节在很宽的范围内对性能的影响微乎其微。简单的方程控制了过拟合与模型/数据集大小之间的依赖关系,以及训练速度与模型大小之间的依赖关系。这些关系使我们能够确定固定计算预算的最优分配策略。更重要的发现是,更大的模型具有显著更高的样本效率,因此最优的计算效率训练涉及在相对适中的数据量上训练非常大的模型,并在显著偏离收敛之前提前停止。

目 录

  1. 引言(Introduction)
  2. 背景与方法(Background and Methods)
  3. 经验结果与基本幂律(Empirical Results)
  4. 探索无限数据极限与过拟合(Infinite Data Limit and Overfitting)
  5. 模型大小与训练时间的缩放定律(Scaling Laws with Model Size and Training Time)
  6. 计算预算的最优分配(Optimal Allocation of the Compute Budget)
  7. 相关工作(Related Work)
  8. 讨论(Discussion)
  9. 附录 A—幂律汇总
  10. 附录 B—计算效率前沿的经验模型
  11. 附录 C—注意事项
  12. 附录 D—补充图表

1 引言(Introduction)

语言为人工智能研究提供了天然的领域,因为绝大多数推理任务都可以用语言高效地表达和评估,而世界上的文本为通过生成建模进行无监督学习提供了丰富的数据。深度学习近年来在语言建模方面取得了快速发展,最先进的模型在许多特定任务上已接近人类水平的表现,包括生成连贯的多段落提示文本样本。

人们可能预期,语言建模性能取决于以下因素:模型架构、神经模型的规模、训练所用的计算能力,以及训练过程可用的数据量。在本工作中,我们将实证研究语言建模损失对所有这些因素的依赖关系,重点关注 Transformer 架构。语言任务的高上限和低下限使我们能够研究超过七个数量级的规模趋势。在整个研究过程中,我们将观察到性能作为训练时间、上下文长度、数据集大小、模型大小和计算预算的函数所呈现的精确幂律缩放关系。

Figure 1

图 1:语言建模性能随着模型大小、数据集大小和训练计算量的增加而平滑改善。对于最优性能,这三个因素必须同时扩展。当不受其他两个因素制约时,经验性能与每个单独的因素都呈幂律关系。注:此处显示的是使用足够小的批量大小时预测的计算量,详见图 13 与纯经验数据的对比。

1.1 主要发现摘要(Summary)

我们关于 Transformer 语言模型的关键发现如下:

性能强烈依赖规模,弱依赖模型形态:模型性能主要强烈依赖于三个规模因素:模型参数量 N(不含嵌入)、数据集大小 D,以及训练所用的计算量 C。在合理范围内,性能对其他架构超参数(如深度 vs. 宽度)依赖非常弱。(第 3 节)

平滑的幂律:当不受其他两个因素制约时,性能与 N、D、C 三个规模因素中的每一个都呈幂律关系,趋势跨越超过六个数量级。我们没有观察到这些趋势在上端有任何偏离的迹象,尽管性能在达到零损失之前最终必然会趋于平坦。(第 3 节)

过拟合的普适性:只要我们同时扩大 N 和 D,性能就会可预测地改善;但如果固定其中一个而只增加另一个,就会进入收益递减的区域。性能损失可预测地取决于比值 N0.74/D——这意味着每当模型大小扩大 8 倍时,只需要增加约 5 倍的数据就可以避免惩罚。(第 4 节)

训练的普适性:训练曲线遵循可预测的幂律,其参数与模型大小大致无关。通过外推训练曲线的早期部分,我们可以大致预测如果训练更长时间将达到的损失。(第 5 节)

迁移随测试性能提升而改善:当我们在与训练分布不同的文本上评估模型时,结果与训练验证集上的结果呈强相关,损失中有一个大致恒定的偏移量——换言之,向不同分布的迁移会产生恒定的惩罚,但其他方面的改善大致与训练集上的性能提升保持一致。(第 3.2.2 节)

样本效率:大模型比小模型更具样本效率,能够用更少的优化步数和使用更少的数据点达到相同的性能水平。(图 2)

收敛是低效的:当在固定计算预算 C 下训练(但对模型大小 N 或可用数据 D 没有任何其他限制)时,我们通过训练非常大的模型并显著提前停止来获得最优性能。因此,最大计算效率的训练比基于将小模型训练到收敛所预期的样本效率要高得多,数据需求的增长非常缓慢,D ~ C0.27。(第 6 节)

最优批量大小:训练这些模型的理想批量大小大约是损失的幂函数,并且仍然可以通过测量梯度噪声尺度来确定;对我们能训练的最大模型,在收敛时大约为 100-200 万个 token。(第 5.1 节)

1.2 缩放定律摘要(Summary of Scaling Laws)

对于训练为自回归语言建模的 Transformer,其测试损失可以通过幂律来预测,当性能仅受以下因素之一限制时:非嵌入参数数量 N、数据集大小 D,或最优分配的计算预算 Cmin(见图 1):

1. 参数受限、收敛训练:
L(N) = (Nc/N)αN
αN ~ 0.076,Nc ~ 8.8 × 1013(非嵌入参数)
2. 大模型、数据集受限、提前停止:
L(D) = (Dc/D)αD
αD ~ 0.095,Dc ~ 5.4 × 1013(token 数)
3. 计算受限、最优批量大小:
L(Cmin) = (Ccmin/Cmin)αCmin
αCmin ~ 0.050,Ccmin ~ 3.1 × 108(PF-day)

这些关系在 Cmin 上跨越八个数量级,N 上跨越六个数量级,D 上跨越两个数量级以上。它们对模型形态和 Transformer 的其他超参数(深度、宽度、自注意力头数量)依赖非常弱,具体的数值与 WebText2 训练集相关联。幂律 αN、αD、αCmin 规定了当我们扩大 N、D 或 Cmin 时预期性能改善的程度;例如,将参数量翻倍会使损失减少约 2N = 0.95 倍(5%)。

临界批量大小(决定数据并行性的速度/效率权衡)也大致服从损失的幂律:

Bcrit(L) = B*/L1/αB
B* ~ 2 × 108 tokens,αB ~ 0.21

公式(1.1)和(1.2)一起表明,随着模型规模的增大,我们应该按 D ∝ NαND ~ N0.74 来亚线性地增加数据集大小。事实上,我们发现有一个将(1.1)和(1.2)结合的单一方程,它控制着 N 和 D 的联合依赖关系以及过拟合的程度:

L(N, D) = [ (Nc/N)αND + Dc/D ]αD

当在无限数据极限下训练给定模型有限步数 S 时,在初始瞬态期之后,学习曲线可以精确地拟合为:

L(N, S) = (Nc/N)αN + (Sc/Smin(S))αS
其中 Sc ≈ 2.1 × 10³,αS ≈ 0.76

当在固定计算预算 C 下训练但没有其他约束时,公式(1.6)导致以下预测:最优模型大小 N、最优批量大小 B、最优步数 S 和数据集大小 D 应该按以下方式增长:

N ∝ CαCminN,B ∝ CαCminB,S ∝ CαCminS,D = B · S

其中 αCmin = 1/(1/αS + 1/αB + 1/αN),这与经验最优结果 N ∝ Cmin0.73、B ∝ Cmin0.24 和 S ∝ Cmin0.03 非常吻合。随着计算预算 C 的增加,大部分应该花在更大的模型上,而训练时间或数据集大小的增加相对较小。这也意味着,随着模型变得更大,它们变得越来越样本高效。

1.3 符号约定(Notation)

符号含义补充说明
L交叉熵损失以 nat 为单位,通常对上下文中的 token 取平均
N模型参数量不含所有词汇和位置嵌入
C ≈ 6NBS总非嵌入训练计算量B 为批量大小,S 为训练步数(参数更新次数);以 PF-day 为单位(1 PF-day ≈ 8.64 × 1019 FLOPs)
D数据集大小以 token 为单位
Bcrit临界批量大小定义见第 5.1 节;在临界批量大小训练能在时间效率和计算效率之间取得大致最优的权衡
Cmin达到给定损失的最小非嵌入计算量相当于在远小于临界批量大小的情况下训练所使用的计算量
Smin达到给定损失所需的最小训练步数相当于在远大于临界批量大小的情况下训练所使用的步数
αX幂律指数L(X) ∝ 1/XαX,X 可为 N、D、C、S、B、Cmin

2 背景与方法(Background and Methods)

我们在 WebText2 数据集上训练语言模型,这是 WebText 数据集的扩展版本,使用字节对编码(BPE)进行分词,词汇表大小 nvocab = 50,257。我们优化对 1024-token 上下文取平均的自回归对数似然(即交叉熵损失),这也是我们的主要性能指标。我们记录 WebText2 测试分布以及一系列其他文本分布上的损失。我们主要训练仅解码器的 Transformer 模型,但也训练 LSTM 和 Universal Transformer 进行对比。

2.1 Transformer 的参数和计算缩放

操作参数量每个 Token 的 FLOPs(前向)
嵌入(nvocab+nctx)·dmodel4dmodel
注意力:QKVnlayer·dmodel·3·dattn2nlayer·dmodel·3·dattn
注意力:Mask2nlayer·nctx·dattn
注意力:投影nlayer·dattn·dmodel2nlayer·dattn·dembd
前馈网络nlayer·2·dmodel·dff2nlayer·2·dmodel·dff
去嵌入2dmodel·nvocab
总计(非嵌入)N = 2dmodel·nlayer·(2dattn+dff)Cforward = 2N + 2nlayer·nctx·dattn

我们使用以下超参数来参数化 Transformer 架构:nlayer(层数)、dmodel(残差流的维度)、dff(前馈层的中间维度)、dattn(注意力输出的维度)和 nheads(每层的注意力头数)。我们在输入上下文中包含 nctx 个 token(除另有说明外,nctx = 1024)。我们使用 N 来表示模型大小,定义为非嵌入参数的数量:

N = 12nlayerdmodel2(使用标准配置 dattn = dff/4 = dmodel

我们的模型在嵌入矩阵中还有 nvocab·dmodel 个参数,并在位置嵌入中使用 nctx·dmodel 个参数,但在讨论"模型大小"N 时不包含这些——我们将会看到,这样做产生了明显更干净的缩放定律。评估 Transformer 的前向传递涉及大约 Cforward ≈ 2N + 2nlayer·nctx·dmodel 次加乘运算。反向传播的计算量约为前向传递的两倍,因此我们将每个训练 token 的估计非嵌入计算量定义为 C ≈ 6N FLOPs。

2.2 训练程序

除非另有说明,我们使用 Adam 优化器训练模型,固定步数为 2.5 × 105 步,批量大小为 512 个序列,每个序列包含 1024 个 token。由于内存限制,我们最大的模型(超过 10 亿参数)使用 Adafactor 训练。我们试验了各种学习率和调度策略,发现收敛时的结果与学习率调度关系不大。除非另有说明,所有训练运行都使用学习率调度,包括 3000 步的线性预热,然后余弦衰减到零。

2.3 数据集

我们在 WebText 数据集的扩展版本(WebText2)上训练模型。原始 WebText 数据集是对 Reddit 外部链接的网络抓取(截至 2017 年 12 月),获得至少 3 个 karma 的链接。在 WebText2 中,我们添加了 2018 年 1 月至 10 月期间的外部 Reddit 链接,同样要求至少 3 个 karma。总数据集包含 20.3M 个文档,96 GB 文本,1.62 × 1010 个词(按 wc 定义)。然后我们应用可逆分词器,得到 2.29 × 1010 个 token。我们预留 6.6 × 108 个 token 作为测试集。

3 经验结果与基本幂律(Empirical Results and Basic Power Laws)

为了表征语言模型缩放,我们训练了各种各样的模型,改变的因素包括:模型大小(从 768 到 15 亿非嵌入参数不等);数据集大小(从 2200 万到 230 亿 token 不等);形态(包括深度、宽度、注意力头数和前馈维度);上下文长度(大多数运行使用 1024,也用更短的上下文进行实验);批量大小(大多数运行使用 219,但也改变它来测量临界批量大小)。

3.1 近似 Transformer 形态和超参数独立性

当我们固定总非嵌入参数数量 N 时,Transformer 性能对形态参数 nlayer、nheads 和 dff 的依赖非常弱。为了建立这些结果,我们固定大小训练模型,同时改变单个超参数。纵横比可以在 40 倍的范围内变化,同时对性能的影响非常小;一个 (nlayer, dmodel) = (6, 4288) 的模型达到的损失与 (48, 1600) 模型相差在 3% 以内。

Figure 4

图 4:左图:当包含嵌入参数时,性能似乎还强烈地依赖于层数以及参数数量。右图:当我们排除嵌入参数时,具有不同深度的模型性能收敛到单一趋势。只有层数少于 2 层或嵌入参数比例过高的模型才显示出明显偏离。

3.2 性能与非嵌入参数数量 N 的关系

当排除了嵌入参数后,不同深度的模型收敛到单一的幂律趋势上(图 6)。我们观察到没有明显的偏离趋势的迹象——性能在我们探索的整个范围内平滑地遵循幂律。这些趋势跨越了超过六个数量级的模型大小。

L(N) = (Nc/N)αN;αN ~ 0.076,Nc ~ 8.8 × 1013

3.3 LSTM 基线

为了评估 Transformer 架构的改进,我们与 LSTM 基线进行了比较。即使控制了参数量,Transformer 也始终显著优于 LSTM,并且随着模型规模的增大,这种优势也在增大。

3.4 跨数据集泛化

我们在各种测试数据集上评估模型,发现性能对训练分布的依赖与对其他分布的依赖之间存在强相关性,损失中有恒定的偏移量。换言之,迁移到一个不同分布会产生恒定的惩罚,但其他方面的改善大致与训练集上的性能提升保持一致。

4 探索无限数据极限与过拟合(Charting the Infinite Data Limit and Overfitting)

当同时考虑 N 和 D 时,损失由以下公式控制:

L(N, D) = [ (Nc/N)αND + Dc/D ]αD

性能损失可预测地取决于比值 N0.74/D——这意味着每当模型大小扩大 8 倍时,只需要增加约 5 倍的数据就可以避免惩罚。我们还发现,当控制 N 和 D 的比例不变时,损失与绝对规模无关。

Figure 5

图 5:样本效率——大模型在看到较少数据时就能达到与小模型相同甚至更低的损失。左图:不同规模模型的损失随 token 数量的变化;右图:达到特定损失水平所需的最小串行步数随模型大小急剧减少。

5 模型大小与训练时间的缩放定律(Scaling Laws with Model Size and Training Time)

5.1 临界批量大小

临界批量大小决定了数据并行性的速度/效率权衡:

Bcrit(L) = B*/L1/αB
B* ~ 2 × 108 tokens,αB ~ 0.21

这意味着在我们能训练的最大模型收敛时,临界批量大小约为 100-200 万个 token。

5.2 学习曲线

当在无限数据极限下训练给定模型有限步数 S 时,在初始瞬态期之后,学习曲线可以精确地拟合为:

L(N, S) = (Nc/N)αN + (Sc/Smin(S))αS
其中 Sc ≈ 2.1 × 10³,αS ≈ 0.76

6 计算预算的最优分配(Optimal Allocation of the Compute Budget)

当在固定计算预算 C 下训练(但对模型大小 N 或可用数据 D 没有任何其他限制)时,我们通过训练非常大的模型并显著提前停止来获得最优性能。当控制 N 和 D 不变时,有一个关于计算最优训练和样本最优训练之间权衡的讨论。

N ∝ C0.73,B ∝ C0.24,S ∝ C0.03,D = B · S
核心结论:随着计算预算 C 的增加,大部分应该花在更大的模型上,而训练时间或数据集大小的增加相对较小。在实践中,由于硬件限制,研究人员通常训练比计算最优方案更小的模型,但训练更长时间。

7 相关工作(Related Work)

之前的工作包括 Hestness et al. (2017) 对深度学习 Scaling 可预测性的经验研究,以及 McCandlish et al. (2018) 关于大批量训练的经验模型。我们是首个系统研究语言模型性能与所有关键因素(N、D、C)之间幂律关系的的工作。

8 讨论(Discussion)

我们的结果表明,只要适当扩展模型大小、数据和计算,语言建模性能会平滑且可预测地改善。我们预期更大的语言模型将表现得更好,并且比当前模型更样本高效。我们的发现对大语言模型训练实践有重要启示。首先,在固定计算预算下,最优策略是训练非常大的模型并在收敛之前提前停止。其次,随着模型变大,它们变得越来越样本高效,因此用更多数据训练更少步数是浪费。第三,批量大小主要由目标损失决定,而不是模型大小。

附录 A — 幂律汇总

关系公式参数
L(N)(Nc/N)αNαN ~ 0.076
L(D)(Dc/D)αDαD ~ 0.095
L(Cmin)(Ccmin/Cmin)αCminαCmin ~ 0.050
Bcrit(L)B*/L1/αBαB ~ 0.21
L(N,D)[(Nc/N)αND + Dc/D]αD联合公式
L(N,S)(Nc/N)αN + (Sc/Smin(S))αSαS ~ 0.76

附录 C — 注意事项(Caveats)

  1. 目前我们对任何提议的缩放定律都没有坚实的理论理解。缩放与模型大小和计算的关系尤为神秘。
  2. 我们对远离已探索范围的损失值下 Bcrit(L) 的预测不是特别有信心。临界批量大小的变化可能对数据并行性和所需串行训练步数之间的权衡产生重大影响。
  3. 我们没有彻底调查小数据 regime,D 最小时 L(N,D) 的拟合较差。此外,我们没有尝试正则化和数据增强。
  4. 我们使用了估计的训练计算量 C ≈ 6NBS,未包括与 nctx 成比例的贡献。
  5. 最优学习率的选择对目标损失敏感。接近收敛时训练可能需要较小的学习率。

参考文献(References)

[1] Kaplan et al. (2020). Scaling Laws for Neural Language Models. arXiv: 2001.08361
[2] Hoffmann et al. (2022). Training Compute-Optimal Large Language Models. arXiv: 2203.15556
[3] Vaswani et al. (2017). Attention Is All You Need. NeurIPS
[4] Brown et al. (2020). Language Models are Few-Shot Learners. NeurIPS
[5] McCandlish et al. (2018). An Empirical Model of Large-Batch Training

本文为 Kaplan et al. (2020) 的完整中文翻译版,保持原文全部内容。
原文:Scaling Laws for Neural Language Models · arXiv: 2001.08361 · OpenAI