LayerNorm 的位置之争：Pre-LN vs Post-LN

在 Transformer 架构的演进过程中，Layer Normalization (LN) 的位置选择是一个看似微小但影响深远的设计细节。原始的 Transformer (Attention is All You Need) 采用了 Post-LN 结构，而随后的 GPT-2、GPT-3 以及当代的 LLaMA 等主流大模型几乎一致转向了 Pre-LN。

这一范式转移背后的动因是什么？Pre-LN 是否是完美的解决方案？DeepNorm 又是如何通过初始化策略调和两者矛盾的？本文将深入探讨这一架构设计中的权衡（Trade-off）。

架构定义

Transformer 的标准模块（Block）包含两个子层：Multi-Head Attention (MHA) 和 Feed-Forward Network (FFN)，并辅以残差连接（Residual Connection）。

Post-LN (Original Design)

在 Post-LN 结构中，Normalization 操作置于残差连接之后。这意味着残差流在进入下一层之前被强制进行了归一化。

$$\begin{aligned} \mathbf{x}' &= \text{LayerNorm}(\mathbf{x}_t + \text{MHA}(\mathbf{x}_t)) \\ \mathbf{x}_{t+1} &= \text{LayerNorm}(\mathbf{x}' + \text{FFN}(\mathbf{x}')) \end{aligned}$$

数据流向：Input -> Sublayer -> Add -> Norm -> Output

Pre-LN (Modern Standard)

在 Pre-LN 结构中，Normalization 操作置于子层输入之前，残差连接则直接相加。这意味着主干（Identity Path）始终保持“纯净”，未受归一化操作的直接干扰。

$$\begin{aligned} \mathbf{x}' &= \mathbf{x}_t + \text{MHA}(\text{LayerNorm}(\mathbf{x}_t)) \\ \mathbf{x}_{t+1} &= \mathbf{x}' + \text{FFN}(\text{LayerNorm}(\mathbf{x}')) \end{aligned}$$

数据流向：Input -> Norm -> Sublayer -> Add -> Output

Post-LN 的优化困境：梯度方差与 Warmup

早期的 BERT 和原始 Transformer 均采用 Post-LN。在训练实践中，它们面临的一个核心挑战是：对超参数高度敏感，且若不使用 Learning Rate Warmup，模型极难收敛，甚至出现梯度爆炸。

梯度范数的理论分析

Xiong 等人（ICML 2020）深入分析了 Post-LN 的梯度传播特性。
在 Post-LN 中，最后一层的梯度需要穿过层层 LayerNorm 才能反向传播到底层。由于 LayerNorm 具有缩放不变性，梯度的数值尺度在反向传播过程中会发生剧烈波动。

理论推导表明，对于 Post-LN，第 $l$ 层的梯度范数 $\|\nabla \mathbf{x}_l\|$ 的期望与 $\sqrt{\frac{L}{l}}$ 成正比（其中 $L$ 是总层数）。
这意味着：

• 输出层附近（$l \approx L$）：梯度较小。
• 输入层附近（$l \approx 1$）：梯度较大，甚至可能发散。

这种梯度尺度的不均匀性导致深层网络的优化极不稳定。底层的参数更新步长往往过大，容易导致模型在训练初期震荡发散。

Warmup 的本质作用

为了缓解这一问题，必须引入 Warmup 策略：在训练初期使用极小的学习率，随后线性增加。
Warmup 的本质是抑制训练初期的梯度方差。由于 Adam 等自适应优化器依赖于梯度的二阶矩估计来调整步长，在训练初期，样本的统计特性尚未稳定，梯度方差极大。Post-LN 结构进一步放大了这种不稳定性。通过 Warmup，我们强制模型在参数处于随机初始化阶段时进行“微小更新”，待参数分布逐渐稳定后，再提升学习率进行正常训练。

Pre-LN 的崛起：稳定压倒一切

Pre-LN 的最大优势在于其卓越的训练稳定性。

梯度高速公路（Gradient Highway）

在 Pre-LN 中，残差连接 $\mathbf{x}_{t+1} = \mathbf{x}_t + F(\text{LayerNorm}(\mathbf{x}_t))$ 形成了一条通畅的梯度通路。
由于 LayerNorm 被置于支路（Sublayer）中，主干通路（Identity Path）未被归一化阻断。反向传播时，梯度可以直接沿着 $\mathbf{x}_t$ 路径无损地传回底层：

$$\frac{\partial L}{\partial \mathbf{x}_t} = \frac{\partial L}{\partial \mathbf{x}_{t+1}} + \frac{\partial L}{\partial \mathbf{x}_{t+1}} \frac{\partial F}{\partial \mathbf{x}_t}$$

公式中的第一项保证了梯度的直接传递。数学推导证明，Pre-LN 使得各层的梯度范数在期望上基本保持一致（与层数 $L$ 无关）。

这种特性使得 Pre-LN 模型可以在没有 Warmup 的情况下，直接使用较大的学习率进行训练，且能够轻松扩展到超深层网络（如 1000 层以上）。因此，GPT-2、GPT-3、PaLM、LLaMA 等大模型纷纷采用了 Pre-LN 架构。

Pre-LN 的隐患：表示能力的瓶颈？

尽管 Pre-LN 在优化方面表现出色，但它并非没有代价。
一些研究指出，Pre-LN 虽然利于训练，但可能会限制模型的表示能力（Representation Capacity），导致在同等参数量下，Pre-LN 的最终性能往往略逊于精细调优后的 Post-LN。

贡献度衰减（Contribution Decay）

在 Pre-LN 中，输出 $\mathbf{y} = \mathbf{x} + F(\text{LayerNorm}(\mathbf{x}))$。
随着层数加深，主干 $\mathbf{x}$ 的数值范数会因不断的残差累加而线性增长。然而，支路 $F(\text{LayerNorm}(\mathbf{x}))$ 的输出经过 LayerNorm 后，其数值范数被限制在常数级别。

这导致了一个结构性问题：在深层网络中，残差分支的相对贡献会被主干淹没。

$$\frac{\| F(\text{LayerNorm}(\mathbf{x})) \|}{\| \mathbf{x} \|} \to 0 \quad \text{as } L \to \infty$$

这意味着深层的 Transformer Layer 实际上退化为了近似恒等映射（Identity Mapping），模型有效深度降低。相比之下，Post-LN 每一层都重置了分布，迫使每一层都必须积极参与特征变换。

破局者：DeepNorm

为了兼得 Post-LN 的高性能和 Pre-LN 的高稳定性，微软提出了 DeepNorm (Wang et al., 2022)。
DeepNorm 的核心思想是：回归 Post-LN 结构，但通过在初始化阶段对参数进行极其精细的缩放（Scaling），来解决梯度不稳的问题。

它引入了一个缩放因子 $\alpha$：

$$\mathbf{x}_{t+1} = \text{LayerNorm}(\alpha \mathbf{x}_t + \text{Sublayer}(\mathbf{x}_t))$$

理论推导与初始化

DeepNorm 证明了，只要将残差连接的权重 $\alpha$ 放大，或者将 Sublayer 的初始化权重缩小，就能限制更新过程中梯度的爆炸。
具体的初始化策略为：

• 对于 Encoder：Sublayer 的权重缩小为 $\frac{0.81}{(N^4 L)^{1/16}}$。
• 对于 Decoder：Sublayer 的权重缩小为 $\frac{1}{\sqrt{2L}}$。

通过这种数学上严谨的初始化控制，DeepNorm 成功训练了 1000 层的 Post-LN Transformer，且在 GLM-130B 等超大模型中证明了其有效性，取得了优于 Pre-LN 的效果。

总结

• Post-LN：
  • 优点：表示能力强，潜力高。
  • 缺点：梯度方差大，训练不稳定，必须依赖 Warmup，难以训练超深层。
• Pre-LN：
  • 优点：拥有梯度高速公路，训练极其稳定，适合超大规模模型（LLaMA 等）。
  • 缺点：深层贡献度衰减，可能有轻微的性能损失。
• DeepNorm：
  • 试图通过初始化技巧复兴 Post-LN，兼顾两者，是当前追求极致性能的一个重要流派。

在当前的开源大模型生态中，Pre-LN + RMSNorm 依然是绝对的王者组合，因为对于大多数从业者来说，“好训练”比“高出 1% 的性能”更重要。但了解这段演变史，有助于我们在面对特定优化需求时做出更理性的选择。

参考文献