北大團隊改造 DeepSeek 注意力機制:速度提升四倍且不損精度

就在大眾正殷切期待 DeepSeek-V4 發布之際,一篇新論文悄然引起了社群關注。

該研究提出了一種新型的稀疏注意力機制 HISA(分層索引稀疏注意力,Hierarchical Indexing Sparse Attention),成功突破了 64K 上下文的索引瓶頸。與 DeepSeek 目前採用的 DSA(DeepSeek Sparse Attention)相比,速度提升了 2 到 4 倍。

HISA 示意圖

HISA 不僅大幅提升速度,且幾乎沒有精度損失,且具備 即插即用、無需重新訓練 的特性。

性能比較

研究團隊直接在 DeepSeek-V3.2 和 GLM-5 上替換索引器,無需進行微調。在關鍵資訊檢索、長文本理解等任務中,其精度與原方法幾乎持平。

精度測試結果

兩步驟消除上下文索引瓶頸

本論文旨在解決一個明確的問題:為大模型的稀疏注意力機制更換一個更高效的「檢索器」。

目前主流的 DSA 等 Token 級稀疏注意力,核心在於僅計算關鍵 Token 的注意力,以降低運算成本。

但此設計存在一個 致命的隱藏問題:為了挑選出相關字符,必須依賴一個「索引器」——它需要 將每個待查詢字符與之前所有字符逐一打分,再選擇分數最高者。

當文本長度 $L$ 增加時,打分的工作量呈 平方級成長($L^2$)。例如,長度翻倍,工作量將增加 4 倍。

在處理超長文本時,索引器的平方級成本反而成為拖慢速度的元兇,甚至比真正的注意力計算更耗時。

基於此,研究團隊思考:能否在不改變最終稀疏注意力結果的前提下,降低索引器的搜尋成本?

於是他們提出了 HISA(分層索引稀疏注意力),核心思路非常簡單:

既然逐一打分太耗時,那就先按區塊篩掉大部分無關內容,再從剩餘的小區塊中精細挑選。

HISA 工作流程

最終挑出的字符與原方法完全一致,後續的注意力計算無需任何修改,相當於「換了一個更高效的篩子,但篩出的東西沒變」。

具體分為兩步,全程複用原模型的打分規則,零學習成本:

第一步:區塊級粗篩(Block-level Coarse Filtering)

  • 將長文本切分為固定大小的「字符塊」(例如每塊 128 個字符),為每個塊計算一個「整體特徵向量」(相當於為每塊貼上總標籤);
  • 利用原索引器的打分方式,僅對這些 塊標籤進行打分
  • 挑選出分數最高的前 $m$ 個塊(例如 64 個),直接捨棄其餘所有塊。由於塊的數量遠少於字符數,此步驟可節省絕大部分工作量。

第二步:塊內精選字符(Intra-block Fine Selection)

僅在第一步選出的 $m$ 個塊中,使用 原索引器的規則對單個字符打分,最後挑出最終需要的 $k$ 個相關字符。

此外還加入了一個優化:首塊與末塊必選,以確保開頭的背景資訊與結尾的最新上下文不被誤篩,同時解決文本拼接的邊界問題。

HISA 詳細機制

HISA 的關鍵優勢在於:複雜度驟降,且能「無縫替換」

HISA 將原索引器每層 $O(L^2)$ 的算力成本降低至 $O(L^2/B + L imes m imes B)$($B$ 為塊大小、$m$ 為選中塊數)。文本越長、區塊選擇越精準,提速效果越顯著。

更重要的是其 工程友好性

  • 輸出與原索引器完全一致,下游注意力計算模組無需更改;
  • 無需重新訓練模型或調整 KV Cache 結構,直接替換索引器即可;
  • 在處理短文本時會自動「退化」回原方法,僅在超長文本時觸發分層篩選,全程自適應。

實測提速驚人,精度幾乎無損

研究團隊在 DeepSeek-V3.2GLM-5 兩大主流模型上進行了全面測試,結果相當亮眼:

速度方面:在 64K 長度的文本下,HISA 比原 DSA 索引器 最高提速 3.75 倍,常規設定也能提速 2 倍以上。

速度提升圖表

索引器耗時從 5.6ms 降至約 1.5ms,徹底解決了索引瓶頸。隨著上下文長度增加,提速效果更為顯著,完全符合 超長上下文(128K/1M) 的實際應用需求。

精度方面:HISA 幾乎完整保留了 DSA 的精度,且顯著優於純塊稀疏(Block-sparse)方法。

「大海撈針」(Needle In A Haystack) 測試中,衡量在超長無關文本中精準檢索指定資訊的能力,HISA 的表現與 DSA 幾乎相同,在所有長度與插入深度下,檢索精度均接近滿分。

大海撈針測試

長文本理解(LongBench 基準測試) 中,HISA 的得分也與 DSA 基本持平。在部分對 Token 篩選精度要求較高的場景(如合成檢索、少樣本學習)中,HISA 甚至實現了小幅反超。

LongBench 測試結果

超參數測試 中,無論塊大小或選塊數量如何變動,HISA 的表現均十分穩定,與 DSA 高度接近,無顯著性能差異

超參數穩定性

這說明 HISA 對超參數不敏感,魯棒性強,在工程落地時無需繁瑣的調參。

不過,作者也指出了目前的不足及後續改進方向:

  • 目前的塊為固定大小,若一塊中同時包含相關與無關內容,整體標籤可能不準。未來可考慮 自適應塊(Adaptive Blocks)、重疊塊 或更優的特徵計算方式。
  • 目前僅在推理階段直接使用,未來可將區塊篩選與模型一同訓練,提升篩選精準度。
  • 目前僅測試索引器速度,未來將整合至完整大模型服務框架中,測試端到端的吞吐量與延遲。

團隊背景

本研究由北京大學 張牧涵 教授團隊完成。

張牧涵為北京大學人工智慧研究院終身教職助理教授及博士生導師。在回國前曾於 Facebook AI(現 Meta AI)擔任研究員,專攻大規模圖學習系統。

張牧涵教授

其 Google Scholar 總引用量超過 13,000 次,多篇第一作者論文引用量達 2,400 至 3,100 次,且連續多年入選 Elsevier 全球前 2% 頂尖科學家(生涯影響力榜單)。

Yufei Xu(徐宇飛)與 Fanxu Meng(孟繁續)為本論文的共同第一作者。

參考連結:https://arxiv.org/abs/2603.28458

相關文章推薦

分享網址
AINews·AI 新聞聚合平台
© 2026 AINews. All rights reserved.