arXiv:2502.20766cs.LGcs.CL2025-02ICLR被引 124

动态调整注意力模式,让大模型长文本推理更快更准。

FlexPrefill: A Context-Aware Sparse Attention Mechanism for Efficient Long-Sequence Inference

论文配图:FlexPrefill: A Context-Aware Sparse Attention Mechanism for Efficient Long-Sequence Inference
图 1 · 摘自论文原文
  • 根据输入内容实时调整注意力稀疏模式和计算量。
  • 在保持精度前提下,推理速度提升显著,优于现有方法。
  • 适合需要高效处理长文本的场景,如文档理解、对话系统。

大语言模型在长序列推理中面临计算挑战,尤其在注意力预填充阶段,复杂度随提示长度呈平方增长。以往方法依赖固定稀疏模式或有限情况下的稀疏模式识别,缺乏对不同输入需求的灵活适应能力。本文提出 FlexPrefill,一种灵活的稀疏预填充机制,可实时动态调整各注意力头的稀疏模式与计算预算,以满足不同输入和注意力头的具体需求。核心创新包括:1)基于查询感知的稀疏模式判定:通过衡量 Jensen-Shannon 散度,自适应切换至查询相关的多样化注意力模式或预设模式;2)累积注意力索引选择:根据不同的注意力模式动态选择需计算的查询-键索引,确保注意力分数总和达到预设阈值。FlexPrefill 能依据提示内容自适应优化每个注意力头的稀疏模式与稀疏比例,显著提升长序列推理效率。实验表明,该方法在速度和准确率上均优于先前方法,为大模型推理提供更灵活高效的解决方案。

原文摘要 · Abstract (English)

Large language models (LLMs) encounter computational challenges during long-sequence inference, especially in the attention pre-filling phase, where the complexity grows quadratically with the prompt length. Previous efforts to mitigate these challenges have relied on fixed sparse attention patterns or identifying sparse attention patterns based on limited cases. However, these methods lacked the flexibility to efficiently adapt to varying input demands. In this paper, we introduce FlexPrefill, a Flexible sparse Pre-filling mechanism that dynamically adjusts sparse attention patterns and computational budget in real-time to meet the specific requirements of each input and attention head. The flexibility of our method is demonstrated through two key innovations: 1) Query-Aware Sparse Pattern Determination: By measuring Jensen-Shannon divergence, this component adaptively switches between query-specific diverse attention patterns and predefined attention patterns. 2) Cumulative-Attention Based Index Selection: This component dynamically selects query-key indexes to be computed based on different attention patterns, ensuring the sum of attention scores meets a predefined threshold. FlexPrefill adaptively optimizes the sparse pattern and sparse ratio of each attention head based on the prompt, enhancing efficiency in long-sequence inference tasks. Experimental results show significant improvements in both speed and accuracy over prior methods, providing a more flexible and efficient solution for LLM inference.

大模型推理稀疏注意力长序列动态优化

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