?PaperWeekly 原創 ·?作者 | 蘇劍林

單位 | 追一科技

研究方向 | NLP、神經網絡

眾所周知，現在的 Transformer 越做越大，但這個“大”通常是“寬”而不是“深”，像 GPT-3 雖然參數有上千億，但也只是一個 96 層的 Transformer 模型，與我們能想象的深度相差甚遠。是什么限制了 Transformer 往“深”發展呢？可能有的讀者認為是算力，但“寬而淺”的模型所需的算力不會比“窄而深”的模型少多少，所以算力并非主要限制，歸根結底還是 Transformer 固有的訓練困難。一般的觀點是，深模型的訓練困難源于梯度消失或者梯度爆炸，然而實踐顯示，哪怕通過各種手段改良了梯度，深模型依然不容易訓練。

近來的一些工作（如 Admin [1]）指出，深模型訓練的根本困難在于“增量爆炸”，即模型越深對輸出的擾動就越大。上周的論文《DeepNet: Scaling Transformers to 1,000 Layers》[2] 則沿著這個思路進行尺度分析，根據分析結果調整了模型的歸一化和初始化方案，最終成功訓練出了 1000 層的 Transformer 模型。整個分析過程頗有參考價值，我們不妨來學習一下。

增量爆炸

原論文的完整分析比較長，而且有些假設或者描述細酌之下是不夠合理的。所以在本文的分享中，筆者會盡量修正這些問題，試圖以一個更合理的方式來得到類似結果。

假設損失函數為 ， 是它的參數，考慮參數由 變為 時損失函數的增量：

對于 SGD 有 ，那么 。設模型有 層，每層的平均參數量為 ，配合 Xavier 初始化以及各種 Normalization 手段，我們可以使得多數參數的梯度是 量級，所以有 。因此，模型每一步的更新量是正比于模型深度 的（寬度不在本文討論范圍），如果模型越深，那么更新量就越大，這意味著初始階段模型越容易進入不大好的局部最優點，然后訓練停滯甚至崩潰，這就是“增量爆炸”問題。

這時候解決方法有兩個，一是初始階段用更小的學習率進行訓練（不超過 量級），然后慢慢增大學習率，這就是 Warmup 技巧；二就是調整初始化方案，使得參數的梯度是 量級，這樣就自動抵消掉模型深度的影響。

量級分析

怎么做到第二種方案呢？我們可以嘗試分析 Transformer 的梯度。然而，精確的梯度求起來比較繁瑣，并且事實上我們也不需要精確的梯度，而只是要對梯度做一個量級分析，所以我們可以用如下的“量級分解”技巧轉化為標量的導數問題。

對于一個矩陣 ，我們將其分解為 的形式，其中

說白了，我們就是要將一個矩陣分解為一個標量 與一個盡可能正交的矩陣 之積。由于 接近正交矩陣，它起到了一個標準參考系的作用，而對應的 則代表了矩陣 的量級。如果 使用 Xavier 初始化，那么 相當于其中的 gain 參數，即在 Xavier 初始化的基礎上還要再乘一個 。這是因為 Xavier 初始化的結果就接近一個正交矩陣，這一點可以參考《從幾何視角來理解模型參數的初始化策略》[3] 。

在此分解之下，我們有

這意味著 跟 在量級上是成正比的，所以對 做量級分析就相當于對 做量級分析。這樣原來的矩陣求導就可以轉化為標量求導，降低了分析難度。

前饋梯度

很多實驗結果都顯示雖然 Pre Norm 比 Post Norm 更容易訓練，但 Post Norm 的最終效果往往更好些，所以原論文保留了 Post Norm 結構，并考慮了更一般的形式：

其中 是一個常數。簡單起見，我們先考慮 FFN 層，此時

這里的 是激活函數，一般為 ReLU 或其變體（Swish、GeLU等），它們（近似）滿足 。使用前一節的量級分解，我們得到

求 的梯度：

我們斷言 、 都是 的，并且由于 、 都接近正交矩陣，所以 也是 的，因此最終有

自注意力

現在考慮自 Self Attention，作為量級分析，我們考慮單頭注意力即可，其形式為

其中 是 softmax 操作的簡寫，這里省略了 Attention 的 scale 操作。對上式進行量級分解后的形式為

現在我們可以對各個 分別求梯度，而由于 softmax 的存在，事實上 的梯度本身會很小，不會明顯影響最終的更新量，所以其實我們考慮 的更新量足矣：

同樣斷言 、 都是 的，并且注意 softmax 出來是一個概率分布，然后對 的各個 token 做加權平均，通常而言，平均前后的向量會在同一數量級，所以我們認為

402 Payment Required

也是 的，因此結果跟 FFN 層的類似：

初步結論

現在不管是 FFN 還是 Self Attention，我們都得到了相似的結論，現在簡單起見，假設每個參數的量級（至少在初始化階段）是一致的，即所有的 取同一個值，那么總的結論是

即梯度的量級是 。另一方面，我們說 層的 Transformer 模型，一般是 層的 Self Attention 加 層的 FFN，所以嚴格來說層數是 。因此，按照“增量爆炸”一節的分析，我們需要將梯度調整到 ，上式告訴我們可以通過讓 來實現。原論文的放縮更為寬松一些，得到的結果是 ，量級上是等價的。

現在我們得到了 與 的一個比例關系，但無法直接得到 和 的具體值。按照論文的說法，是從對稱角度出發，讓 ，從而可以解得

然而，單純對稱的解釋顯然是不夠說服力的，我們需要搞清楚不同的選擇究竟有什么不同的結果。為此，我們可以比較另外兩組解：

另解一：，此時參數的初始化縮小到原來的 倍，梯度也被縮小到原來的 倍，根據 SGD 的 得出每步的更新量也是原來的 倍，也就是說，調整前后的相對學習幅度是沒有變化的，因此有可能剛開始 級別，但訓練集幾步后就脫離了這個量級了。

另解二：，此時參數的初始化沒有縮小，但梯度也被縮小到原來的 倍，根據 SGD 的 得出每步的更新量也是原來的 倍，調整前后的相對學習幅度是明顯縮小了，因此有可能出現學習得非常慢的情況。

這兩種情況看上去都各有缺點，因此介乎兩者之間的式（14）似乎就能解釋得通了，它就是保持梯度縮放到原來的 倍的同時，讓初始學習步伐稍微慢一些，但又不至于太慢，隱式地起到了 Warmup 的作用。

多種優化

上面的分析都是基于 SGD 進行的，但事實上我們很少直接用 SGD 去訓練 NLP 模型，我們更多是自適應學習率優化器，主要有兩大類：一是用二階矩來校正學習率，Adam、AdamW 等都屬此類；另一類是通過參數模長進一步校正學習率，比如 LAMB [4]、AdaFactor [5]。原論文的說法是“我們在 SGD 上進行推導，然后在Adam上驗證發現也還可以”，但從理論上來講，它們并不完全通用，這一節我們就來針對性地做一下分析。

對于 Adam 類優化器來說，每一步的更新量大約為 ，所以 ，它是正比于梯度的 1 次方而不是 2 次方，因此要過要讓更新量跟層數無關，那么梯度應該縮小到原來的 倍才對，即應該有 ，如果同樣讓 ，那么有

對于 LAMB 類優化器來說，每一步更新量大約為 ，所以 ，注意到參數的縮放比例是 、梯度的縮放比例是 ，所以，從而是 。注意這類優化器每步的相對更新量是一樣的（等于學習率 ），所以不管怎么調整相對更新大小都不會變化，所以我們可以直接取 。

結果匯總對比如下：

事后分析

前面的兩節推導過程都用到了斷言“、 都是 的”，那么它是否成立呢？這里我們事后分析一下。

其實也很簡單，經過前述調整后，不管是 FFN 層（6）還是 Self Attention 層 （10），初始階段每個殘差分支的權重被縮放到原來的 倍，不管是哪種優化器的結果， 都是一個比較小的數字，這意味著初始階段整個模型其實接近一個恒等函數，因此 、 自然都是 的，所以結論和斷言是自洽的。

另外，可能有讀者想問同樣的分析是否可以用到 Pre Norm 結構上呢？答案是可以的，并且結論是基本一致的，只是因為 Norm 放在了殘差分支之前，所以就沒必要設置 參數了，所以結論就是上述關于 Post Norm 的結果中所有的 都等于為 1，然后重新計算相應的 。

最后，讀者可能有疑問的是花了那么多功夫討論把模型做深，那么模型深度真有那么重要嗎？有，原論文給出了一個漂亮的實驗結果，用一個 200 層的“深而窄”的模型（32 億參數），戰勝了之前48層“淺而寬”的 SOTA 模型（120 億參數）：

▲ “深而窄”的模型勝于“淺而寬”的模型

文章小結

本文分析了將 Transformer 做“深”的瓶頸所在并給出了相應的解決方案，文章的主要思路源于微軟新出的 DeepNet，并對原論文的分析過程做了一定的簡化和完善。

參考文獻

[1] https://arxiv.org/abs/2004.08249

[2] https://arxiv.org/abs/2203.00555

[3] https://kexue.fm/archives/7180

[4] https://kexue.fm/archives/7094

[5] https://kexue.fm/archives/7302

更多閱讀

#投 稿?通 道#

?讓你的文字被更多人看到?

如何才能讓更多的優質內容以更短路徑到達讀者群體，縮短讀者尋找優質內容的成本呢？答案就是：你不認識的人。

總有一些你不認識的人，知道你想知道的東西。PaperWeekly 或許可以成為一座橋梁，促使不同背景、不同方向的學者和學術靈感相互碰撞，迸發出更多的可能性。?

PaperWeekly 鼓勵高校實驗室或個人，在我們的平臺上分享各類優質內容，可以是最新論文解讀，也可以是學術熱點剖析、科研心得或競賽經驗講解等。我們的目的只有一個，讓知識真正流動起來。

📝?稿件基本要求：

? 文章確系個人原創作品，未曾在公開渠道發表，如為其他平臺已發表或待發表的文章，請明確標注?

? 稿件建議以?markdown?格式撰寫，文中配圖以附件形式發送，要求圖片清晰，無版權問題

? PaperWeekly 尊重原作者署名權，并將為每篇被采納的原創首發稿件，提供業內具有競爭力稿酬，具體依據文章閱讀量和文章質量階梯制結算

📬?投稿通道：

? 投稿郵箱：hr@paperweekly.site?

? 來稿請備注即時聯系方式（微信），以便我們在稿件選用的第一時間聯系作者

? 您也可以直接添加小編微信（pwbot02）快速投稿，備注：姓名-投稿

△長按添加PaperWeekly小編

🔍

現在，在「知乎」也能找到我們了

進入知乎首頁搜索「PaperWeekly」

點擊「關注」訂閱我們的專欄吧

·

總結

以上是生活随笔為你收集整理的训练1000层的Transformer究竟有什么困难？的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。