0. TL; DR

本文提出了大语言模型（LLMs）中的“深度诅咒”（Curse of Depth）现象，即现代LLMs中近一半的深度层（Transformer块）比预期的贡献小。研究通过理论和实证分析，发现这一现象的根本原因是广泛使用的Pre-Layer Normalization（Pre-LN）。Pre-LN虽然稳定了Transformer LLMs的训练，但其输出方差随模型深度指数增长，导致深度Transformer块的导数接近单位矩阵，几乎不贡献于训练。

为解决这一训练困境，文章提出了LayerNorm Scaling技术，通过按深度的平方根对Layer Normalization的输出进行缩放，有效抑制了深度层输出方差的爆炸性增长，提高了深度层的贡献。实验结果表明，LayerNorm Scaling显著提升了LLMs的预训练性能，并且这种提升无缝地延续到监督微调阶段。

1. 背景介绍

近年来，大语言模型（LLMs）在自然语言处理领域取得了巨大成功，但研究发现这些模型的深度层往往不如预期有效。这一现象不仅为模型压缩提供了机会，更揭示了训练过程中的不足。现代LLMs训练资源密集，通常需要数千个GPU训练数月，而大量层未能有效训练意味着资源的浪费。因此，解决这一问题对于提高LLMs的效率和性能至关重要。

2. LayerNorm Scaling 技术

深度诅咒指的是LLMs中深度层在学习和表示上的贡献显著低于早期层的现象。这些深度层对剪枝和扰动表现出显著的鲁棒性，表明它们未能执行有意义的转换，导致资源效率低下。

Pre-LN在训练Transformer LLMs时虽然稳定，但其输出方差随模型深度指数增长。这导致深度Pre-LN层的导数接近单位矩阵，使得这些层在训练中几乎不贡献于表示学习。具体来说，Pre-LN的输出方差在深度层中显著累积，导致深度层的梯度几乎为恒等映射，限制了模型的表达能力和学习能力。

Pre-LN的形式如下：

\[x_{t+1} = x_t + F_t(\text{LayerNorm}(x_t))\]

Pre-LN可以展开为：

\[\begin{aligned} x_{t+1} &= x_t + F_t(\text{LayerNorm}(x_t)) \\ &= x_{t-1} + F_{t-1}(\text{LayerNorm}(x_{t-1})) + F_t(\text{LayerNorm}(x_t))\\ &= \cdots \\ &= x_0 + \sum_{i=0}^t F_i(\text{LayerNorm}(x_i)) \\ \end{aligned}\]

假设$x_0,…,x_i$的方差都为$1$，且相互独立；此时$x_l$的方差为$l$。

LayerNorm Scaling通过按深度的平方根对Layer Normalization的输出进行缩放，有效控制了深度层输出方差的增长。具体来说，对于一个具有$L$层的Transformer模型，LayerNorm Scaling将每一层$l$的输出进行缩放:

\[x_l = \text{LayerNorm}(x_l) \cdot \frac{1}{\sqrt{l}}\]

这种缩放方式不仅抑制了输出方差的爆炸性增长，还提高了深度层在训练中的贡献，确保了所有层都能有效地参与学习。

3. 实验分析

实验使用了基于LLaMA架构的不同大小模型（从130M到1B参数），比较了Post-LN、DeepNorm、Pre-LN和LayerNorm Scaling的性能。结果表明，LayerNorm Scaling在所有模型大小上均优于其他归一化技术，显著降低了困惑度，表明其在预训练阶段就展现出了更好的性能。

在监督微调阶段，LayerNorm Scaling预训练的模型在多个下游任务上的表现优于使用Pre-LN和Post-LN的模型。这表明LayerNorm Scaling训练的深度层能够捕捉到更丰富和多样化的特征，从而在复杂任务中实现更好的泛化能力。

通过层剪枝实验，LayerNorm Scaling显著提高了深度层的有效性。实验表明，移除深度层会导致更大的性能下降，表明这些层在LayerNorm Scaling下对模型的贡献更大。