深度學習基礎架構革新？透過梯度近似尋找Normalization的替代品

©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。