Linear Attention/zh

Article
Topic area	Deep Learning
Prerequisites	Transformer, Softmax Function

概述

線性注意力是一類注意力機制，其時間和內存開銷隨序列長度線性擴展，與原始 Transformer 架構中所用標準 softmax 注意力的二次開銷形成對比。其核心思想是用一個可寫成特徵映射內積的核替代 softmax 注意力中的指數相似度函數，再利用矩陣乘法的結合律重排計算順序。其結果是一種注意力層，對於自回歸解碼而言，可表示為具有固定大小隱藏狀態的線性循環神經網絡，因而在長上下文語言建模、流式推理和端側部署中頗具吸引力。

線性注意力以犧牲部分 softmax 注意力的表達能力為代價換取漸近效率。這種取捨是否划算取決於工作負載：對於很長的序列與常數內存推理，它通常是明顯的勝利；而對於短到中等長度的上下文，常數因子和質量差距可能抵消其理論上的優勢。諸如門控線性注意力以及狀態空間式模型等現代變體已顯著縮小了質量差距，並構成了 2020 年代初期若干高效序列模型的基礎。

背景與動機

標準的縮放點積注意力對查詢 $ Q \in \mathbb{R}^{N \times d} $、鍵 $ K \in \mathbb{R}^{N \times d} $ 和值 $ V \in \mathbb{R}^{N \times d} $ 計算輸出

$ {\displaystyle \mathrm{Attention}(Q, K, V) = \mathrm{softmax}\!\left(\tfrac{Q K^\top}{\sqrt{d}}\right) V.} $

矩陣 $ Q K^\top $ 的形狀為 $ N \times N $，因此其時間和內存開銷均為 $ \mathcal{O}(N^2 d) $。對於長文檔、高解像度圖像或音頻波形，這種二次擴展成為訓練和推理的主要開銷，催生了大量關於高效注意力的文獻。線性注意力是其中最簡單也最具影響力的方法之一。

第二個動機來自自回歸解碼。標準 Transformer 解碼器必須在每一步關注所有先前的詞元，這就需要存儲一份隨上下文長度增長的鍵值緩存。線性注意力解碼器維護對過去的固定大小摘要，因此每詞元的解碼開銷在序列長度上是常數，且內存不會增長。

核重構

線性注意力的出發點是把（未歸一化的）softmax 相似度寫成核 $ k(q, k) = \exp(q^\top k / \sqrt{d}) $，然後用一個允許顯式特徵映射 $ \phi : \mathbb{R}^d \to \mathbb{R}^{d'} $ 的正定核加以替換，使得

$ {\displaystyle k(q, k) = \phi(q)^\top \phi(k).} $

代入並寫出對單個查詢 $ q_i $ 的注意力輸出，可得

$ {\displaystyle y_i = \frac{\sum_{j=1}^{N} \phi(q_i)^\top \phi(k_j)\, v_j}{\sum_{j=1}^{N} \phi(q_i)^\top \phi(k_j)} = \frac{\phi(q_i)^\top \sum_{j=1}^{N} \phi(k_j)\, v_j^\top}{\phi(q_i)^\top \sum_{j=1}^{N} \phi(k_j)}.} $

關鍵的一步是第二個等號，它利用矩陣乘法的結合律把 $ \phi(q_i) $ 提到求和號外。兩個求和

$ {\displaystyle S = \sum_{j=1}^{N} \phi(k_j)\, v_j^\top \in \mathbb{R}^{d' \times d}, \qquad z = \sum_{j=1}^{N} \phi(k_j) \in \mathbb{R}^{d'}} $

並不依賴於查詢，因此可在 $ \mathcal{O}(N d' d) $ 時間內預先計算並在所有查詢間重用。完整的注意力輸出此時按 $ \mathcal{O}(N d' d) $ 而非 $ \mathcal{O}(N^2 d) $ 計算，從而獲得名副其實的線性擴展。

同樣的技巧也適用於 softmax 的分母，它轉化為對 $ \phi(q_i)^\top z $ 的歸一化。一些實現完全省略分母，轉而依靠層歸一化來穩定輸出幅度。

特徵映射

$ \phi $ 的選擇同時決定了該層的表達能力與開銷。已有若干家族被提出：

• 逐元素的正映射。 Katharopoulos 等人引入了 $ \phi(x) = \mathrm{elu}(x) + 1 $，這是一種與輸入維度相同的廉價非負映射，是規範的「線性 Transformer」基線。
• 隨機特徵映射。 Performer 模型用正隨機特徵 $ \phi(x) \propto \exp(W x - \|x\|^2 / 2) $ 近似 softmax 核，其中 $ W $ 包含隨機正交投影。這在保持線性開銷分解的同時，在期望意義上恢復了 softmax 注意力。
• 多項式特徵映射。 取 $ \phi(x) = (1, x, x \otimes x, \dots) $ 得到一個多項式核；截斷到低次給出一種可處理的、表達能力受控的線性注意力變體。
• 恆等映射。 令 $ \phi(x) = x $ 把該層退化為普通雙線性注意力。這種做法很快，但除非配合歸一化或門控，否則通常表現不佳。

在實踐中，特徵映射很少是質量的瓶頸；更重要的設計選擇是歸一化、門控以及循環狀態如何更新。

循環形式與自回歸推理

對於因果（自回歸）建模，累積和 $ S_t $ 與 $ z_t $ 可增量更新：

$ {\displaystyle S_t = S_{t-1} + \phi(k_t)\, v_t^\top, \qquad z_t = z_{t-1} + \phi(k_t).} $

每個新詞元都對矩陣值狀態 $ S_t $ 施加一次秩一更新。在時刻 $ t $ 的輸出為

$ {\displaystyle y_t = \frac{\phi(q_t)^\top S_t}{\phi(q_t)^\top z_t}.} $

這正好是線性 RNN 的形式，其固定大小的隱藏狀態維度為 $ d' \times d $。每詞元的解碼開銷為 $ \mathcal{O}(d' d) $，與序列長度無關，也不存在不斷增長的鍵值緩存。推理時的內存消耗為常數，正是這一性質推動了對長上下文語言模型中線性注意力的重新關注。

訓練：並行形式

若依賴循環形式進行訓練會很慢，因為通過長時序循環進行反向傳播難以沿時間軸並行化。所幸，給出循環形式的同一套代數也給出了並行形式：累積和可以用前綴掃描計算，或者更常見地，把整個 $ N \times N $ 的注意力矩陣以適中大小的分塊物化，以便讓 GPU 保持忙碌。諸如 FLA（Flash Linear Attention 庫）的分塊並行形式以及 RetNet 中所用的核函數等現代實現，正是利用這一點在長序列上以可與 softmax Transformer相媲美的速度訓練，同時保持線性的漸近開銷。

因果掩碼帶來一個微妙之處：累積和必須遵守詞元的順序，因此先對所有 $ j $ 求和再施加掩碼的樸素向量化做法是錯誤的。正確的實現要麼使用前綴掃描，要麼把序列切分為分塊，並把塊內的 softmax 風格注意力與塊間的線性更新結合起來。

變體與擴展

大量後續模型都建立在線性注意力的基本思路之上：

• Performer。 使用隨機特徵對 softmax 核進行逼近，提供無偏估計以及對逼近方差的理論保證。
• 帶 elu+1 的線性 Transformer。 Katharopoulos 等人的原始公式，被廣泛用作基線。
• RetNet。 用衰減因子 $ \gamma \in (0, 1) $ 取代無界的累積和，得到 $ S_t = \gamma S_{t-1} + \phi(k_t) v_t^\top $。該衰減賦予模型多尺度保持的性質，並被證明等價於一種分塊計算。
• 門控線性注意力（GLA）。 用依賴於數據的門控替代標量衰減，在保留線性開銷的同時恢復了更多 softmax 注意力的選擇性行為。
• 選擇性狀態空間模型。 諸如 Mamba 的架構雖然嚴格來說並非線性注意力模型，卻共享線性循環結構，並可在一個密切相關的框架中加以表述。在近期文獻中，這兩個家族已大幅趨同。

與 softmax 注意力的比較

相對於標準的 softmax 注意力，線性注意力提供：

• 漸近開銷。 時間和內存上從 $ \mathcal{O}(N^2 d) $ 降為 $ \mathcal{O}(N d' d) $。這一優勢在大約 $ N \approx 2{,}000 $ 到 $ 8{,}000 $ 詞元處變得決定性，具體取決於硬件和常數因子。
• 常數推理內存。 沒有不斷增長的 KV 緩存；狀態形狀固定為 $ d' \times d $。
• 便於流式處理。 新詞元可通過秩一更新加入。

代價是：

• 表達能力降低。 Softmax 注意力可以以非常尖銳的分佈選擇性地聚焦於少量詞元；而線性注意力的有界秩狀態無法重現任意尖銳的模式。這在經驗上表現為更弱的關聯式召回和複製能力。
• 對特徵映射選擇的敏感性。 質量隨 $ \phi $、歸一化和門控顯著變化，因此樸素的替換往往表現不佳。

在實踐中，把 softmax 與線性注意力層交錯排列的混合架構是一種常見折中。

局限性與開放問題

線性注意力有界秩的循環狀態既是其定義性特徵，也是其主要局限。需要精確檢索任意過去詞元的任務，例如使用長示例集的上下文學習，最清楚地暴露了這一差距。若干工作試圖彌補缺失的容量：門控、多尺度衰減、更大的頭數，以及使用諸如 delta 規則等非線性更新規則。

第二個開放問題與硬件有關。循環形式緊湊但本質上是順序的，而並行形式則需要精心的核函數工程才能匹配融合的 softmax 注意力核函數的吞吐量。Flash Linear Attention 等庫已經大幅縮小了差距，但其實現層面的成熟度仍遜於標準注意力。

最後，線性注意力有時被批評為向 RNN 的倒退，確實，循環形式承襲了訓練深層循環網絡的全部經典困難。現代變體通過精心的初始化、層歸一化和衰減參數化來緩解這一問題，但該問題真實存在，在採用線性注意力構建新系統之前值得認真理解。