arXiv:2506.13059cs.CLcs.LG2025-06NeurIPS被引 7

用多极注意力只精确计算关键步骤,让长思维链推理快4.5倍。

Multipole Attention for Efficient Long Context Reasoning

  • 聚类关键向量,仅对重要项算精确注意力
  • 在复杂任务上保持高准确率,支持激进稀疏设置
  • 适合需要高效长文本推理的模型部署

大型推理模型(LRMs)在复杂问题求解任务中展现出显著的准确率提升。这些模型通过测试时增加计算量,生成数千个词的思维链来思考后作答,导致自回归生成过程中需处理超长上下文。现有稀疏注意力方法虽能缓解键值缓存压力,但可能引入错误干扰推理过程。此外,以往方法依赖预处理识别重要提示词,难以在线处理新生成的推理内容。本文提出多极注意力机制,仅对最关键令牌计算精确注意力,其余使用近似表示。该方法首先对语义相似的关键向量进行聚类,利用聚类中心既识别重要向量,又近似其余向量,以保持高精度。设计快速聚类更新机制,实时重聚输入与已生成的令牌,加速对历史输出的注意力计算。我们在Qwen-8B等新兴LRM上评估,结果表明在极端注意力稀疏设置下仍能维持高准确率。同时提供内核实现,验证实际效率提升,长上下文推理中的注意力计算最高提速4.5倍。代码已开源。

原文摘要 · Abstract (English)

Large Reasoning Models (LRMs) have shown promising accuracy improvements on complex problem-solving tasks. While these models have attained high accuracy by leveraging additional computation at test time, they need to generate long chain-of-thought reasoning in order to think before answering, which requires generating thousands of tokens. While sparse attention methods can help reduce the KV cache pressure induced by this long autoregressive reasoning, these methods can introduce errors which disrupt the reasoning process. Additionally, prior methods often pre-process the input to make it easier to identify the important prompt tokens when computing attention during generation, and this pre-processing is challenging to perform online for newly generated reasoning tokens. Our work addresses these challenges by introducing Multipole Attention, which accelerates autoregressive reasoning by only computing exact attention for the most important tokens, while maintaining approximate representations for the remaining tokens. Our method first performs clustering to group together semantically similar key vectors, and then uses the cluster centroids both to identify important key vectors and to approximate the remaining key vectors in order to retain high accuracy. We design a fast cluster update process to quickly re-cluster the input and previously generated tokens, thereby allowing for accelerating attention to the previous output tokens. We evaluate our method using emerging LRMs such as Qwen-8B, demonstrating that our approach can maintain accuracy on complex reasoning tasks even with aggressive attention sparsity settings. We also provide kernel implementations to demonstrate the practical efficiency gains from our method, achieving up to 4.5$\times$ speedup for attention in long-context reasoning applications. Our code is available at https://github.com/SqueezeAILab/MultipoleAttention.

注意力机制长序列推理模型加速稀疏计算

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。