经典题解析:RMSNorm 比 LayerNorm 少做了什么

图中比较的是对同一个隐藏向量的处理路径,归一化范围仍需由模型定义。

一句“RMSNorm 是简化版 LayerNorm”,容易让人以为两者只是少写一行代码。真正要问的是:它省掉了什么统计量,因此失去了什么性质,又保留了什么性质?把这三件事说清,比背一个速度结论更有用。

先把归一化的范围圈出来

假设一个 token 的隐藏向量只有两个数:[3,4]。以下都沿这个向量的特征维计算,不跨不同样本凑均值。实际模型可能对最后若干维做归一化,必须看配置,不能把所有 normalization 都理解成对整个批次求统计量。

LayerNorm 先求均值,再求各分量与均值之差的平方均值。这个例子均值为 3.5,方差为 0.25。暂时忽略稳定项和可学习参数,两个分量变为 [-1,1]。它既调整中心,也调整尺度。

RMSNorm 不先减去均值,而是求各分量平方的平均值,再开平方。这个例子的均方根是 sqrt(12.5),约为 3.5355。因此输出约为 [0.8485,1.1314]。输出依旧都是正数,均值也没有变成零。

一个不减均值的公式

常见写法是:y_i = gamma_i × x_i / sqrt(mean(x_j²) + epsilon)。gamma 是可学习的逐特征缩放参数;这里采用常见的无偏置形式。LayerNorm 则在分子里使用 x_i 减均值,分母使用方差加稳定项,并常配合缩放与偏置。

不要把“没有居中”翻译成“没有归一化”。RMSNorm 仍根据当前输入的整体幅度缩放各分量,只是不主动抹去共同偏移。也不要看到 gamma 就以为最终输出的均方根一定严格为一:逐维缩放会改变结果,epsilon 也会影响这个等式。

不减均值,保留幅度归一化

缩放和偏移,是两种不同的检验

先忽略 epsilon,并令 gamma 为一。如果把输入乘一个正数 c,均方根也乘 c,两者相除后结果不变。这是它保留的正向缩放不变性。若 c 为负数,结果会翻转符号,不能不加条件地说任意缩放都完全不变。

再把输入所有分量同时加一个数。LayerNorm 的减均值步骤会消去这次共同平移,而 RMSNorm 通常不会。例如 [3,4] 变成 [13,14],两个方向的比例已经改变。RMSNorm 不是先算均值再偷偷放回去,它从设计上就没有执行这一操作。

实际存在非零 epsilon 时,正向缩放不变性也只能在稳定项相对较小的范围近似成立。极小输入会让稳定项影响变大;全零输入则要依赖有效的稳定处理,避免分母为零。数学性质要和实现条件放在一起回答。

少做一些运算,不等于固定快多少

省掉均值与居中操作,会改变归一化计算量和内存访问。但在真实设备上,算子融合、数据类型、内存带宽、张量形状,以及整个网络中该算子的占比,都会影响最终收益。不能拿一篇论文的特定任务实验,换算成所有模型端到端的固定加速比例。

同样,减少一个统计步骤不代表在每个任务上效果都更好。模型通常连同归一化结构一起训练;把训练好的 LayerNorm 直接替换成 RMSNorm,不会自动保留原模型输出。即使张量形状一致,表示的中心与幅度也可能发生变化,需要重新验证。

还可以问一个反例:如果每个分量都等于三,LayerNorm 的居中结果是全零,而没有居中的 RMSNorm 在忽略稳定项时会得到全一。两个算子对共同分量的处理从根本上不同,不能仅凭参数数量相近就认定可以无损替换。

面试时用一组数字收尾

可以按“范围、公式、性质、边界”四步回答:沿特征维计算均方根,不减均值;保留对正向幅度变化的适应,失去一般的平移不变性;稳定项、可学习缩放与数值精度会改变理想公式的表现;速度和质量依赖具体实现与训练。

检查理解还有个简单办法:对 [3,4]、[6,8] 和 [13,14] 分别计算。前两组在忽略稳定项时对应相同归一化结果,第三组一般不同。这是性质验证,不是训练效果评测。

如果有人说“RMSNorm 的输出一定均值为零、方差为一”,你会用哪一步计算指出问题?