RMSNorm与SwiGLU：改变Transformer的两个小改动

你好，我是 Shrijith Venkatramana。我正在构建 git-lrc，一款在每次提交时运行的 AI 代码审查工具。给我们点个 Star 来帮助开发者发现这个项目。一定要试试，并分享你的反馈来改进这个产品。

如果你打开一个现代 LLM 的实现，有趣的部分往往隐藏在这几行代码中。

有获得最多关注的注意力机制。然后还有两个看起来简单到有些"令人尴尬"的组件：

x = x * rsqrt(mean(x^2) + eps)

y = W2 * (Swish(W1 * x) * W3 * x)

第一个是 RMSNorm。

第二个是 SwiGLU。

它们共同构成了 LLaMA 及其衍生模型的标志性架构。Meta 的 Llama 2 技术报告描述了一个使用 RMSNorm 进行预归一化并采用 SwiGLU 前馈层的纯解码器 Transformer。

这两个想法都不复杂。

有趣的工程故事是：为什么这些简单操作如此有效？为什么它们取代了 LayerNorm + ReLU/GELU 这些旧选择？当把它们扩展到数十亿参数和数万亿 token 时，它们的算术表现又是怎样的？

本文先建立直观理解，再深入方程、参数经济性和实现细节。

1. RMSNorm 和 SwiGLU 实际上在哪里？

一个 Transformer block 大致长这样：

             ┌──────────────────────┐
             │      Attention       │
             └──────────────────────┘
                       │
                       v
                  RMSNorm
                       │
                       v
             ┌──────────────────────┐
             │       SwiGLU         │
             │        FFN           │
             └──────────────────────┘
                       │
                       v
                    output

一个现代 pre-norm Transformer 更精确地说是这样的：

h1 = x + Attention(RMSNorm(x))

h2 = h1 + SwiGLU(RMSNorm(h1))

这两个组件解决的是完全不同的问题。

RMSNorm 控制表示的尺度。

当信息通过数十甚至数百个 Transformer block 传递时，激活值的量级可能会变得差异巨大。RMSNorm 将它们的整体尺度控制在合理范围内。

SwiGLU 提供非线性计算。

注意力机制让 token 之间相互通信。FFN 独立处理每个 token 的表示，并为网络提供大量计算能力来转换该表示。

这个区分很有用：

Attention  -> "我应该检索哪些信息？"
SwiGLU     -> "我应该如何转换这些信息？"
RMSNorm    -> "保持表示的数值稳定性。"

因此，现代 LLM 技术栈不仅仅是"注意力"。

模型中相当大一部分计算发生在 FFN 中。

这使得 FFN 的设计变得极其重要。

让我们从更早、更常见的想法开始。

假设一个隐藏向量是：

x = [2, 4, 6, 8]

LayerNorm 计算均值：

mean(x) = 5

variance(x)
    = ((2-5)^2 + (4-5)^2 + (6-5)^2 + (8-5)^2) / 4
    = 5

然后它大致按以下方式转换每个值：

x_normalized = (x - mean) / sqrt(variance + eps)

所以 LayerNorm 做两件事：

1. 减去均值
2. 除以标准差

RMSNorm 提出了一个出人意料的挑衅性问题：

我们真的需要第一个操作吗？

Biao Zhang 和 Rico Sennrich 在他们 2019 年的论文《Root Mean Square Layer Normalization》中基本上提出了这个问题。他们的论点是：将激活值重新居中到零可能是不必要的。

与其计算相对于均值的方差，RMSNorm 计算的是均方根：

RMS(x) = sqrt( (x1^2 + x2^2 + ... + xd^2) / d )

RMSNorm(x) = gamma * x / sqrt(mean(x^2) + eps)

其中 gamma 是一个学习到的缩放向量。

mean = sum(x) / d
x - mean

这在大规模时很重要，因为归一化是持续执行的。

原始 RMSNorm 论文报告了与 LayerNorm 相当的性能，同时在他们测试的模型上减少了 7% 到 64% 的运行时间。

确切的加速很大程度上取决于周围架构和实现，但基本的工程吸引力是显而易见的：

LayerNorm:

均值
  ↓
减去均值
  ↓
平方
  ↓
方差
  ↓
开方
  ↓
缩放

RMSNorm:

平方
  ↓
均值
  ↓
开方
  ↓
缩放

一个统计操作消失了。

考虑两个向量：

A = [1, 2, 3, 4]

B = [101, 102, 103, 104]

它们形状相同但偏移差异很大。

LayerNorm 将偏移视为需要移除的东西。

这个向量整体有多大？

mean(square(A))
    = (1 + 4 + 9 + 16) / 4
    = 7.5

RMS(A) = sqrt(7.5)
       ~= 2.74

A / RMS(A)
 ~= [0.36, 0.73, 1.09, 1.46]

该操作在控制幅度的同时保留了向量的方向。

这提供了一个有用的几何解释。

把一个高维向量想象成一支箭。

RMSNorm 大致在说：

          x
          |
          |       x
          |     /
          |   /
          | /
----------+----------------

归一化长度。

然后学习到的 gamma 允许网络决定每个特征最终应该接收多少尺度。

这在深度残差网络中特别有吸引力。

Transformer 重复执行如下操作：

x -> x + f(x)

如果 f(x) 的幅度相对于 x 不断变化，优化可能变得困难。

RMSNorm 提供了一个持续的尺度控制机制。

原始论文中还有一个有趣的理论解释：RMSNorm 具有一种隐式的自适应学习率效果，因为它的归一化依赖于当前激活值的幅度。

原始 RMSNorm 论文早于当前的 LLM 爆发。

这个想法提出于 2019 年，当时大型 Transformer 语言模型已经出现，但现代 LLM 生态还需要数年才能到来。

然后架构开始收敛。

Transformer 于 2017 年出现。

T5 探索了大规模 Transformer 语言建模，并深刻影响了后来的架构。

Noam Shazeer 探索了门控 FFN。

Zhang 和 Sennrich 探索了 RMSNorm。

Meta 的 LLaMA 架构随后将几个这样的想法整合到一个紧凑的纯解码器设计中。

例如，Llama 2 在其 Transformer block 中使用了 RMSNorm 和 SwiGLU。

这是深度学习架构中一个有趣的模式：

研究论文
      ↓
小型架构修改
      ↓
实证验证
      ↓
被主流模型采用
      ↓
成为"显而易见"的架构
      ↓
最终成为样板代码

当你遇到以下代码时：

self.norm = RMSNorm(dim)

在一个 LLaMA 实现中，你看到的是数年实验的终点。

现在来看另一半。

原始 Transformer FFN 概念上很简单：

FFN(x) = W2 * ReLU(W1 * x)

有两个线性投影。

第一个扩展表示。

激活函数引入非线性。

第二个投影回来。

hidden size = 4096
FFN size    = 16384

表示被扩展了：

4096 -> 16384 -> 4096

因为中间层提供了一个大的计算工作空间。

一个有用的类比是临时草稿纸。

模型取一个 4096 维的表示，扩展到一个更大的空间，在那里执行非线性计算，然后压缩回来。

但传统的激活函数（如 ReLU）有一个问题。

ReLU(x) = max(0, x)

它做了一个硬决策：

负数 -> 0
正数 -> x

Shazeer 2020 年的论文《GLU Variants Improve Transformer》探索了一个不同的想法：门控。

不是简单地对一个投影应用激活函数，而是构建两个投影：

a = W1 * x

b = W3 * x

对一个应用非线性函数：

g = activation(a)

h = g * b

第二个投影就像一个学习到的门。

这创造了一个强大的机制：

输入
  |
  +----> 投影 A ---> 激活函数 ---+
  |                                       |
  +----> 投影 B --------------------*----> 输出投影

每个维度可以有效地调制另一个维度。

这就是 GLU 的核心思想。

Shazeer 测试了几种 GLU 变体，包括：

GLU
ReGLU
GEGLU
SwiGLU

SwiGLU 中的激活函数是 SiLU/Swish 函数：

SiLU(x) = x * sigmoid(x)

sigmoid(x) = 1 / (1 + exp(-x))

SwiGLU(x)
    = SiLU(W1 * x) * (W3 * x)

= [W1*x * sigmoid(W1*x)] * [W3*x]

然后是输出投影：

FFN(x) = W2 * (SiLU(W1*x) * (W3*x))

注意一个重要的事情。

一个传统的 FFN 有两个矩阵：

W1
W2

W1
W2
W3

这听起来像是参数增加了 50%。

如果保持完全相同的中间维度，确实如此。

Shazeer 的实验转而缩减了中间维度，使总参数计数大致保持不变。

这是最重要的实现细节之一。

假设一个常规 FFN 有：

d_model = 4096
d_ff    = 16384

忽略偏置，其参数计数约为：

W1: 4096 * 16384
W2: 16384 * 4096

total = 2 * 4096 * 16384
      ~= 1.342 亿参数

现在假设我们天真地引入相同 d_ff 的 SwiGLU。

我们有三个矩阵：

W1: 4096 * 16384
W2: 16384 * 4096
W3: 4096 * 16384

3 * 4096 * 16384
~= 2.013 亿参数

这多了大约 50%。

d_ff_swiglu ~= (2/3) * d_ff

d_ff_swiglu ~= 10923

3 * 4096 * 10923
~= 1.342 亿

参数预算大致相同。

这就是为什么你会经常看到现代架构有着看似奇怪的 FFN 维度。

例如，不要这样想：

4096 -> 16384 -> 4096

你可能会遇到更接近：

4096 -> 11008 -> 4096

这个数字不是任意的。

它还经常被四舍五入为硬件友好的值，如 256 的倍数。

这是机器学习架构与系统工程交汇的地方。

一个理论上完美的维度，如果产生丑陋的 GPU kernel，就毫无意义。

参数数量
FLOPs
内存带宽
tensor-core 利用率
kernel 形状
对齐
batch size

一个 10,922 宽的矩阵与 11,008 宽的矩阵可能有几乎相同的数学代价，而后者可能更方便地映射到硬件上。

让我们把各个部分拼在一起。

一个简化的现代 Transformer block 可以这样：

def block(x):
    h = x + attention(rmsnorm(x))
    h = h + swiglu(rmsnorm(h))
    return h

RMSNorm 大致是：

def rmsnorm(x, weight, eps=1e-6):
    rms = sqrt(mean(x * x) + eps)
    return weight * x / rms

def swiglu(x):
    a = W1 @ x
    b = W3 @ x
    return W2 @ (silu(a) * b)

有几个实现细节值得注意。

首先，规范的 LLaMA 风格 FFN 中没有偏置。

其次，向中间维度的两个投影是并行发生的。

第三，elementwise 乘法相对于矩阵乘法非常便宜。

d_model = 4096
d_ff    = 11008

一个 token 大约需要：

W1: 4096 * 11008
W3: 4096 * 11008
W2: 11008 * 4096

乘加运算。

3 * 45.1M
~= 1.35 亿 MACs

270M FLOPs

如果一个乘加被算作两个浮点运算的话。

而这只是每层 FFN 每个 token 的计算量。

对于一个 32 层的模型：

270M * 32
~= 86 亿 FLOPs/token

注意力机制占据了概念上的焦点，但 FFN 正在做大量的数值计算。

有两种方法可以很好地理解这两个操作。

RMSNorm 控制数值范围

没有归一化，连续的变换会改变表示的尺度。

RMSNorm 给网络一个稳定的参考：

"无论这个表示当前有多大，
在下次变换之前将其缩放到可预测的范围。"

SwiGLU 增加条件计算

一个标准的 FFN 本质上是说：

transform -> activate -> transform

transform A -> 非线性门
                   \
                    multiply -> transform
                   /
transform B -------

所以模型获得了一个依赖于输入的机制来控制流经 FFN 的信息流。

这个门不是一个像 MoE（Mixture-of-Experts）模型那样的独立高级路由系统。

它是每个 token FFN 内一个廉价的 elementwise 交互。

你可以把它想象成给每个中间特征一个学习到的、依赖于输入的音量旋钮。

LLM 架构最终受算术约束。

假设你训练一个模型：

1 万亿 tokens

而你的架构因为一个架构选择额外执行了：

1 亿 FLOPs/token

100M * 1T
= 1e20 FLOPs

一个表面上微小的 per-token 变化在训练规模上变得巨大。

推理也是一样。

如果一个模型每天生成：

10 亿 tokens

100M FLOPs/token

1e17 FLOPs/天

所以架构细节重要的原因是经济上的，不仅仅是数学上的。

一行归一化的改变会影响：

训练稳定性
kernel 复杂度
内存带宽
延迟
功耗
GPU 小时数

而对 FFN 的修改影响每个 Transformer 层最大的计算组件之一。

这就是为什么 Shazeer 的论文从工程角度来看很有趣。

他不仅仅是在发明另一种激活函数。他是在探索在相同的计算和参数预算下，不同的 FFN 结构是否能提高模型质量。

这种约束更接近于生产级 ML 的实际工作方式。

这里有一个有趣的历史反讽。

当你第一次学习 Transformer 时，架构看起来几乎是必然的：

Attention
+
Feed Forward
+
Normalization

现代 LLM 让这个历史过程变得可见。

架构不是一次性设计出来的。

研究人员不断提出小问题：

我们真的需要均值中心化吗？

FFN 能使用更好的非线性吗？

一个门能改善表示吗？

我们能保持相同的参数预算吗？

我们能让操作更便宜吗？

这个改变能在大规模实验中存活下来吗？

Zhang 和 Sennrich 在 2019 年提出了 RMSNorm。

Noam Shazeer 在 2020 年探索了 GLU 变体，发现像 SwiGLU 和 GEGLU 这样的变体在他的实验中可以相对于常规 ReLU/GELU FFN 改善 Transformer 结果。

几年后，这些想法一起出现在 LLaMA 风格的架构中。Meta 在 2023 年发布的 Llama 2（7B 到 70B 参数规模）明确记录了这个架构家族。

这是在阅读新架构论文时一个有用的模式。

重要的创新常常隐藏在一个看起来太小以至于不值得拥有自己章节的东西里面。

一个新的归一化规则。

一个不同的激活函数。

一个改变了的投影。

一个不同的 tensor 布局。

一种不同的花销相同 FLOPs 的方式。

在十亿参数规模上，这些细节会累积。

如果你剥离掉巨大的矩阵和数十亿参数，两个想法非常紧凑。

RMS(x) = sqrt(mean(x^2) + eps)

y = gamma * x / RMS(x)

它的作用是在避免 LayerNorm 显式均值中心化的同时控制表示尺度。

SwiGLU(x)
    = W2 * (SiLU(W1*x) * (W3*x))

它的作用是给 FFN 一个依赖于输入的 gating 机制，同时通过缩小中间维度保持大致相同的参数/FLOP 预算。

它们一起说明了一个关于 LLM 工程的重要现象：

小的数学改变
        +
大量层
        +
数万亿 tokens
        =
重大系统后果

也许最有趣的部分是 RMSNorm 和 SwiGLU 都不是概念上的异类。

困难的工程问题往往是：

一台 700 亿参数的机器应该执行哪种简单的操作数万亿次？

如果你今天从头设计一个 Transformer，你会先研究哪个"无聊"的组件：归一化、激活函数、注意力，还是 FFN 本身？

*AI 智能体写代码很快。它们也会静默删除逻辑、改变行为、引入 bug——而不告诉你。你经常在生产环境中才发现。

git-lrc 修复了这个问题。它 hook 到 git commit 并在每个 diff 落地之前审查它们。60 秒配置。完全免费。*

任何反馈或贡献者都欢迎！它已上线、开源、可供任何人使用。

免费、微型 AI 代码审查，在 Git Commit 时运行

| 🇩🇰 丹麦语 | 🇪🇸 西班牙语 | 🇮🇷 波斯语 | 🇫🇮 芬兰语 | 🇯🇵 日语 | 🇳🇴 挪威语 | 🇵🇹 葡萄牙语 | 🇷🇺 俄语 | 🇦🇱 阿尔巴尼亚语 | 🇨🇳 中文 | 🇮🇳 印地语 |

免费、在提交时运行的微型 AI 代码审查

今天的 GenAI 是一辆没有刹车的赛车。它加速很快——你描述一个东西，大块代码就瞬间出现了。但 AI 智能体静默地破坏东西：它们删除逻辑、松弛约束、引入昂贵的云调用、泄露凭证、改变行为——而不告诉你。你经常在生产环境中才发现。

git-lrc 就是你的刹车系统。它挂载到 git commit 钩子上，在每次 diff 合并之前运行 AI 审查。60 秒配置完成，完全免费。

简而言之，git-lrc 帮助在停机、安全漏洞和技术债务发生之前就加以阻止

一目了然：10 个风险类别 · 追踪 100+ 种故障模式 · 每次提交……

如需进一步行动，你可以考虑屏蔽此人或举报滥用行为