论文

LISA:面向高效长上下文推理的线性索引稀疏注意力

LISA: Linear-Indexed Sparse Attention for Efficient Long-Context Reasoning

模型架构Transformer

摘要

以DeepSeek-R1为代表的长思维链推理模型的进展,使测试时扩展范式下的推理上下文长度不断攀升。但标准自注意力O(n²)的计算复杂度使推理成本随长序列急剧增长,限制长CoT推理在生产环境的部署。为此我们提出LISA(线性索引稀疏注意力),一个即插即用的注意力替换模块,无需从零预训练。LISA在原模型内并行集成两个轻量组件:(1)线性注意力模块,以O(n)时间复杂度提供长程记忆;(2)闪电索引器,从全上下文中选出top-M重要token送入稀疏自注意力。两支路经门控融合,把生成n个token的推理复杂度从O(n²)降到O(nM)(M≪n)。我们设计两阶段训练管线:第一阶段整合线性注意力以捕捉长程依赖,辅以经知识蒸馏优化的滑窗注意力来逼近冻结教师模型的完整自注意力分布;第二阶段引入索引器替换静态滑窗,实现从更广上下文的动态token选择——索引器以新颖的逐头KL散度损失训练,使其选择行为对齐教师模型的注意力模式。在DeepSeek蒸馏的Qwen模型上的实验显示:LISA在16K上下文下取得50%推理加速,同时在AIME与MATH-500等推理基准上把平均表现提升5.6%。

LISA:面向高效长上下文推理的线性索引稀疏注意力:论文配图
图 3:索引器跨层选择的前 256 个标记的可视化。