LISA让长文本推理快50%且更准,无需重训练。
LISA: Linear-Indexed Sparse Attention for Efficient Long-Context Reasoning

- 用线性注意力+动态选重要词,把计算量从平方级降到线性级。
- 16K上下文下推理速度提升50%,数学推理准确率提高5.6%。
- 可直接替换原模型,适合部署长链条逻辑推理任务的工程师。
近年来,如DeepSeek-R1等长链思维模型在测试时缩放范式下推理上下文长度不断增长。然而,标准自注意力机制的O(n²)复杂度导致长序列推理成本急剧上升,限制了其在生产环境中的应用。为此,我们提出LISA(Linear-Indexed Sparse Attention),一个无需从头预训练的即插即用注意力模块。LISA在原始模型中并行集成两个轻量组件:(1) 线性注意力模块,以O(n)时间复杂度实现长程记忆;(2) Lightning Indexer,从完整上下文中选出前M个重要标记输入稀疏自注意力。两分支通过门控机制融合,将生成n个标记的推理复杂度从O(n²)降至O(nM)(M << n)。我们设计了两阶段训练流程:第一阶段通过知识蒸馏优化滑动窗口注意力,使其逼近冻结教师模型的全自注意力分布;第二阶段引入Indexer替代静态滑动窗口,实现从更广上下文中动态选择关键标记。Indexer采用新型每头KL散度损失进行训练,使其选择行为与教师模型注意力模式对齐。在DeepSeek-distilled-Qwen模型上的实验表明,LISA在16K令牌上下文中实现50%的推理加速,同时在AIME和MATH-500等推理基准上平均性能提升5.6%。
原文摘要 · Abstract (English)
Recent advances in long chain-of-thought reasoning models such as DeepSeek-R1 have led to increasingly longer inference context lengths under the test-time scaling paradigm. However, the O(n^2) computational complexity of standard self-attention causes inference costs to grow sharply with long sequences, limiting the deployment of long-CoT reasoning in production settings. To address this, we propose LISA (Linear-Indexed Sparse Attention), a plug-and-play attention replacement module that requires no pretraining from scratch. LISA integrates two lightweight components in parallel within the original model: (1) a Linear Attention module that provides long-range memory with O(n) time complexity; (2) a Lightning Indexer that selects the top-M important tokens from the full context to feed into a Sparse Self-Attention. The two branches are fused via a gating mechanism, reducing inference complexity from O(n^2) to O(nM) (M << n) for generating n tokens. We design a two-stage training pipeline: Stage 1 initializes the model by integrating the linear attention to capture long-range dependencies, complemented by a sliding-window attention mechanism that is optimized via knowledge distillation to approximate the full self-attention distribution of a frozen teacher model. In Stage 2, we further introduce the Indexer to replace the static sliding-window mechanism, enabling dynamic token selection from broader contexts. The Indexer is trained using a novel per-head KL divergence loss, which aligns its selection behavior with the attention patterns of the teacher model. Experiments on DeepSeek-distilled-Qwen models demonstrate that LISA achieves a 50% inference speedup under 16K-token context, while improving average performance by 5.6% on reasoning benchmarks including AIME and MATH-500.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。