[轉] Transformer中warm-up和LayerNorm的重要性探究

idea很棒，實驗結果也很棒

from： https://zhuanlan.zhihu.com/p/84614490

Transformer中的warm-up與LayerNorm

之前知乎上有一個問題：神經網絡中 warmup 策略為什么有效；有什么理論解釋么？在這個問題下，由於理論解釋的缺乏，通過現有的一些文章，我將（可能的）原因歸結為兩點：（1）減緩模型在初始階段對mini-batch的提前過擬合現象，保持分布的平穩；（2）保持模型深層的穩定性。本文從一個新的觀點解釋warm-up：和Layer Norm有關。而且在另一個問題：為什么Transformer 需要進行 Multi-head Attention？我也提到，Transformer中的初始化其實相當有探索空間。本文正是將這兩個問題聯系了起來。

我們知道，在原始的Transformer中，Layer Norm在跟在Residual之后的，我們把這個稱為Post-LN Transformer；而且用Transformer調過參的同學也知道，Post-LN Transformer對參數非常敏感，需要很仔細地調參才能取得好的結果，比如必備的warm-up學習率策略，這會非常耗時間。

所以現在問題來了，為什么warm-up是必須的？能不能把它去掉？本文的出發點是：既然warm-up是訓練的初始階段使用的，那肯定是訓練的初始階段優化有問題，包括模型的初始化。從而，作者發現，Post-LN Transformer在訓練的初始階段，輸出層附近的期望梯度非常大，所以，如果沒有warm-up，模型優化過程就會炸裂，非常不穩定。這一發現之前總結的2點原因完全符合，使用warm-up既可以保持分布的平穩，也可以保持深層的穩定。既然如此，本文作者嘗試把Layer Norm換個位置，比如放在Residual的過程之中（稱為Pre-LN Transformer），再觀察訓練初始階段的梯度變化，發現比Post-LN Transformer不知道好到哪里去了，甚至不需要warm-up，從而進一步減少訓練時間，這一結果的確令人震驚。

下圖是兩種結構的示意圖，相當直觀：

Post-LN Transformer與Pre-LN Transformer

為了之后敘述方便，我們還是先約定一些記號。我們還是用表示參數，用表示大家熟知的操作。用表示LayerNorm操作。Post-LN Transformer和Pre-LN Transformer還是參考上圖，可以用這些記號表示，就不多贅述了。

Post-LN Transformer中warm-up的重要性

Transformer中的warm-up可以看作學習率隨迭代數的函數：

之后，學習率會以某種方式（線性或其他）遞減，這里我們只關心warm-up階段。可以看到，學習率從0開始增長，經過次迭代達到最大。

下面，我們在IWSLT14 De-En數據集上研究兩個問題：（1）是否warm-up必要，（2）是否的設置很重要。對第一個問題，我們分別用Adam和SGD訓練模型，並且分別看保留和移除warm-up的效果。對Adam，設置，對SGD，設置。當使用warm-up時，。對第二個問題，分別設置，並且只使用Adam及其設置。

下圖是相關結果。可以看到，無論對Adam還是SGD，warm-up都相當必要，沒有warm-up，結果瞬間爆炸。而且，實際上還是有影響（至少對Adam）。

從上面的實驗可以得出：（1）warm-up的使用增加了訓練時間；（2）在訓練的開始階段，loss很大，一開始使用大學習率對模型是毀滅性的打擊。（Liu et al. 2019a）認為warm-up對Adam而言非常重要，因此提出了RAdam，但是這里的結果表明，warm-up對SGD同樣非常重要，並不是Adam的“寵妃”。

一窺Transformer訓練的初始階段

上面的發現啟發我們去探究模型訓練的初始階段發生了什么。這一部分大多是理論證明，我們略去表述和證明，只講與之相關的結論，有興趣的同學可以參考原文學習一個（其實是公式太多人有點懶）。注意到這里使用的是Xavier Gaussian分布，而在Transformer原文里使用的是Xavier Uniform。

定理1（Transformer最后一個FFN層的梯度） 告訴我們，Post-LN Transformer最后一個FFN層的梯度是的，和無關，而對Pre-LN Transformer，是的，小得多。

然后引理1，引理2，引理3 又告訴我們，LayerNorm的確會normalize梯度。在Post-LN Transformer中，輸入到最后一層LN的scale是與無關的，因此最后一層的梯度也是與無關的；而在Pre-LN Transformer中，輸入到最后一層的LN的scale是隨線性增長的，因此梯度將會以的比例normalized。

現在我們要擴展到每一層。這里的主要結論是：Post-LN Transformer的梯度范數在輸出層附近很大，因此很可能隨着BP的進行梯度越來越小。相反，Pre-LN Transformer在每層的梯度范數都近似不變（證明在Appendix F）。

下面我們又以IWSLT14 De-En數據集為例實驗，實驗設置同上。我們計算不同階段、不同參數的梯度期望值，結果如下。

 

可以看到，實驗結果驗證了上述的推論。而且，注意圖中的綠色部分（真的有，不是忘了畫！），非常小，這時候可以用大學習率訓練。這說明了，warm-up的確有助於訓練的穩定性。那么，用Pre-LN Transformer能不能去掉warm-up呢，在接下來一節將用實驗證明。

在Pre-LN Transformer中去掉warm-up

我們在兩個數據集上實驗：IWSLT14 De-En和WMT14 En-De，使用Transformer base結構。對Pre-LN Transformer，我們去掉warm-up，在IWSLT14 De-En中保持然后從第8個epoch開始下降；在WMT14 En-De上，分別實驗，都在第6個epoch下降。用帶warm-up的Post-LN Transformer作為基線。對所有實驗，使用Adam，並且也在IWSLT14 De-En上實驗了RAdam，參數設置與原論文保持一致。

對於非翻譯任務，我們使用預訓練的BERT base模型，只不過是把其中的Transformer相應地替換為Post-LN與Pre-LN而已。

下圖是機器翻譯的結果。對Pre-LN Transformer，warm-up不再必要，而且收斂更快。此外，相比用RAdam替換Adam，似乎改變LN的結構更顯著。

下圖是無監督預訓練的結果，指標有Validation Loss，MRPC和RTE上的Accuracy。可以看到，Pre-LN Transformer收斂更快，效果更好。總而言之一句話，Pre-LN Transformer似乎不再需要warm-up了。

小結

本文別出心裁，用實驗和理論驗證了Pre-LN Transformer結構不需要使用warm-up的可能性，其根源是LN層的位置導致層次梯度范數的增長，進而導致了Post-LN Transformer訓練的不穩定性。本文很好進一步follow，比如Residual和LN的其他位置關系，如何進一步設計初始化方法，使得可以完全拋棄warm-up等。比如Improving Deep Transformer with Depth-Scaled Initialization and Merged Attention這篇文章（第二次說到了），其實也提到了Transformer底層多頭的方差過大的問題，然后提出使用depth-scaled初始化參數緩解。我們不得不去猜測：Residual、LN、Initialization、Gradient這四者之間，肯定有千絲萬縷的聯系。

 

另外其實在Kaiming大神2016年的這篇文章中已經指出了BN、activation和Residual的關系了，並且提供了簡單的數學證明，有興趣的同學可以看看。（再次膜拜Kaiming大神）

https://arxiv.org/pdf/1603.05027.pdf​arxiv.org