深度学习基础架构革新？通过梯度近似寻找Normalization的替代品-CSDN博客

本文链接：https://blog.csdn.net/c9Yv2cf9I06K2A9E/article/details/147809926

©PaperWeekly 原创 · 作者 | 苏剑林

单位 | 科学空间

研究方向 | NLP、神经网络

不知道大家有没有留意到前段时间的《Transformers without Normalization》[1]？这篇论文试图将 Transformer 模型中的 Normalization 层用一个 Element-wise 的运算 DyT 替代，以期能提高速度并保持效果。

这种基础架构的主题本身自带一点吸引力，加之 Kaiming He 和 Yann LeCun 两位大佬挂名，所以这篇论文发布之时就引起了不少围观，评价也是有褒有贬。

无独有偶，近期的一篇新论文《The Mathematical Relationship Between Layer Normalization and Dynamic Activation Functions》[2] 从梯度分析和微分方程的视角解读了 DyT，并提出了新的替代品。个人感觉这个理解角度非常本质，遂学习和分享一波。

写在前面

DyT 全称是 Dynamic Tanh，它通过如下运算来替代 Normalization 层：

其中都是可学参数，是 Normalization 层本来就有的，所以这里的关键是用替代了 Normalize 运算。是逐元素的运算，免除了均值、方差这两个统计量的计算。

关于 DyT，笔者曾在知乎《如何评价 Meta 新论文 Transformers without Normalization？》[3] 发表过一些看法，简单来说就是不大看好。

理由是 Normalization 无脑地稳定了模型的前向传播，那么就留了更多的自由度和可能性给模型的其他方面（比如效果），所以笔者不认为比有 Normalization 更简化的通用操作能实现更好的效果（No Free Lunch）。

事实上早在 2021 年的《浅谈 Transformer 的初始化、参数化与标准化》[4] 我们就讨论过去掉 Normalization 这个话题，相关工作有 SkipInit [5]、ReZero [6]、Fixup [7] 等。

当时笔者试了一些方案，发现它们即便在某些方面能够追平 Normalization，但仍会存在另一些方面的不足，比如预训练效果尚可，但微调效果较差等，所以就没再深究下去了。

因此，笔者现在对类似工作都只视为简化维度上的极限探索来欣赏，正如《nGPT: Normalized Transformer with Representation Learning on the Hypersphere》[8] 几乎将每一处能 Normalize 的地方都加上 Normalize 一样，都属于某个方向的极限探索。

梯度计算

当然，不看好归不看好，不妨碍我们的学习和分析。要想寻找 Normalization 的替代或者说近似，最直接的思路就是从梯度入手，因为深度学习说到底也就是前向传播和反向传播那点事，反向传播也就是求梯度，往往扮演着比较本质的角色。

接下来我们只考虑 RMS Norm，它的关键运算是：

其中，以及：

所以要求的梯度，等价于求的梯度，我们可以通过如下方式计算：

比较复杂的地方是展开：

代入式（4）得：

最后代回式（2）得：

DyT现！

注意都是一个向量，所以是一个矩阵（雅可比矩阵）。现在我们考虑给 RMS Norm 找一个 Element-wise 近似，即每个分量是独立运算的：

这个独立性意味着它的雅可比矩阵一定是对角阵！我们希望这个近似能尽可能保留 RMS Norm 的梯度，所以我们考虑保留式（7）的对角线部分：

如果我们进一步假设是常数，那么可以直接求解上述微分方程得到：

这样我们就得到了 DyT 的 T（），其中求解过程选择的初值条件为。

DyT 相当于将前面的吸收到参数里，然后将括号内的视为训练参数，缓解“ 是常数”这一假设带来的限制。不过在笔者看来，显式保留可能会更有价值，只要将部分视为可训练参数就好。

DyISRU

不知道大家有没有留意到，对于 RMS Norm 我们恒有，所以方程（9）的我们可以换成，从而得到：

这是一个只有的方程，免除了对的近似处理。求解该方程得：

其中C是任意常数。这种形式有个名字叫做 ISRU（Inverse Square Root Unit，我们之前也叫过 SoftSign），出自论文《Improving Deep Learning by Inverse Square Root Linear Units (ISRLUs)》 [9]。如果将 C 视为可训练参数，那么就可以类比 DyT 称为 DyISRU（Dynamic ISRU）。

从梯度（7）到方程（9）再到（11）来看，DyISRU 是我们用 Element-wise 函数能做到的最好结果，因为除对角线假设外没有再加额外近似了。从形式上看，DyISRU 其实也比 DyT 更直观，因为即，既然要寻求 Element-wise 的运算，只好将换成了，最后加 C 乘算是平滑操作：

相关工作

和 ISRU 都可以视为符号函数的光滑近似，而基于它们，我们可以构建运算的光滑近似，例如：

由此，我们也可以将 DyT 理解为引入（光滑的）操作来防止前向传播的爆炸，从而稳定模型。

提出自 Google 的 Gemma2 [10]，当时的用途是加在 Softmax 前的 Attention Logits 矩阵上，防止出现过大的 Logits 值。然而，我们实测中发现，尽管之后的 Logits 不会爆炸，但之前的 Logits 仍有爆炸风险，所以用防止 Logits 爆炸纯粹是将问题换了个出处，治标不治本。

不知道是否 Google 后来也意识到了这个问题，他们在最新的 Gemma3 [11] 中，选择去掉而改用 QK-norm。我们自己的实验也显示，QK-norm 可以更好地抑制 Attention Logits 的增长。这个改动和结论实际上再次间接传递了一个悲观信号：DyT 等类操作在实践中很难完全取代 Normalization。