論文標題：
Thomas Bachlechner, Bodhisattwa Prasad Majumder, Huanru Henry Mao, Garrison W. Cottrell, Julian McAuley
論文作者：
Thomas Bachlechner, Bodhisattwa Prasad Majumder, Huanru Henry Mao, Garrison W. Cottrell, Julian McAuley
論文連結：
https://arxiv.org/abs/2003.04887
代碼連結：
https://github.com/majumderb/rezero
在NLP領域，如Transformer的深度模型由於梯度消失/爆炸難以訓練，往往需要花費很長的時間才能收斂。
為此，本文提出ReZero：在殘差連接前增加一個權重，使模型能夠更好接受到梯度信號，加快收斂速度。
這種方法能在上百層的Transformer上收斂，並在常見深層模型上大大縮短訓練時間，同時取得相近的結果。
一般來說，模型越深效果越好。但是同時，模型越深也更難訓練——即無法收斂到訓練集上一個好的值。
普遍認為，這是深度模型的梯度消失/梯度爆炸現象導致的：梯度的指數級爆炸使得訓練極其不穩定，指數級消失使得訓練非常緩慢。
拿Transformer來說，模型越深，底層收到的梯度就越低，如下圖所示（來自Improving Deep Transformer with Depth-Scaled Initialization and Merged Attention）：
比如看圖中紅色曲線，在高層的時候，梯度的範數尚且較大，但當來到底層的時候，梯度的範數就非常小，特別是在Decoder端，幾乎趨於零。這說明了深層模型仍然有比較嚴重的梯度消失問題。
那麼，這個問題的解決方案有什麼呢？在計算機視覺上，Normalization和殘差連接有很大幫助，但在NLP領域，它們的作用就相當有限，因為Transformer本身就內置了殘差連接和Normalization技術。
為了能夠更好地解決NLP模型上的訓練問題，本文提出了一種非常簡單的方法：ReZero——對每個殘差連接加上一個可學習的係數——使得模型能在訓練初期更加穩定，進而促進整個訓練過程。
本方法在12層的Transformer上取得約56%的訓練加速，在32層前饋網絡上取得15倍的加速，在ResNet54上實現了32%的加速。
更重要的是，該方法能夠直接在128層的Transformer上取得收斂，同時還能丟掉Transformer中的Normalization和訓練時的warm-up環節。
在閱讀完本文後，讀者可以思考如下問題：
首先來簡要回顧一下殘差連接。對一個層的模型，第層的輸入（也即第層的輸出）可以表示為：
而ReZero做的很簡單，它只是在殘差連接前加了一個可學習的參數：
並且所有的都初始為0。我們通過一個簡單的實驗說明並分析該方法的有效性。
設現在有個模型有層，每一層只有一個神經元，而且所有層都共享這一個神經元，那麼，模型的輸出就可以表示為：
顯然，當時就是殘差網絡，當初始化時就是ReZero。此時，輸入對輸出的Jacobian矩陣就是。
此時，如果初始化且，那麼就趨近，這就會導致輸入很小的擾動造成輸出的巨大變化，從而梯度非常不穩定。
而當時，這種不穩定就會大大緩解。下面我們用反向傳播加以解釋。設學習率為，損失函數為，那麼參數的更新就是：
在這個式子裡，我們主要看學習率，參數和。當，上式右邊第二項（忽略損失項和輸入）成為，這個時候為了讓訓練更加穩定，學習率就需要正比於。
但是對於過大的和接近-1的參數，單純依靠學習率的手動調整就難以解決這個問題了。
當初始化的時候，上述問題就可以得到解決。因為初始化，從而第一輪的梯度更新沒有更新參數，但是是可以更新的：
這樣，在下一輪梯度更新的時候，由於此時不為0，則模型參數就可以更新。把當前的帶入到上面更新的式子中，我們有：
比較兩式，其最大的不同在於括號內的損失函數。換言之，如果損失函數是合理的，那麼當前的更新就不會導致過分的梯度波動（），每次梯度下降使得模型參數既可以更新，又不會不穩定。
下圖是模型擬合的對數等高線示意圖，左圖是使用二次損失的損失圖，右圖是對應的梯度範數。紅線是初始化和不同的初始化在訓練中的變化軌跡。我們注意到以下事實：
訓練全連接模型
首先我們在簡單的全連接模型上實驗。該全連接模型激活函數為ReLU，層數為32，數據集為CIFAR-10，比較的模型有：
FC：單純的全連接
FC+Res：全連接和殘差結合
FC+Norm：全連接和normalization
FC+Res+Pre-Norm: 全連接和殘差和pre-normalization
FC+Res+Post-Norm：全連接和殘差和post-normalization
FC+ReZero：
注意到，ReZero沒有使用Normalization。下圖是各個模型的收斂情況:
可以看到，ReZero比其他三個模型能夠快7~15倍收斂，而且，當我們把全連接層加到10K層，ReZero依然可以收斂，而其他模型卻無法收斂。
訓練Transformer
我們的重點是訓練Transformer，一個被普遍認為非常難以訓練的模型。我們把Transformer的結構修正為：
同樣，我們丟掉了Transformer本來的Normalization，而且也不需要warm-up。下圖呈現了Transformer輸出-輸入的Jacobian矩陣中的奇異值在不同層的模型下的分布。
先看左圖，當層數較少（4）時，保持在附近（圖中取了對數），這就表明輸入的變化既不會造成輸出的太大變化，也能使得模型得以訓練。
而當層數增加，尤其是增加到64層時，該值非常小，這說明無論輸入怎麼改變，模型的輸出都八九不離十，模型難以訓練。
再看右圖是使用ReZero對64層的Transformer訓練的情況。在一開始的時候，由於被初始化為0，模型參數無法更新，值都是1，而當模型開始訓練後，也能保持以1為中心的分布，使得模型能夠訓練下去。
下面來具體看ReZero的速度提升和最終的效果。下表是12層的Transformer使用各方法在enwiki8數據集上的收斂情況。可以看到，ReZero只需要8800次迭代，而原始Transformer需要13690次迭代，速度方面加快了56%。
下表是各方法在enwiki8測試集上的結果。128層的Transformer能夠在ReZero下收斂，並且取得了和Char-level的Transformer非常相似的結果，而參數量只有後者的一半。
最後我們來可視化64層Transformer在訓練過程中，各層的係數參數的變化情況，如下圖所示。
由於我們初始化，故在訓練的開始階段所有層的係數都比較小，這可以使得各層訓練穩定。
當模型訓練到一定階段之後，輸入的擾動不會造成輸出的劇烈改變，係數開始增大，這是模型加速訓練、訓練的階段。
在此之後，係數再一次減小，使得模型訓練進入第三個階段，直到最終收斂。
訓練ResNet
最後，我們也將ReZero應用到了ResNet56上。在CIFAR-10上的結果顯示，使用ReZero可以降低基線模型的Error rate，同時也加快了約30%的收斂速度。
本文提出了一個非常簡單但是非常有效的殘差網絡變體——ReZero。和殘差網絡不同，它是在殘差連接前增加一個可學習的係數，並初始化為0，使得模型在訓練初期更加穩定。
實驗表明，運用這種方法可以直接在非常深的Transformer上收斂，而且丟掉了原有的Normalization和warmup。如果它能經得住大量實驗的檢驗，那麼ReZero非常有希望成為今後Transformer的標配。
ReZero和Highway的區別；Highway的公式是，其中表示各部分前饋的信息量。
在開始的時候，初始化為負值，這時候，就趨近於0，也就是上式趨近於，這實際上是和ReZero的想法是一致的——在訓練開始的時候保持穩定有助於深層模型的收斂。
不同的是，ReZero更加簡單粗暴——直接使用一個可學習的係數而不是所謂的門控，降低了模型學習的難度，簡化了複雜度，從而更有利於模型控制訓練過程。
和Transformer中warmup訓練的聯繫。在我們的回答神經網絡中 warmup 策略為什麼有效；有什麼理論解釋麼？中，我們說：「剛開始模型對數據的「分布」理解為零，在第一輪訓練的時候，每個數據點對模型來說都是新的，模型會很快地進行數據分布修正，如果這時候學習率就很大，極有可能導致開始的時候就對該數據「過擬合」，後面要通過多輪訓練才能拉回來，浪費時間。
當訓練了一段時間（比如兩輪、三輪）後，模型已經對每個數據點看過幾遍了，或者說對當前的batch而言有了一些正確的先驗，較大的學習率就不那麼容易會使模型學偏，所以可以適當調大學習率。
這個過程就可以看做是warmup。那麼為什麼之後還要decay呢？當模型訓到一定階段後（比如十個epoch），模型的分布就已經比較固定了，或者說能學到的新東西就比較少了。
如果還沿用較大的學習率，就會破壞這種穩定性，用我們通常的話說，就是已經接近loss的local optimal了，為了靠近這個point，我們就要慢慢來。」
從本文最後的可視化圖中可見，我們的解釋是完全符合這種直覺的。所以說，ReZero和warmup的作用都是為了訓練的穩定，只是一個選擇了自動學習，一個選擇手動調整。
🔍
現在，在「知乎」也能找到我們了
進入知乎首頁搜索「PaperWeekly」
點擊「關注」訂閱我們的專欄吧
關於PaperWeekly
PaperWeekly 是一個推薦、解讀、討論、報導人工智慧前沿論文成果的學術平臺。如果你研究或從事 AI 領域，歡迎在公眾號後臺點擊「交流群」，小助手將把你帶入 PaperWeekly 的交流群裡。