解决Transformer训练难题，微软亚研直接把Transformer干到了1000层

昨日出炉的论文《DeepNet: Scaling Transformers to 1,000 Layers》在钻研社区引起了热议，作者来自微软亚洲钻研院。

该钻研直接把Transformer深度提升到1000层！

下面让我们看下这篇钻研说了什么。

近年来，大规模 Transformer模型出现了这样一种趋势：随着模型参数从数百万增加至数十亿甚至数万亿，性能相应地实现了显著提升。大规模模型在一系列任务上都取得了SOTA性能，并在小样本和零样本学习设置下展现出了令人瞩目的能力。如下图1所示，尽管参数量已经很大了，但Transformer模型的深度（depth）却受到了训练不稳定的限制。

Nguyen和Salazar (2019)发现，基于post-norm连接（Post-LN），pre-norm 残差连接（Pre-LN）能够提升 Transformer的稳定性。但是，Pre-LN在底层的梯度往往大于顶层，因而导致与 Post-LN相比性能下降。为了缓解这一问题，钻研人员一直努力通过更好的初始化或更好的架构来改进深度Transformer的优化。这些要领可以使多达数百层的Transformer模型实现稳定化，然而以往的要领没有能够成功地扩张至1000层。

微软亚研在一篇新论文《DeepNet: Scaling Transformers to 1,000 Layers》中终于将Transformer的深度扩张到了1000层。

论文地址：https://arxiv.org/pdf/2203.00555.pdf

钻研者的目标是提升 Transformer 模型的训练稳定性，并将模型深度进行数量级的扩张。为此，他们钻研了不稳定优化的原因，并且发现爆炸式模型革新是造成不稳定的罪魁祸首。基于这些观察，钻研者在残差连接处引入了一个新的归一化函数 —— DEEPNORM，它在将模型革新限制为常数时具有理论上的合理性。这一要领简单但高效，只需要改变几行代码即可。最终，该要领提升了Transformer模型的稳定性，并实现了将模型深度扩张到了1000多层。

此外，实验结果表明，DEEPNORM 能够将 Post-LN 的良好性能和Pre-LN的稳定训练高效结合起来。钻研者提出的要领可以成为Transformers的首选替代方案，不仅适用于极其深（多于1000层）的模型，也适用于现有大规模模型。值得指出的是，在大规模多言语机器翻译基准上，文中 32 亿参数量的 200 层模型（DeepNet）比120亿参数量的48层SOTA模型（即 Facebook AI的M2M模型）实现了 5 BLEU值提升。

有知乎网友疑问：就实现效果来说，1000层是不是有必要？论文作者之一董力（Li Dong）表示，1000层更多地是为了探究上限，实际跑的过程中并非一定要上千层。此外，训练代码很快就会公开。

DEEPNORM要领

如下图2所示，使用 PostLN 实现基于 Transformer 的要领很简单。与 Post-LN 相比，DEEPNORM 在执行层归一化之前up-scale了残差连接。

（图注）图2：(a) DEEPNORM 的伪代码，例如可以用其他标准初始化代替 Xavier 初始化 (Glorot and Bengio, 2010) ，其中α是一个常数。(b) 不同架构的 DEEPNORM 参数（N 层编码器，M 层解码器）。

此外，该钻研还在初始化期间down-scale了参数。值得注意的是，该钻研只扩张了前馈网络的权重，以及注意力层的值投影和输出投影。此外，残差连接和初始化的规模取决于图2中不同的架构。

深度Transformer的不稳定性

该钻研分析了深度Transformer不稳定的原因。

首先，钻研者观察发现：更好的初始化要领可以让 Transformer 的训练更稳定。之前的工作（Zhang et al., 2019a; Huang et al., 2020; Xu et al., 2021）也证实了这一点。

因此，钻研者分析了有无适当初始化的 Post-LN 的训练过程。通过更好的初始化，在执行 Xavier 初始化后通过down-scale第l层的权重。例如，第l层 FFN 的输出投影被初始化为

其中d’是输入和输出维度的平均值。钻研者将此模型命名为 Post-LN-init。请注意，与之前的工作（Zhang et al., 2019a）不同， Post-LN-init是缩窄了较低层的扩张而不是较高层。钻研者相信这种要领有助于将梯度扩张的影响与模型革新区分开来。此外，Post-LN-init 与 Post-LN 具有相同的架构，从而消除了架构的影响。

该钻研在 IWSLT-14 De-En 机器翻译数据集上训练了 18L-18L Post-LN 和 18L-18L Post-LN-init。图 3 可视化了它们的梯度和验证损失曲线。如图 3(c) 所示，Post-LN-init 收敛，而 Post-LN 没有。 Post-LN-init 在最后几层中具有更大的梯度范数，尽管其权重已按比例缩小。此外，钻研者可视化最后一个解码器层的梯度范数，模型深度从 6L-6L 到 24L-24L。

下图 3 显示，无论模型深度如何，最后一层 Post-LN-init 的梯度范数仍远大于 Post-LN 的梯度范数。得出的结论是，深层梯度爆炸不应该是 Post-LN 不稳定的根本原因，而模型革新的扩张往往可以解释这一点。

然后钻研者证明 Post-LN 的不稳定性来自一系列问题，包括梯度消失以及太大的模型革新。如图 4(a) 所示，他们首先可视化模型革新的范数 ||ΔF||在训练的早期阶段：

其中x和θ_i分别代表输入和第i次革新后的模型参数。Post-LN在训练一开始就有爆炸式的革新，然后很快就几乎没有革新了。这表明该模型已陷入虚假的局部最优。

warm-up和更好的初始化都有助于缓解这个问题，使模型能够顺利革新。当革新爆炸时，LN 的输入会变大（见图 4(b) 和图 4(c)）。根据Xiong等人(2020)的理论分析，通过 LN 的梯度大小与其输入的大小成反比：

相比于没有warm-up或正确初始化的情况，图 4(b) 和图 4(c) 表明 ||x||的明显大于。这解释了 Post-LN 训练中出现的梯度消失问题（见图 4(d)）。

最重要的是，不稳定性始于训练开始时的大型模型革新。它使模型陷入糟糕的局部最优状态，这反过来又增加了每个 LN 的输入量。随着训练的继续，通过 LN 的梯度变得越来越小，从而导致严重的梯度消失，使得难以摆脱局部最优，并进一步破坏了优化的稳定性。相反，Post-LN-init 的革新相对较小，对 LN 的输入是稳定的。这减轻了梯度消失的问题，使优化更加稳定。

DeepNet：极深的Transformer模型

钻研者首先介绍了极深的Transformer模型——DeepNet，该模型可以通过缓解爆炸式模型革新问题来稳定优化过程。

DeepNet基于Transformer架构。与原版Transformer相比，DeepNet在每个子层使用了新要领DEEPNORM，而不是以往的Post-LN。DEEPNORM的公式如下所示。

其中，α是一个常数，G_l(x_l , θ_l)是参数为θ_l的第l个Transformer子层（即注意力或前馈网络）的函数。DeepNet还将残差内部的权重θ_l扩张了β。

接着，钻研者提供了对DeepNet模型革新预期大小（expected magnitude）的估计。

他们可视化了IWSLT-14 De-En翻译数据集上，Post-LN和DeepNet在早期训练阶段的模型革新情况，如下图5所示。可以看到，相较于Post-LN，DeepNet的模型革新几乎保持恒定。

最后，钻研者提供理论分析，以表明 DeepNet的革新受到了 DEEPNORM 的常数限制。具体地，他们展示了 DeepNet的预期模型革新受到了适当参数α 和 β的常数限制。钻研者的分析基于 SGD 革新，并通过实证证明对 Adam 优化器效果很好。

钻研者提供了对编码器-解码器架构的分析，它能够以相同的方式自然地扩张到仅编码器和仅解码器的模型。具体如下图所示，他们将模型革新的目标设定如下：

仅编码器（例如 BERT）和仅解码器（例如 GPT）架构的推导能够以相同的方式进行。钻研者将步骤总结如下：

神经机器翻译

该钻研验证了DeepNet 在流行的机器翻译基准上的有效性，包括 IWSLT-14 德语-英语 (De-En) 数据集和 WMT-17 英语-德语 (En-De) 数据集。该钻研将DeepNet 与多个SOTA深度 Transformer 模型进行比较，包括 DLCL 、NormFormer 、ReZero 、R- Fixup 、T-Fixup 、DS-init 和 Admin。

下表 1 报告了 WMT-17 En-De 翻译数据集上的基线和DeepNet 的结果：

下图 6 显示了 IWSLT-14 数据集的结果

下图 7 报告了 WMT-17 验证集的损失曲线

大规模多言语神经机器翻译

该钻研首先使用 OPUS-100 语料库来评价模型。OPUS100 是一个以英语为中心的多言语语料库，涵盖 100 种言语，是从 OPUS 集合中随机抽取的。该钻研将 DeepNet   扩张到 1,000 层，该模型有一个 500 层的编码器、 500 层的解码器、512 个隐藏大小、8 个注意力头和 2,048 维度的前馈层。

下表2总结了 DeepNet 和基线的结果。结果表明，增加网络深度可以显着提高 NMT 的翻译质量：48 层的模型比 12 层的模型平均获得 3.2 点的提高。 DeepNet  可以成功地将深度扩张到 1,000 层，比基线提高4.4 BLEU。值得注意的是，DeepNet 只训练了 4 个 epoch，并且在计算预算更多的情况下，性能可以进一步提高。

深度扩张规律：该钻研在OPUS100数据集上训练具有{12,20,100,200,1000}层的DeepNet，图8显示了深度扩张曲线。与双语NMT相比，多语NMT从扩张模型深度受益更多。可以观察到多语 NMT 的 BLEU 值呈对数增长，规律可以写成：L(d) = A log(d) + B，其中d是深度，A, B是关于其他超参数的常数。

更多数据和言语说明：为了探索DeepNet在多语NMT上的局限性，该钻研随后使用Schwenk等人提出的CCMatrix扩张训练数据。此外，该钻研还扩张了CCAligned 、OPUS 和Tatoeba的数据，以涵盖Flores101评价集的所有言语。最终的数据由102种言语、1932个方向和12B对句子组成。利用这些数据，该钻研用100层编码器、100层解码器、1024个隐藏维度、16个头、4096个前馈层中间维度对DeepNet进行训练。

该钻研将 DeepNet 与SOTA多语 NMT 模型 M2M-100进行了比较。M2M-100 有一个 24 层的编码器、一个 24 层的解码器和 4,096 个隐藏大小，从而产生高达 12B 的参数。与M2M-100相比，DeepNet深而窄，参数只有3.2B。

在 M2M-100 之后，该钻研在几个多言语翻译评价数据集上评价模型，包括 WMT、OPUS 、TED、 Flores。WMT的言语对是以英语为中心的。包括英语在内的10种言语，其中大部分是高资源言语。对于 OPUS 数据集，该钻研从包含 30 个评价对的测试集中选择非英语方向。TED评价集有28种言语和756个方向，数据来自口语领域。 Flores 数据集包含 102 种言语之间的所有翻译对。该钻研使用涵盖 M2M-100 和  DeepNet  支持的言语的子集，产生 87 种言语和 7,482 个翻译方向。

下表 3 报告了结果，为了公平比较，该钻研使用与基线相同的评价要领。结果表明 DeepNet 在所有评价数据集上的性能都明显优于 M2M-100，表明深化模型是提高 NMT 模型质量的一个非常有前景的方向。

感兴趣的读者可阅读论文原文，了解更多细节内容。