深度学习进阶(二十六)现代 LLM 的核心架构设计其一:RMSNorm

上一篇我们介绍了旋转位置编码 RoPE ,它通过在 QK 点积中注入旋转矩阵,让注意力自然而然地感知到相对位置。如今 RoPE 已经成为几乎所有主流开源大模型的位置编码标准。

但位置编码只是现代大模型改造的一个环节。 如果把标准 Transformer 和今天的大模型放在一起对比,就会发现归一化方式、激活函数、注意力结构......几乎每个模块都被动过。

因此,近几篇会正式进入现代 LLM 模型架构,展开其相关技术。

本篇展开第一个改动:归一化

1. 归一化的发展历程

我们在很早之前就展开过归一化的内容:归一化它的基本逻辑是通过调整数据尺度来加速模型训练,帮助优化,从而让模型更快的收敛,这里再简单复述一下:

一个深层网络在训练时,每层的输入分布都会随着前一层参数的变化而不断变化。这意味着后一层需要持续适应前一层输出的分布漂移,训练也就变得更困难,甚至可能导致梯度现象。

当时的研究者将这一问题归因于内部协变量偏移(Internal Covariate Shift,ICS),并尝试通过归一化稳定中间层分布。

最初的方案是 15 年的 Batch Normalization:对每一层的激活值做归一化,减去 batch 均值,除以 batch 标准差。这个操作让网络可以使用更大的学习率、更少的 epoch 就收敛,也因此迅速成了 CNN 的标配。

但 Batch Normalization 的局限在于依赖 batch 维度的统计量。 在序列长度变化、 batch size 不一定大的 NLP 任务中,其效果会大打折扣。

而且,在 Transformer 这种结构中,BN 的 "竖着统计"和与自注意力的 "横着注意" 在同时应用时,也存在逻辑上的不协调。

于是 16 年的 Layer Normalization 换了一个思路:不沿 batch 维度归一化,而是沿特征维度归一化。

这种方式针对单个样本,计算它所有特征的均值和方差,然后做归一化。既不依赖 batch size,也不受序列长度影响,天然适合序列模型。

此后,LayerNorm 就成了 Transformer 的标准配置,一直到今天的少部分 大模型也仍在使用。

因此,在进入 RMSNorm 前,我们再回顾一下 LN 本身的逻辑。

2. LayerNorm 有什么问题?

我们在之前的 Transformer Block 里也介绍过 LN,先看一个标准的 LayerNorm 公式:

\\\text{LayerNorm}(\\mathbf{x}) = \\frac{\\mathbf{x} - \\mu}{\\sqrt{\\sigma\^2 + \\epsilon}} \\odot \\boldsymbol{\\gamma} + \\boldsymbol{\\beta} \\

其中:

\\\mu = \\frac{1}{d}\\sum_{i=1}\^d x_i \\

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

\(\boldsymbol{\gamma}\) 和 \(\boldsymbol{\beta}\) 是可学习的参数,维度与 \(\mathbf{x}\) 相同。\(\epsilon\) 是防止除零的极小常数。

总结来看,LN 做了三件事:

  1. 居中(centering) :减去均值 \(\mu\),让分布中心归零。
  2. 缩放(scaling) :除以标准差 \(\sigma\),把方差归一化到 1 。
  3. 恢复(affine transformation) :用 \(\boldsymbol{\gamma}\) 和 \(\boldsymbol{\beta}\) 恢复模型需要的尺度和偏移。

前两步是固定的归一化操作,第三步是把"自由度"还给模型。因为强制把所有激活值变成零均值单位方差不一定是最优的,模型可能希望某些层保留特定的分布特征。

看这个结构,感觉每一步都很合理。但科研往往打破常规,19 年的论文 Root Mean Square Layer Normalization 提出了问题:

这三步里,真的每一步都是必要的吗?特别是第一步:居中。

以此,我们来展开 RMSNorm 的思路:

3. RMSNorm 的内容

RMSNorm 提出了一个很直接的实验:把 LayerNorm 的居中部分去掉,只保留缩放,看看效果如何。去掉之后的公式长这样:

\\\text{RMSNorm}(\\mathbf{x}) = \\frac{\\mathbf{x}}{\\sqrt{\\frac{1}{d}\\sum_{i=1}\^d x_i\^2 + \\epsilon}} \\odot \\boldsymbol{\\gamma} \\

对比一下:

操作 LayerNorm RMSNorm
计算均值 \(\mu = \frac{1}{d}\sum x_i\)
减去均值 \(\mathbf{x} - \mu\)
除以标准差 \(\sqrt{\frac{1}{d}\sum (x_i-\mu)^2 + \epsilon}\) \(\sqrt{\frac{1}{d}\sum x_i^2 + \epsilon}\)
可学习偏置 \(\boldsymbol{\beta}\)
可学习增益 \(\boldsymbol{\gamma}\) \(\boldsymbol{\gamma}\)

很明显,RMSNorm 简化了两处内容:

  1. 取消中心化,直接使用原始输入的均方根来做归一化的分母,而不是用标准差。
  2. 去掉偏置参数 \(\boldsymbol{\beta}\) ,这是因为不做中心化后的分布中心不再是零,就不需要专门用一个可学习参数去专门调整了。

4. 为什么可以去掉中心化?

值得一提的是,这一问题的答案并不在 RMSNorm 论文本身,而是在后续大模型实践中逐渐清晰的,其关键在于我们之前提到的 Pre-Norm 结构

我们在 Transformer Block那里展开过 Post-Norm 和 Pre-Norm 两种结构:

  1. Post-Norm :归一化放在子层之后(\(\text{LayerNorm}(\mathbf{x} + \text{Sublayer}(\mathbf{x}))\))
  2. Pre-Norm :归一化放在子层之前(\(\mathbf{x} + \text{Sublayer}(\text{LayerNorm}(\mathbf{x}))\))

已知 Post-Norm 是原版 Transformer 的设计,但深层训练不稳定。

于是 Pre-Norm 把归一化放在残差分支的入口处,保证了主干的信号流通更加顺畅,训练更稳定,从而成为事实标准。

后续大量实践表明,Transformer 的训练稳定性主要依赖于尺度控制而非严格的零均值约束。

所以 Pre-Norm 结构改善了梯度传播,而 RMSNorm 保留了最关键的尺度归一化功能。两者结合后,即使不再执行中心化操作,模型依然能够稳定训练并保持性能。

但毕竟是黑箱,还有一种说法是 Pre-Norm 结构和残差的组合可以特征值保持较小的均值,让中心化不再那么必要。

总之,在主流认知中的共同点都有 Pre-Norm 的重要性,但验证其效果后,剩下的其实也就是怎么讲故事的问题了。

我们来重点看看其效果:

5. RMSNorm 的表现

首先,RMSNorm 论文在多个任务上做了对比。核心结论是:去掉均值中心化之后,效果与 LayerNorm 基本持平,甚至在某些设定下略有提升。

但更重要的是计算效率。RMSNorm 省去了两个操作:

  1. 前向 :不需要计算均值 \(\mu\),均方根的计算也比方差计算少一次减法。
  2. 反向 :不需要对 \(\boldsymbol{\beta}\) 求梯度。

原论文在多种网络结构上进行了实验:在保持性能接近的情况下,整体运行时间可减少约 7%~64%。对于 Transformer 模型,加速通常在 7%~15% 左右。

在单层上这点差异微乎其微,但在 70B、上百层的模型中,累积起来就是可观的训练加速。

但同样要说明的是:今天的大模型层数更深,使用的技术更广,Norm 层的占比变小,实际上的收益已经不会像上面一样那么可观了。

从 LLaMA 开始,RMSNorm 几乎成了所有现代大模型的默认选择:

模型系列 归一化方式
LLaMA 1/2/3 RMSNorm
Mistral / Mixtral RMSNorm
Qwen 系列 RMSNorm
Gemma RMSNorm
DeepSeek RMSNorm
GPT-NeoX / Pythia LayerNorm
原始 Transformer LayerNorm

2023 年之后发布的大模型,几乎全都换成了 RMSNorm。

这便是 RMSNorm 的全部逻辑,就是这样的一个个改进,支撑起了庞大的现代 LLM 架构。

相关推荐
小七-七牛开发者3 小时前
周一上线 | SpaceX 收购 Cursor、支付宝进入 AI 时代、DeepSeek 完成 500 亿元融资
ai·agent·token·glm·智谱·claudecode·ai coding·周一上线
doiito21 小时前
【Agent Harness】为什么我把 JSON‑LD “编译成 DAG” 后,整个 Agent 平台立刻聪明了
ai·rust·架构设计·系统设计·ai agent
xiezhr1 天前
折腾半小时,终于让AI 能直接帮我写飞书文档了
ai·飞书·ai agent·飞书cli·飞书文档
岳小哥AI1 天前
Claude Fable和Claude Mythos 5同时发布:注意力机制下愈加强大的AI大模型
ai·ai基础
Artech1 天前
[MAF预定义的AIContextProvider-04]Mem0Provider——长期记忆基于的云端解决方案
ai·agent·maf·aicontextprovider·chathistorymemoryprovider·mem0provider
哥不是小萝莉2 天前
一文读懂 OpenAI Codex 源码的原理、架构与未来
ai
AlfredZhao2 天前
AI 编程工作总结:从体验问题到模块能力建设
ai·codex
cup113 天前
[技术复盘] Windows Python 打包实战:Nuitka 环境踩坑总结与 CI 自动化构建全指南
python·ai·环境变量·ci·nuitka·skill
IT王师傅3 天前
从 豆包 到 Codex CLI:一名普通开发者的 AI 工具进化路线
ai·codex cli·openclaw