返回首頁

AI 研究

WISE 重用遞迴模型注意力範圍,注意力流程測得 1.36 倍加速

方法保留全部遞迴步驟,透過早期探索縮小後續注意力運算範圍。效能數字僅涵蓋注意力負載,4K 上下文的品質下降與公開重現材料仍須核對。

Idawriter · CC BY-SA 3.0 · Image source
zh-Hant

維吉尼亞大學研究者於 9 月 23 日公開 WISE,嘗試降低遞迴語言模型反覆計算全域注意力的成本。方法毋須重新訓練,利用早期運算找出重要上下文區塊,再於後續步驟重用選取範圍;模型仍持續更新內部表徵。[論文摘要](https://arxiv.org/abs/2609.27373)

公開實作預設執行 32 次遞迴,前 12 次保留全域注意力。系統以 32 個 token 為區塊,在第 9 至 12 步分別挑選累積注意力權重達 95% 的區塊,再合併成後續可存取的集合。第 13 步起固定的是區塊範圍,Q、K、V 與區塊內注意力權重仍重新計算,因此保留了後期細化表徵的能力。[程式庫說明](https://github.com/tbn5pj/WISE_code)

這項設計的工程價值,在於把不規則的注意力稀疏性轉成 GPU 可執行的工作排程。作者提供 Triton 核心,略過未選取的 key 區塊,並融合分數計算、因果遮罩、softmax 與數值聚合,避免只加遮罩卻仍掃過所有區塊。[實作文件](https://github.com/tbn5pj/WISE_code/blob/main/docs/REPRODUCING.md)

作者在 RTX A6000、Huginn 模型的 30 組注意力負載上測量:4K 上下文的後 20 步,相對原生 FlashAttention 達約 1.76 倍加速;計入前 12 步探索後,完整 32 步注意力流程約為 1.36 倍。數字包含相應的排程建立成本,但未涵蓋整個模型推論。[論文測速表](https://arxiv.org/html/2609.27373v1#S7.SS4)

品質代價不能略過。論文的配對測試在最高 2K 上下文未偵測到平均 F1 下降,4K 則下降約 3.3 分。這使最高加速數字同時成為需要檢視品質取捨的測試點,不能概括為所有長文本任務都能維持原有表現。[品質與效率評估](https://arxiv.org/html/2609.27373v1#S7.SS5)

此外,論文將上下文品質測試描述為一百個例子,公開重現文件卻列每種長度三十個,共一百二十筆提示。兩者的樣本對應仍須釐清;重新繪製已保存的結果表,也不能代替重新執行原始評測。[重現文件](https://github.com/tbn5pj/WISE_code/blob/main/docs/REPRODUCING.md)

重現材料也有缺口:程式庫提供核心、設定、資料識別碼及提示雜湊,但沒有附上完整測速張量;文件並說明,打包公開版本時未重新執行大型模型評測與 GPU 測速。對部署者而言,下一步應是取得可核對的測試輸入,在相同軟硬體環境重測端到端延遲,並另驗證中文、多輪對話與更長上下文的品質。[公開版本查核說明](https://github.com/tbn5pj/WISE_code/blob/main/PAPER_CODE_CONSISTENCY.md)

來源

  1. Attention Routing Stabilizes Early: Working-Set Inference for Recurrent Language Models
  2. WISE 論文全文與實驗附錄
  3. WISE_code:官方參考實作
  4. WISE 重現指南
  5. WISE 論文與程式一致性查核