盡管Transformer架構(gòu)已經(jīng)主宰了當(dāng)今幾乎所有的大模型，但我們依舊對(duì)它的工作原理知之甚少。

而且，基于Transformer的預(yù)訓(xùn)練LLM動(dòng)輒有幾十億參數(shù)，很難直接對(duì)模型進(jìn)行可解釋性分析。

同時(shí)，模型中間層由N個(gè)相同的塊堆疊在一起，它們之間唯一的區(qū)別只有層次位置和權(quán)重值，這就讓理解中間層更加困難。

然而，最近發(fā)表的一篇論文卻給出了一個(gè)十分通俗易懂的比喻——「畫(huà)家流水線」。

論文地址：https://arxiv.org/pdf/2407.09298v1

有著「東京AI夢(mèng)之隊(duì)」之稱(chēng)的Sakana AI，聯(lián)合IBM前AI負(fù)責(zé)人Satya Nitta創(chuàng)始的Emergence AI，兩個(gè)團(tuán)隊(duì)的研究人員用一種新的「打開(kāi)方式」來(lái)解釋Transformer架構(gòu)的中間層。

值得一提的是，這篇論文作者之一Llion Jones同樣也是當(dāng)年Transformer架構(gòu)的共同創(chuàng)建者之一。

那么，「畫(huà)家流水線」這個(gè)比喻該如何理解呢？

首先，輸入被看作是一張畫(huà)布，輸入通過(guò)N個(gè)組成中間層的塊的過(guò)程，就像是畫(huà)布在「畫(huà)家流水線」上進(jìn)行傳遞的過(guò)程。

有些畫(huà)家擅長(zhǎng)畫(huà)鳥(niǎo)，而有些畫(huà)家則更擅長(zhǎng)畫(huà)魚(yú)。每個(gè)畫(huà)家從前面的畫(huà)家手中接過(guò)畫(huà)布，然后決定是在畫(huà)上添幾筆，還是直接傳給后面的畫(huà)家。

在這個(gè)類(lèi)比中，非常重要的一點(diǎn)是，每個(gè)畫(huà)家都使用相同的「詞匯」來(lái)理解畫(huà)作，因此一個(gè)畫(huà)家可以在流水線上從前一個(gè)畫(huà)家手中接過(guò)畫(huà)作，但不會(huì)因?yàn)閷?duì)畫(huà)面理解不同而造成災(zāi)難。

畫(huà)家們也可以重新排序（調(diào)整圖層的前后順序），甚至可以同時(shí)添加筆觸，就像N個(gè)塊可以并行運(yùn)行。

這個(gè)類(lèi)比并不是一個(gè)嚴(yán)謹(jǐn)?shù)睦碚摚梢蕴峁┮粋€(gè)幫助我們思考Transformer層的有趣視角。

在這個(gè)類(lèi)比的啟發(fā)下，研究人員提出了一些假設(shè)，并通過(guò)實(shí)驗(yàn)來(lái)驗(yàn)證這些假設(shè)是否成立——

• 不同層是否使用相同的表征空間？
• 所有的層都是有必要的嗎？
• 中間層是否都在執(zhí)行相同的功能？
• 層的順序重要嗎？
• 我們能并行運(yùn)行各層嗎？
• 順序是否對(duì)與某些特定任務(wù)而言更重要
• 循環(huán)是否有助于并行層？
• 哪些變體對(duì)性能的損害最小？

實(shí)驗(yàn)

主要用于實(shí)驗(yàn)包括兩種預(yù)訓(xùn)練LLM，分別是decoder-only架構(gòu)的Llama2-7B，以及encoder-only架構(gòu)的BERT。Llama2-7B有70億個(gè)參數(shù)和32層（每層含2.02億個(gè)參數(shù)），BERT僅有24層和3.4億個(gè)參數(shù)。

在下述所有實(shí)驗(yàn)過(guò)程中，模型都是凍結(jié)的。除了對(duì)BERT進(jìn)行GLUE基準(zhǔn)測(cè)試時(shí)進(jìn)行了標(biāo)準(zhǔn)的微調(diào)步驟，參數(shù)沒(méi)有經(jīng)過(guò)任何修改。

評(píng)估過(guò)程采用了ARC（科學(xué)考試題）、HellaSwag（常識(shí)）、GSM8K（數(shù)學(xué)應(yīng)用題）、LAMBADA（單詞預(yù)測(cè)）等常用基準(zhǔn)。

其中LAMBADA任務(wù)可以衡量模型困惑度（perplexity），任務(wù)最接近預(yù)訓(xùn)練時(shí)的原始token預(yù)測(cè)。

結(jié)果發(fā)現(xiàn)，Transformer的中間層有一定程度的一致性，但不冗余，而且對(duì)數(shù)學(xué)、推理任務(wù)而言，各層的運(yùn)行順序比在語(yǔ)義任務(wù)中有更重要的影響。

各層「說(shuō)同一種語(yǔ)言」？

Transformer中的不同層是否共享相同的表示空間？

為了回答這個(gè)問(wèn)題，論文采用的方法是讓模型跳過(guò)特定層或調(diào)換相鄰層的順序，觀察會(huì)不會(huì)出現(xiàn)災(zāi)難性后果。

圖2中展示了Llama 2 7B在跳過(guò)或調(diào)換一些層后，模型整體在Open-LAMADA基準(zhǔn)上的表現(xiàn)。

可以看到，除了起始和末端的幾層，模型對(duì)這兩種架構(gòu)修改都表現(xiàn)出了相當(dāng)強(qiáng)的魯棒性。

因此可以得出初步結(jié)論：1)中間層共享同一個(gè)表示空間，2)表示空間與「外層」（第一層和最后幾層）不同。

為了進(jìn)一步驗(yàn)證，論文還進(jìn)入模型內(nèi)部，測(cè)量了不同層中隱藏狀態(tài)內(nèi)激活函數(shù)的余弦相似度（圖3），表明這種一致性在三個(gè)模型的所有中間層都成立。

上圖還可以很清晰看到，模型各層自然形成了4～5個(gè)不同的相似組，比如Llama 2 13B模型中分別是：第0層，1-3層、中間層，以及最后的1層或2層。

據(jù)此，Transformer中的所有層可以被大致分為三類(lèi)：起始層、中間層和結(jié)束層。

此外，圖3中的矩陣也能和圖2中的模型分?jǐn)?shù)相對(duì)應(yīng)，更能有力證明，中間層之間共享語(yǔ)義表達(dá)空間。

所有層都必要？

為了進(jìn)一步檢驗(yàn)中間層的重定向空間是否真正共享（除了具有接近的余弦相似性），研究人員嘗試跳過(guò)多個(gè)層。

也就是說(shuō)，將第N層的輸出直接送入第N+M層的輸入（其中M>1），從而「跳過(guò)」M-1層。

在不進(jìn)行任何微調(diào)的情況下，這個(gè)實(shí)驗(yàn)是要看看N+M層能否理解來(lái)自N層的激活，盡管它在訓(xùn)練中只接受了來(lái)自N+M-1層的輸入。

結(jié)果顯示，Llama2-7B和BERT-Large的許多基準(zhǔn)性能都出現(xiàn)了一定程度的下降。

那么，所有層都有必要嗎？這一問(wèn)題已經(jīng)有了答案。

No! 并非所有層都是必要的，至少有幾個(gè)中間層可以跳過(guò)，而不會(huì)發(fā)生災(zāi)難性故障。

左圖：Llama2-7B跳過(guò)N層～32-N層的基準(zhǔn)測(cè)試結(jié)果（歸一化）；右圖：BERT跳過(guò)N層～24-N 層的基準(zhǔn)測(cè)試結(jié)果（未歸一化）

中間層功能相同嗎？

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

為了驗(yàn)證這一點(diǎn)，研究人員重新進(jìn)行了上一小節(jié)的「跳過(guò)」實(shí)驗(yàn)。

但不同的是，這次不是直接跳過(guò)M個(gè)中間層，而是用模型最中心的的一層代替全部M個(gè)層（Llama是第16層，BERT是第12層），相當(dāng)于在這一層上循環(huán)T-2N+1次，其中T是層的總數(shù)。

結(jié)果表明，隨著被替換層數(shù)M的增加，基準(zhǔn)測(cè)試結(jié)果迅速下降。

在研究人員所嘗試的所有測(cè)試中，這一項(xiàng)測(cè)試的變化是最嚴(yán)重的，比直接跳過(guò)一些層還要嚴(yán)重得多。

因此，中間層功能相同嗎？這一問(wèn)題的答案是——

No! 在中間層之間共享權(quán)重是災(zāi)難性的，這表明中間層在執(zhí)行不同的功能。

用中心層替換M個(gè)中間層（左側(cè)經(jīng)過(guò)歸一化，右側(cè)未經(jīng)歸一化）

順序重要嗎？

之前的實(shí)驗(yàn)表明，中間層共享一個(gè)表征空間，但對(duì)這個(gè)空間執(zhí)行不同的操作。

那么另一個(gè)問(wèn)題來(lái)了——這些操作的執(zhí)行順序有多重要？

論文進(jìn)行了兩組實(shí)驗(yàn)來(lái)檢驗(yàn)這個(gè)問(wèn)題。首先，以與預(yù)訓(xùn)練完全相反的順序運(yùn)行中間層，如下圖所示：

第二組則是以隨機(jī)順序運(yùn)行中間層，最終結(jié)果是取10個(gè)隨機(jī)種子進(jìn)行實(shí)驗(yàn)后的均值。

圖6和圖7分別展示了中間層完全翻轉(zhuǎn)和隨機(jī)順序的結(jié)果，雖然都出現(xiàn)了一定程度的性能下降，但兩者的結(jié)果都優(yōu)于直接跳過(guò)的情況。

所以，中間層順序重要嗎？這一問(wèn)題的答案是——

比較重要。改變中間層的執(zhí)行順序，無(wú)論是隨機(jī)打亂或者完全翻轉(zhuǎn)，都會(huì)導(dǎo)致模型性能退化。

并行運(yùn)行

如果層本身的存在比它們的執(zhí)行順序更重要，那么我們是否可以獨(dú)立運(yùn)行各層，最后合并它們的結(jié)果呢？

比如像下圖中，將原本堆疊在一起的中間層展開(kāi)，并行運(yùn)行后取各層輸出的平均值，傳遞給最后的N個(gè)層。

實(shí)驗(yàn)結(jié)果顯示，GSM8K（數(shù)學(xué)應(yīng)用題）基準(zhǔn)中，模型性能有劇烈的變化，直線下降，其他基準(zhǔn)分?jǐn)?shù)的下滑則平緩得多。

我們暫且可以下這樣一個(gè)結(jié)論：并行運(yùn)行是可行的，但解決數(shù)學(xué)問(wèn)題除外。

要理解這種性能下降，可以用我們的「畫(huà)家流水線」進(jìn)行類(lèi)比：某些中間層只有在看到合適輸入時(shí)，才能對(duì)結(jié)果有所貢獻(xiàn)，就像一個(gè)擅長(zhǎng)畫(huà)車(chē)輪的畫(huà)家，只有在畫(huà)面上看到汽車(chē)車(chē)身時(shí)，才更有可能畫(huà)出輪子。

如果是這種情況，將中間層并行運(yùn)行的過(guò)程迭代多次應(yīng)該會(huì)提高性能。

如下圖所示，論文將多個(gè)并行層的平均輸出再作為輸入反饋回去，如此進(jìn)行一定次數(shù)的循環(huán)。

圖9顯示了循環(huán)3次的結(jié)果，與圖8中沒(méi)有循環(huán)的方案相比，性能曲線的確相對(duì)平緩，尤其是在圖右BERT模型未經(jīng)歸一化的分?jǐn)?shù)上更加明顯。

圖10更清楚直觀地展示了，并行的中間層數(shù)和循環(huán)次數(shù)如何影響性能，其中紅框圈出了每列上的最高值。

除了29層和31層（接近Llama 2 7B的總層數(shù)32）得出例外的結(jié)果，從5層到27層都呈現(xiàn)出一致的趨勢(shì)：最佳迭代次數(shù)大致與并行化層數(shù)呈線性比例。

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

將上述所有實(shí)驗(yàn)結(jié)果放到同一張圖中（圖11），我們就能比較不同變體對(duì)模型性能的影響程度。

左圖（Llama2）取各基準(zhǔn)的中值，右圖（BERT）取各基準(zhǔn)的平均值

「隨機(jī)化層順序」和「循環(huán)并行」分別在Llama2和BERT-Large上造成了最少的性能下降，「中間重復(fù)」方案（用中心層運(yùn)行多次代替整個(gè)中間層）則在兩個(gè)模型上都造成了最嚴(yán)重的滑坡。

討論

自從Transformer發(fā)布后，大多數(shù)工作都在關(guān)注架構(gòu)的修改和優(yōu)化，以達(dá)到性能提升或參數(shù)減少。這篇論文則提供了另一種視角，調(diào)查了層并行化和重用的影響。

基于「Transformer層即畫(huà)家」這個(gè)類(lèi)比，我們開(kāi)頭提出的幾個(gè)問(wèn)題都通過(guò)實(shí)驗(yàn)得到了答案，最后得到了3個(gè)有趣的發(fā)現(xiàn)：

• 所有Transformer層可以大致分為三類(lèi)：起始層、中間層和結(jié)束層，其中中間層占比最大；
• 中間層具有一定程度的一致性，但并不冗余；
• 與語(yǔ)義任務(wù)相比，各層的執(zhí)行順序?qū)?shù)學(xué)和推理任務(wù)更為重要。

為什么Transformer架構(gòu)面對(duì)各種架構(gòu)修改時(shí)能表現(xiàn)出如此強(qiáng)大的魯棒性？作者表示將在之后的工作中再深入研究。

一個(gè)可能的假設(shè)是，訓(xùn)練過(guò)程中的殘差連接是各層共享相同表征的必要條件。

我們已經(jīng)知道，殘差連接有助于解決梯度消失問(wèn)題，然而相比沒(méi)有殘差連接的Transformer，加上殘差會(huì)降低性能。

如果能在沒(méi)有殘差的Transformer上重新運(yùn)行上述架構(gòu)的變體，看看是否會(huì)破壞完全無(wú)殘差模型所取得的微薄收益，那將會(huì)非常有趣。

對(duì)于未來(lái)的其他工作，研究人員還計(jì)劃「解凍」模型，并研究Transformer是否需要（以及需要多長(zhǎng)時(shí)間）通過(guò)微調(diào)來(lái)適應(yīng)上述的架構(gòu)變化。

雖然本文的目的是更好地理解Transformer的中間層，而非引入新模型，但根據(jù)實(shí)驗(yàn)結(jié)果，中間層并行或者干脆跳過(guò)都可以用適度的準(zhǔn)確性損失換取更低的推理延遲。

作者團(tuán)隊(duì)

本文作者分別來(lái)自?xún)杉褹I初創(chuàng)公司：Sakana AI和Emergence AI。

Sakana AI在今年年初剛剛獲得3000萬(wàn)美元的種子輪融資，由Lux Capital領(lǐng)投，并得到了硅谷頂級(jí)風(fēng)投公司Khosla Ventures以及Jeaf Dean、Alexandr Wang等大佬的支持。

公司研發(fā)的重點(diǎn)是基于自然啟發(fā)的新型基礎(chǔ)模型，創(chuàng)始團(tuán)隊(duì)也是星光熠熠，一半成員來(lái)自「AI黃埔軍校」——谷歌大腦和DeepMind。

相比于關(guān)注基礎(chǔ)研究的Sakana，Emergence AI更關(guān)注應(yīng)用，專(zhuān)門(mén)從事LLM驅(qū)動(dòng)的multi-agent系統(tǒng)研發(fā)。

公司聯(lián)合創(chuàng)始Satya Nitta曾擔(dān)任IBM研究院「AI解決方案」領(lǐng)域的全球主管，其中的許多研究人員和工程師也同樣來(lái)自谷歌、Meta、微軟、亞馬遜和Allen AI等頂尖機(jī)構(gòu)。

Emergence上個(gè)月剛剛從Learn Capital獲得9720萬(wàn)美元的資金，以及額外的總計(jì)超過(guò)一億美元的信貸額度，未來(lái)的發(fā)展也是前途可期。