打亂/跳過Transformer層會(huì)怎樣？最新研究揭開其信息流動(dòng)機(jī)制，一口氣解答8大問題

西風(fēng) 2024-07-26 13:30:21 來源：量子位

西風(fēng) 發(fā)自 凹非寺

量子位 | 公眾號(hào) QbitAI

Transformer中的信息流動(dòng)機(jī)制，被最新研究揭開了：

所有層都是必要的嗎？中間層在做同樣的事嗎？層的順序重要嗎？

如果跳過一些層，比如第4層輸出接到第6層會(huì)怎樣。隨機(jī)打亂層的順序，比如4-6-5-7又會(huì)怎樣。

最近一項(xiàng)名為“Transformer Layers as Painters”的研究火了，由來自AI初創(chuàng)公司Sakana AI、Emergence AI的研究團(tuán)隊(duì)完成。

他們從Transformer內(nèi)部工作原理出發(fā)，經(jīng)過一系列實(shí)驗(yàn)對(duì)以上問題得出了結(jié)論。團(tuán)隊(duì)表示深入理解這些原理不僅能提高現(xiàn)有模型利用效率，還能幫助改進(jìn)架構(gòu)開發(fā)新的變體。

谷歌DeepMind研究員、ViT作者Lucas Beyer看過后直接點(diǎn)了個(gè)贊：

很棒的總結(jié)！盡管一些實(shí)驗(yàn)在之前的研究中已經(jīng)被展示過了，但我喜歡你添加的新細(xì)節(jié)，特別是強(qiáng)調(diào)了“推理”類任務(wù)比其他任務(wù)受影響更大！

還有不少學(xué)者、工程師也表示強(qiáng)烈推薦。

敢打賭，其中一些見解最終將會(huì)用于改進(jìn)Transformer。

其中的實(shí)驗(yàn)再次證實(shí)了：復(fù)制層對(duì)創(chuàng)造性任務(wù)有幫助，但對(duì)推理任務(wù)通常無效；改變層的順序行不通；剪枝在中間層效果最佳，但仍需要進(jìn)行修復(fù)調(diào)整。

所以，在這項(xiàng)研究中，研究團(tuán)隊(duì)都進(jìn)行了哪些實(shí)驗(yàn)？回答了哪些問題？

實(shí)驗(yàn)?zāi)Ｐ瓦x擇和基準(zhǔn)

先來看一下實(shí)驗(yàn)配置～

實(shí)驗(yàn)在decoder-only和encoder-only模型上進(jìn)行。

其中decoder-only模型選擇的是Llama2，主要研究32層、70億參數(shù)的Llama2-7B，擴(kuò)展實(shí)驗(yàn)中也包含13B（40層）和70B（80層）模型。

encoder-only模型選擇的是BERT，有24層、3.4億參數(shù)。

研究人員使用這些模型的標(biāo)準(zhǔn)預(yù)訓(xùn)練checkpoints。在所有實(shí)驗(yàn)中，模型都是凍結(jié)的，除BERT的評(píng)估中包含一個(gè)標(biāo)準(zhǔn)的微調(diào)步驟，其它情況未通過微調(diào)等方法修改模型參數(shù)。

基準(zhǔn)測(cè)試方面，Llama2使用以下標(biāo)準(zhǔn)基準(zhǔn)：ARC（科學(xué)考試問題）、HellaSwag（常識(shí)問題）、GSM8K（數(shù)學(xué)題），WinoGrande（常識(shí)推理）、LAMBADA（詞匯預(yù)測(cè)）。其中LAMBADA用于測(cè)困惑度，與訓(xùn)練期間使用的原始token預(yù)測(cè)最接近。

對(duì)于Llama2的性能評(píng)估，提供了基準(zhǔn)測(cè)試的標(biāo)準(zhǔn)化中位數(shù)，將性能從0到1（模型最優(yōu)性能）進(jìn)行量化。

對(duì)于BERT，采用GLUE基準(zhǔn)并遵循其評(píng)估指標(biāo)，包括基準(zhǔn)的未標(biāo)準(zhǔn)化平均分。注意，標(biāo)準(zhǔn)的BERT評(píng)估包括一個(gè)微調(diào)步驟，因此對(duì)模型進(jìn)行了適應(yīng)性調(diào)整。在附錄中研究人員也展示了一個(gè)只有模型頭部可以調(diào)整的評(píng)估結(jié)果。

實(shí)驗(yàn)動(dòng)機(jī)最初源于這樣一個(gè)問題：

是否可以將多個(gè)層以某種方式合并成一個(gè)可能更大的單一層？假設(shè)可能由于訓(xùn)練過程中使用了殘差連接，神經(jīng)網(wǎng)絡(luò)的中間層可能使用了一個(gè)共同的表征空間。（對(duì)于標(biāo)準(zhǔn)的多層感知機(jī)來說不成立，它們之間沒有促使共同表征或?qū)娱g排列一致性的機(jī)制）

如果層能共享一個(gè)表征空間，將對(duì)后續(xù)條件計(jì)算或向預(yù)訓(xùn)練Transformer模型動(dòng)態(tài)添加新知識(shí)及下游應(yīng)用產(chǎn)生重要影響。

關(guān)于Transformer的8大問題

層是否使用相同的表征空間？

為確定不同層是否共享相同的表征空間，研究人員檢驗(yàn)了Transformer對(duì)于跳過特定層或更改相鄰層順序的魯棒性。

例如，在Llama2-7B模型中將輸出流從“第4層->第5層->第6層”的正常順序，改為“第4層->第6層”，跳過第5層，會(huì)怎樣？

又或者將第4層的輸出送到第6層，然后將第6層的輸出送到第5層，再送到第7層，會(huì)怎樣？

如下圖所示，實(shí)驗(yàn)發(fā)現(xiàn)除了最前面的和最末尾的幾層，Llama2-7B跳過或改變層序表現(xiàn)出很好的魯棒性。

也就是說，中間層共享一個(gè)表征空間，中間層與“外層”（最前面的和最末尾的幾層）具有獨(dú)立的表征空間。

為了進(jìn)一步證實(shí)這一假設(shè)，研究人員測(cè)量了不同模型（Llama2-7B、Llama2-13B和BERT-Large）中不同層的隱藏狀態(tài)激活之間的平均余弦相似度，并跨基準(zhǔn)測(cè)試進(jìn)行了比較。

下圖3展示了所有中間層之間的一致性。例如，底部第四層的激活與頂部第四層的激活高度相似。對(duì)于40層的Llama2-13B，可以看到這些層按相似性可劃分成4-5個(gè)組：第0層，1-3層，中間層，然后是最后一兩層。

這表明模型可能對(duì)“開始”、“中間”和“結(jié)束”層具有三個(gè)不同的表征空間。研究人員還發(fā)現(xiàn)，“開始層”的數(shù)量似乎隨著模型總層數(shù)的增加而增加。

此外，高余弦相似度可能證明有一個(gè)共享的表征空間，低相似度更能表明這些空間不是共享的。而上圖3中Llama2-7B的數(shù)據(jù)與圖2所示的性能結(jié)果高度一致，這進(jìn)一步證明了：

至少中間層的表征空間是共享的。

所有層都是必要的嗎？

為了進(jìn)一步驗(yàn)證中間層的表征空間真正共享，研究人員還進(jìn)行了層跳過實(shí)驗(yàn)（實(shí)驗(yàn)中未進(jìn)行任何微調(diào)）。

具體來說，將第N層的輸出直接傳遞為第N+M層的輸入（M>1），從而“跳過”了M-1層，如下圖所示。

原本第N+M層僅針對(duì)來自第N+M-1層的輸入進(jìn)行訓(xùn)練，那么現(xiàn)在它能否理解第N層的激活？

此類實(shí)驗(yàn)中，研究人員正常執(zhí)行第一層和最后N-1層，而跳過或修改第N+1到第T-N層（T是模型總層數(shù)）。

如下圖4，在多個(gè)基準(zhǔn)測(cè)試中，Llama2-7B和BERT-Large的性能均逐漸下降（圖從左至右展示了跳過層數(shù)逐漸遞增的變化）。這一結(jié)果揭示了：

不是所有層都是必要的，至少省略部分中間層不會(huì)對(duì)整體性能造成嚴(yán)重影響。

中間層是否都執(zhí)行相同的功能？

如果中間層共享一個(gè)共同的表征空間，這些層是否多余？

為了解答該問題，研究人員重新進(jìn)行了前面的“跳過”實(shí)驗(yàn)，但這次不是跳過中間層，而是用最中心層的權(quán)重替換了這些所有中間層的權(quán)重，如下圖所示。

實(shí)際上就是在最中心層上循環(huán)執(zhí)行了T-2N+1次，其中T是模型總層數(shù)（Llama2-7B為32層，BERT-Large為24層）。

結(jié)果基準(zhǔn)測(cè)試中，隨著被替換的層數(shù)增加，模型性能迅速下降。而且性能下降速度比僅僅跳過某些層要嚴(yán)重得多，這種權(quán)重替換極具破壞性。

因此，中間層各執(zhí)行不同的功能并非多余，中間層之間共享權(quán)重會(huì)產(chǎn)生災(zāi)難性后果。

層的順序重要嗎？

上面實(shí)驗(yàn)表明中間層雖共享表征空間，卻在該空間上執(zhí)行不同操作。那么這些操作順序重要嗎？研究人員進(jìn)行了兩組實(shí)驗(yàn)。

首先，將中間層按照與其訓(xùn)練順序相反的順序（逆序）執(zhí)行。將第T-N層的輸出傳遞給第T-N-1層，依此類推，直至第N層，然后將該層的輸出傳至最后的T-N層。

如下圖：

第二個(gè)實(shí)驗(yàn)，隨機(jī)排列中間層順序，并平均10個(gè)隨機(jī)種子結(jié)果。

結(jié)果如下圖，兩種情況模型都呈現(xiàn)出緩慢的性能下降。

這里劇透一下下面的一項(xiàng)實(shí)驗(yàn)結(jié)果，無論是逆序還是隨機(jī)順序，模型表現(xiàn)均優(yōu)于直接跳過這些層，說明即使層在非訓(xùn)練順序的輸入上運(yùn)行，依然能進(jìn)行有效輸出。

因此，層順序重要嗎？結(jié)論是：

層順序調(diào)整對(duì)性能有一定影響，隨機(jī)順序和逆序都表現(xiàn)出一定的性能退化。

值得注意的是，隨機(jī)順序性能優(yōu)于逆序?？赡苁且?yàn)槟嫘蚺c訓(xùn)練時(shí)的順序完全相反，而任何隨機(jī)順序都至少保持了一些順序上的連貫性（即某層i總在另一層j之后，其中i>j）。

可以并行運(yùn)行這些層嗎？

如果層的存在，即沒有被跳過，比它們執(zhí)行的順序更重要，那么是否可以考慮獨(dú)立地運(yùn)行這些層，然后將它們的結(jié)果合并？如下圖所示。

研究人員進(jìn)行了一個(gè)實(shí)驗(yàn)，不是跳過第N到第T-N層，而是并行運(yùn)行這些中間層，然后將它們的平均結(jié)果傳遞到最后的N層。

結(jié)果如下圖所示，除了GSM8K數(shù)學(xué)題基準(zhǔn)外，所有基準(zhǔn)測(cè)試都表現(xiàn)出緩慢的性能退化。

有趣的是，并行層的表現(xiàn)優(yōu)于跳過層，但不如逆序運(yùn)行層。

總之，可以并行運(yùn)行這些層嗎？答案是：可以，數(shù)學(xué)為主的基準(zhǔn)測(cè)試除外。

對(duì)于某些任務(wù)，順序是否更重要？

大多數(shù)變體（包括逆序、跳過和并行）在抽象推理ARC或數(shù)學(xué)推理GSM8K基準(zhǔn)中，表現(xiàn)出最快速的性能下降。

可以解釋為逐步推理任務(wù)比“語(yǔ)義”任務(wù)（如Winogrande或HellaSwag）對(duì)層順序的變化更為敏感。

這是因?yàn)橥评砣蝿?wù)需要結(jié)合結(jié)構(gòu)和語(yǔ)義雙重信息，而HellaSwag這類任務(wù)僅需語(yǔ)義即可完成。

通過實(shí)驗(yàn)，研究人員得出結(jié)論：數(shù)學(xué)和推理任務(wù)比“語(yǔ)義”任務(wù)更依賴順序。

迭代對(duì)并行層有幫助嗎？

如果把Transformer內(nèi)部運(yùn)行機(jī)制比作是畫一幅畫的過程：畫布（輸入）在一些畫家之間傳遞，一些畫家專門畫鳥，一些則更擅長(zhǎng)畫輪子……每個(gè)畫家都依次從另一位畫家手里接過畫布，然后決定對(duì)這幅畫進(jìn)行補(bǔ)充，還是將其直接傳遞給下一位畫家（使用殘差連接）。

可以想象，某些層在收到適當(dāng)?shù)妮斎霑r(shí)才會(huì)對(duì)畫作進(jìn)行“補(bǔ)充”。例如，如果“畫輪子”的畫家先看到汽車的車身，才更有可能畫上輪子。

在Transformer中，某些層可能只有在接收到適當(dāng)?shù)妮斎霑r(shí)才會(huì)對(duì)前向傳遞發(fā)揮作用，而不是通過殘差連接將輸入直接傳遞出去。

這么來看的話，那么相比于僅執(zhí)行一次并行層，迭代執(zhí)行并行層應(yīng)該會(huì)提高性能。

研究人員通過將并行層的平均輸出回饋到同一層并固定迭代次數(shù)來進(jìn)行測(cè)試，如下圖：

下圖9中，研究人員展示了并行層迭代3次的結(jié)果，這種方法顯著優(yōu)于僅執(zhí)行一次并行層。

唯一的例外是在Llama2-7B的起始層N為15或BERT的起始層N為11時(shí)。在這種情況下，循環(huán)并行3次的效果相當(dāng)于僅重復(fù)中間層3次，此時(shí)的并行層等同于完整模型。

研究人員還用不同的迭代次數(shù)重復(fù)進(jìn)行了實(shí)驗(yàn)。

下圖展示了Llama2-7B的性能隨并行層數(shù)M和迭代次數(shù)的變化而變化。

每個(gè)M的最佳迭代次數(shù)用紅框表示。除了M=29和M=31（幾乎并行所有層）外，最佳迭代次數(shù)大致與并行層數(shù)成線性比例。

因此結(jié)論是：迭代對(duì)并行層有幫助，最佳迭代次數(shù)與并行層數(shù)成比例。

哪些變體對(duì)性能損害最?。?/h1>

最后，研究人員將實(shí)驗(yàn)中的所有不同變體在同一圖表上進(jìn)行了比較。

結(jié)果顯示，重復(fù)單一層（如上面提到的用同等數(shù)量的最中心的層替換中間層）效果最差，性能迅速退化至隨機(jī)基準(zhǔn)線。

迭代并行和隨機(jī)層順序性能退化最小，其中迭代并行在BERT和Llama2-7B中表現(xiàn)最好。

論文附錄中還補(bǔ)充了更多實(shí)驗(yàn)結(jié)果，感興趣的家人們可以查看原論文。

論文鏈接：https://arxiv.org/abs/2407.09298v1
參考鏈接：https://x.com/A_K_Nain/status/1812684597248831912

版權(quán)所有，未經(jīng)授權(quán)不得以任何形式轉(zhuǎn)載及使用，違者必究。