Layer Normalization/zh
| Article | |
|---|---|
| Topic area | Deep Learning |
| Prerequisites | Backpropagation, Batch Normalization, Stochastic Gradient Descent |
概述
层归一化是一种通过在单个训练样本的特征维度上(而非跨样本批次)计算均值和方差来标准化神经网络层激活值的技术。该方法由 Jimmy Lei Ba、Jamie Ryan Kiros 和 Geoffrey Hinton 于 2016 年提出,旨在解决 Batch Normalization 在批次统计量不可靠或不可用场景下的局限性,例如循环网络与在线学习。[1] 如今它已成为几乎所有现代 Transformer 架构的标准组件,支撑着大型语言模型、视觉transformer以及许多序列模型。
其动机是通过降低每一层输出对输入幅度变化的敏感度来稳定训练并加速收敛。通过将激活值重新归一化为零均值和单位方差,层归一化将信号在深层网络中传播的尺度保持在受控范围内,进而在 Backpropagation 过程中维持梯度的良好行为。
直觉
深层网络难以训练,部分原因在于随着前面各层权重的更新,每一层激活值的分布会发生偏移。每一层必须不断适应新的输入统计量,这会减慢学习速度,并可能以病态的方式放大或衰减梯度。归一化技术通过显式固定激活值的某些统计量来应对这一问题。
Batch Normalization 在小批次中的样本之间计算统计量,而层归一化则在单个样本的特征之间计算统计量。这一差异虽属概念性,却影响深远:统计量不再依赖于批次中的其他样本,因此训练和推理时操作完全相同,行为与批次大小无关,甚至可以在每次只处理一个样本时应用。这种与批次组成无关的特性使其非常适合循环计算和自回归生成,因为这些场景中每个 token 都是按序处理的。
一种有用的可视化方式是:想象网络中某一位置上的激活值向量,对该向量进行重新中心化和重新缩放,使其各分量具有零均值和单位方差。然后,一个可学习的增益和偏置恢复了模型表示任意所需尺度和偏移的自由度。
公式化
设 $ x \in \mathbb{R}^d $ 是某一特定层在单个样本上的激活值向量。层归一化首先在 $ d $ 个分量上计算均值和方差:
$ {\displaystyle \mu = \frac{1}{d} \sum_{i=1}^{d} x_i, \qquad \sigma^2 = \frac{1}{d} \sum_{i=1}^{d} (x_i - \mu)^2.} $
接着对每个分量进行标准化,并应用一个由可学习向量 $ \gamma, \beta \in \mathbb{R}^d $ 参数化的逐元素仿射变换:
$ {\displaystyle \hat{x}_i = \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}}, \qquad y_i = \gamma_i \hat{x}_i + \beta_i.} $
常数 $ \epsilon $(通常为 $ 10^{-5} $ 或 $ 10^{-6} $)可在激活值方差极小时防止除零。可学习的 $ \gamma $ 和 $ \beta $ 使网络能够恢复标准化后激活值的任意仿射变换(包括恒等变换),因此该操作不会限制该层的表示容量。
关键的是,相同的过程在训练和推理时都运行。无需跟踪任何滑动平均,批次大小为一时也不会出现行为变化。
在现代架构中的位置
在 Transformer 块中,层归一化应用于每个子层(多头注意力和位置式前馈网络)的两侧。常见两种约定。
原始的表述方式,常称为 Post-LN,将归一化应用于残差相加之后:$ \mathrm{LN}(x + \mathrm{Sublayer}(x)) $。这正是原始Transformer论文中所采用的方式。[2] 与之相对,Pre-LN 则在残差相加之前对每个子层的输入应用归一化:$ x + \mathrm{Sublayer}(\mathrm{LN}(x)) $。研究表明,Pre-LN 在初始化时可产生条件更好的梯度,并且无需在深层模型中使用学习率预热,因此它已成为当今大多数大型语言模型训练时的默认选择。[3]
除 Transformers 之外,层归一化也是循环单元内部(在每个时间步稳定循环过程)、图神经网络内部(由于图大小不一,批次统计量可能产生误导)以及强化学习代理(数据以非独立同分布流的形式到来)的标准选择。
训练与推理
由于归一化统计量仅依赖于当前样本的激活值,层归一化不会引入训练与测试之间的差异。在典型实现中,前向传播会计算 $ \mu $ 和 $ \sigma^2 $,对 $ x $ 进行标准化,应用仿射变换,并存储中间张量以用于反向传播。反向传播通过一个闭式表达式将梯度传过归一化,该表达式考虑了每个输出分量通过共享的均值和方差对每个输入分量的依赖关系。
开销适中。对于具有 $ d $ 个特征的激活张量,前向和反向均需 $ O(d) $ 的工作量以及一次跨特征维度的同步。在现代加速器上,这通常受限于内存带宽而非计算能力,因此将标准化与仿射变换融合为单个内核的实现被广泛使用。
初始化约定很简单。增益 $ \gamma $ 初始化为一,偏置 $ \beta $ 初始化为零,因此该层初始时(在标准化之后)为恒等映射。此后,参数会在训练过程中漂移到模型认为有用的尺度和偏移上。
变体
若干改进对原始表述的某些部分进行了删减或重构。
RMSNorm(均方根层归一化)放弃了均值中心化步骤和偏置 $ \beta $,仅按激活值的均方根进行归一化:$ \hat{x}_i = x_i / \sqrt{\tfrac{1}{d}\sum_j x_j^2 + \epsilon} $。[4] 它的计算开销更小,在许多语言建模任务上实证表现达到或超过层归一化;LLaMA、PaLM 以及若干其他大型模型都采用了它。
ScaleNorm 用单个标量替代了逐元素增益,将每个向量归一化到一个可学习的范数上。FixNorm 及类似变体进一步约束归一化,使其位于超球面上。DeepNorm 重新缩放 Post-LN Transformer中的残差分支,从而支持训练数千层的网络。[5]
Group Normalization 介于批归一化和层归一化之间:它将特征维度划分为若干组,并在每个组内进行归一化。当只有一个组时,它退化为层归一化;当组数等于通道数时,它退化为实例归一化。当小批次大小使有效批归一化变得不可行时,它在计算机视觉中被广泛使用。
与其他归一化方法的比较
这四种经典归一化方案的差异仅在于计算均值和方差所沿的轴。
Batch Normalization 在每个特征通道上沿批次维度计算统计量。它在批次较大的卷积视觉模型中表现极佳,但当批次较小、训练与推理之间统计量发生漂移或序列长度可变时,性能会下降。
层归一化对每个样本沿特征维度计算统计量。它不受批次大小影响,并且在训练和推理时行为完全相同,这使其成为序列模型和Transformer的主导选择。
实例归一化按样本和通道分别计算统计量,仅在空间维度上进行归一化。它常见于风格迁移和生成式图像模型。前面提到的组归一化通过参数化归一化前通道的分组方式,将这一族进行了推广。
一条实用的经验法则:对于具有大批次和稳定输入分布的视觉模型,优先使用Batch Normalization;对于序列模型、Transformer,以及任何批次组成不可靠的场景,优先使用层归一化(或 RMSNorm)。
为什么有效
原始论文将层归一化的有效性归因于与Batch Normalization相同的假设,即减少了内部协变量偏移。后续研究对该解释提出了质疑。一项有影响力的分析表明,批归一化通过改善损失及其梯度的 Lipschitz 性质来平滑优化景观,而该效应并不需要协变量偏移的减少。[6] 类似的论证也被用于层归一化,并附有额外证据表明该操作充当了梯度范数的隐式正则化器,并将激活值的幅度和方向解耦。
无论具体机制为何,实证效应都很稳健:层归一化使训练更加稳定,对学习率选择不那么敏感,对深层架构也更具容忍度。
局限性
层归一化并非没有代价。沿特征维度的归约引入了一种比纯逐点操作更难并行化的串行依赖,在带宽受限的硬件上,它可能在 Transformer 块的开销中占据相当可观的一部分。RMSNorm 及类似的简化方案部分地缓解了这一问题。
该技术还假设被归一化的特征共享一种应被标准化的共同尺度。在某些架构中,某些特征携带着归一化会抹去的有意义的幅度信息;可学习的 $ \gamma $ 和 $ \beta $ 原则上可以恢复它,但优化器可能难以恢复那些从未以标准化形式出现过的尺度。
最后,对方差极低的 token 进行层归一化即使加入了 $ \epsilon $ 项也可能产生数值不稳定的输出,尤其是在低精度训练中。混合精度实现通常将归一化本身保留在更高精度下以避免这种失败模式。
另请参阅
- Batch Normalization
- Group Normalization
- RMSNorm
- Transformer
- Backpropagation
- Stochastic Gradient Descent
参考文献
- ↑ Template:Cite arxiv
- ↑ Template:Cite arxiv
- ↑ Xiong, Yang, et al. On Layer Normalization in the Transformer Architecture, 2020.
- ↑ Template:Cite arxiv
- ↑ Wang et al., DeepNet: Scaling Transformers to 1,000 Layers, 2022.
- ↑ Santurkar et al., How Does Batch Normalization Help Optimization?, NeurIPS 2018.