解決訓(xùn)練難題，1000層的Transformer來了，訓(xùn)練代碼很快公開

來源：機(jī)器之心

機(jī)器之心編輯部

1000 層的 Transformer，深得嚇人。

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

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

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

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

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

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

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

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

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

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

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

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

DEEPNORM 方法

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

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

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

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

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

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

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

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

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

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

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

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

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

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

DeepNet：極深的 Transformer 模型

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

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

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

接著，研究者提供了對(duì) DeepNet 模型更新預(yù)期大?。╡xpected magnitude）的估計(jì)。

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

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

研究者提供了對(duì)編碼器 - 解碼器架構(gòu)的分析，它能夠以相同的方式自然地?cái)U(kuò)展到僅編碼器和僅解碼器的模型。具體如下圖所示，他們將模型更新的目標(biāo)設(shè)定如下：

僅編碼器（例如 BERT）和僅解碼器（例如 GPT）架構(gòu)的推導(dǎo)能夠以相同的方式進(jìn)行。研究者將步驟總結(jié)如下：

神經(jīng)機(jī)器翻譯

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

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

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

下圖 7 報(bào)告了 WMT-17 驗(yàn)證集的損失曲線

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

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

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

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

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

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

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

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

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

——The ?End——