深层网络一层接一层计算,激活数值可能越走越大,训练随之变得难控制,LayerNorm 会先减去一组激活的均值,再用标准差缩放。RMSNorm 留下缩放,去掉了减均值这一步,只用平方平均后的开方值控制整组激活的大小。

我对照了 RMSNorm 原始论文的公式、机器翻译实验和公开实现。作者的判断很明确,LayerNorm 提供的重新缩放能力承担了主要作用,重新居中未必不可缺;RMSNorm 因此保留一个可学习的缩放参数,也保留一个防止除零的小量,但无需计算均值和中心化后的方差。

少一步计算也能稳住训练

论文把这种性质称为缩放不变性。输入或权重整体放大以后,均方根也会一起放大,归一化后的结果仍处在相近尺度;参数越大,归一化带来的有效更新会相对变小,作者把它解释为一种隐式的学习率调节。

这套说法经过了一个很实在的对照。论文在 WMT14 英德翻译上训练基础规格 Transformer,模型维度为 512,前馈层为 2048,使用 8 个注意力头。完全去掉归一化的版本没有训成。LayerNorm 在两个测试集上的 BLEU 为 26.6 和 27.7,RMSNorm 得到 26.8 和 27.7,差距很小。

同一实验还记录了运行时间。Tesla V100 上每一千个训练步,LayerNorm 用时约 248 秒,RMSNorm 约 231 秒,快了 6.9%。只从前 6.25% 激活估计均方根的 pRMSNorm 约为 225 秒。作者在翻译、阅读理解和图像任务等不同架构上汇总的加速范围更宽,从 7% 到 64% 都有。

这些数字不能直接搬到今天的大模型训练账单里,论文也发现,pRMSNorm 理论计算更少,某些框架里却会被张量切片的低效实现拖慢。现代训练常把归一化、残差相加等操作融合进一个 GPU 内核,内存读写和算子启动次数会改变最终差距。

论文还测了循环网络、阅读理解和图像任务。RMSNorm 在不同任务里大致守住了 LayerNorm 的效果,速度差异却随架构变化很大。循环网络每一步都重复归一化,省掉均值计算更容易积出可见差距;Transformer 能把许多位置并行处理,论文里的相对收益便收窄到个位数。这也解释了 7% 到 64% 为何不能压成一个固定加速比例。

换归一化方式需要重新训练

RMSNorm 属于模型结构的一部分,一个已经用 LayerNorm 训练好的模型,直接删掉减均值步骤,权重并不会自动适应。新模型可以在同样数据和预算下做对照,现成模型则应遵守原架构。

评估时至少要同时记录训练损失、任务分数、每步时间和数值稳定性。只看公式少了一次均值计算,很容易高估收益;只看最终分数,又会漏掉长训练中反复发生的那点开销。RMSNorm 的意义正在这处取舍里,它用更简单的尺度控制换到接近的训练表现,省下多少时间仍由具体实现回答。

参考资料