一種神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu)改進(jìn)方法「ReZero」

深度學(xué)習(xí)在眾多領(lǐng)域都取得了顯著進(jìn)展，但與此同時也存在一個問題：深層網(wǎng)絡(luò)的訓(xùn)練常常面臨梯度消失或梯度爆炸的阻礙，尤其是像 Transformer 這樣的大型網(wǎng)絡(luò)?，F(xiàn)在，加州大學(xué)圣迭戈分校的研究者提出了一種名為 ReZero 的神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu)改進(jìn)方法，并使用 ReZero 訓(xùn)練了具有一萬層的全連接網(wǎng)絡(luò)，以及首次訓(xùn)練了超過 100 層的 Tansformer，效果都十分驚艷。

深度學(xué)習(xí)在計算機(jī)視覺、自然語言處理等領(lǐng)域取得了很多重大突破。神經(jīng)網(wǎng)絡(luò)的表達(dá)能力通常隨著其網(wǎng)絡(luò)深度呈指數(shù)增長，這一特性賦予了它很強(qiáng)的泛化能力。然而深層的網(wǎng)絡(luò)也產(chǎn)生了梯度消失或梯度爆炸，以及模型中的信息傳遞變差等一系列問題。研究人員使用精心設(shè)計的權(quán)值初始化方法、BatchNorm 或 LayerNorm 這類標(biāo)準(zhǔn)化技術(shù)來緩解以上問題，然而這些技術(shù)往往會耗費更多計算資源，或者存在其自身的局限。

近日，來自加州大學(xué)圣迭戈分校（UCSD）的研究者提出一種神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu)改進(jìn)方法「ReZero」，它能夠動態(tài)地加快優(yōu)質(zhì)梯度和任意深層信號的傳播。

論文地址：https://arxiv.org/abs/2003.04887v1

代碼地址：https://github.com/majumderb/rezero

這個想法其實非常簡單：ReZero 將所有網(wǎng)絡(luò)層均初始化為恒等映射。在每一層中，研究者引入了一個關(guān)于輸入信號 x 的殘差連接和一個用于調(diào)節(jié)當(dāng)前網(wǎng)絡(luò)層輸出 F（x） 的可訓(xùn)練參數(shù)α，即：

在剛開始訓(xùn)練的時候?qū)ⅵ猎O(shè)置為零。這使得在神經(jīng)網(wǎng)絡(luò)訓(xùn)練初期，所有組成變換 F 的參數(shù)所對應(yīng)的梯度均消失了，之后這些參數(shù)在訓(xùn)練過程中動態(tài)地產(chǎn)生合適的值。改進(jìn)的網(wǎng)絡(luò)結(jié)構(gòu)如下圖所示：

圖 1：ReZero 結(jié)構(gòu)示意圖

ReZero 主要帶來了以下兩個益處：

1. 能夠訓(xùn)練更深層神經(jīng)網(wǎng)絡(luò)

學(xué)習(xí)信號能夠有效地在深層神經(jīng)網(wǎng)絡(luò)中傳遞，這使得我們能夠訓(xùn)練一些之前所無法訓(xùn)練的網(wǎng)絡(luò)。研究者使用 ReZero 成功訓(xùn)練了具有一萬層的全連接網(wǎng)絡(luò)，首次訓(xùn)練了超過 100 層的 Tansformer 并且沒有使用學(xué)習(xí)速率熱身和 LayerNorm 這些奇技淫巧。

2. 更快的收斂速度

與帶有標(biāo)準(zhǔn)化操作的常規(guī)殘差網(wǎng)絡(luò)相比，ReZero 的收斂速度明顯更快。當(dāng) ReZero 應(yīng)用于 Transformer 時，在 enwiki8 語言建?；鶞?zhǔn)上，其收斂速度比一般的 Transformer 快 56%，達(dá)到 1.2BPB。當(dāng) ReZero 應(yīng)用于 ResNet，在 CIFAR 10 上可實現(xiàn) 32% 的加速和 85% 的精度。

ReZero （residual with zero initialization）

ReZero 對深度殘差網(wǎng)絡(luò)的結(jié)構(gòu)進(jìn)行了簡單的更改，可促進(jìn)動態(tài)等距（dynamical isometry）并實現(xiàn)對極深網(wǎng)絡(luò)的有效訓(xùn)練。研究者在初始階段沒有使用那些非平凡函數(shù) F［W_i］ 傳遞信號，而是添加了一個殘差連接并通過初始為零的 L 個可學(xué)習(xí)參數(shù)α_i（作者稱其為殘差權(quán)重）來重新縮放該函數(shù)。目前，信號根據(jù)以下方式進(jìn)行傳遞：

在初始階段，該網(wǎng)絡(luò)表示為恒等函數(shù)并且普遍滿足動態(tài)等距關(guān)系。在該架構(gòu)修改中，即使某一層的 Jacobian 值消失，也可以訓(xùn)練深度網(wǎng)絡(luò)（正如 ReLU 激活函數(shù)或自注意力機(jī)制出現(xiàn)這樣的狀況）。這一技術(shù)還可以在現(xiàn)有的已訓(xùn)練網(wǎng)絡(luò)上添加新層。

實驗結(jié)果

更快的深層全連接網(wǎng)絡(luò)訓(xùn)練

圖 3 展示了訓(xùn)練損失的演變過程。在一個簡單實驗中，一個使用了 ReZero 的 32 層網(wǎng)絡(luò)，擬合訓(xùn)練數(shù)據(jù)的收斂速度相比其他技術(shù)快了 7 到 15 倍。值得注意的是，與常規(guī)的全連接網(wǎng)絡(luò)相比，殘差連接在沒有額外的標(biāo)準(zhǔn)化層時會降低收斂速度。這可能是因為初始化階段信號的方差并不獨立于網(wǎng)絡(luò)深度。

隨著深度的增加，ReZero 架構(gòu)的優(yōu)勢更加明顯。為了驗證該架構(gòu)可用于深度網(wǎng)絡(luò)訓(xùn)練，研究者在一臺配備 GPU 的筆記本電腦上成功訓(xùn)練了多達(dá) 1 萬層的全連接 ReZero 網(wǎng)絡(luò)，使其在訓(xùn)練數(shù)據(jù)集上過擬合。

圖 3：256 寬度和 ReLU 激活的 32 層全連接網(wǎng)絡(luò)四種變體，在訓(xùn)練過程中的交叉熵?fù)p失。

更快的深層 Transformer 訓(xùn)練

研究者提出，常規(guī)的 Transformer 會抑制深層信號傳遞，他們在輸入序列 x 的 n x d 個 entry 的無窮小變化下評估其變化，獲得注意力處理的輸入-輸出 Jacobian，從而驗證了之前的觀點。

圖 5a 展示了不同深度中使用 Xavier 統(tǒng)一初始化權(quán)重的 Transformer 編碼層的輸入-輸出 Jacobian 值。淺層的 Transformer 表現(xiàn)出峰值在零點附近的單峰分布，可以發(fā)現(xiàn)，深層結(jié)構(gòu)中 Jacobian 出現(xiàn)了大量超出計算精度的峰值。雖然這些分布取決于不同初始化方法，但以上量化的結(jié)論在很大范圍內(nèi)是成立的。這些結(jié)果與普遍認(rèn)為的相一致，也就是深層 Transformer 很難訓(xùn)練。

圖 5：多個輸入-輸出 Jacobian 矩陣中對數(shù)奇異值λ_io 的直方圖。（a）層數(shù)分別為 4、12、64 層的 Transformer 編碼器網(wǎng)絡(luò)；（b）是 64 層時訓(xùn)練前和訓(xùn)練中的 ReZero Transformer 編碼器網(wǎng)絡(luò)。深層 Transformer 距離動態(tài)等距很遠(yuǎn)，即λ_io 1，而 ReZero Transformer 更接近動態(tài)等距，平均奇異值 λ_io ≈ 1。

能夠在多項 NLP 任務(wù)中實現(xiàn) SOTA 的 Transformer 模型通常是小于 24 層的，這項研究中，最深層模型最多使用了 78 層，并且需要 256 個 GPU 來訓(xùn)練。研究者又將這一模型擴(kuò)展至數(shù)百個 Transformer 層，并且仍然可以在臺式機(jī)上訓(xùn)練。為了檢查該方法是否可以擴(kuò)展至更深層的 Transformer 模型之上，研究者將 ReZero Transformer 拓展到了 64 及 128 層，并與普通 Transformer 進(jìn)行了對比。

結(jié)果顯示，收斂之后，12 層的 ReZero Transformer 與常規(guī)的 Transformer 取得了相同的 BPB。也就是說，用 ReZero 來替代 LayerNorm 不會失去任何模型表現(xiàn)。訓(xùn)練普通的 Transformer 模型會導(dǎo)致收斂困難或訓(xùn)練緩慢。當(dāng)達(dá)到 64 層時，普通的 Transformer 模型即使用了 warm-up 也無法收斂。ReZero Transformer 在α初始化為 1 時發(fā)散，從而支持了α = 0 的初始化理論。深層的 ReZero Transformer 比淺層的 Transformer 表現(xiàn)出了更優(yōu)越的性能。

表 3：在 enwiki8 測試集上的 Transformers （TX） 對比。

收斂速度比較

選擇 enwiki8 上的語言建模作為基準(zhǔn)，因為較難的語言模型是 NLP 任務(wù)性能的良好指標(biāo)。在實驗中，其目標(biāo)是通過測量 12 層的 Transformer 在 enwiki8 上達(dá)到 1.2 位每字節(jié)（BPB）所需的迭代次數(shù)，由此來衡量所提出的每種方法的收斂速度。

表二：針對 ReZero 的 12 層 Transformers 歸一化后與 enwiki8 驗證集上達(dá)到 1.2 BPB 時所需的訓(xùn)練迭代比較。

更快的殘差網(wǎng)絡(luò)訓(xùn)練

通過前述部分，看到了 ReZero 的連接是如何使深層網(wǎng)絡(luò)的訓(xùn)練成為可能的，并且這些深層網(wǎng)絡(luò)都包含會消失的 Jacobian 奇異值，例如 ReLU 激活或自我注意力。但是，如果沒有 ReZero 的連接或者是其他架構(gòu)的更改，其中某些架構(gòu)將無法執(zhí)行訓(xùn)練。在本節(jié)中，會將 ReZero 連接應(yīng)用于深層殘差網(wǎng)絡(luò)從而進(jìn)行圖像識別。

雖然這些網(wǎng)絡(luò)并不需要 ReZero 連接便可以進(jìn)行訓(xùn)練，但通過觀察發(fā)現(xiàn)，在 CIFAR-10 數(shù)據(jù)集上訓(xùn)練的 ResNet56 model4（最多 200 個 epochs）的驗證誤差得到了非常明顯的提升：從（7.37±0.06）％到（6.46±0.05）％。這一效果是將模型中的所有殘差連接轉(zhuǎn)換為 ReZero 連接之后得到的。在實施 ReZero 之后，驗證誤差降低到 15％以下的次數(shù)也減少了（32±14）％。盡管目前這些結(jié)果只提供了有限的信息，但它們?nèi)灾赋隽?ReZero 連接擁有更廣泛的適用性，從而也推進(jìn)了進(jìn)一步的研究。

上手實操

項目地址：

https://github.com/majumderb/rezero

在此提供了自定義的 ReZero Transformer 層（RZTX），比如以下操作將會創(chuàng)建一個 Transformer 編碼器：

import torchimport torch.nn as nnfrom rezero.transformer import RZTXEncoderLayerencoder_layer = RZTXEncoderLayer（d_model=512， nhead=8）transformer_encoder = nn.TransformerEncoder（encoder_layer， num_layers=6）src = torch.rand（10， 32， 512）out = transformer_encoder（src）

創(chuàng)建一個 Transformer 解碼器：

import torchimport torch.nn as nnfrom rezero.transformer import RZTXDecoderLayerdecoder_layer = RZTXDecoderLayer（d_model=512， nhead=8）transformer_decoder = nn.TransformerDecoder（decoder_layer， num_layers=6）memory = torch.rand（10， 32， 512）tgt = torch.rand（20， 32， 512）out = transformer_decoder（tgt， memory）

注意確保 norm 參數(shù)保留為 None，以免在 Transformer 中用到 LayerNorm。