LISA：线性索引稀疏注意力助力高效长上下文推理
针对长链思维推理模型在测试时缩放中面临的自注意力二次复杂度问题，本文提出LISA（线性索引稀疏注意力），一种即插即用的注意力替换模块，无需从头预训练。LISA并行集成线性注意力和闪电索引器，通过门控机制融合，将推理复杂度从O(n²)降至O(nM)。在DeepSeek蒸馏Qwen模型上的实验表明，在16K上下文下实现50%推理加速，并在AIME和MATH-500等基准上平均提升5.6%的性能。
近年来，以DeepSeek-R1为代表的长链思维推理模型在测试时缩放范式下不断推长推理上下文长度，使得模型能够处理更加复杂的推理任务。然而，标准自注意力机制具有O(n²)的计算复杂度，导致长序列推理成本急剧上升，严重限制了长CoT推理在生产环境中的部署。为了解决这一瓶颈，研究人员提出了LISA（Linear-Indexed Sparse Attention，线性索引稀疏注意力），这是一种即插即用的注意力替换模块，无需从头预训练即可集成到现有模型中。LISA的设计精巧而高效。它在原始模型内部并行集成了两个轻量级组件：线性注意力模块和闪电索引器。线性注意力以O(n)的时间复杂度提供长距离记忆能力，能够捕捉上下文中的全局依赖关系；而闪电索引器则从完整上下文中选出最重要的前M个令牌，并将其输入稀疏自注意力机制。这两个分支通过门控机制融合，从而将生成n个令牌的推理复杂度从O(n²)降低到O(nM)，其中M远小于n。训练过程采用两阶段流水线。第一阶段通过集成线性注意力来捕获长距离依赖，并结合滑动窗口注意力机制，利用知识蒸馏优化以逼近冻结教师模型的全自注意力分布。第二阶段引入索引器替代静态滑动窗口机制，实现从更广上下文中的动态令牌选择。索引器通过一种新颖的逐头KL散度损失进行训练，使其选择行为与教师模型的注意力模式对齐。在DeepSeek蒸馏后的Qwen模型上进行的实验表明，LISA在16K令牌上下文下实现了50%的推理加速，同时在包括AIME和MATH-500在内的推理基准上平均性能提升5.6%。这一工作为长上下文推理的高效部署提供了全新的解决方案，有望推动长链思维推理模型在实际应用中的广泛落地。