Transformer 架構(gòu)可以說(shuō)是近期深度學(xué)習(xí)領(lǐng)域許多成功案例背后的主力軍。構(gòu)建深度 Transformer 架構(gòu)的一種簡(jiǎn)單方法是將多個(gè)相同的 Transformer 「塊」（block）依次堆疊起來(lái)，但每個(gè)「塊」都比較復(fù)雜，由許多不同的組件組成，需要以特定的排列組合才能實(shí)現(xiàn)良好的性能。

自從 2017 年 Transformer 架構(gòu)誕生以來(lái)，研究者們基于其推出了大量衍生研究，但幾乎沒(méi)有改動(dòng)過(guò) Transformer 「塊」。

那么問(wèn)題來(lái)了，標(biāo)準(zhǔn) Transformer 塊是否可以簡(jiǎn)化？

在最近的一篇論文中，來(lái)自 ETH Zurich 的研究者討論了如何在不影響收斂特性和下游任務(wù)性能的情況下簡(jiǎn)化 LLM 所必需的標(biāo)準(zhǔn) Transformer 塊?；谛盘?hào)傳播理論和經(jīng)驗(yàn)證據(jù)，他們發(fā)現(xiàn)可以移除一些部分，比如殘差連接、歸一化層（LayerNorm）、投影和值參數(shù)以及 MLP 序列化子塊（有利于并行布局），以簡(jiǎn)化類似 GPT 的解碼器架構(gòu)以及編碼器式 BERT 模型。

對(duì)于每個(gè)涉及的組件，研究者都探討了是否可以在不降低訓(xùn)練速度的情況下將其移除（包括每次更新步驟和運(yùn)行時(shí)間），以及為此需要 Transformer 塊進(jìn)行哪些架構(gòu)修改。

論文鏈接：https://arxiv.org/pdf/2311.01906.pdf

Lightning AI 創(chuàng)始人、機(jī)器學(xué)習(xí)研究者 Sebastian Raschka 將這項(xiàng)研究稱為自己的「年度最愛(ài)論文之一」：

但也有研究者質(zhì)疑：「這很難評(píng)，除非我看過(guò)完整的訓(xùn)練過(guò)程。如果沒(méi)有歸一化層，也沒(méi)有殘差連接，如何能在大于 1 億參數(shù)的網(wǎng)絡(luò)中進(jìn)行擴(kuò)展？」

Sebastian Raschka 表示贊同：「是的，他們?cè)囼?yàn)的架構(gòu)相對(duì)較小，這是否能推廣到數(shù)十億參數(shù)的 Transformer 上還有待觀察?！沟匀槐硎具@項(xiàng)工作令人印象深刻，并認(rèn)為成功移除殘差連接是完全合理的（考慮到其初始化方案）。

對(duì)此，圖靈獎(jiǎng)得主 Yann LeCun 的評(píng)價(jià)是：「我們僅僅觸及了深度學(xué)習(xí)架構(gòu)領(lǐng)域的皮毛。這是一個(gè)高維空間，因此體積幾乎完全包含在表面中，但我們只觸及了表面的一小部分。」

為什么需要簡(jiǎn)化 Transformer 塊？

研究者表示，在不影響訓(xùn)練速度的前提下簡(jiǎn)化 Transformer 塊是一個(gè)有趣的研究問(wèn)題。

首先，現(xiàn)代神經(jīng)網(wǎng)絡(luò)架構(gòu)設(shè)計(jì)復(fù)雜，包含許多組件，而這些不同組件在神經(jīng)網(wǎng)絡(luò)訓(xùn)練動(dòng)態(tài)中所扮演的角色，以及它們之間如何相互作用，人們對(duì)此尚不清楚。這個(gè)問(wèn)題事關(guān)深度學(xué)習(xí)理論與實(shí)踐之間存在的差距，因此非常重要。

信號(hào)傳播理論（Signal propagation）已被證明具有影響力，因?yàn)樗軌蚣?lì)深度神經(jīng)網(wǎng)絡(luò)架構(gòu)中的實(shí)際設(shè)計(jì)選擇。信號(hào)傳播研究了初始化時(shí)神經(jīng)網(wǎng)絡(luò)中幾何信息的演化，通過(guò)跨輸入的分層表征的內(nèi)積來(lái)捕捉，在訓(xùn)練深度神經(jīng)網(wǎng)絡(luò)方面取得了許多令人印象深刻的成果。

然而，目前該理論只考慮初始化時(shí)的模型，而且往往只考慮初始前向傳遞，因此無(wú)法揭示深度神經(jīng)網(wǎng)絡(luò)訓(xùn)練動(dòng)態(tài)的許多復(fù)雜問(wèn)題，例如殘差連接對(duì)訓(xùn)練速度的助益。雖然信號(hào)傳播對(duì)修改動(dòng)機(jī)至關(guān)重要，但研究者表示，他們不能僅從理論上就得出簡(jiǎn)化的 Transformer 模塊，還要依靠經(jīng)驗(yàn)見(jiàn)解。

在實(shí)際應(yīng)用方面，考慮到目前訓(xùn)練和部署大型 Transformer 模型的高昂成本，Transformer 架構(gòu)的訓(xùn)練和推理流水線的任何效率提升都代表著巨大的潛在節(jié)約意義。如果能夠通過(guò)移除非必要組件來(lái)簡(jiǎn)化 Transformer 模塊，既能減少參數(shù)數(shù)量，又能提高模型的吞吐量。

這篇論文也提到，移除殘差連接、值參數(shù)、投影參數(shù)和序列化子塊之后，可以同時(shí)做到在訓(xùn)練速度和下游任務(wù)性能方面與標(biāo)準(zhǔn) Transformer 相匹配。最終，研究者將參數(shù)量減少了 16%，并觀察到訓(xùn)練和推理時(shí)間的吞吐量增加了 16%。

如何簡(jiǎn)化 Transformer 塊？

研究者結(jié)合信號(hào)傳播理論和經(jīng)驗(yàn)觀察，介紹了如何從 Pre-LN 模塊出發(fā)，生成最簡(jiǎn)單的 Transformer 塊（如下圖）。

在論文第四章的每一個(gè)小節(jié)，作者分別介紹了如何在不影響訓(xùn)練速度的情況下每次刪除一個(gè)塊組件。

這一部分的所有實(shí)驗(yàn)都在 CodeParrot 數(shù)據(jù)集上使用了一個(gè) 18-block 768-width 的因果僅解碼器類 GPT 模型，這個(gè)數(shù)據(jù)集足夠大，因此當(dāng)作者處于單個(gè)訓(xùn)練 epoch 模式時(shí)，泛化差距非常?。ㄒ?jiàn)圖 2），這使得他們可以專注于訓(xùn)練速度。

刪除殘差連接

研究者首先考慮刪除注意力子塊中的殘差連接。在公式（1）的符號(hào)中，這相當(dāng)于將 α_SA 固定為 0。簡(jiǎn)單地移除注意力殘差連接會(huì)導(dǎo)致信號(hào)退化，即秩崩潰（rank collapse），從而導(dǎo)致可訓(xùn)練性差。在論文 4.1 部分，研究者詳細(xì)解釋了他們的方法。

刪除投影 / 值參數(shù)

從圖 3 中可以得出結(jié)論，完全移除值和投影參數(shù) W^V、W^P 是可能的，而且每次更新的訓(xùn)練速度損失最小。也就是說(shuō)，當(dāng) β_V = β_P = 0 和 identity 初始化的

時(shí)，在相同的訓(xùn)練步數(shù)后，本研究基本上能達(dá)到 Pre-LN 塊的性能。在這種情況下，在整個(gè)訓(xùn)練過(guò)程中都有 W^V = W^P = I，即值和投影參數(shù)是一致的。作者在 4.2 節(jié)介紹了詳細(xì)方法。

刪除 MLP 子塊殘差連接

與上述幾個(gè)模塊相比，刪除 MLP 子塊殘差連接要更具挑戰(zhàn)性。與之前的研究一樣，作者發(fā)現(xiàn)，在使用 Adam 時(shí)，如果沒(méi)有 MLP 殘差連接，通過(guò)信號(hào)傳播使激活更加線性仍會(huì)導(dǎo)致每次更新訓(xùn)練速度的顯著下降，如圖 22 所示。

他們還嘗試了 Looks Linear 初始化的各種變體，包括高斯權(quán)重、正交權(quán)重或恒等權(quán)重，但都無(wú)濟(jì)于事。因此，他們?cè)谡麄€(gè)工作中使用標(biāo)準(zhǔn)激活（例如 ReLU）和 MLP 子塊中的初始化。

作者轉(zhuǎn)向并行 MHA 和 MLP 子塊的概念，這在幾個(gè)近期的大型 transformer 模型中已被證明很受歡迎，例如 PALM 和 ViT-22B。并行 transformer 塊如下圖所示。

作者在論文 4.3 節(jié)詳細(xì)介紹了移除 MLP 子塊殘差連接的具體操作。

刪除歸一化層

最后一個(gè)被刪除的是歸一化層，這樣就得到了圖 1 右上角的最簡(jiǎn)塊。從信號(hào)傳播初始化的角度來(lái)看，作者可以在本節(jié)簡(jiǎn)化的任何階段移除歸一化層。他們的想法是，Pre-LN 塊中的歸一化會(huì)隱式地降低殘差分支的權(quán)重，而這種有利的效果可以通過(guò)另一種機(jī)制在沒(méi)有歸一化層的情況下復(fù)制：要么在使用殘差連接時(shí)明確降低殘差分支的權(quán)重，要么將注意力矩陣偏向 identity / 將 MLP 非線性轉(zhuǎn)化為「更」線性。

由于作者在修改過(guò)程中考慮到了這些機(jī)制（如降低 MLP β_FF 和 Shaped Attention 的權(quán)重），因此無(wú)需進(jìn)行歸一化處理。作者在第 4.4 節(jié)介紹了更多信息。

實(shí)驗(yàn)結(jié)果

深度擴(kuò)展

鑒于信號(hào)傳播理論通常關(guān)注很大的深度，而這種情況下通常會(huì)出現(xiàn)信號(hào)退化。因此一個(gè)很自然的問(wèn)題就是，本文的簡(jiǎn)化 transformer 塊所提高的訓(xùn)練速度是否也能擴(kuò)展到更大的深度？

從圖 6 中可以觀察到，將深度從 18 個(gè)塊擴(kuò)展到 72 個(gè)塊后，本研究的模型和 Pre-LN transformer 的性能都得到了提高，這表明本研究中的簡(jiǎn)化模型不僅訓(xùn)練速度更快，而且還能利用更大的深度所提供的額外能力。事實(shí)上，在使用歸一化時(shí)，本研究中的簡(jiǎn)化塊和 Pre-LN 的每次更新軌跡在不同深度下幾乎沒(méi)有區(qū)別。

BERT

接下來(lái)，作者展示了他們的簡(jiǎn)化塊性能除了適用于自回歸解碼器之外，還適用于不同的數(shù)據(jù)集和架構(gòu)，以及下游任務(wù)。他們選擇了雙向僅編碼器 BERT 模型的流行設(shè)置，用于掩蔽語(yǔ)言建模，并采用下游 GLUE 基準(zhǔn)。

如圖 7 所示，在 24 小時(shí)運(yùn)行時(shí)內(nèi)，與（Crammed）Pre-LN 基線相比，本研究的簡(jiǎn)化塊可以媲美掩蔽語(yǔ)言建模任務(wù)的預(yù)訓(xùn)練速度。另一方面，在不修改值和投影的情況下刪除殘差連接再次導(dǎo)致訓(xùn)練速度的顯著下降。在圖 24 中，作者提供了 microbatch 步驟的等效圖。

此外，在表 1 中，研究者發(fā)現(xiàn)他們的方法在 GLUE 基準(zhǔn)上經(jīng)過(guò)微調(diào)后，性能與 Crammed BERT 基準(zhǔn)相當(dāng)。

他們?cè)诒?2 中對(duì)下游任務(wù)進(jìn)行了細(xì)分。為了進(jìn)行公平比較，他們使用了與 Geiping & Goldstein (2023) 相同的微調(diào)協(xié)議（5 個(gè) epoch、各任務(wù)超參數(shù)恒定、dropout regularisation）。

效率提升

在表 1 中，研究者還詳細(xì)列出了使用不同 Transformer 塊的模型在掩蔽語(yǔ)言建模任務(wù)中的參數(shù)數(shù)量和訓(xùn)練速度。他們以預(yù)訓(xùn)練 24 小時(shí)內(nèi)所采取的 microbatch 步驟數(shù)與基線 Pre-LN Crammed BERT 的比率計(jì)算了速度。結(jié)論是，模型使用的參數(shù)減少了 16%，SAS-P 和 SAS 的每次迭代速度分別比 Pre-LN 塊快 16% 和 9%。

可以注意到，在這里的實(shí)現(xiàn)中，并行塊只比 Pre-LN 塊快 5%，而 Chowdhery et al.（2022 ）觀察到的訓(xùn)練速度則快 15%，這表明通過(guò)更優(yōu)化的實(shí)現(xiàn)，整個(gè)訓(xùn)練速度有可能進(jìn)一步提高。與 Geiping & Goldstein（2023 年）一樣，此處實(shí)現(xiàn)也使用了 PyTorch 中的自動(dòng)算子融合技術(shù) (Sarofeen et al., 2022)。

更長(zhǎng)的訓(xùn)練

最后，考慮到當(dāng)前在更多數(shù)據(jù)上長(zhǎng)時(shí)間訓(xùn)練較小模型的趨勢(shì)，研究者討論了簡(jiǎn)化塊在長(zhǎng)時(shí)間訓(xùn)練后是否仍能達(dá)到 Pre-LN 塊的訓(xùn)練速度。為此，他們?cè)?CodeParrot 上使用圖 5 中的模型，并使用 3 倍 token 進(jìn)行訓(xùn)練。準(zhǔn)確地說(shuō)，是在批大小為 128、序列長(zhǎng)度為 128 的情況下進(jìn)行了約 120K 步（而不是 40K 步）的訓(xùn)練，這將導(dǎo)致約 2B 個(gè) token。

從圖 8 可以看出，當(dāng)使用更多的 token 進(jìn)行訓(xùn)練時(shí)，簡(jiǎn)化的 SAS 和 SAS-P 代碼塊的訓(xùn)練速度仍然與 PreLN 代碼塊相當(dāng)，甚至優(yōu)于 PreLN 代碼塊。

更多研究細(xì)節(jié)，可參考原論文。