1000 层的 Transformer,深得吓人。
昨日出炉的论文《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 © 所示,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 ©)。根据 Xiong 等人 (2020) 的理论分析,通过 LN 的梯度大小与其输入的大小成反比:
相比于没有 warm-up 或正确初始化的情况,图 4 (b) 和图 4 © 表明 ||x|| 的明显大于。这解释了 Post-LN 训练中出现的梯度消失问题(见图 4 (d))。
最重要的是,不稳定性始于训练开始时的大型模型更新。它使模型陷入糟糕的局部最优状态,这反过来又增加了每个 LN 的输入量。随着训练的继续,通过 LN 的梯度变得越来越小,从而导致严重的梯度消失,使得难以摆脱局部最优,并进一步破坏了优化的稳定性。相反,Post-LN-init 的更新相对较小,对 LN 的输入是稳定的。这减轻了梯度消失的问题,使优化更加稳定。