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

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

ENHANCING LINEAR ATTENTION WITH RESIDUALLEARNING

https://arxiv.org/pdf/2509.25223

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

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

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

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

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

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

這種循環形式在推理時每步保持恒定的時間和內存復雜度，并通過分塊并行算法促進了高效訓練 (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）這樣的模型能夠實現細粒度的記憶控制。Gated Delta Net（Yang et al., 2024a）則通過在學習過程中引入權重衰減，進一步增強了這種方法。

2.3 分解為預測與校正

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

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

3 方法

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

3.1 顯式殘差擬合

利用第2節中線性注意力的在線學習視角，我們對輔助狀態應用了類似的更新規則。這產生了以下循環過程：

3.2 自適應門控與校正因子

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

3.3 歸一化與殘差裁剪

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

3.4 最終公式

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

4 實驗 4.1 設置

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

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

4.2 主要結果

核效率：我們針對線性注意力基線和 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) 進行評估。我們的主要結果總結于表2，顯示我們提出的殘差學習變體 RLA 和 RDN 在各自的基線（sGLA 和 GDN）上，困惑度取得了持續改進。此外，我們的模型在多個基準測試中均優于其他領先的線性注意力方法，并能與標準 Transformer 的性能相競爭。

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

4.3 消融研究

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

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

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

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

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

5 相關工作

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

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

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

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

6 結論

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

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