Press "Enter" to skip to content

解决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值提升 。

 

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 模型质量的一个非常有前景的方向。

 

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

Be First to Comment

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注