首次！無殘差連接或歸一化層，也能成功訓(xùn)練深度transformer。

盡管取得了很多顯著的成就，但訓(xùn)練深度神經(jīng)網(wǎng)絡(luò)（DNN）的實踐進展在很大程度上獨立于理論依據(jù)。大多數(shù)成功的現(xiàn)代 DNN 依賴殘差連接和歸一化層的特定排列，但如何在新架構(gòu)中使用這些組件的一般原則仍然未知，并且它們在現(xiàn)有架構(gòu)中的作用也依然未能完全搞清楚。

殘差架構(gòu)是最流行和成功的，最初是在卷積神經(jīng)網(wǎng)絡(luò)（CNN）的背景下開發(fā)的，后來自注意力網(wǎng)絡(luò)中產(chǎn)生了無處不在的 transformer 架構(gòu)。殘差架構(gòu)之所以取得成功，一種原因是與普通 DNN 相比具有更好的信號傳播能力，其中信號傳播指的是幾何信息通過 DNN 層的傳輸，并由內(nèi)核函數(shù)表示。

最近，使用信號傳播原則來訓(xùn)練更深度的 DNN 并且殘差架構(gòu)中沒有殘差連接和 / 或歸一化層的參與，成為了社區(qū)感興趣的領(lǐng)域。原因有兩個：首先驗證了殘差架構(gòu)有效性的信號傳播假設(shè)，從而闡明對 DNN 可解釋性的理解；其次這可能會實現(xiàn)超越殘差范式的 DNN 可訓(xùn)練性的一般原則和方法。

對于 CNN，Xiao et al. (2018)的工作表明，通過更好初始化提升的信號傳播能夠高效地訓(xùn)練普通深度網(wǎng)絡(luò)，盡管與殘差網(wǎng)絡(luò)比速度顯著降低。Martens et al. (2021) 的工作提出了 Deep Kernel Shaping (DKS)，使用激活函數(shù)轉(zhuǎn)換來控制信號傳播，使用 K-FAC 等強二階優(yōu)化器在 ImageNet 上實現(xiàn)了普通網(wǎng)絡(luò)和殘差網(wǎng)絡(luò)的訓(xùn)練速度相等。Zhang et al. (2022) 的工作將 DKS 擴展到了更大類的激活函數(shù)，在泛化方面也實現(xiàn)了接近相等。

信號傳播中需要分析的關(guān)鍵量是 DNN 的初始化時間內(nèi)核，或者更準(zhǔn)確地說，是無限寬度限制下的近似內(nèi)核。對于多層感知機（MLP）以及使用 Delta 初始化的 CNN，該內(nèi)核可以編寫為僅包含 2D 函數(shù)的簡單層遞歸，以便于進行直接分析。跨層 transformer 的內(nèi)核演化更加復(fù)雜，因此 DKS 等現(xiàn)有方法不適用 transformer 或?qū)嶋H上任何包含自注意力層的架構(gòu)。

在 MLP 中，信號傳播是通過查看（一維）內(nèi)核的行為來判斷的，而 transformer 中的信號傳播可以通過查看（高維）內(nèi)核矩陣在網(wǎng)絡(luò)層中的演化來判斷。

該研究必須避免一種情況：對角線元素隨深度增加快速增長或收縮，這與不受控制的激活范數(shù)有關(guān)，可能導(dǎo)致飽和損失或數(shù)值問題。避免秩崩潰（rank collapse）對于深度 transformer 的可訓(xùn)練性是必要的，而是否可以訓(xùn)練深度無殘差 transformer 仍是一個懸而未決的問題。

ICLR 2023 盲審階段的這篇論文解決了這個問題，首次證明了無需殘差連接或歸一化層時也可能成功訓(xùn)練深度 transformer。為此，他們研究了深度無殘差 transformer 中的信號傳播和秩崩潰問題，并推導(dǎo)出三種方法來阻止它們。具體而言，方法中使用了以下組合：參數(shù)初始化、偏置矩陣和位置相關(guān)的重縮放，并強調(diào)了 transformer 中信號傳播特有的幾種復(fù)雜性，包括與位置編碼和因果掩蔽的交互。研究者實證證明了他們的方法可以生成可訓(xùn)練的深度無殘差 transformer。

在實驗部分，在 WikiText-103 和 C4 數(shù)據(jù)集上，研究者展示了使用他們主要的方法——指數(shù)信號保持注意力（Exponential Signal Preserving Attention, E-SPA），可以通過延長大約五倍的訓(xùn)練時間使得標(biāo)準(zhǔn) transformer 與文中無殘差 transformer 的訓(xùn)練損失相當(dāng)。此外通過將這一方法與殘差連接結(jié)合，研究者還表明無歸一化層的 transformer 能夠?qū)崿F(xiàn)與標(biāo)準(zhǔn) transformer 相當(dāng)?shù)挠?xùn)練速度。

論文地址：https://openreview.net/pdf?id=NPrsUQgMjKK

對于這篇論文，Google AI 首席工程師 Rohan Anil 認為是 Transformer 架構(gòu)向前邁出的一大步，還是一個基礎(chǔ)性的改進。

構(gòu)造無捷徑可訓(xùn)練的深層 Transformer

迄今為止，糾正 Transformer 秩崩潰（rank collapse）的唯一策略依賴于殘差連接，該方式跳過了自注意力層固有的可訓(xùn)練性問題。與此相反，該研究直接解決這個問題。首先通過注意力層更好地理解信號傳播，然后根據(jù)見解（insights）進行修改，以在深度 transformer 中實現(xiàn)對忠實信號的傳輸，無論是否使用殘差連接，都可以對信號進行訓(xùn)練。

具體而言，首先，該研究對僅存在注意力的深度 vanilla transformer 進行了一下簡單設(shè)置，之后他們假設(shè)該 transformer 具有單一頭（h = 1）設(shè)置或具有多頭設(shè)置，其中注意力矩陣 A 在不同頭之間不會變化。如果塊 l≤L 初始化時有注意力矩陣 A_l，則最終塊的表示形式為 X_L：

對于上式而言，如果和采用正交初始化，那么就可以在初始化時正交。

在上述假設(shè)下，如果采用表示跨位置輸入核矩陣，經(jīng)過一些簡化處理后，可以得到如下公式：

從這個簡化公式（深度僅注意力 transformer 中的核矩陣）中，可以確定對 (A_l)_l 的三個要求：

必須在每個塊中表現(xiàn)良好，避免退化情況，如秩崩潰和爆炸 / 消失的對角線值；

A_l 必須是元素非負 ?l；

A_l 應(yīng)該是下三角?l，以便與因果掩碼注意力兼容。

在接下來的 3.1 和 3.2 節(jié)中，該研究專注于尋找滿足上述需求的注意力矩陣，他們提出了 3 種方法 E-SPA、U-SPA 和 Value-Skipinit，每種方法都用來控制 transformer 的注意力矩陣，即使在很深的深度也能實現(xiàn)忠實的信號傳播。此外，3.3 節(jié)演示了如何修改 softmax 注意力以實現(xiàn)這些注意力矩陣。

下圖中，該研究對提出的兩個 SPA 方案進行了驗證，U-SPA 和 E-SPA，結(jié)果顯示即使在網(wǎng)絡(luò)較深時也能成功地避免僅注意力 vanilla transformers 中的秩崩潰現(xiàn)象。

實驗

WikiText-103 基線：首先，該研究驗證了沒有殘差連接的標(biāo)準(zhǔn)深度 transformer 是不可訓(xùn)練的，即使它們有歸一化層 (LN) 和 transformed 激活，但本文的方法可以解決這個問題。如圖 2 所示，可以清楚地看到，從標(biāo)準(zhǔn) transformer 中移除殘差連接使其不可訓(xùn)練，訓(xùn)練損失穩(wěn)定在 7.5 左右。正如圖 1 所示，標(biāo)準(zhǔn) transformer 遭受了秩崩潰。

另一方面，該研究提出的 E-SPA 方法優(yōu)于 U-SPA 和 Value-Skipinit。然而，與本文無殘差方法相比，帶有殘差和 LN 的默認 transformer 仍然保持訓(xùn)練速度優(yōu)勢。

在表 1 中，該研究使用提出的方法評估了 MLP 塊中不同激活函數(shù)的影響，以及 LN 在無殘差 transformer 的使用。可以看到在深度為 36 處，本文方法針對一系列激活實現(xiàn)了良好的訓(xùn)練性能：DKS-transformed GeLU、TAT-transformed Leaky ReLU 以及 untransformed GeLU ，但不是 untransformed Sigmoid。通過實驗還看到，層歸一化對于訓(xùn)練速度而言相對不重要，甚至在使用 SPA 時對 transformed activation 的激活有害，因為 SPA 已經(jīng)具有控制激活規(guī)范的內(nèi)置機制。

在圖 3 中，我們看到一種不需要更多迭代就能匹配默認 transformer 訓(xùn)練損失的方法是使用歸一化殘差連接。

表 2 顯示帶有歸一化殘差和 LN 的 E-SPA 優(yōu)于默認的 PreLN transformer。

下圖 4(a)表明 E-SPA 再次優(yōu)于其他方法；4(b)表明訓(xùn)練損失差距可以通過簡單地增加訓(xùn)練時間來消除。

審核編輯 ：李倩