层归一化 LayerNorm：Transformer 中的标准归一化

什么是层归一化

层归一化（Layer Normalization，LayerNorm）是一种用于稳定和加速深度神经网络训练的技术。与批归一化（Batch Normalization）在批次维度上进行标准化不同，层归一化在特征维度上进行标准化，使得每个样本的所有特征均值为 0、标准差为 1。这种特性使其非常适合处理序列数据，成为 Transformer 架构中的标准归一化方式。

为什么需要层归一化

深度网络在训练时，每一层的输入分布会随着前一层参数更新而改变，这种现象叫做 内部协变量偏移（Internal Covariate Shift）。层归一化通过固定每一层输入的均值和方差来缓解这一问题，从而带来以下好处：

• 加速收敛：使损失曲面更平滑，允许使用更大的学习率。
• 降低对初始化的敏感性：即使权重初始值不够理想，逐层标准化也能稳定训练。
• 缓解梯度消失/爆炸：将激活值控制在合理范围内。
• 适应小批次或动态输入：不像批归一化那样依赖批次统计量，在批次大小变化或仅为 1 时依然稳定。

层归一化的计算过程

对于一个输入向量 ( \mathbf{x} = (x_1, x_2, ..., x_H) ) ，层归一化的计算分为三步。

1. 计算统计量

在特征维度上计算均值和方差：

[ \mu = \frac{1}{H}\sum_{i=1}^{H} x_i ]

[ \sigma^2 = \frac{1}{H}\sum_{i=1}^{H} (x_i - \mu)^2 ]

这里 ( H ) 是隐藏层的神经元数量（即特征数）。

2. 标准化

用得到的均值和方差将输入标准化：

[ \hat{x}_i = \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}} ]

其中 ( \epsilon ) 是一个很小的常数（例如 ( 10^{-5} ) ），防止除零。

3. 缩放和平移

为了让网络能够学习到更有表达能力的分布，引入可学习的参数 ( \gamma ) （缩放因子）和 ( \beta ) （偏移因子）：

[ y_i = \gamma , \hat{x}_i + \beta ]

这两个参数与原始输入 ( \mathbf{x} ) 的维度相同，通过反向传播更新。

最终输出 ( \mathbf{y} = (y_1, y_2, ..., y_H) ) 即为层归一化的结果。

层归一化与批归一化的对比

层归一化在序列建模中优势明显，因为序列长度可能动态变化，批归一化很难跨时间步维护稳定统计量。

在 Transformer 中的应用

Transformer 模型大量使用层归一化，主要出现在两个位置：残差连接之后 和 每个子层之前（或之后，取决于具体变体）。

原始 Transformer 的 Post-LN 结构

在 “Attention Is All You Need” 论文中，每个子层（自注意力和前馈网络）的计算顺序为：

输出 = LayerNorm( x + Sublayer(x) )

这种在残差加法之后施加归一化的方式称为 Post-LN。

改进的 Pre-LN 结构

后来的研究发现，将层归一化放在子层之前（Pre-LN）有助于训练更深的 Transformer，使梯度流动更顺畅：

输出 = x + Sublayer( LayerNorm(x) )

Pre-LN 已经成为许多现代 Transformer（如 GPT 系列、ViT 等）的默认选择。

代码实现示例（PyTorch 风格）

import torch
import torch.nn as nn

class LayerNorm(nn.Module):
    def __init__(self, features, eps=1e-6):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(features))
        self.beta  = nn.Parameter(torch.zeros(features))
        self.eps = eps

    def forward(self, x):
        # x shape: (batch_size, ..., features)
        mean = x.mean(-1, keepdim=True)
        std  = x.std(-1, keepdim=True, unbiased=False)
        return self.gamma * (x - mean) / (std + self.eps) + self.beta

在实际框架中，直接调用 nn.LayerNorm 即可，内部实现会针对特定硬件优化。

层归一化的变体与扩展

• RMSNorm：只使用均方根（Root Mean Square）进行缩放，移除均值中心化，计算更高效，常用于 LLaMA 等大模型。
• DeepNorm：结合深度缩放的层归一化，专门为训练非常深的 Transformer 设计。
• PowerNorm：在标准化时引入幂项，以应对不同分布的数据。

所有这些变体都保留了层归一化“每个样本独立归一化”的核心思想，只是在统计量的选择和缩放策略上有所不同。

总结

层归一化是深度学习中不可或缺的基础组件。它通过单样本内跨特征的标准化，解决了小批次动态序列场景下的训练稳定性问题。尤其在 Transformer 架构中，LayerNorm 配合残差连接，支撑起了从机器翻译到大规模语言模型的各类成功应用。理解其原理和计算过程，能够帮助你在设计或调试网络时做出恰当选择。

熟练使用层归一化，是深入掌握现代深度学习模型的必修课。