LISA：線性索引稀疏注意力機制用於高效長文本推理

針對 DeepSeek-R1 等長鏈式思維（long chain-of-thought reasoning）模型在測試時間擴展（test-time scaling）下導致推理上下文長度增加，進而使標準自注意力機制（standard self-attention）的 O(n^2) 計算複雜度引發推理成本急劇上升的問題，研究團隊提出了 LISA（Linear-Indexed Sparse Attention）。這是一個即插即用的注意力替換模組，無需重新預訓練即可整合到現有模型中。

LISA 在原始模型內並行整合了兩個輕量級組件：一是提供 O(n) 時間複雜度長程記憶的線性注意力模組（Linear Attention module）；二是從完整上下文中選取前 M 個重要 token 並輸入稀疏自注意力（Sparse Self-Attention）的閃電索引器（Lightning Indexer）。透過閾值機制（gating mechanism）融合這兩個分支，將生成 n 個 token 的推理複雜度從 O(n^2) 降低至 O(nM)，其中 M 遠小於 n。

研究設計了兩階段訓練流程：第一階段透過知識蒸餾（knowledge distillation）優化滑動窗口注意力機制（sliding-window attention mechanism），以擬合凍結教師模型的完整自注意力分佈，並整合線性注意力以捕捉長程依賴；第二階段引入索引器取代靜態滑動窗口機制，實現從更廣闊上下文中的動態 token 選取。索引器採用一種新的每頭 KL 散度損失（per-head KL divergence loss）進行訓練，使其選取行為與教師模型的注意力模式對齊。

在 DeepSeek-distilled-Qwen 模型上的實驗顯示，在 16K-token 上下文下，LISA 實現了 50% 的推理速度提升，並在包含 AIME 和 MATH-500 的推理基準測試中將平均性能提升了 5.6%。