Abstract
Loss spikes 经常出现在大型语言模型的预训练过程中。The spikes 会降低大型语言模型的性能,有时还会破坏预训练。由于预训练需要大量的计算预算,我们应该避免这样的 spikes。为了研究 loss spikes 的原因,我们重点研究了internal layers 的梯度 。通过理论分析,揭示了产生爆炸梯度 的两种原因,并提出了防止爆炸的要求。此外,我们提出了一种将初始化方法 与对 embeddigs 进行简单修改 相结合的方法来满足要求。我们进行各种实验,从经验上验证我们的理论分析。实验结果表明,该组合可以有效地防止预训练过程中的 spikes。
Introduction
大型语言模型(llm)已经成为各种应用的基础资产。增加(神经)语言模型中的参数数量和训练数据的数量通常会导致更好的llm 。因此,预训练需要大量的预算,因此,最小化预训练失败的风险是最重要的问题。
尽管它们被广泛用作LLM的基础架构,但尚未实现对trasnformer的全面理论理解。一个关键的未解决的问题是,在基于transformer的llm中,由于损失值的spikes(loss spike)可能导致灾难性的分歧,导致频繁发生预训练失败的原因 ,如图1中的Vanilla所示。虽然已经提出了一些经验策略来缓解这一问题,这些方法缺乏理论依据,使得它们在其他情况下的可泛化性不明确,例如模型参数的大小变化。
Figure 1:Training loss values of Transformers, whose dimensions and the number of layers are the same as the 1.7 billion parameters configuration in Narayanan et al. (2021). In Vanilla, some spikes occur at the beginning of the training, and its loss value exploded at about 13000 steps.
在本研究中,我们针对LLM预训练过程中的 loss spike 问题 进行了理论分析。我们确定了导致该问题的两个主要因素:(1) 前向传播过程中 residual connection 中 shortcut 部分 norm 的快速放大 ;(2 ) 每层的层归一化(LNs)前后梯度的增强 。我们的研究还表明,每个因素的影响与所选择的 参数初始化技术密切相关 。我们提供的证据表明,使用 typical 方法初始化参数会导致归因于 shortcut 的梯度爆炸 。相反,LLM预训练中流行的初始化方法将某些参数减小到更小的值,虽然减轻了 shortcut 引起的梯度爆炸,但却促进了 LN 引起的梯度爆炸 。为了防止两个梯度爆炸,我们描述了一个需求,并引入了一个简单的修改来满足需求。
为了证实我们的理论分析,我们进行了一系列的实证实验。结果表明,满足要求的方法抑制了loss spikes 和 explosions 的发生 。这与不能满足要求的方法形成对比,即使以前主张作为经验解决方案,这些方法仍然容易受到爆炸梯度的影响。此外,我们证明了满足要求的方法使llm能够以相对较大的学习率进行预训练,最终获得优异的性能结果 。
Preliminary
Pre-LN Transformer
本文主要关注GPT系列中使用的神经架构。他们使用 Pre-LN Transformer ,这是最近Transformer实现中的事实上的标准架构,因为当我们堆叠多层时,该架构比原始Transformer架构更稳定 。设
是Transformer某一层的输入,其中
表示该层的维数。该层输出
的方程如下:其中LN是层归一化函数。我们称式(1)和(2)中的第一项,即
和
为 shotcut 。此外,FFN和Attention的定义如下:
在这里插入图片描述
我们讨论了与原始Transformer的区别,即附录D中的 Post-LN Transformer,并在第5.3节中使用 Root Mean Square layer normalization(RMSNorm)代替 LN 。
Initialization
在这里插入图片描述
Analyses on Exploding Gradients
本节讨论爆炸梯度的两个原因以及防止它们的方法。我们指出,当我们使用2.2节中描述的使用plain初始化时,由于 shortcut 而发生的爆炸 。此外,我们指出 scaled 初始化减轻由shortcut 而引起的爆炸,但引起了LNs的爆炸 。为了防止这两种爆炸,我们描述了一个需求,并在 scaled 初始化中引入了一个修改。
