智東西
作者｜陳駿達(dá)
編輯｜云鵬

“一切都需要被重新思考，深度學(xué)習(xí)2.0時(shí)代即將來(lái)臨?！?/strong>在讀完下方這篇來(lái)自月之暗面的最新論文后，前OpenAI大牛、“推理模型之父”Jerry Tworek發(fā)出感嘆。

智東西3月16日?qǐng)?bào)道，今天，月之暗面發(fā)布論文，提前預(yù)覽了下一代模型的關(guān)鍵模塊——注意力殘差（Attention Residuals，簡(jiǎn)稱AttnRes）。論文的核心亮點(diǎn)在于對(duì)大模型中最基礎(chǔ)、但長(zhǎng)期被忽視的結(jié)構(gòu)——殘差連接（Residual Connection）進(jìn)行了重新設(shè)計(jì)。

在傳統(tǒng)Transformer架構(gòu)中，每一層都會(huì)將前一層的輸出與當(dāng)前層的計(jì)算結(jié)果直接相加，這種“等權(quán)累加”的方式雖然穩(wěn)定了深度網(wǎng)絡(luò)訓(xùn)練，但隨著模型層數(shù)不斷增加，所有歷史層的表示被簡(jiǎn)單疊加，容易導(dǎo)致信息被逐漸稀釋，深層網(wǎng)絡(luò)的表達(dá)能力受到限制。

AttnRes用注意力機(jī)制替代固定的殘差相加方式，讓每一層能夠從所有歷史層表示中動(dòng)態(tài)選擇最有用的信息。每一層相當(dāng)于一個(gè)查詢（Query），在之前層的表示中進(jìn)行檢索并分配權(quán)重，使模型能更靈活地利用不同層級(jí)的特征。

在實(shí)現(xiàn)層面，論文還提出了塊級(jí)注意力殘差（Block AttnRes）來(lái)解決大模型訓(xùn)練時(shí)的計(jì)算和通信開(kāi)銷問(wèn)題。

實(shí)驗(yàn)中，月之暗面將AttnRes集成到Kimi Linear架構(gòu)（總參數(shù)量48B/激活參數(shù)量3B）中，在1.4T token上進(jìn)行預(yù)訓(xùn)練。AttnRes緩解了預(yù)歸一化的稀釋問(wèn)題，產(chǎn)生更均勻的跨深度輸出幅度和梯度分布，從而緩解深層模型訓(xùn)練中的不穩(wěn)定問(wèn)題。

在相同計(jì)算預(yù)算下，引入AttnRes的模型在多個(gè)基準(zhǔn)測(cè)試中均取得穩(wěn)定提升，尤其是在數(shù)學(xué)推理、科學(xué)問(wèn)答以及代碼生成等需要復(fù)雜推理能力的任務(wù)上。

此外，Block AttnRes的驗(yàn)證損失為1.692，而基線為1.714，相當(dāng)于約1.25倍的計(jì)算效率優(yōu)勢(shì)。

這篇論文在大模型圈引發(fā)廣泛討論，月之暗面官方在X平臺(tái)發(fā)布的帖子閱讀量已逼近60萬(wàn)次。除了Jerry Tworek之外，OpenAI聯(lián)合創(chuàng)始人Andrej Karpathy也評(píng)論道：“隨機(jī)梯度下降（常用的優(yōu)化算法之一）也是一種ResNet，殘差流就是權(quán)重……看來(lái)我們還沒(méi)把‘Attention is All You Need’ 這句話完全按字面意思理解透。”

月之暗面的數(shù)十名研究員參與了這項(xiàng)研究，對(duì)項(xiàng)目貢獻(xiàn)最大的三位研究員包括Guangyu Chen、Yu Zhang（Kimi-Linear一作）、Jialin Su。月之暗面的三位聯(lián)合創(chuàng)始人楊植麟、吳育昕、周昕宇帶領(lǐng)團(tuán)隊(duì)完成了這項(xiàng)研究。

一、標(biāo)準(zhǔn)殘差連接，面臨“等權(quán)累加”之困

由知名華人AI學(xué)者何愷明等人提出的標(biāo)準(zhǔn)殘差連接，最早在ResNet中系統(tǒng)化提出，并迅速成為現(xiàn)代深度神經(jīng)網(wǎng)絡(luò)的核心結(jié)構(gòu)之一。

隨著深度學(xué)習(xí)的發(fā)展，這一設(shè)計(jì)不僅在計(jì)算機(jī)視覺(jué)中取得巨大成功，也逐漸成為大型語(yǔ)言模型和各類Transformer系統(tǒng)中的基礎(chǔ)構(gòu)建模塊。

其核心思想非常直觀：讓網(wǎng)絡(luò)學(xué)習(xí)“殘差”（即輸入與輸出之間的差值），而不是直接擬合完整映射。標(biāo)準(zhǔn)殘差連接中，每一層的輸入等于所有之前層輸出的累積和。殘差連接實(shí)際上是一種跨深度的信息聚合機(jī)制，所有層的貢獻(xiàn)權(quán)重都是1，也就是均勻加權(quán)。

在現(xiàn)代大模型中，殘差結(jié)構(gòu)通常還會(huì)與PreNorm（預(yù)歸一化）結(jié)合使用。具體做法是在每個(gè)子層計(jì)算之前先進(jìn)行層歸一化，然后再進(jìn)入注意力或前饋網(wǎng)絡(luò)模塊，最后再通過(guò)殘差連接將輸出加回原始表示。

這種結(jié)構(gòu)在Transformer模型中被廣泛采用，因?yàn)樗軌蜻M(jìn)一步改善訓(xùn)練穩(wěn)定性，使梯度在深層網(wǎng)絡(luò)中更容易傳播，并避免訓(xùn)練后期出現(xiàn)數(shù)值不穩(wěn)定的問(wèn)題。

標(biāo)準(zhǔn)殘差連接之所以能獲得廣泛采用，在于其解決了深層網(wǎng)絡(luò)訓(xùn)練困難的問(wèn)題，允許梯度繞過(guò)復(fù)雜變換，緩解了梯度消失的問(wèn)題。同時(shí)，殘差使每層對(duì)表示的修改是漸進(jìn)的，不會(huì)破壞已有表示。模型訓(xùn)練更穩(wěn)定、更容易優(yōu)化。

然而，這一方式也存在一些局限性。標(biāo)準(zhǔn)殘差連接中所有層的權(quán)重一致，但現(xiàn)實(shí)情況是不同層信息的重要性不同。理想情況下應(yīng)該采用可學(xué)習(xí)的權(quán)重。

此外，學(xué)習(xí)時(shí)，每層網(wǎng)絡(luò)只能看到壓縮后的狀態(tài)，導(dǎo)致許多早期信息被混合壓縮，無(wú)法單獨(dú)訪問(wèn)某一層輸出，導(dǎo)致信息不可逆的丟失。同時(shí)，后期層在模型的貢獻(xiàn)中變大，早期層被逐漸淹沒(méi)。

模型剪枝（prune）等方式對(duì)模型性能影響不大，這一現(xiàn)象也從側(cè)面證明，當(dāng)前廣泛采用的殘差結(jié)構(gòu)，使很多層貢獻(xiàn)很小。

正是在這樣的背景下，月之暗面的研究團(tuán)隊(duì)提出了一種新的思路：用習(xí)得的softmax注意力機(jī)制來(lái)替代固定的殘差累積，也就是AttnRes。

二、用注意力替代固定權(quán)重，讓每一層學(xué)會(huì)“檢索”歷史

在AttnRes中，每一層在計(jì)算時(shí)不再僅僅依賴上一層的輸出，而是可以對(duì)之前多個(gè)層的表示進(jìn)行一次注意力檢索。模型通過(guò)學(xué)習(xí)得到的一組權(quán)重，對(duì)不同層的輸出進(jìn)行加權(quán)組合，從而生成當(dāng)前層的輸入表示。

由于這些權(quán)重通過(guò)softmax歸一化，它們能夠反映不同層在當(dāng)前上下文中的相對(duì)重要性，使信息聚合過(guò)程更加靈活和自適應(yīng)。

這種機(jī)制帶來(lái)的一個(gè)直接好處是，模型能夠更有效地利用早期層的信息。由于每一層都可以直接訪問(wèn)歷史層的表示，早期信息不再被完全壓縮到單一狀態(tài)中，而是可以在需要時(shí)被重新調(diào)用。

這種“跨層檢索”的能力，使得深層網(wǎng)絡(luò)在信息利用上更加接近于Transformer在時(shí)間維度上的注意力機(jī)制。

對(duì)于每個(gè)token，完整AttnRes需要O(L2d)的算術(shù)運(yùn)算和O(Ld)的內(nèi)存來(lái)存儲(chǔ)層輸出。由于深度遠(yuǎn)小于序列長(zhǎng)度，算術(shù)成本適中。

當(dāng)然，如果讓每一層都對(duì)所有歷史層進(jìn)行注意力計(jì)算，也就是使用全注意力殘差，計(jì)算成本會(huì)迅速增加。為了解決這一問(wèn)題，研究團(tuán)隊(duì)進(jìn)一步提出了Block AttnRes（塊注意力殘差）的結(jié)構(gòu)設(shè)計(jì)。

該方法將網(wǎng)絡(luò)層劃分為若干塊，每個(gè)塊內(nèi)部保持常規(guī)計(jì)算方式，而跨塊之間則使用壓縮后的表示進(jìn)行注意力交互。

通過(guò)這種方式，模型在保持靈活信息檢索能力的同時(shí)，也能有效控制計(jì)算復(fù)雜度，使這一機(jī)制在大規(guī)模模型中仍然具有可行性。

▲從左到右：傳統(tǒng)殘差、全注意力殘差、塊注意力殘差

三、基礎(chǔ)設(shè)施協(xié)同優(yōu)化，計(jì)算效率達(dá)1.25倍

不過(guò)，Block AttnRes也給基礎(chǔ)設(shè)施帶來(lái)新的挑戰(zhàn)。

在訓(xùn)練階段，Block AttnRes帶來(lái)更多的流水線并行中通信開(kāi)銷。標(biāo)準(zhǔn)殘差結(jié)構(gòu)在流水線階段之間只需傳輸固定大小的隱藏狀態(tài)，而B(niǎo)lock AttnRes需要在每個(gè)階段訪問(wèn)此前所有塊的表征。

如果采用樸素實(shí)現(xiàn)，就必須在階段轉(zhuǎn)換時(shí)重復(fù)傳輸完整的歷史塊表示，通信量會(huì)隨著塊數(shù)量累積呈平方增長(zhǎng)。為了解決這一問(wèn)題，月之暗面引入跨階段緩存（cross-stage cache）機(jī)制，大幅減少冗余通信。

同時(shí)，每個(gè)塊只需在虛擬階段中存儲(chǔ)一次，結(jié)合激活檢查點(diǎn)機(jī)制后，每層激活內(nèi)存占用與標(biāo)準(zhǔn)Transformer基本一致，因此整體訓(xùn)練的額外開(kāi)銷很小，在實(shí)際系統(tǒng)中端到端訓(xùn)練時(shí)間增加不到4%。

在推理階段，Block AttnRes的主要挑戰(zhàn)來(lái)自跨層塊表征的重復(fù)訪問(wèn)和長(zhǎng)上下文緩存帶來(lái)的內(nèi)存壓力。為此他們采用了兩階段計(jì)算策略，首先在階段一中對(duì)緩存的塊表征執(zhí)行一次批量查詢，為同一塊內(nèi)的所有層同時(shí)計(jì)算塊間注意力，從而將原本每層都需要訪問(wèn)歷史塊的操作合并為一次矩陣計(jì)算，顯著降低內(nèi)存讀取次數(shù)。

隨后在階段二中按順序計(jì)算塊內(nèi)各層的注意力，并通過(guò)在線softmax將結(jié)果與階段一的輸出合并，這一階段可以與周圍算子進(jìn)行內(nèi)核融合，從而進(jìn)一步減少I/O開(kāi)銷。

此外，為緩解長(zhǎng)上下文預(yù)填充時(shí)塊表征緩存過(guò)大的問(wèn)題，系統(tǒng)將這些表征沿序列維度在張量并行設(shè)備之間進(jìn)行分片，使每個(gè)設(shè)備只存儲(chǔ)部分序列。通過(guò)兩階段計(jì)算和序列分片的結(jié)合，BlockAttnRes在推理時(shí)僅帶來(lái)不到2%的延遲開(kāi)銷，同時(shí)顯著降低了長(zhǎng)上下文場(chǎng)景下的顯存需求。

在實(shí)驗(yàn)方面，研究團(tuán)隊(duì)將這一結(jié)構(gòu)集成到Kimi Linear架構(gòu)中進(jìn)行驗(yàn)證。

實(shí)驗(yàn)表明，在相同參數(shù)規(guī)模下，引入Block AttnRes后，模型在多個(gè)基準(zhǔn)測(cè)試中都取得了性能提升。其中提升最多的是多步推理基準(zhǔn)測(cè)試GPQA-Diamond，采用AttnRes的模型得分提升幅度超過(guò)了20%。

與此同時(shí)，該方法還表現(xiàn)出良好的計(jì)算效率，Block AttnRes在損失函數(shù)上與使用1.25倍計(jì)算量訓(xùn)練的基線模型相當(dāng)。

結(jié)語(yǔ)：將注意力引入“網(wǎng)絡(luò)深度維度”，有望成為下一代模型關(guān)鍵架構(gòu)

通過(guò)將注意力機(jī)制引入深度維度，Attention Residuals為深層網(wǎng)絡(luò)提供了一種更加靈活的信息聚合方式，使模型能夠動(dòng)態(tài)選擇歷史表示，而不是被動(dòng)接受均勻累積的結(jié)果。這一思路有望進(jìn)一步提升深層模型的信息利用效率。

從架構(gòu)演進(jìn)的角度來(lái)看，這項(xiàng)工作也提出了一個(gè)值得關(guān)注的思路：在過(guò)去，注意力機(jī)制主要用于解決序列維度上的信息交互問(wèn)題，而AttnRes則將注意力引入到了“網(wǎng)絡(luò)深度維度”，讓模型能夠在不同層之間建立更加靈活的信息通路，可能為下一代大模型架構(gòu)演進(jìn)提供新的方向。