當 DeepSeek 的 NSA 與月之暗面的 MoBA 以稀疏注意力掀起長序列技術(shù)熱潮，行業(yè)對 “效率革命” 的追逐迎來關(guān)鍵一躍 —— 華為諾亞方舟實驗室正式發(fā)布全新 ESA 算法（Efficient Selective Attention）。

論文地址：https://arxiv.org/pdf/2502.14477

通過稀疏化注意力的創(chuàng)新設(shè)計，ESA 突破了大模型在長文本處理中的瓶頸。ESA 不僅實現(xiàn)了數(shù)倍序列長度的拓展，還引入獨創(chuàng)的動態(tài)計算范式，結(jié)合鄰域影響力有效避免了單純選擇 top-ranked token 所帶來的性能損失。通過對關(guān)鍵 token 的精確選擇，ESA 在優(yōu)化長序列處理效率的同時，提升了計算性能，為大模型在長序列任務(wù)中的應(yīng)用帶來了新的可能性。

在大語言模型的推理過程中，長序列模型的訓練需要極高的算力和海量數(shù)據(jù)支持，理想的解決方案是通過短序列的訓練成果外推到長序列。然而，隨著序列長度的增加，注意力計算的復(fù)雜度呈平方級增長，這使得高效且準確的長序列推理成為了一大挑戰(zhàn)。為此，研究人員提出了多種方法，以應(yīng)對這一挑戰(zhàn)。

ESA 方案正是在這一背景下提出的創(chuàng)新外推解決方案。ESA 通過對 query 和 key 的低維壓縮，有效減少了 token 選擇的計算復(fù)雜度。該方案通過靈活高效地選擇關(guān)鍵 token 進行注意力計算，大幅度降低了 LLMs 在處理長文本時的計算負擔，且在性能上與全注意力外推方法相當，甚至在高倍外推場景下優(yōu)于全注意力算法，實現(xiàn)了上下文長度的有效擴展。

當大模型訓練長度有限，隨著序列長度的增長，一方面會出現(xiàn) OOD (out-of-distribution) 的問題，另一方面注意力計算量會迅速增大?，F(xiàn)有的研究表明，注意力矩陣具有稀疏性，對于長序列而言，稀疏程度進一步擴大。選擇性注意力（Selective Attention）利用了稀疏性這一特性，選擇部分 token 來計算注意力，結(jié)合外推的位置編碼能將短序列模型應(yīng)用到長序列任務(wù)上的同時，顯著降低計算量。在計算稀疏注意力時細粒度的 token 選擇方法能夠更加靈活、精準地定位到關(guān)鍵信息。然而，token 粒度選擇會引入巨大的計算開銷。這引出了一個核心的問題：如何在選擇性注意力方法中平衡靈活性與效率。針對這一挑戰(zhàn)，ESA 方法通過將 query 和 key 進行低維壓縮，顯著降低 token 選擇的計算復(fù)雜度，在外推場景下實現(xiàn) token 粒度動態(tài)稀疏注意力機制。

具體而言，ESA 包括以下兩個核心步驟：

高效選擇：ESA 引入了一種基于 query 感知的 token 粒度選擇機制，基于壓縮后的 query 和 key 計算 token 的重要性分數(shù)，同時考慮周圍 token 的影響（鄰距影響力），以避免直接選擇 top-ranked token 導(dǎo)致的性能下降。

注意力計算：在選擇關(guān)鍵 token 后，ESA 使用被選中的 token 的完整的 query 和 key 進行注意力計算，而非對所有前序 token 進行計算，從而大幅降低復(fù)雜度。

ESA 的主要創(chuàng)新點在于通過 token 粒度選擇性注意力機制，在保持模型準確率的同時顯著降低計算復(fù)雜度。具體來說，與現(xiàn)有的長序列外推方法不同，ESA 提出了一種基于 token 的細粒度選擇注意力，能夠在 prefilling 和 decoding 階段動態(tài)選擇最關(guān)鍵的少量 token，而不是固定 block 選擇或者永久丟棄不重要的 token。首先，ESA 將 query 和 key 經(jīng)過簡單的一層 MLP 壓縮到原有維度的大約 3.2%，在低維空間計算重要性分數(shù)，顯著降低計算復(fù)雜度；其次，根據(jù)重要性分數(shù)選擇 topk 的 token，控制 key 的長度是固定的，這樣將注意力計算由原有的平方復(fù)雜度降低為線性復(fù)雜度。雖然選擇 token 是平方復(fù)雜度，但是由于將 query 和 key 壓縮到了更低維的空間，使得對于算力要求大大降低。

ESA 的具體實現(xiàn)方式如下：輸入序列的 token 被分為 4 部分，注意力包括全局注意力和 window 的局部注意力，初始 token 和 ESA 選擇的 topk 中間 token 拼接起來計算全局注意力，localtoken 用于計算 window 的注意力，兩部分注意力進行融合計算最終的注意力。ESA 按照 chunked-prefill 緩存 key 和 value，即基于當前 chunk 的 query 選擇重要的中間 tokens，計算 token 的重要性時兼顧當前的所有 query；在解碼階段，只需要考慮當前的一個 token 的 query 即可。如果計算中間某個 token 重要性，需要計算和當前所有 token 的重要性，其中單個 token 的重要性用 query 和 key 的點積表示：

這里 H 是 head 的數(shù)量，為了降低復(fù)雜度 ESA 整合了所有的 head。為了進一步降低計算復(fù)雜度，不要求準確計算重要性分數(shù)，而是更關(guān)注相對大小，ESA 將 query 和 key 分別通過一層 MLP 進行壓縮。ESA 采取 offline 的方式學習 MLP 的權(quán)重：

ESA 使用一個小的校準數(shù)據(jù)集用模型進行推理，保存中間的 query、key 和 value，用于訓練降維 MLP，只增加了極少量的降低 query 和 key 大小的網(wǎng)絡(luò)權(quán)重，且無需對模型微調(diào)。

為了確保分數(shù)的相對大小，避免某個 token 在重要性分數(shù)中占據(jù)主導(dǎo)地位，ESA 對分數(shù)進行修正：

進一步的，作者發(fā)現(xiàn)僅選擇 topk 的 token 模型在大海撈針任務(wù)中只能檢索到部分信息，提出了鄰距影響力的概念，即對于某個中間的 token，其重要性分數(shù)不僅取決于自身的分數(shù)，還受到周圍 token 的影響，更新后的分數(shù)為：

在選擇完重要 token 后，ESA 使用完整的 query、key 和 value 計算注意力，最終的注意力輸出如下所示：

ESA 的計算復(fù)雜度降低主要來源于低維的 query 和 key 計算重要性分數(shù)以及選擇完成以后的線性注意力計算復(fù)雜度，經(jīng)過理論計算，一步 attention 計算在長序列場景下能降低為原有的：

實際實驗中我們將 query 和 key 壓縮為原有的 3.2%，一步 attention 計算量在輸入序列足夠長時理論能降低至 1.6% 左右。

論文選擇開源訓練集 Pile 的 2 條 Books3 樣本收集用于訓練降維 MLP 的 qk 樣本，query 和 key 從 4096 壓縮為 128，壓縮比例約為 l3.2%，注意力計算的窗口長度約為 6k。為了將開源的短序列模型應(yīng)用到長序列中，ESA 沿用了 Infllm 的外推位置編碼設(shè)置，使用 Llama-3-8B-Instruct 和 Mistral-7B-Instruct-v0.2，在多個公開的長序列基準測試中驗證了 ESA 的性能，包括 Longbench、InfiniteBench、NeedleBench 等。作者對比了 full attention 的外推方法和同類型的基于 window 的外推方法，且同類型方法的 window 長度一致。實驗結(jié)果表明，ESA 通過高效靈活選擇重要的 token，總體性能在外推倍數(shù)足夠大時候優(yōu)于 full attention 的方法，且均明顯優(yōu)于同類型的方法，尤其在 multi needles 檢索場景下例如數(shù)星星和 NeedleBench，在其他同類型方法失效的時候，ESA 仍然有較高的準確率。

ESA 不對每個 head 單獨選擇 token，而是將所有 head 整合到一起計算重要性分數(shù)，有利于降低計算復(fù)雜度，提升效率，為了驗證這一操作對算法的影響，作者做的對比實驗如下所示，可以看出這樣的整合對于算法影響有限。

論文研究了鄰距影響力的超參數(shù)影響，結(jié)果如下所示，對不同的測評集該參數(shù)的影響不同，取值較小有利于 multi needles 類型的檢索任務(wù)，取值較大則有利于 single needle 類型任務(wù)，這可能是由于單針檢索任務(wù)只需要關(guān)注 ground truth 所在的片段即可，增大鄰距影響力有利于 attention 集中到較長的片段上。

ESA 有效平衡了長序列外推場景下的選擇性注意力中的靈活性和計算效率，用于在不進行模型參數(shù)增量微調(diào)的情況下擴展上下文長度。ESA 的核心思想是在每個步驟中選擇固定數(shù)量的最重要 token 來計算注意力，利用注意力矩陣的稀疏性。當輸入序列足夠長時，ESA 通過將 query 和 key 壓縮為低維表征，有效降低選擇 token 的計算復(fù)雜度。實驗評估表明，ESA 能夠有效處理長度為訓練長度 4 倍甚至 25 倍的各種長序列任務(wù)。未來的研究需要探索更準確、更高效的選擇重要 token 的方法，以及軟硬件協(xié)同的高效外推方案。