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

所以現在問題來了，為什麼warm-up是必須的？能不能把它去掉？本文的出發點是：既然warm-up是訓練的初始階段使用的，那肯定是訓練的初始階段優化有問題，包括模型的初始化。從而，作者發現，Post-LN Transformer

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

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大神）