从 π₀.₅ 的 adaRMS 理解 Transformer 归一化

原始问题

π₀.₅ 的 adarms 参数控制是否启用 adaptive RMSNorm。启用后，归一化层接受条件向量（通常来自 timestep），用一个小 Dense 生成 scale、shift、gate 三组参数，逐层调制特征。

这个设计的核心差别在于：π₀ 只在输入端把 action 和 timestep 拼接一次，π₀.₅ 则在每一层 Transformer 都注入 timestep 信息。

为什么要逐层注入

Flow matching 的去噪是一个随时间变化的任务。t 接近 1 时网络需要恢复粗轮廓，t 接近 0 时需要修细节。只在输入注入一次，深层必须自己传递这个信息；逐层注入则让每一层直接读到当前处于哪个去噪阶段。

adaRMS 的三个参数各有职责：

scale(t)：控制归一化后各特征的幅值，网络可以根据 timestep 强调或压制不同维度
shift(t)：加法偏移，让网络在不同阶段调整特征的基准
gate(t)：控制这一层计算结果往 residual stream 写入多少

完整形式可以写成：

x_{l+1} = x_l + g_l(t) \odot F_l((1+s_l(t)) \odot \text{RMSNorm}(x_l) + b_l(t))

其中 $g_l(t)$ 是 gate， $s_l(t)$ 是 scale， $b_l(t)$ 是 shift， $F_l$ 是 Attention 或 FFN。

RMSNorm 做了什么

在理解 adaRMS 之前，需要先理解 RMSNorm 本身在干什么。

Transformer 层层累积 residual，hidden state 的 magnitude 容易漂移。RMSNorm 的作用是：把 hidden state 的整体尺度稳定下来，但不改变它的方向。

\text{RMS}(x) = \sqrt{\frac{1}{d}\sum_{i=1}^d x_i^2}

x_{\text{norm}} = \frac{x}{\text{RMS}(x)}

归一化后有一个重要性质： $\text{RMS}(x_{\text{norm}}) \approx 1$ 。所以 RMSNorm 实际在做：去掉整体能量的不可控漂移，保留特征的相对模式。

然后通常会有一个可学习参数 $\gamma$ ：

y = \gamma \odot \text{RMSNorm}(x)

网络可以学习不同维度的重要性。实现中常写成 (1 + scale)，这样初始化时 scale = 0 就是 identity，训练更稳定。

为什么不像 LayerNorm 那样减 mean

LayerNorm 认为 hidden vector 的两类整体变化都没意义：

整体平移（所有维度加常数）
整体缩放（所有维度乘常数）

所以它做 减 mean + 除标准差：

x_{\text{LN}} = \frac{x - \mu}{\sigma}

RMSNorm 做了更激进的简化：只控制整体尺度，不删除 mean 方向的信息。

从几何上看，RMSNorm 相当于保持向量方向、只调整向量长度。LayerNorm 则先把向量沿 all-ones 方向的分量删掉，再调整长度。

例如 $x = [10, 11, 12]$ 和 $x' = [101, 102, 103]$ ：

LayerNorm 后：两者完全相同（都变成同一个零均值、单位方差的形状）
RMSNorm 后：两者仍然不同

RMSNorm 的设计哲学是：**不要主动删除 hidden state 的一个方向，只把它的 magnitude 控制住。**而且工程上更简单：

\frac{x}{\sqrt{E[x^2] + \epsilon}}

比 LayerNorm 少一次 reduction。

Pre-Norm 为什么更稳定

现代 Transformer 普遍用 Pre-Norm：

x_{l+1} = x_l + F(\text{Norm}(x_l))

而不是 Post-Norm：

x_{l+1} = \text{Norm}(x_l + F(x_l))

核心原因是：Pre-Norm 给残差流保留了一条完全不受阻碍的信息和梯度通道。

Pre-Norm 对 $x_l$ 求导：

\frac{\partial x_{l+1}}{\partial x_l} = I + \frac{\partial F(\text{Norm}(x_l))}{\partial x_l}

注意这里有一个 identity $I$ 。即使 Attention/FFN 那部分梯度很小，我们依然有：

\frac{\partial x_{l+1}}{\partial x_l} \approx I

梯度可以直接往前传。

而 Post-Norm 的 Jacobian 是：

\frac{\partial x_{l+1}}{\partial x_l} = J_{\text{Norm}}(I + J_F)

$I$ 被包在 $J_{\text{Norm}}$ 里面，梯度必须经过 Norm。几十层叠起来更容易不稳定。

Pre-Norm 还有一个特性：假设某层 Attention 刚初始化时几乎什么都不做（ $F(x) \approx 0$ ），那么：

x_{l+1} \approx x_l

这一层天然接近 identity。100 层叠起来，一开始也可以 $x_{100} \approx x_0$ ，网络至少是"通的"。

而 Post-Norm 即使 $F(x) = 0$ ，每一层仍然会 Norm 一次，深网络并不天然接近 identity。

Pre-Norm 中 Norm 的位置

Norm 不在残差主干上，它只负责：某个模块准备读取 residual stream 时，先把输入尺度整理一下。

可以理解成：

Attention 本身非常怕输入尺度变化。如果 hidden state 整体放大 10 倍：

Q \rightarrow 10Q, \quad K \rightarrow 10K

那么：

QK^T \rightarrow 100 QK^T

Attention logits 暴涨 100 倍，softmax 可能从均匀分布变成极端的 one-hot。

所以在 Attention 之前做 RMSNorm 可以保证它读取的是尺度稳定的输入。

回到 π₀.₅ 的 adaRMS

现在可以完整理解 π₀.₅ 的设计了：

Residual stream 保持稳定的 identity highway： $x$
RMSNorm 保证 Attention/FFN 每次读取的是尺度稳定的特征： $\text{RMSNorm}(x)$
scale/shift 告诉这一层：当前 timestep 下，你应该如何理解这些 feature
gate 告诉这一层：当前 timestep 下，你算出来的信息应该往 residual stream 写多少

完整流程：

π₀.₅ 做的事情很漂亮：

residual stream 始终保持稳定的 identity path
RMSNorm 稳定子模块的输入尺度
adaRMS 让 timestep 控制"怎么看"和"写多少"
但不破坏 residual highway

这就是为什么 π₀.₅ 要把 timestep 注入每一层，而不是只在输入注入一次。