Press "Enter" to skip to content

1000层 Transformer 也能稳定训练?详解微软亚研院提出新工作DeepNet

本站内容均来自兴趣收集,如不慎侵害的您的相关权益,请留言告知,我们将尽快删除.谢谢.

作者丨科技猛兽

 

本文目录

 

43 解决 Transformer 训练难题,1000层 Transformer 也能稳定训练

 

(来自微软亚洲研究院)

 

43.1 DEEPNET 原理分析

 

Transformer 是 Google 的团队在 2017 年提出的一种 NLP 经典模型,现在比较火热的 Bert 也是基于 Transformer。Transformer 模型使用了 Self-Attention 机制, 不采用 RNN 的 顺序结构 ,使得模型 可以并行化训练 ,而且能够 拥有全局信息。

 

视觉 Transformer 模型在计算机视觉领域取得了巨大的成功,甚至大有取代 CNN 之势。但是虽然 Transformer 模型的参数从数百万增加至数十亿甚至数万亿,大规模模型也在一系列任务上都取得了 SOTA 的性能,并在小样本和零样本学习设置下展现出了令人瞩目的能力,但是视觉 Transformer 模型的深度 (depth) 却受到了训练不稳定的限制,使得 Transformer 模型不容易像 CNN 那样做 Deep。

 

本文就是为了解决这个问题,所提出的 DeepNet 模型成功地将 Transformer 的深度扩展到了1000层。

 

43 解决 Transformer 训练难题,1000层 Transformer 也能稳定训练

 

论文名称:DeepNet: Scaling Transformers to 1,000 Layers

 

论文地址:

 

https://arxiv.org/pdf/2203.00555.pdf

 

43.1 DeepNet 原理分析:

 

在 Toan Q. Nguyen and Julian Salazar 的这篇文章里面:

 

Transformers without tears: Improving the normalization of self-attention

 

Toan Q. Nguyen and Julian Salazar 发现把每个 Block 里面的 Layer Normalization 层从 Self-Attention 层的后面 (Post-LN) 换到 Self-Attention 层之前 (Pre-LN),能够提升 Transformer 训练的稳定性。但是后人的研究又发现 Pre-LN 层的梯度存在这样的现象,即:浅层的梯度高于深层,这个现象会使 Pre-LN 的性能不如 Post-LN。之后,又有几个工作试图将 Transformer 模型做得更深,如下:

 

从初始化角度:Improving deep transformer with depth-scaled initialization and merged attention (EMNLP-IJCNLP 2019)

 

从初始化角度:Improving transformer optimization through better initialization (ICML 2020)

 

从结构的角度:Learning deep transformer models for machine translation (ACL 2019)

 

从结构的角度:Understanding the difficulty of training transformers (EMNLP 2020)

 

从结构的角度:Rezero is all you need: Fast convergence at large depth

 

从结构的角度:Normformer: Improved transformer pretraining with extra normalization

 

这些工作能够把 Transformer 训练到数百层的深度,但是训练到1000层,是前所未有的。

 

本文作者的目标是提升 Transformer 模型的训练稳定性,并将模型深度进行数量级的扩展。

 

43.1.1 Motivation

 

论文里面并没有很清楚地说明把 Transformer 做深的动机是什幺,一开始读这篇论文时我首先想到的是为什幺一定要把 Transformer 模型做得这幺深 (除去发论文的目的)。

 

后来结合 MSRA 的一系列工作 (Swin, Swin V2),和这个知乎的回答:如何评价微软亚洲研究院的Swin Transformer V2:在4个数据集上达到SOTA?:https://www.zhihu.com/question/500004483/answer/2236028403。个人觉得目前大模型是诸如 MSRA 这类大厂的研究趋势,从 Swin V2 强行扩展就能够看得出,SwinV2 这个工作本身, 不是开创性的 ,也未必会对整个领域产生深远的影响, 但这个工作本身对 Swin 是重要的 ,因为在大规模无监督数据加大模型的故事里, 一个新结构,必须要证明自己能有效的训练大模型。

 

那幺在这个 无监督+大模型 的故事背景下, 简单地通过增加 Block 的 channel 来把 Swin Transformer 从 Base 扩展到 G 是可以的,但如果想再大,就要从 Depth 的角度来扩展了 。

 

Swin V2 搞懂 Vision Transformer 原理和代码,看这篇技术综述就够了(二十) 已经研究了几种解决训练中的不稳定性问题的方法,如:

 

Post Normalization 技术。

 

Scaled Cosine Attention 技术。

 

和对数连续位置编码技术等等。

 

但是还是那句话,这些方法只是辅助 Transformer 在 channel 维度增加的训练方法,对于 Depth 维度的增加,需要探索新的稳定训练的方式,这也是本文的价值所在。

 

为此,作者们研究了不稳定优化的原因,并且发现爆炸式模型更新 (exploding model update) 是造成不稳定的罪魁祸首。基于这些观察,研究者在残差连接处引入了一个新的归一化函数 —— DEEPNORM,它从理论上保证了把模型更新过程限制为常数。这一方法简单但高效,只需要改变几行代码即可。最终,该方法提升了 Transformer 模型的稳定性,并实现了将模型深度扩展到了1000多层。

 

结果显示,本文的方法能够将 Post-LN 性能的优势和 Pre-LN 训练稳定的优势结合起来,且对于目前多个大型 Transformer 模型也是适用的。

 

43.1.2 太长不看版

 

如果没时间仔细阅读,只需要看本小节这一点点就行。

 

DeepNorm 伪代码:

 

def deepnorm(x):
   return LayerNorm(x∗α + f(x))
def deepnorm_init(w):
   if w is ['ffn', 'v_proj', 'out_proj']:
      nn.init.xavier_normal_(w, gain=β)
   elif w is ['q_proj', 'k_proj']:
      nn.init.xavier_normal_(w, gain=1)

 

1 使用 DeepNorm 时,拿它来替换 Post-LN。

 

2 DeepNorm 其实就是 LN,只是在执行层归一化之前 up-scale 了残差连接。x∗α + f(x) 里面的 f(x) 代表 Self-Attention 等等的 Token-mixer,x 代表 Token-mixer 的输出,α 是常数。

 

3 torch.nn.init.xavier_normal_(tensor, gain=1) 是 xavier 高斯初始化,参数由0均值,标准差为 gain × sqrt(2 / (fan_in + fan_out)) 的正态分布产生,其中 fan_in 和 fan_out 是分别权值张量的输入和输出元素数目。这种初始化同样是为了保证 输入输出的方差不变。

 

4 DeepNorm 还在初始化期间 down-scale 了参数。值得注意的是,对于 ffn,v_proj,out_proj 和 q_proj,k_proj 的初始化,是不一样的。

 

5 不同架构的 α 和 β 值不一样,如下图所示。

 

图1:不同架构的 α 和 β 值不一样

43.1.3 深度 Transformer 模型训练的不稳定问题

 

根据前人的研究,更好的初始化方法可以让 Transformer 的训练更稳定。因此,研究者分析了有无适当初始化的 Post-LN 的训练过程。通过更好的初始化,在执行 Xavier 初始化后对第层的权重进行 downscale,下降的比例是: 。比如,第 层的 FFN 的权重初始化为:

 

式中,是输入和输出维度的均值。

 

研究者将此模型命名为 Post-LN-init。注意到这个值是随着层数的增加而越来越小的,也就意味着对于浅层而言,权重进行 downscale 的比例较大,对于深层而言,权重进行 downscale 的比例较小。此外,Post-LN-init 与 Post-LN 具有相同的架构,从而消除了架构的影响。

 

作者接下来在 IWSLT-14 De-En 机器翻译数据集上面训练了 18L-18L Post-LN 和 18L-18L Post-LN-init。如下图1所示是它们的梯度和 validation loss 的曲线。如图1(c) 所示,Post-LN-init 的 validation loss 是收敛的,但是 Post-LN 却没有。如图1(a) 所示,Post-LN-init 在最后几层中具有更大的梯度范数,尽管其权重已按比例缩小。这说明梯度爆炸不是导致 Transformer 训练不稳定的原因。如图1(b) 所示,作者可视化最后一个解码器层的梯度范数,模型深度从 6L-6L 到 24L-24L,发现 Post-LN-init 的梯度范数大于 Post-LN 的梯度范数,但是初始化时 Post-LN-init 的范数是更小的。

 

图1:IWSLT-14 De-En 机器翻译数据集上面训练了 18L-18L Post-LN 和 18L-18L Post-LN-init 的实验结果

接下来作者通过一系列实验,从实验的角度验证 Post-LN 的不稳定性来自一系列的问题。

 

首先可视化在训练的早期阶段,模型更新的范数,如下图2(a)所示:

 

式中,和 代表第 次更新的输入和模型参数。Post-LN 在训练一开始就有爆炸式的更新,然后很快就几乎没有更新了。这表明该模型已陷入虚假的局部最优。warm-up,或者更好的初始化都有助于缓解这个问题,使模型能够更加顺利地更新。图2(b)和(c)分别是从 FFN 层和 Attention 层到 LN 层的输入,作者发现当更新爆炸时,LN 的输入会变大。对于这种现象,之前已有理论分析:

 

On layer normalization in the transformer architecture (ICML 2020)

 

通过 LN 的梯度大小与其输入的大小成反比:

 

当不使用 warm-up 以及合适的初始化的情况下,图2(b)和(c) 表明明显大于 ,这也就解释了 Post-LN 训练中出现的梯度消失问题,如图2(d)所示。

 

图2:18L-18L 模型在训练时的一些可视化结果。(a) 累计参数更新。(b) (c) 从 FFN 层和 Attention 层到 LN 层的输入。(d) decoder 层的梯度范数

本节的结论是:

 

大的 Transformer 模型训练的不稳定的原因是训练的 起始阶段模型参数更新非常快 ,它使模型陷入一个坏的局部最优,这反过来增加了每个 LN 的输入。随着训练的继续进行,通过 LN 的梯度变得越来越小,从而导致严重的梯度消失。消失的梯度使得难以脱离局部最优,并且进一步破坏了优化的稳定性。所以最关键的问题是 不要让起始阶段模型参数更新得太快 。相反,Post-LN-init 的更新相对较少,LN 的输入是稳定的。这减轻了梯度消失的问题,使优化更加稳定。

 

43.1.4 DeepNet 的架构

 

作者提出了一种非常深的 Transformer 模型,设计的初衷就是 不要让起始阶段模型参数更新得太快 。

 

作者首先估计了 DeepNet 模型更新的预期幅度,然后作者提供理论分析来证明它的更新可以被一个常数所限制。

 

DeepNet 的架构基于 Transformer,把 Post-LN 替换成 DeepNorm。DeepNorm 的表达式可以写成:

 

式中,是常数, 是第 层 Transformer 的输出 (Attention 层或者 FFN)。DeepNet 还将残差内部的权重 扩展了 倍。

 

43.1.5 DeepNet 模型更新预期大小的估计

 

在下面这一小节里面,作者提供了对 DeepNet 模型更新预期大小 (Expected Magnitude) 的估计。不是一般性,研究一个 head 为1的 Attention 层。

 

Attention 模块的表达式可以写成:

 

式中,依次代表了 query,key 和 value。 代表相应的映射矩阵。

 

接下来作者通过3个定理展示了 Attention 模块的一些性质:

 

[引理 4.1] 给定 ,其中 。对所有的 都有 ,则满足下式:

 

其中,代表左右两侧的幅值是一样的 (Equal bound of magnitude)。

 

证明:把 softmax 展开:

 

这样,上式就变成了:

 

因为,所以有 。所以上式就变成了:

 

就等价于5式。 得证。

 

引理 4.1说明,attention 层的输出幅度其实只取决于 Value 和输出的映射矩阵,即:

 

作者这里只考虑 hidden dimension的情况,所以此时 就变成了标量 ,则意味着 。

 

同样的道理有:,这里 分别代表 FFN 网络的参数。

 

定义模型 某一层的输出的变化量 为 ,基于以上分析,对于一个 层的 DeepNet 而言,对于 范围的估计有如下定理:

 

[定理 4.2] 给定一个 层的 DeepNet: ,其中 和 分别代表第 层的 Self-attention 和 FFN 的参数,且每层都使用 DeepNorm 去替换 LN,即: ,则 满足:

 

**证明:**作者的目的是研究模型更新的幅度,作出以下假设以简化推导:

 

 

Hidden dimension 。

 

相关权重 都是小于1的正数, 都是大于1的常数。

 

 

考虑到引理4.1,如果是参数为 的 FFN 网络,则 。如果 是参数为 的 Self-attention 网络,则 。如果对各个投影矩阵使用 Xavier 初始化,则输出可以保持输入方差,这相当于 。根据假设2,有:

 

根据上式,和 的范围界限是:

 

最后一层的输出的变化量满足:

 

对上式使用泰勒展开,有:

 

上式表明了与 的关系,因此有:

 

对于原始的 Post-LN,有 ,所以有 。

 

得证。

 

定理4.2说明,Post-LN 其实相当于 DeepNorm 的一种 的特殊情况, 说明 Post-LN 会使得每个层的参数更新累加起来,导致模型输出更新非常大。

 

Warm-ups 通过减小 来减小 的幅度,而 Xavier 初始化通过减小 来达到同样的效果。

 

同样的方法,作者也研究了 Encoder-Decoder 架构,有 层的 Encoder 和 层的 Decoder,模型为 。定义模型 某一层的输出的变化量 为 ,基于以上分析,对于一个有 层的 Encoder 和 层的 Decoder 的 DeepNet 而言,对于 范围的估计有如下定理:

 

[定理 4.3] 给定一个 层 Encoder 和   层 Decoder 的 DeepNet:   ,

 

其中 Encoder 的归一化方式是 ,Decoder 的归一化方式是 , 则 满足:

 

原始的 encoder-decoder 模型满足 ,所以有: 。说明 Post-LN 会使得每个层的参数更新累加起来,导致模型输出更新非常大。此外,cross-attention 将幅值从编码器传播到解码器,这解释了为什幺解码器比编码器更不稳定。

 

43.1.6 DeepNet 的更新受到了 DeepNorm 的常数限制

 

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

 

如图3所示是 IWSLT-14 De-En 翻译数据集上,Post-LN 和 DeepNet 在早期训练阶段的模型更新情况,相比于Post-LN,DeepNet 的模型更新几乎保持恒定。

 

 

图3:IWSLT-14 De-En 翻译数据集上,Post-LN 和 DeepNet 在早期训练阶段的模型更新情况

 

如下图4和图5分别是 Encoder-Decoder 架构和只有 Encoder 的架构的 DeepNorm 流程,与上文的伪代码一致。

 

 

图4:Encoder-Decoder 架构的 DeepNorm 流程

 

 

图5:只有 Encoder 的架构的 DeepNorm 流程

 

43.1.7 实验结果

 

作者首先在神经机器翻译任务上验证了 DeepNet 的有效性,包括 IWSLT-14 德语-英语 (De-En) 数据集和 WMT-17 英语-德语 (En-De) 数据集。该研究将 DeepNet 与多个 SOTA 的深度 Transformer 模型进行比较,包括 DLCL 、NormFormer 、ReZero 、R- Fixup 、T-Fixup 、DS-init 和 Admin。

 

如下图6所示为 WMT-17 En-De 翻译数据集上的基线和 DeepNet 的结果,以 BLEU 为评价指标。

 

 

图6:WMT-17 En-De 翻译数据集上的基线和 DeepNet 的结果

 

DeepNet 相比 Post-LN 的训练更加稳定,可以成功扩展到 100L-100L,在测试集上达到 28.9 BLEU。相比之下,使用 Post-LN 的基线模型当深度达到 50L-50L 时就出现了优化不稳定的问题。此外,当模型较浅时,DeepNet 可以获得与这些基线相当的性能。

 

与 Post-LN 相比,使用 Pre-LN 训练的模型更稳定。原始的 Pre-LN 和 DLCL 都可以缩放到 100L-100L,50L-50L NormFormer 也可以训练成功。但是与 Post-LN 相比,Pre-LN 会导致 0.5-1.0 的 BLEU 下降。作者推测这应该是由于 Pre-LN 的浅层梯度往往大于深层的梯度。相比之下,DeepNet 通过使用 Post-LN 缓解了这个问题,并且优于所有 Pre-LN 的基线。

 

如下图7所示是不同深度的模型 IWSLT-14 数据集的结果,作者改变模型的深度,从 10L-10L 到 100L-100L。除了ReZero3之外,所有实验都采用混合精度训练。作者对模型进行8000个 steps 的训练,大多数分歧发生在优化的开始。总体来说,DeepNet 从浅到深都是稳定的。它收敛很快,仅用8000个 steps 就能达到30个BLEU以上,而大多数基线却不能。而且,随着模型的深入,性能不断提高。图8是 WMT-17 验证集的损失曲线,DeepNet 可以毫无障碍进行训练。

 

 

图7:IWSLT-14 数据集的结果

 

 

图8:WMT-17 验证集的损失曲线

 

大规模多语言神经机器翻译实验结果

 

作者首先使用 OPUS-100 语料库来评估模型。OPUS100 是一个以英语为中心的多语言语料库,涵盖 100 种语言,是从 OPUS 集合中随机抽取的。该研究将 DeepNet 扩展到 1,000 层,该模型有一个 500 层的编码器、 500 层的解码器、512 的 hidden dimension、8 个注意力头和 2,048 维度的前馈层。下图9所示是 DeepNet 和基线的结果。研究表明,增加深度可以显着提高 NMT 的翻译质量:48层的基线比12层模型平均提高3.2分。DeepNet 可以成功地将深度扩展到1000层,比基线提高了4.4 的 BLEU。

 

 

图9:DeepNet 和基线在大规模多语言神经机器翻译的实验结果

 

作者在 OPUS100 数据集上训练具有 {12,20,100,200,1000} 层的 DeepNet,图7展示了 BLEU 值与深度的变化规律。与双语 NMT 相比,多语 NMT 从扩展模型深度受益更多。可以观察到多语 NMT 的 BLEU 值呈对数增长,规律可以写成 ,其中 是深度, 是关于其他超参数的常数。

 

 

图10:BLEU 值与深度的变化规律

 

总结

 

本文作者提高了深度 Transformer 的训练稳定性,并成功将其扩展到1000层。这是通过 DeepNorm 这个新颖的归一化层实现的。它从理论上保证了把模型更新过程限制为常数。这一方法简单但高效,只需要改变几行代码即可。最终,该方法提升了 Transformer 模型的稳定性,实验结果验证了本文的方法在各种基准测试中的有效性。

 

参考:

 

https://zhuanlan.zhihu.com/p/474844668

 

https://zhuanlan.zhihu.com/p/436381997

 

https://www.zhihu.com/question/500004483/answer/2236028403

 

Be First to Comment

发表回复

您的电子邮箱地址不会被公开。