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

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

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

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

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

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

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

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

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

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

在實(shí)驗(yàn)部分，在 WikiText-103 和 C4 數(shù)據(jù)集上，研究者展示了使用他們主要的方法——指數(shù)信號(hào)保持注意力（Exponential Signal Preserving Attention, E-SPA），可以通過延長(zhǎng)大約五倍的訓(xùn)練時(shí)間使得標(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

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

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

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

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

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

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

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

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

A_l 必須是元素非負(fù) ?l；

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

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

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

實(shí)驗(yàn)

WikiText-103 基線：首先，該研究驗(yàn)證了沒有殘差連接的標(biāo)準(zhǔn)深度 transformer 是不可訓(xùn)練的，即使它們有歸一化層 (LN) 和 transformed 激活，但本文的方法可以解決這個(gè)問題。如圖 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 的默認(rèn) transformer 仍然保持訓(xùn)練速度優(yōu)勢(shì)。

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

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

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

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

審核編輯 ：李倩