稀疏注意力再添一員，華為諾亞推出高效選擇注意力架構ESA

為大模型在長序列任務中的應用帶來了新的可能性。

原標題：稀疏注意力再添一員，華為諾亞推出高效選擇注意力架構ESA
文章來源：智猩猩GenAI
內容字數：6412字

華為諾亞方舟實驗室提出高效選擇性注意力算法ESA

本文介紹了華為諾亞方舟實驗室提出的高效選擇性注意力算法ESA (Efficient Selective Attention)，該算法旨在解決大語言模型(LLMs)處理長文本序列時的效率瓶頸。ESA通過創新性的稀疏注意力機制，在顯著降低計算復雜度的同時，保持了模型的準確性，實現了上下文長度的有效擴展。

1. 長序列處理的挑戰

隨著序列長度的增加，注意力計算的復雜度呈平方級增長，這使得高效且準確的長序列推理成為一大挑戰。現有方法試圖通過短序列訓練成果外推到長序列，但面臨著OOD(out-of-distribution)問題和計算量迅速增大的難題。選擇性注意力（Selective Attention）利用注意力矩陣的稀疏性，選擇部分token進行計算，但token粒度選擇會引入巨大的計算開銷。

2. ESA算法的核心思想

ESA算法的核心在于平衡選擇性注意力方法中的靈活性與效率。它通過以下兩個核心步驟實現高效的長序列處理：

1. 高效選擇：ESA采用基于query感知的token粒度選擇機制。通過對query和key進行低維壓縮(約3.2%)，計算token的重要性分數，并考慮周圍token的影響(鄰距影響力)，避免僅選擇top-ranked token導致的性能下降。
2. 注意力計算：在選擇關鍵token后，ESA使用被選中token的完整query和key進行注意力計算，而非對所有前序token進行計算，從而大幅降低復雜度。

ESA的具體實現包括將輸入序列的token分為四部分：全局注意力和局部注意力。初始token和ESA選擇的topk中間token拼接起來計算全局注意力，局部token用于計算窗口注意力，兩部分注意力融合計算最終的注意力。ESA采用chunked-prefill緩存key和value，動態選擇關鍵token。

3. ESA算法的創新點

ESA的主要創新在于其基于token的細粒度選擇性注意力機制。它能夠在prefilling和decoding階段動態選擇最關鍵的少量token，而不是固定block選擇或永久丟棄不重要的token。通過低維壓縮query和key，將注意力計算復雜度從平方級降低到線性級。此外，ESA引入了鄰距影響力，考慮周圍token的影響，進一步提升了模型性能。

4. 實驗結果與分析

論文在Longbench、InfiniteBench、NeedleBench等公開的長序列基準測試中驗證了ESA的性能。結果表明，ESA在高倍外推場景下優于全注意力算法，尤其在multi needles檢索場景下表現突出。實驗還分析了鄰距影響力超參數的影響，以及將所有head整合計算重要性分數對算法的影響。

5. 總結與展望

ESA有效平衡了長序列外推場景下的選擇性注意力中的靈活性和計算效率，無需模型參數增量微調即可擴展上下文長度。ESA通過選擇固定數量的最重要token計算注意力，并利用注意力矩陣的稀疏性，在輸入序列足夠長時，通過低維壓縮query和key有效降低計算復雜度。未來的研究方向包括探索更準確高效的選擇重要token的方法，以及軟硬件協同的高效外推方案。

聯系作者

文章來源：智猩猩GenAI
作者微信：
作者簡介：智猩猩旗下賬號，專注于生成式人工智能，主要分享技術文章、論文成果與產品信息。