解決訓練難題，1000層的Transformer來了，訓練代碼很快公開

來源：機器之心

機器之心編輯部

1000 層的 Transformer，深得嚇人。

昨日出爐的論文《DeepNet: Scaling Transformers to 1,000 Layers》在研究社區(qū)引起了熱議，作者來自微軟亞洲研究院。

該研究直接把 Transformer 深度提升到 1000 層！

下面讓我們看下這篇研究說了什么。

近年來，大規(guī)模 Transformer 模型出現(xiàn)了這樣一種趨勢：隨著模型參數(shù)從數(shù)百萬增加至數(shù)十億甚至數(shù)萬億，性能相應地實現(xiàn)了顯著提升。大規(guī)模模型在一系列任務上都取得了 SOTA 性能，并在小樣本和零樣本學習設置下展現(xiàn)出了令人矚目的能力。如下圖 1 所示，盡管參數(shù)量已經(jīng)很大了，但 Transformer 模型的深度（depth）卻受到了訓練不穩(wěn)定的限制。

Nguyen 和 Salazar (2019) 發(fā)現(xiàn)，基于 post-norm 連接（Post-LN），pre-norm 殘差連接（Pre-LN）能夠提升 Transformer 的穩(wěn)定性。但是，Pre-LN 在底層的梯度往往大于頂層，因而導致與 Post-LN 相比性能下降。為了緩解這一問題，研究人員一直努力通過更好的初始化或更好的架構來改進深度 Transformer 的優(yōu)化。這些方法可以使多達數(shù)百層的 Transformer 模型實現(xiàn)穩(wěn)定化，然而以往的方法沒有能夠成功地擴展至 1000 層。

微軟亞研在一篇新論文《DeepNet: Scaling Transformers to 1,000 Layers》中終于將 Transformer 的深度擴展到了 1000 層。

論文地址：https://arxiv.org/pdf/2203.00555.pdf

研究者的目標是提升 Transformer 模型的訓練穩(wěn)定性，并將模型深度進行數(shù)量級的擴展。為此，他們研究了不穩(wěn)定優(yōu)化的原因，并且發(fā)現(xiàn)爆炸式模型更新是造成不穩(wěn)定的罪魁禍首?；谶@些觀察，研究者在殘差連接處引入了一個新的歸一化函數(shù) —— DEEPNORM，它在將模型更新限制為常數(shù)時具有理論上的合理性。

這一方法簡單但高效，只需要改變幾行代碼即可。最終，該方法提升了 Transformer 模型的穩(wěn)定性，并實現(xiàn)了將模型深度擴展到了 1000 多層。

此外，實驗結果表明，DEEPNORM 能夠將 Post-LN 的良好性能和 Pre-LN 的穩(wěn)定訓練高效結合起來。研究者提出的方法可以成為 Transformers 的首選替代方案，不僅適用于極其深（多于 1000 層）的模型，也適用于現(xiàn)有大規(guī)模模型。

值得指出的是，在大規(guī)模多語言機器翻譯基準上，文中 32 億參數(shù)量的 200 層模型（DeepNet）比 120 億參數(shù)量的 48 層 SOTA 模型（即 Facebook AI 的 M2M 模型）實現(xiàn)了 5 BLEU 值提升。

有知乎網(wǎng)友疑問：就實現(xiàn)效果來說，1000 層是否有必要？論文作者之一董力（Li Dong）表示，1000 層更多地是為了探究上限，實際跑的過程中并非一定要上千層。此外，訓練代碼很快就會公開。

DEEPNORM 方法

如下圖 2 所示，使用 PostLN 實現(xiàn)基于 Transformer 的方法很簡單。與 Post-LN 相比，DEEPNORM 在執(zhí)行層歸一化之前 up-scale 了殘差連接。

圖 2：(a) DEEPNORM 的偽代碼，例如可以用其他標準初始化代替 Xavier 初始化 (Glorot and Bengio, 2010) ，其中 α 是一個常數(shù)。(b) 不同架構的 DEEPNORM 參數(shù)（N 層編碼器，M 層解碼器）。

此外，該研究還在初始化期間 down-scale 了參數(shù)。值得注意的是，該研究只擴展了前饋網(wǎng)絡的權重，以及注意力層的值投影和輸出投影。此外，殘差連接和初始化的規(guī)模取決于圖 2 中不同的架構。

深度 Transformer 的不穩(wěn)定性

該研究分析了深度 Transformer 不穩(wěn)定的原因。

首先，研究者觀察發(fā)現(xiàn)：更好的初始化方法可以讓 Transformer 的訓練更穩(wěn)定。之前的工作（Zhang et al., 2019a; Huang et al., 2020; Xu et al., 2021）也證實了這一點。

因此，研究者分析了有無適當初始化的 Post-LN 的訓練過程。通過更好的初始化，在執(zhí)行 Xavier 初始化后通過down-scale 第 l 層的權重。例如，第 l 層 FFN 的輸出投影被初始化為其中 d’是輸入和輸出維度的平均值。研究者將此模型命名為 Post-LN-init。請注意，與之前的工作（Zhang et al., 2019a）不同， Post-LN-init 是縮窄了較低層的擴展而不是較高層。研究者相信這種方法有助于將梯度擴展的影響與模型更新區(qū)分開來。此外，Post-LN-init 與 Post-LN 具有相同的架構，從而消除了架構的影響。

該研究在 IWSLT-14 De-En 機器翻譯數(shù)據(jù)集上訓練了 18L-18L Post-LN 和 18L-18L Post-LN-init。圖 3 可視化了它們的梯度和驗證損失曲線。如圖 3 (c) 所示，Post-LN-init 收斂，而 Post-LN 沒有。Post-LN-init 在最后幾層中具有更大的梯度范數(shù)，盡管其權重已按比例縮小。此外，研究者可視化最后一個解碼器層的梯度范數(shù)，模型深度從 6L-6L 到 24L-24L。

下圖 3 顯示，無論模型深度如何，最后一層 Post-LN-init 的梯度范數(shù)仍遠大于 Post-LN 的梯度范數(shù)。得出的結論是，深層梯度爆炸不應該是 Post-LN 不穩(wěn)定的根本原因，而模型更新的擴展往往可以解釋這一點。

然后研究者證明 Post-LN 的不穩(wěn)定性來自一系列問題，包括梯度消失以及太大的模型更新。如圖 4 (a) 所示，他們首先可視化模型更新的范數(shù) ||ΔF|| 在訓練的早期階段：

其中 x 和 θ_i 分別代表輸入和第 i 次更新后的模型參數(shù)。Post-LN 在訓練一開始就有爆炸式的更新，然后很快就幾乎沒有更新了。這表明該模型已陷入虛假的局部最優(yōu)。

warm-up 和更好的初始化都有助于緩解這個問題，使模型能夠順利更新。當更新爆炸時，LN 的輸入會變大（見圖 4 (b) 和圖 4 (c)）。根據(jù) Xiong 等人 (2020) 的理論分析，通過 LN 的梯度大小與其輸入的大小成反比：

相比于沒有 warm-up 或正確初始化的情況，圖 4 (b) 和圖 4 (c) 表明 ||x|| 的明顯大于。這解釋了 Post-LN 訓練中出現(xiàn)的梯度消失問題（見圖 4 (d)）。

最重要的是，不穩(wěn)定性始于訓練開始時的大型模型更新。它使模型陷入糟糕的局部最優(yōu)狀態(tài)，這反過來又增加了每個 LN 的輸入量。隨著訓練的繼續(xù)，通過 LN 的梯度變得越來越小，從而導致嚴重的梯度消失，使得難以擺脫局部最優(yōu)，并進一步破壞了優(yōu)化的穩(wěn)定性。相反，Post-LN-init 的更新相對較小，對 LN 的輸入是穩(wěn)定的。這減輕了梯度消失的問題，使優(yōu)化更加穩(wěn)定。

DeepNet：極深的 Transformer 模型

研究者首先介紹了極深的 Transformer 模型 ——DeepNet，該模型可以通過緩解爆炸式模型更新問題來穩(wěn)定優(yōu)化過程。

DeepNet 基于 Transformer 架構。與原版 Transformer 相比，DeepNet 在每個子層使用了新方法 DEEPNORM，而不是以往的 Post-LN。DEEPNORM 的公式如下所示。

其中，α 是一個常數(shù)，G_l (x_l , θ_l) 是參數(shù)為 θ_l 的第 l 個 Transformer 子層（即注意力或前饋網(wǎng)絡）的函數(shù)。DeepNet 還將殘差內部的權重 θ_l 擴展了 β。

接著，研究者提供了對 DeepNet 模型更新預期大?。╡xpected magnitude）的估計。

他們可視化了 IWSLT-14 De-En 翻譯數(shù)據(jù)集上，Post-LN 和 DeepNet 在早期訓練階段的模型更新情況，如下圖 5 所示?？梢钥吹?，相較于 Post-LN，DeepNet 的模型更新幾乎保持恒定。

最后，研究者提供理論分析，以表明 DeepNet 的更新受到了 DEEPNORM 的常數(shù)限制。具體地，他們展示了 DeepNet 的預期模型更新受到了適當參數(shù) α 和 β 的常數(shù)限制。研究者的分析基于 SGD 更新，并通過實證證明對 Adam 優(yōu)化器效果很好。

研究者提供了對編碼器 - 解碼器架構的分析，它能夠以相同的方式自然地擴展到僅編碼器和僅解碼器的模型。具體如下圖所示，他們將模型更新的目標設定如下：

僅編碼器（例如 BERT）和僅解碼器（例如 GPT）架構的推導能夠以相同的方式進行。研究者將步驟總結如下：

神經(jīng)機器翻譯

該研究驗證了 DeepNet 在流行的機器翻譯基準上的有效性，包括 IWSLT-14 德語 - 英語 (De-En) 數(shù)據(jù)集和 WMT-17 英語 - 德語 (En-De) 數(shù)據(jù)集。該研究將 DeepNet 與多個 SOTA 深度 Transformer 模型進行比較，包括 DLCL 、NormFormer 、ReZero 、R- Fixup 、T-Fixup 、DS-init 和 Admin。

下表 1 報告了 WMT-17 En-De 翻譯數(shù)據(jù)集上的基線和 DeepNet 的結果：

下圖 6 顯示了 IWSLT-14 數(shù)據(jù)集的結果

下圖 7 報告了 WMT-17 驗證集的損失曲線

大規(guī)模多語言神經(jīng)機器翻譯

該研究首先使用 OPUS-100 語料庫來評估模型。OPUS100 是一個以英語為中心的多語言語料庫，涵蓋 100 種語言，是從 OPUS 集合中隨機抽取的。該研究將 DeepNet ??擴展到 1,000 層，該模型有一個 500 層的編碼器、 500 層的解碼器、512 個隱藏大小、8 個注意力頭和 2,048 維度的前饋層。

下表 2 總結了 DeepNet 和基線的結果。結果表明，增加網(wǎng)絡深度可以顯著提高 NMT 的翻譯質量：48 層的模型比 12 層的模型平均獲得 3.2 點的提高。DeepNet ?可以成功地將深度擴展到 1,000 層，比基線提高 4.4 BLEU。值得注意的是，DeepNet 只訓練了 4 個 epoch，并且在計算預算更多的情況下，性能可以進一步提高。

深度擴展規(guī)律：該研究在 OPUS100 數(shù)據(jù)集上訓練具有 {12,20,100,200,1000} 層的 DeepNet，下圖 8 顯示了深度擴展曲線。與雙語 NMT 相比，多語 NMT 從擴展模型深度受益更多?？梢杂^察到多語 NMT 的 BLEU 值呈對數(shù)增長，規(guī)律可以寫成：L (d) = A log (d) + B，其中 d 是深度，A, B 是關于其他超參數(shù)的常數(shù)。

更多數(shù)據(jù)和語言說明：為了探索 DeepNet 在多語 NMT 上的局限性，該研究隨后使用 Schwenk 等人提出的 CCMatrix 擴展訓練數(shù)據(jù)。此外，該研究還擴展了 CCAligned 、OPUS 和 Tatoeba 的數(shù)據(jù)，以涵蓋 Flores101 評估集的所有語言。最終的數(shù)據(jù)由 102 種語言、1932 個方向和 12B 對句子組成。利用這些數(shù)據(jù)，該研究用 100 層編碼器、100 層解碼器、1024 個隱藏維度、16 個頭、4096 個前饋層中間維度對 DeepNet 進行訓練。

該研究將 DeepNet 與 SOTA 多語 NMT 模型 M2M-100 進行了比較。M2M-100 有一個 24 層的編碼器、一個 24 層的解碼器和 4,096 個隱藏大小，從而產(chǎn)生高達 12B 的參數(shù)。與 M2M-100 相比，DeepNet 深而窄，參數(shù)只有 3.2B。

在 M2M-100 之后，該研究在幾個多語言翻譯評估數(shù)據(jù)集上評估模型，包括 WMT、OPUS 、TED、 Flores。WMT 的語言對是以英語為中心的。包括英語在內的 10 種語言，其中大部分是高資源語言。對于 OPUS 數(shù)據(jù)集，該研究從包含 30 個評估對的測試集中選擇非英語方向。TED 評估集有 28 種語言和 756 個方向，數(shù)據(jù)來自口語領域。Flores 數(shù)據(jù)集包含 102 種語言之間的所有翻譯對。該研究使用涵蓋 M2M-100 和? DeepNet ?支持的語言的子集，產(chǎn)生 87 種語言和 7,482 個翻譯方向。

下表 3 報告了結果，為了公平比較，該研究使用與基線相同的評估方法。結果表明 DeepNet 在所有評估數(shù)據(jù)集上的性能都明顯優(yōu)于 M2M-100，表明深化模型是提高 NMT 模型質量的一個非常有前景的方向。

感興趣的讀者可閱讀論文原文，了解更多細節(jié)內容。

——The ?End——