殘差學習的線性注意力增強方法

殘差學習的線性注意力增強方法

ENHANCING LINEAR ATTENTION WITH RESIDUALLEARNING

https://arxiv.org/pdf/2509.25223

摘要
線性注意力以線性時間替代自注意力，卻常難以捕獲長程模式。本文以“預測–校正”視角重新審視線性注意力，證明主流變體均可寫成“歷史預測 + 單 token 校正”之和，從而造成表達能力瓶頸。為此，我們提出殘差線性注意力（RLA），通過顯式殘差擬合機制增強線性注意力：維護一條輔助遞推狀態(tài)，隨時間累積殘差誤差并修正基礎預測。我們進一步實例化出基于 delta 規(guī)則的版本——殘差 Delta 網(wǎng)絡（RDN），引入自適應門控與殘差裁剪以增強校正控制與穩(wěn)定性。實現(xiàn)上，我們復用高度優(yōu)化的線性注意力核，保持線性時間與內(nèi)存。在語言建模與強回憶任務中，RLA 與 RDN 均穩(wěn)定超越對應基線及現(xiàn)代線性注意力方法，縮小與標準 Transformer 的差距，同時維持線性擴展。

1 引言
Transformer 架構(gòu)已成為大語言模型標配，但其自注意力的二次復雜度仍是處理長序列的關(guān)鍵瓶頸。線性注意力通過將注意力改寫為遞推過程，實現(xiàn)線性時間訓練與推理，成為高效替代方案。RetNet、Mamba 等已展現(xiàn)競爭力；GLA、DeltaNet 等通過數(shù)據(jù)依賴門控與狀態(tài)更新規(guī)則進一步提升單狀態(tài)矩陣的信息流。現(xiàn)代線性注意力可被統(tǒng)一為“鍵到值的直接映射”，等價于測試時訓練。例如，delta 更新規(guī)則可由一步在線梯度下降導出，這為改進提供多條路徑：探索不同在線損失、設計更強映射函數(shù)、修改梯度更新機制等。TTT-MLP、Titans 等用 MLP 作深層記憶模塊，卻犧牲線性遞推，增加并行訓練難度。

本文提出對注意力輸出的新解讀：主流線性注意力的輸出可分解為“歷史狀態(tài)生成的基礎項”與“僅由當前 token 得到的校正項”（見 2.3 節(jié)）。依賴單 token 進行系統(tǒng)性校正造成瓶頸，損害模型表達力。為此，我們提出殘差線性注意力（RLA），通過輔助狀態(tài)矩陣顯式建模并修正基礎線性注意力的系統(tǒng)預測誤差，最終輸出為“基礎預測 + 學習到的誤差校正”。該框架可推廣至任意線性注意力方法，為構(gòu)建更強序列模型提供高效策略。

基于現(xiàn)有線性注意力，我們實現(xiàn)兩種殘差擬合變體：RLA 與 RDN，并在語言建模與強回憶任務上評估。實驗表明，二者均優(yōu)于對應基線及現(xiàn)代線性注意力方法；消融分析驗證了框架內(nèi)各關(guān)鍵設計的重要性。

2 預備知識
2.1 線性注意力作為遞推模型

softmax 注意力機制的計算復雜度與序列長度呈二次關(guān)系，在處理長序列時構(gòu)成顯著瓶頸。線性注意力（Katharopoulos 等，2020）通過移除 softmax 函數(shù)，允許對計算順序進行重排，從而解決這一問題。

這種循環(huán)形式在推理時每步保持恒定的時間和內(nèi)存復雜度，并通過分塊并行算法促進了高效訓練 (Yang et al., 2023)。此外，門控機制的應用催生了更多變體，例如 RetNet (Sun et al., 2023)、Lightning Attention (Qin et al., 2024a) 和 Mamba-2 (Dao & Gu, 2024)。

2.2 在線學習視角

該公式使得像Delta Net（Yang et al., 2024b; Schlag et al., 2021）這樣的模型能夠?qū)崿F(xiàn)細粒度的記憶控制。Gated Delta Net（Yang et al., 2024a）則通過在學習過程中引入權(quán)重衰減，進一步增強了這種方法。

2.3 分解為預測與校正

我們通過預測-校正的視角來解釋線性注意力。標準的線性注意力輸出 可以被視為來自過去狀態(tài)的基礎預測與基于當前令牌的校正之和：

基于預測-校正的視角，我們引入了一種殘差擬合框架來增強線性注意力。該框架通過學習一個更具表達力的校正項，來顯式擬合超出當前令牌的上下文信息。

3 方法

本節(jié)介紹我們提出的方法，該方法通過殘差擬合過程來增強線性注意力。我們首先描述支撐本方法的基礎性殘差學習框架。接著，引入自適應校正因子以增強建模能力，以及裁剪方法以穩(wěn)定殘差擬合過程。最后，我們給出該方法的兩種最終變體。

3.1 顯式殘差擬合

利用第2節(jié)中線性注意力的在線學習視角，我們對輔助狀態(tài)應用了類似的更新規(guī)則。這產(chǎn)生了以下循環(huán)過程：

3.2 自適應門控與校正因子

該公式使用衰減因子和校正因子，分別對從基礎狀態(tài)和輔助狀態(tài)的檢索進行動態(tài)門控。

3.3 歸一化與殘差裁剪

這確保了即使基礎模型產(chǎn)生暫時性的大預測誤差時，誤差校正狀態(tài) r t
也能保持穩(wěn)定的學習軌跡。該裁剪方法的詳細推導見附錄B。

3.4 最終公式

殘差擬合原理是一種通用技術(shù)，可與多種線性注意力骨干網(wǎng)絡相結(jié)合。通過將我們的殘差機制應用于標準的加法更新規(guī)則和 delta 更新規(guī)則，我們推導出了兩種強大的變體。這構(gòu)成了我們的最終模型：

4 實驗 4.1 設置

實現(xiàn)： 為了最大化效率，我們在 flash-linear-attention 庫（Yang & Zhang, 2024）的基礎上，使用 Triton（Tillet et al., 2019）實現(xiàn)了自定義的注意力核。我們利用了一個事實，即我們的狀態(tài)更新規(guī)則與線性注意力的規(guī)則相同，只需對其內(nèi)核進行微小修改：我們增強了它，使其同時返回注意力結(jié)果和中間殘差。這一設計使得同一個高度優(yōu)化的內(nèi)核可以在所有殘差擬合階段重復使用，從而確保了高吞吐量。

模型設置：我們通過與多個近期線性注意力架構(gòu)進行比較來評估我們的模型，包括 Retentive Network (RetNet) (Sun et al., 2023)、Mamba2 (Dao & Gu, 2024) 和 Gated Delta Net (GDN) (Yang et al., 2024a)。此外，我們通過評估標量門控線性注意力（sGLA），一個配備了查詢-鍵歸一化和標量門（α 和 β）的線性注意力變體，為 RLA 建立了一個基線。在我們的主要實驗中，我們將裁剪閾值設為 c = 1。所有模型都包含大約 15 億參數(shù)，并在相同條件下使用 1000 億令牌進行訓練，以確保公平比較。訓練配置的更多細節(jié)見附錄C。

4.2 主要結(jié)果

核效率：我們針對線性注意力基線和 FlashAttention（Dao et al., 2022; Dao, 2023）對我們的核的運行時間進行了基準測試，如圖2所示。盡管殘差擬合過程增加了計算開銷，但我們方法的運行時間隨序列長度呈線性增長。這使得其在更長的序列上比呈二次方增長的 FlashAttention 快得多。在吞吐量方面，我們的方法與其他線性注意力機制類似，保持近乎恒定的高吞吐量。相反，受計算限制的 FlashAttention 的吞吐量則隨著序列長度增加而迅速下降。

語言建模與常識推理：我們在 WikiText (Merity et al., 2016) 困惑度和一系列評估推理與常識理解的基準測試上評估了 RLA 和 RDN。推理任務包括 ARC-Easy、ARC-Challenge (Clark et al., 2018)、PIQA (Bisk et al., 2020) 和 MMLU (Hendrycks et al., 2020)，而常識理解則通過 HellaSwag (Zellers et al., 2019)、Winogrande (Sakaguchi et al., 2021)、SocialIQA (Sap et al., 2019) 和 LAMBADA (Paperno et al., 2016) 進行評估。我們的主要結(jié)果總結(jié)于表2，顯示我們提出的殘差學習變體 RLA 和 RDN 在各自的基線（sGLA 和 GDN）上，困惑度取得了持續(xù)改進。此外，我們的模型在多個基準測試中均優(yōu)于其他領(lǐng)先的線性注意力方法，并能與標準 Transformer 的性能相競爭。

召回密集型任務：為了評估記憶容量，我們在 Arora 等人 (2024) 的召回密集型任務上對我們的模型進行了基準測試。此外，我們還使用"大海撈針"任務 (NIAH) (gkamradt, 2023) 直接評估模型的檢索能力，該任務需要從長文檔中不同深度插入的鍵值對中進行檢索。如表3所示，這些基準對線性注意力模型具有挑戰(zhàn)性，因為其有限的狀態(tài)空間構(gòu)成了信息瓶頸。結(jié)果表明，我們提出的 RLA 和 RDN 始終優(yōu)于其對應的基線，尤其在 DROP 和 FDA 基準測試上提升顯著。此外，它們在 NIAH 任務上大幅優(yōu)于其他模型，突顯了其增強的信息召回能力。

4.3 消融研究

在本節(jié)中，我們進行了一系列消融研究以驗證關(guān)鍵組件的貢獻。我們首先量化了學習式殘差擬合方法相對于預定義校正的優(yōu)勢。接著，我們研究了使用專用校正因子的重要性，然后分析了使用門控機制來結(jié)合基礎預測與校正的必要性。最后，我們考察了歸一化和殘差裁剪的效果。

殘差擬合
為了驗證累積過去誤差的重要性，我們測試了一個使用更簡單的預定義校正項的變體。在此消融中，我們將我們持久的輔助狀態(tài)替換為僅從當前殘差推導出的無狀態(tài)校正項：

如表4所示，缺少顯式殘差擬合的變體性能不如我們的完整方法。盡管這個消融變體在某些基準測試上保持了有競爭力的性能，但其在訓練集和評估集上的困惑度均大幅增加。這種性能下降擴展到專業(yè)領(lǐng)域，其數(shù)學和代碼能力嚴重退化，這由 GSM8k (Cobbe et al., 2021) 和 HumanEval (Chen et al., 2021) 的困惑度衡量指標所證實。這表明輔助狀態(tài)在累積過去殘差以有效精煉模型輸出方面起著關(guān)鍵作用。

專用校正因子：通過比較我們的完整模型與校正因子 γ 受更新因子 β 約束的變體，我們分析了使用專用校正因子 γ 的優(yōu)勢。在圖3a中，具有獨立 γ 的模型始終獲得更低的評估損失，其中 RDN 變體顯示出更大的改進。這一趨勢延續(xù)到了下游性能，如圖3b的結(jié)果所示，該結(jié)果也表明專用校正因子能在多個基準測試中帶來性能提升。值得注意的是，我們無需額外 γ 的基礎架構(gòu)，相較于基線線性注意力方法，仍標志著顯著的改進。

歸一化與殘差裁剪：最后，我們研究了歸一化和殘差裁剪的重要性。我們在 RLA 上進行了消融研究，移除了歸一化和裁剪。如圖4所示，這兩個組件對于穩(wěn)定訓練都至關(guān)重要；移除它們會導致激活值無界和性能下降。相比之下，RDN 模型對殘差裁剪基本不敏感。這種穩(wěn)健性歸因于其 delta 規(guī)則更新固有的穩(wěn)定性，該更新在沒有殘差裁剪的情況下也能保持一致的損失曲線（圖4b）。

5 相關(guān)工作

序列建模歷史上由循環(huán)神經(jīng)網(wǎng)絡（RNNs）主導（Lipton et al., 2015），包括長短期記憶網(wǎng)絡（LSTM）（Hochreiter & Schmidhuber, 1997）和門控循環(huán)單元（GRU）（Cho et al., 2014）等變體。雖然有效，但其固有的順序性阻礙了訓練的并行化。Transformer 架構(gòu)（Vaswani et al., 2017）克服了這一限制，成為序列建模的事實標準。然而，其自注意力機制的計算復雜度與序列長度成二次方，對長上下文應用構(gòu)成了顯著瓶頸。

為了應對這些挑戰(zhàn)，近期研究重新審視了線性 RNNs，將其作為高效 Transformer 替代方案的基礎。通過將序列處理公式化為線性遞歸，這些模型實現(xiàn)了可并行訓練和線性時間推理。該領(lǐng)域的早期探索，如 S4 (Gu et al., 2021)、LRU (Orvieto et al., 2023) 和 RetNet (Sun et al., 2023)，使用了結(jié)構(gòu)化的狀態(tài)轉(zhuǎn)移矩陣。隨后通過引入數(shù)據(jù)相關(guān)的動態(tài)實現(xiàn)了性能飛躍。像 Mamba (Gu & Dao, 2023; Dao & Gu, 2024)、HGRN (Qin et al., 2023; 2024b) 和門控線性注意力 (Yang et al., 2023) 等模型利用依賴于輸入的動態(tài)門控來控制狀態(tài)轉(zhuǎn)移，從而增強了表達能力。

更先進的方法引入了 delta 學習規(guī)則，將狀態(tài)更新從簡單的門控衰減重新定義為細粒度的記憶校正。這種方法以 DeltaNet (Yang et al., 2024b; Schlag et al., 2021) 和 Gated DeltaNet (Yang et al., 2024a) 為代表，實現(xiàn)了更精確的動態(tài)記憶修改。該機制可以從在線學習的角度解釋，其中狀態(tài)更新被視為一個優(yōu)化過程，正如 TTT (Sun et al., 2024) 中所探索的那樣。這一觀點啟發(fā)了旨在發(fā)現(xiàn)和改進序列模型中內(nèi)在學習算法的進一步工作（von Oswald et al., 2023; 2025）。

并行的研究集中在增強狀態(tài)轉(zhuǎn)移的表達能力上。例如，RWKV-7 (Peng et al., 2025) 采用了對角線加低秩結(jié)構(gòu)，而 DeltaProduct (Siems et al., 2025) 通過為每個令牌執(zhí)行多個更新步驟來泛化 DeltaNet。為了進一步提升容量，近期的架構(gòu)如 Titans (Behrouz et al., 2024) 和 Miras (Behrouz et al., 2025) 引入了非線性深度記憶，用 MLP 對狀態(tài)進行參數(shù)化。

6 結(jié)論

本文介紹了殘差線性注意力，這是一個通過顯式殘差擬合過程來增強線性注意力模型的框架。我們的方法利用一個輔助狀態(tài)來校正基礎模型的預測誤差，從而構(gòu)建更穩(wěn)健、更準確的上下文表示。該框架具有高度適應性，可應用于多種線性注意力方法。我們的實驗證明了這種通用性，表明我們的方法在各項基準上均持續(xù)優(yōu)于其相應的基線。盡管這一改進是以額外計算開銷為代價的，但平衡這一權(quán)衡為未來研究提供了一個有前景的方向。

原文鏈接：https://arxiv.org/pdf/2509.25223