arXiv cs.CL· Zhenxu Tian, Kebin Liu, Zhengwu Yang, Yi Su, Qingqing Dang, Yanlin Sha, Kaipeng Deng, Dianhai Yu, Yanjun Ma, Juntao Li, Min Zhang·· 11 小时前AI 评分34
Prism Attention:面向高效 LLM 推理的提案精炼索引共享机制
Prism Attention: Proposal-Refined Index Sharing Mechanism for Efficient LLMs Inference
AI 导读
Prism Attention 提出一种“提案精炼”的稀疏注意力框架,用粗到细的策略替代跨层稀疏索引的静态复用,在保留各层自适应性的同时维持跨层共享的效率。在推理与长上下文基准上,该方法取得主流稀疏注意力方法中的 SOTA,并在 128K 上下文长度下相比 FlashAttention 实现 3.4 倍端到端解码加速。该论文已被 NeurIPS 2026 收录为主会议论文。
正文 · AI 翻译
作者:Zhenxu Tian、Kebin Liu、Zhengwu Yang、Yi Su、Qingqing Dang、Yanlin Sha、Kaipeng Deng、Dianhai Yu、Yanjun Ma、Juntao Li、Min Zhang
摘要:大语言模型(LLM)中的长上下文解码日益受到瓶颈制约,原因在于每个自回归步骤都需要对不断增长的键值(KV)缓存进行注意力计算,从而带来巨大的内存与计算开销。尽管 token 级稀疏注意力通过将计算限制在一小部分相关的历史 token 上减轻了这一负担,但其实际收益往往被随上下文增长而反复重新计算逐层稀疏索引所抵消。一种常见的补救措施是仅在指定的锚点层计算稀疏索引,并在各层之间共享,但这种静态复用忽略了注意力模式的逐层演化,导致复用的索引逐渐偏离各层的最优值,并造成显著的注意力损失。通过从注意力动态的角度重新审视跨层稀疏性,我们发现层间显著 token 的偏移高度局部化:尽管精确的 top-$\kappa$ 索引在各层之间有所不同,但核心高注意力 token 保持稳定,变化集中在先前选定区域的扩展边界附近。受此启发,我们提出 Prism Attention,一种提议-精炼的稀疏注意力框架。它用由粗到细的策略取代僵化的跨层索引复用,在保留逐层适应性的同时保留跨层共享的效率优势。大量实验表明,Prism Attention 在推理和长上下文基准上均达到主流稀疏注意力方法中的最先进性能,同时在 128K 上下文长度下相比 FlashAttention 实现 $3.4\times$ 的端到端解码加速。
| 评论: | 作为 NeurIPS 2026 主会议论文发表 |
| 主题: | 计算与语言(cs.CL);人工智能(cs.AI) |
| 引用为: | arXiv:2610.22985 [cs.CL] |
| (或 arXiv:2610.22985v1 [cs.CL] 用于此版本) | |
| https://doi.org/10.48550/arXiv.2610.22985 由 DataCite 颁发的 arXiv DOI(待注册) |
提交历史
来自:Zhenxu Tian [查看邮箱]
[v1]
2026年10月16日星期五 06:11:20 UTC(898 KB)
来源:arXiv cs.CL · arxiv.org