🚑DeepNet: Scaling Transformers to 1,000 Layers
2024-7-18
| 2024-7-18
0  |  阅读时长 0 分钟
type
status
password
date
slug
summary
category
URL
tags
icon

文章简介

这篇文章针对了当前经典的transformer模型在深度很深的情况下训练往往不稳定,容易发散的现象进行了一定的研究,对这个现象的原因进行了比较深入的分析,并基于此提出了一种deepnorm的layernorm方法,该方法结合了Post - LN的良好性能和Pre - LN的稳定训练的优点,并在数学上可以确保训练的稳定性。基于此,文中直接把transformer的最大训练层数推到了1000层,并且在各项指标上吊打其他模型。
使用DeepNorm的千层transformer在3.2B参数下的性能会超过12B模型5个点。
使用DeepNorm的千层transformer在3.2B参数下的性能会超过12B模型5个点。

TL;DR for Practitioners(面向从业者的简短总结)

与 Post-LN 相比,DEEPNORM 在执行层归一化之前放大了残差连接。此外,在初始化期间缩小了参数。值得注意的是,只缩放了前馈网络的权重,以及注意力层的值投影和输出投影。而且,残差连接和初始化的缩放比例取决于架构(图 2)
(a) 部分是 DEEPNORM 的伪代码。以 Xavier 初始化(Glorot 和 Bengio,2010)为例,也可以用其他标准初始化来替代。注意是一个常数。(b) 部分展示了 DEEPNORM 在不同架构(N 层编码器、M 层解码器)中的参数。
(a) 部分是 DEEPNORM 的伪代码。以 Xavier 初始化(Glorot 和 Bengio,2010)为例,也可以用其他标准初始化来替代。注意是一个常数。(b) 部分展示了 DEEPNORM 在不同架构(N 层编码器、M 层解码器)中的参数。

Xavier Normal初始化

先说具体怎么操作,很简单。假设有 层全连接层,激活函数为 ,第 层输入为 维向量 ,状态值为 维向量 ,满足:
那么参数初始化时, 置为 每个元素 均采样自正态分布 ,方差 ,即输入维度的倒数。
💡
注:xavier normal初始化一般的标准做法是方差 ,为输入和输出维度均值的倒数。这是同时考虑了前向传播输入输出的方差和后向传播梯度的方差都要基本不变,做的折衷方案。对于这篇paper,主要考虑前向传播,所以我认为不采用这种反而更好。

为什么会出现训练不稳定

作者研究了深度 Transformer 不稳定的原因。发现,更好的初始化方法能稳定 Transformer 的训练,这也被之前的工作(Zhang 等人,2019a;Huang 等人,2020;Xu 等人,2021)所证实。因此,研究了有或没有适当初始化的 Post - LN 的训练过程。有更好的初始化时,在执行 Xavier 初始化后,将第 层的权重按 )进行缩小。例如,第 层 FFN 的输出投影 的初始化如下:
其中,是输入和输出维度的平均值。我们将这个模型命名为 Post-LN-init。需要注意的是,与之前的工作(Zhang 等人,2019a)不同,我们缩小的是较低层的参数,而非较高层。我们认为这有助于将梯度大小的影响与模型更新的影响区分开。此外,Post-LN-initPost-LN 具有相同的架构,这消除了架构带来的影响。
如图 3(c) 所示,Post-LN-init 收敛了,而 Post-LN 没有。如图 3(a) 所示Post-LN-init 在最后几层的gradient norm甚至更大,尽管其权重已经缩小。图 3(c) 可视化了从 6L - 6L 到 24L - 24L 不同模型深度的最后一个解码器层的gradient norm;发现,无论模型深度如何,Post-LN-init 在最后一层的gradient norm仍然比 Post-LN 大得多。由此得出结论,深层的梯度爆炸不应该是 Post-LN 不稳定的根本原因,而模型更新的规模更可能是其原因。
如图 3(c) 所示,Post-LN-init 收敛了,而 Post-LN 没有。如图 3(a) 所示Post-LN-init 在最后几层的gradient norm甚至更大,尽管其权重已经缩小。图 3(c) 可视化了从 6L - 6L 到 24L - 24L 不同模型深度的最后一个解码器层的gradient norm;发现,无论模型深度如何,Post-LN-init 在最后一层的gradient norm仍然比 Post-LN 大得多。由此得出结论,深层的梯度爆炸不应该是 Post-LN 不稳定的根本原因,而模型更新的规模更可能是其原因。
notion image
然后我们证明 Post - LN 的不稳定性源于一系列问题,包括梯度消失以及模型更新过大。如图 4(a)所示,我们首先在训练的早期阶段可视化模型更新的范数
notion image
其中, 分别表示输入和第 次更新后的模型参数。Post-LN 在训练刚开始时更新量爆炸,然后很快几乎不再更新。这表明模型陷入了一个虚假的局部最优。Post-LN+4k warmupPost-LN-init + 4k warmup都有助于缓解这个问题,使模型能够平稳更新。当更新量爆炸时,输入到 LN(可能是层归一化 Layer Normalization)的值会变大(见图 4(b)和图 4(c))。根据 Xiong 等人(2020)的理论分析,通过 LN 的梯度大小与其输入的大小成反比:
图 4(b) 和图 4(c) 表明,在没有warmup或适当初始化的情况下,显著大于 。这解释了在 Post-LN 的训练中出现的梯度消失问题(见图 4(d))。
💡
综上所述,不稳定性始于训练开始时的大规模模型更新。这使得模型陷入不良的局部最优值,进而增加了每个 LN 输入的量值。随着训练的继续,通过 LN 的梯度变得越来越小,从而导致严重的梯度消失。消失的梯度使得模型难以从局部最优值中逃脱,进一步破坏了优化的稳定性。相反,Post-LN-init 的更新相对较小,输入到 LN 的值是稳定的。这减轻了梯度消失的问题,使优化更稳定。
 
Nguyen和Salazar(2019)发现基于后归一化连接(Post - LN)的Transformer,其前归一化残差连接(Pre - LN)能提高稳定性。然而,Pre - LN在底层的梯度往往比顶层更大(Shleifer等人,2021),导致与Post - LN相比性能下降。

DeepNet原理

模型结构

DEEPNET 基于 Transformer 架构。与原始的 Transformer 相比,它在每个子层中使用我们新的 DEEPNORM,而不是 Post-LN。DEEPNORM 的公式可以写成
其中的 为常数, 是代表第 个Transformer的子层(如attention层或FFN层)的函数, 为子层的模型参数。公式中看不出来的是,attention层和FFN层中的参数初始化很有讲究,对于attention中的 、FFN中的 ,都用nn.init.xavier_normal_(w, gain=)初始化;而对于attention中的 ,用 nn.init.xavier_normal_(w, gain=1)初始化。

影响模型更新的因素

💡
Attention中 不会影响注意力输出的大小,只有会影响注意力输出的大小。FFN中 会影响注意力输出的大小
  • Attention
    • 不会影响注意力输出的大小,只有会影响注意力输出的大小
  • FFN
    • 会影响注意力输出的大小

推导DEEPNORM参数

接下来我们让这个 在SGD的每一步迭代中,和学习率 在同一个数量级,而与网络层数无关,即当 时,有 ,这样就可以使得训练初始阶段,输出 的变化幅度受控,跟网络层数无关,达到我们初衷。这里 表示lost function对模型参数 的导数。因此有
常见的任务都是输出 之后接一个softmax之类的,假设 有界是合理的,且一般与模型参数和层数都无关,所以有
按照paper里的说法,根据参考文献[2]的结论有

代码实战

参考

  • LLM
  • GLU 和 SwiGLUGLM论文解读
    Loading...
    目录