论文
独特:用于 无需训练 推理和稀疏性感知训练的通用 Top-k 稀疏注意力
UNIQUE: Universal Top-k Sparse Attention for Training-free Inference and Sparsity-aware Training
摘要
大语言模型 (LLM) 中的长上下文推理受到自注意力键值 (KV) 缓存的线性增长的瓶颈。 Top-k 稀疏注意力通过仅加载 KV 缓存的一小部分来缓解这一问题,但对于 无需训练 使用和稀疏感知训练来说,准确且廉价地估计缓存重要性仍然具有挑战性。本文提出了 UNIQUE,这是一种通用的 top-k 稀疏注意力框架,可以满足这两个要求,并在 LLM 模式中保持一致的有效性。 UNIQUE 以 KV 页面的粒度进行操作,并通过简单而准确的分数来估计每页的重要性,该分数将页面键的平均值作为代表向量与其标准差作为偏移项相结合。为了进一步缩小训练推理差距,本文引入了一种软掩码稀疏感知训练方案,该方案使用 top-k 分数边界作为每个查询的阈值,并在其周围使用 sigmoid 软掩码,既不需要辅助损失,也不需要架构更改。文本和语音 LLM 的实验表明,UNIQUE 在 LongBench Pro 等长上下文基准测试和长格式语音识别上保留了任务性能,同时与 FlashInfer 密集注意力相比,注意力内核加速高达 11.4 倍,与基于 vLLM 的密集模型相比,端到端解码加速至少提高 5.3 倍。