提出UNIQUE框架,实现无需训练的高效长文本推理与稀疏感知训练。
UNIQUE: Universal Top-k Sparse Attention for Training-free Inference and Sparsity-aware Training

- 按键值页粒度计算重要性得分,结合均值与标准差快速筛选关键缓存。
- 在长文本任务中保持性能,推理速度比FlashInfer快11.4倍,端到端提速5.3倍。
- 无需额外损失或结构修改,适配多模态LLM,适合追求效率的部署场景。
大语言模型(LLMs)在长上下文推理中受自注意力键值(KV)缓存线性增长的限制。顶-k稀疏注意力通过仅加载少量缓存缓解此问题,但准确且低成本地估计缓存重要性,对无训练使用和稀疏感知训练仍具挑战。本文提出UNIQUE,一种通用的顶-k稀疏注意力框架,可同时满足两类需求,并在多种LLM模态下保持稳定有效性。UNIQUE以键值页为粒度,通过结合页面键的均值(代表向量)与标准差(偏移项)构建简单而精确的重要性评分。为缩小训练-推理差距,引入软掩码稀疏感知训练方案:以顶-k得分边界作为每查询阈值,并在其周围使用sigmoid软掩码,无需辅助损失或架构改动。在文本与语音LLM上的实验表明,UNIQUE在LongBench Pro等长上下文基准及长语音识别任务中保持任务性能,相比FlashInfer密集注意力最高实现11.4倍注意力内核加速,相比vLLM基线密集模型至少获得5.3倍端到端解码加速。
原文摘要 · Abstract (English)
Long-context inference in large language models (LLMs) is bottlenecked by the linear growth of the self-attention key-value (KV) cache. Top-k sparse attention alleviates this by loading only a small fraction of the KV cache, but accurately and cheaply estimating cache importance, for both training-free use and sparsity-aware training, remains challenging. This paper proposes UNIQUE, a universal top-k sparse attention framework that addresses both requirements and stays consistently effective across LLM modalities. UNIQUE operates at the granularity of KV pages and estimates per-page importance with a simple yet accurate score combining the mean of the page's keys as a representative vector with their standard deviation as an offset term. To further close the train-inference gap, this paper introduces a soft-mask sparsity-aware training scheme that uses the top-k score boundary as a per-query threshold and a sigmoid soft mask around it, requiring neither auxiliary losses nor architectural changes. Experiments on text and speech LLMs show that UNIQUE preserves task performance on long-context benchmarks such as LongBench Pro and on long-form speech recognition, while delivering up to 11.4x attention-kernel speedup over FlashInfer dense attention and at least 5.3x end-to-end decoding speedup over a vLLM-based dense model.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。