智東西
作者｜陳駿達
編輯｜云鵬

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

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

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

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

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

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

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

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

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

月之暗面的數(shù)十名研究員參與了這項研究，對項目貢獻最大的三位研究員包括Guangyu Chen、Yu Zhang（Kimi-Linear一作）、Jialin Su。月之暗面的三位聯(lián)合創(chuàng)始人楊植麟、吳育昕、周昕宇帶領(lǐ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è)計不僅在計算機視覺中取得巨大成功，也逐漸成為大型語言模型和各類Transformer系統(tǒng)中的基礎(chǔ)構(gòu)建模塊。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

在實驗方面，研究團隊將這一結(jié)構(gòu)集成到Kimi Linear架構(gòu)中進行驗證。

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

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

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

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

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