归一化与残差
RMSNorm 为何取代 LayerNorm、pre-norm 怎么解掉深层训练不稳,以及残差作为 attention 与 FFN 通信通道的几何意义
核心要点:
- LayerNorm 沿 hidden 维归一,$2h$ 参数 ($\gamma, \beta$)
- RMSNorm 去掉均值减法和 bias,参数减半,加速 7-64%
- re-centering 是多余的,re-scaling 才是稳定训练核心
- Post-norm (Vaswani) 深层梯度爆炸,必须 warmup; pre-norm 解掉
- Residual 不只是数值技巧,是 attention 与 FFN 通信的主干带 (Anthropic residual stream)
- DeepNorm 用 post-norm 加缩放训出 1000 层,极深方向唯一成熟方案
名词定义
本篇共享名词在 5.1 总览 已定义 (LayerNorm / RMSNorm / Pre-norm / Post-norm / Residual stream)。本篇新引入:
| 名词 | 定义 |
|---|---|
| Re-centering invariance | LayerNorm 减均值带来的"对输入加常数不影响输出"的不变性 |
| Re-scaling invariance | LayerNorm 除标准差带来的"对输入乘正常数不影响输出"的不变性;RMSNorm 仅保留这条 |
| Read-write framework | Anthropic Transformer Circuits 视角:attention 和 FFN 都是从 residual stream 读、计算、再写回的操作 |
| DeepNorm | Wang 2022 提出的 post-norm 变体,残差用 $\alpha > 1$ 缩放,权重用 $\beta < 1$ 初始化,让 1000 层 Transformer 稳定 |
@tbl-norm-glossary 本篇新引入名词
LayerNorm 与 RMSNorm:差一个减法的工业级提速
核心问题:LayerNorm 在 Transformer 里几乎随处都是。Llama / Qwen / DeepSeek 全改用 RMSNorm,差异只在去掉均值减法和 bias,凭什么这么小的改动值得改名换姓?
RMSNorm 用 7-64% 的实测加速换 0% 的性能损失,配合大模型训练的总开销,这是值得的工程权衡;同时 Zhang 2019 证明 re-centering 在 LayerNorm 里其实是多余的,数学上也站得住。
LayerNorm 数学
Ba, Kiros, Hinton 2016 提出 LayerNorm[1] 解决 BatchNorm 不适合可变序列长度 RNN 的问题:
$$\begin{equation} \mathrm{LN}(\mathbf{x}) = \gamma \odot \frac{\mathbf{x} - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta \label{eq:norm-ln} \end{equation}$$- $\mu = \frac{1}{h} \sum_{i=1}^h x_i$, $\sigma^2 = \frac{1}{h} \sum_{i=1}^h (x_i - \mu)^2$:沿 hidden 维度算
- $\gamma \in \mathbb{R}^h$:可学习 gain
- $\beta \in \mathbb{R}^h$:可学习 bias
- 每层 LayerNorm 总参数 $2h$
- $\epsilon$ 防 0 除,典型 $10^{-5}$
两个不变性:
- Re-centering invariance: $\mathrm{LN}(\mathbf{x} + c \cdot \mathbf{1}) = \mathrm{LN}(\mathbf{x})$ (输入加常数不变)
- Re-scaling invariance: $\mathrm{LN}(\alpha \mathbf{x}) = \mathrm{LN}(\mathbf{x})$ (输入乘正常数不变)
RMSNorm:去掉均值减法和 bias
Zhang, Sennrich NeurIPS 2019[2] 提出 RMSNorm:
$$\begin{equation} \mathrm{RMSNorm}(\mathbf{x}) = \gamma \odot \frac{\mathbf{x}}{\sqrt{\frac{1}{h} \sum_{i=1}^h x_i^2 + \epsilon}} \label{eq:norm-rms} \end{equation}$$- 去掉均值减法 $\mu$ → 失去 re-centering invariance
- 去掉 $\beta$ bias → 参数从 $2h$ 减到 $h$
- 保留 $\gamma$ gain 和 re-scaling invariance
论文核心论断:LayerNorm 的两个不变性中,re-centering 是多余的,re-scaling 才是稳定训练的真正核心。论文实验把 LayerNorm 改成"仅 re-centering"或"仅 re-scaling",仅 re-scaling 就够。
实测加速 7-64% 来自哪里
Zhang 2019 论文 abstract 原文实测 RMSNorm 在 RNN + Transformer 上 7%-64% 加速。现代 GPU 上 Transformer 实测约 15%。
加速的核心不是省了几个 FLOP,而是内存访问模式:
- LayerNorm 需要两遍扫描 hidden 维:第 1 遍算均值,第 2 遍算方差 (或用 Welford 单遍算法,但 GPU 上不友好),引入两次全局归约
- RMSNorm 单遍扫描:直接算 $\sum x_i^2$,一次归约搞定
GPU 上多一次同步屏障 (sync barrier) 就是几十微秒的代价,大模型里这种归一化层每个 block 出现两次 ($L$ 层就是 $2L$ 次),累积起来非常可观。
HuggingFace LlamaRMSNorm.forward 实现:
def forward(self, hidden_states):
variance = hidden_states.pow(2).mean(-1, keepdim=True)
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
return self.weight * hidden_states
仅一个 pow + mean + rsqrt,无均值减法。
Llama / Qwen / DeepSeek 全线 RMSNorm
| 模型 | 归一化选择 |
|---|---|
| GPT-2 / GPT-3 | LayerNorm |
| Llama 1 / 2 / 3 | RMSNorm |
| Qwen 2 / 2.5 / 3 | RMSNorm |
| DeepSeek-V3 | RMSNorm (含 MLA 内部压缩向量) |
| Gemma / PaLM | RMSNorm |
@tbl-norm-models 主流 LLM 归一化选择 (2023 年后基本全转 RMSNorm)
Llama 1 原话:"To improve the training stability, we normalize the input of each transformer sub-layer... We use the RMSNorm normalizing function, introduced by Zhang and Sennrich (2019)."[3]
唯一已知弱点:量化场景 (FP8 / INT8) 下非零均值激活会引入额外误差,DeepSeek-V3 对归一化操作保持 BF16/FP32 不下沉到低精度。
Pre-norm vs Post-norm:一行代码挪位置的革命
核心问题:Vaswani 2017 原版 Transformer 把 LayerNorm 放在残差之后 (post-norm),后人发现这样深层训练不稳,改成放残差之前 (pre-norm)。具体差在哪?
Post-norm 让深层模型的梯度量级随层数爆炸,训练时必须用精细 warmup 才能稳;pre-norm 让梯度量级在初始化时就有界,训练直接稳,这是现代 LLM 全线 pre-norm 的根本原因。
两种顺序的公式
设 $\mathrm{Sublayer}(\mathbf{x})$ 是 attention 或 FFN 子层。
Post-norm (Vaswani 2017 §3.1 原版):
$$\begin{equation} \mathbf{x}_{\text{out}} = \mathrm{LN}(\mathbf{x} + \mathrm{Sublayer}(\mathbf{x})) \label{eq:norm-post} \end{equation}$$残差加完再归一化。
Pre-norm (GPT-2 起的现代主流):
$$\begin{equation} \mathbf{x}_{\text{out}} = \mathbf{x} + \mathrm{Sublayer}(\mathrm{LN}(\mathbf{x})) \label{eq:norm-pre} \end{equation}$$归一化先做,再算子层,最后加残差。残差通道里的信号永远不被 LayerNorm 截断。
Post-norm 的训练痛点
Vaswani 2017 用 post-norm,训练时必须配精细的 warmup + cosine schedule,否则训练崩。当时这一约束没被深入分析,大家以为是 Transformer 这种架构的固有特性。
直到 Xiong et al. ICML 2020[4] 用 mean field theory 给出严格分析:
- Post-norm 初始化时,靠近输出层的参数期望梯度量级随层数 $L$ 增大而指数爆炸
- Pre-norm 初始化时,各层梯度量级有界,与 $L$ 无关
- 论文实证:pre-norm 可以不用 warmup 达到同等性能,训练时间显著缩短
这是工程上"挪一行 LayerNorm 代码"带来训练稳定性的根本理论解释:不是 Transformer 难训,是 post-norm 难训。
GPT-2 改 pre-norm:一段标志性原话
GPT-2 (Radford 2019) 报告里只用一句话说明了这次架构改动:
"Layer normalization was moved to the input of each sub-block, similar to a pre-activation residual network, and an additional layer normalization was added after the final self-attention block."
两个关键改动:
- 子层内:
LN(x + Sublayer(x))→x + Sublayer(LN(x)) - 顶层末尾:加一个最终 LayerNorm (
ln_f或类似命名),让 residual stream 输出在送进 LM head 前再归一化一次
这条改动后续被所有 GPT 系 / Llama / Qwen / Mistral / Gemma 沿用。
Llama 系列在 GPT-2 pre-norm 基础上再换 RMSNorm:现代 LLM block 内归一化已经收敛到 "pre-RMSNorm" 这一标准。
极深方向:DeepNorm 让 post-norm 训出 1000 层
Wang et al. 2022 DeepNet (DeepNorm)[5] 在 post-norm 上加缩放因子,解决极深 Transformer 训练:
$$\begin{equation} \mathrm{DeepNorm}: \quad \mathbf{x}_{\text{out}} = \mathrm{LN}(\alpha \cdot \mathbf{x} + G(\mathbf{x}, \theta)) \label{eq:norm-deepnet} \end{equation}$$- $\alpha > 1$:放大残差 → 让 Sublayer 输出成"小扰动",而非主导信号
- 权重 $\theta$ 用 $\beta < 1$ 初始化,$\beta \approx (8N)^{-1/4}$ ($N$ 是层数)
- $\alpha, \beta$ 的具体值由层数 $N$ 和 $M$ (encoder/decoder layer 数) 的幂次公式给出
实证:Microsoft 用此方案训出 1000 层 Transformer 稳定收敛,是目前把 post-norm 推到极深的唯一成熟方案。
但实际工业大模型 (Llama 405B 仅 126 层,DeepSeek-V3 仅 61 层) 远远不需要 1000 层,所以 DeepNorm 仍是学术上的好奇心,工业主流仍是 pre-norm。
残差:不只是数值技巧,是通信通道
核心问题:残差连接 $\mathbf{x} + \mathrm{Sublayer}(\mathbf{x})$ 看起来只是个"加上原值防梯度消失" 的技巧。这就是它的全部意义吗?
残差是 attention head 和 FFN 之间的通信通道,整个 Transformer 可以看成一个"读-计算-写" 的迭代过程,residual stream 是主干带。
ResNet 起源:解决深网络退化
残差连接由 He et al. 2015 ResNet 引入 CNN,经典动机:
- 数学上 $f(\mathbf{x}) + \mathbf{x}$ 让梯度通过恒等路径直接反传,缓解梯度消失
- 直觉上让网络学"残差 / 增量" $f(\mathbf{x})$ 而非完整映射
Transformer 直接继承这条设计,在 attention 和 FFN 各加一次残差。
Anthropic Residual Stream: read-write framework
更深的视角来自 Anthropic Transformer Circuits 工作。Elhage et al. 2021[6] 把每个 block 看成对residual stream 的读写操作:
"Both the attention and MLP layers each 'read' their input from the residual stream (by performing a linear projection), and then 'write' their result to the residual stream by adding a linear projection back in."
几个关键直觉:
- Residual stream 是"主干带":每个 block 不替换 stream,而是从 stream 读、做计算、把结果写回 stream
- Attention 和 FFN 各操作不同子空间:它们的 $W_{\text{out}}$ 投影到 residual stream 的不同方向,互不干扰
- 跨层信息流通:早期 layer 写入的信息可以被晚期 layer 读出,中间层若不动这些信息就保持不变 (恒等映射)
- mechanistic interpretability 的基础:induction head 这类跨层电路依赖 residual stream 作为通信通道
Pre-norm 的"几何收益":pre-norm 让残差通道 $\mathbf{x}$ 永远不被 LayerNorm 截断,residual stream 保持几何连续性;post-norm 每层都 LN 截断,residual stream 视角下信号被反复重整,不利于跨层信息流通。
Residual 的工程性质
| 性质 | 描述 |
|---|---|
| 加法不引入参数 | Residual 本身无可学习参数,只是 $+x$ |
| 梯度无衰减直通 | 反传时残差路径梯度 = 1,不经过任何变换 |
| 几何稳定 | $\| \mathbf{x} + \delta \mathbf{x} \| \approx \| \mathbf{x} \|$ 当 $\delta \mathbf{x}$ 量级远小于 $\mathbf{x}$,是 pre-norm 让模型稳定的关键 |
| 加深友好 | 加层数不破坏现有信息流通,新 layer 学增量即可 |
@tbl-norm-residual-properties 残差连接的四个工程性质
Takeaway
| 知识点 | 核心结论 |
|---|---|
| LayerNorm 公式 | $\gamma \cdot (\mathbf{x} - \mu) / \sqrt{\sigma^2 + \epsilon} + \beta$,沿 hidden 维,$2h$ 参数 |
| RMSNorm 公式 | $\gamma \cdot \mathbf{x} / \sqrt{\mathrm{mean}(x^2) + \epsilon}$,去均值减法和 bias, $h$ 参数 |
| RMSNorm 核心论断 | re-centering 多余,re-scaling 才是核心 |
| RMSNorm 加速 | 7-64% (论文实测),Transformer 上约 15%;主因省去一次 GPU 同步屏障 |
| 主流采用 | 2023 年后 Llama / Qwen / DeepSeek / Gemma / PaLM 全线 RMSNorm |
| Post-norm 公式 | $\mathrm{LN}(\mathbf{x} + \mathrm{Sublayer}(\mathbf{x}))$ |
| Pre-norm 公式 | $\mathbf{x} + \mathrm{Sublayer}(\mathrm{LN}(\mathbf{x}))$ |
| Xiong 2020 理论 | post-norm 梯度量级随层数指数爆炸,pre-norm 有界,是 pre-norm 训练稳的根本原因 |
| GPT-2 关键改动 | "Layer normalization was moved to the input of each sub-block" + 顶层末尾加 final LN |
| DeepNorm | post-norm 加缩放,训出 1000 层,工业不用 |
| Residual stream | Anthropic 视角:attention 和 FFN 是 read-write on residual stream,不是 stream 替换 |
| Pre-norm 几何收益 | residual stream 不被 LayerNorm 截断,几何连续,利于跨层信息流通 |
开放问题
- RMSNorm 是否还有改进空间:当前实现已经很简洁,是否还有更省的归一化形式?RMSNorm-free (无归一化) 方案有人测过但实证收益小
- Pre-norm 是否真的不可替代:DeepNorm 证明 post-norm 加策略也能稳,极深方向是否会有新设计影响主流?
- Residual stream 的容量极限:模型大了 stream 维度也大,但是否存在"信息饱和" 问题?Anthropic 工作显示 residual stream 是 mechanistic interp 的根基,但理论分析有限
- Parallel block (PaLM) 是否值得:共享一个 LN 同时算 attention + FFN 训练快 15%,但精度无显著优势,工业为何不广泛采用仍开放
- 量化下的 norm 行为:FP8/INT8 训练里归一化层的精度损失最严重,RMSNorm 是否需要专门变种
延伸阅读
- 上一步:激活与 FFN → 5.2 激活与 FFN
- 下一步:Block 与堆叠 → 5.4 Block 与堆叠
- Anthropic Transformer Circuits residual stream → https://transformer-circuits.pub/2021/framework/index.html
- DeepNorm 1000 层 Transformer → https://arxiv.org/abs/2203.00555
- 完整训练稳定性 (warmup / lr schedule) → 不在本知识库,见各模型 tech report
参考资料
- Ba, Kiros, Hinton. Layer Normalization. 2016. https://arxiv.org/abs/1607.06450
- Zhang, Sennrich. Root Mean Square Layer Normalization. NeurIPS 2019. https://arxiv.org/abs/1910.07467
- Touvron et al. LLaMA: Open and Efficient Foundation Language Models. 2023. https://arxiv.org/abs/2302.13971
- Xiong et al. On Layer Normalization in the Transformer Architecture. ICML 2020. https://arxiv.org/abs/2002.04745
- Wang et al. DeepNet: Scaling Transformers to 1000 Layers. 2022. https://arxiv.org/abs/2203.00555
- Elhage et al. A Mathematical Framework for Transformer Circuits. Anthropic, 2021. https://transformer-circuits.pub/2021/framework/index.html