让大模型自己学着在注意力中删减冗余块,提速长文本处理。
SeerAttention: Learning Intrinsic Sparse Attention in Your LLMs
- 通过可学习的门控机制动态识别重要注意力块
- 长文本预填充时延迟降低,精度优于现有方法
- 轻量训练,适配主流大模型,易集成
注意力是现代大语言模型的核心,但其二次复杂度限制了效率与可扩展性,尤其在长序列处理时。现有稀疏化方法多依赖预定义模式或启发式规则,难以自适应不同上下文。我们提出SeerAttention,一种直接从大模型内部学习块级注意力稀疏性的简单有效机制。受混合专家(MoE)门控启发,SeerAttention在传统注意力基础上增加可学习门控,通过池化查询(Q)和键(K)并经可学习线性层处理,生成门控分数以预测块级稀疏结构。结合块稀疏FlashAttention内核,可在GPU上实现显著加速。应用于预训练模型时,仅需轻量级自蒸馏方式训练门控参数,收敛迅速。实验表明,相较于先前方法,SeerAttention在长序列预填充任务中实现了更高精度与更低延迟。代码已开源:https://github.com/microsoft/SeerAttention
原文摘要 · Abstract (English)
Attention is the cornerstone of modern Large Language Models (LLMs). Yet its quadratic complexity hinders efficiency and scalability, especially for long-context processing. A promising approach is to leverage sparsity in attention. However, existing sparsity-based solutions predominantly rely on predefined patterns or heuristics at the attention head level, struggling to adapt dynamically to different contexts efficiently. We propose SeerAttention, a simple yet effective attention mechanism that directly learns the block-level attention sparsity from the LLM itself. Inspired by the gating mechanism in Mixture of Experts (MoE), SeerAttention augments the conventional attention with a learnable gate that selectively activates important blocks within the attention map. Specifically, the gate first pools the query (Q) and key (K) tensors along the sequence dimension and processes them through learnable linear layers. The resulting matrices are then multiplied together to produce the gating scores, which are used to predict block-level attention sparsity. Combined with our block-sparse FlashAttention kernel, SeerAttention can achieve significant speedup on GPUs. When applied to pre-trained LLMs, SeerAttention only requires training the gate parameters in a lightweight self-distillation manner, allowing rapid convergence. Our evaluation results demonstrate that SeerAttention achieves better model accuracy and lower latency for long-context pre-filling compared to prior methods. Code is available at: https://github.com/microsoft/SeerAttention
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。