用指数衰减记忆增强注意力，提升长文本查询的稀疏推理精度
Augmenting Attention with Exponentially Decaying Memory Improves Query-Aware KV Sparsity
- 引入指数衰减记忆模块，动态调整注意力权重
- 在8个长文本找针任务中，稀疏预算下准确率全面超越标准注意力
- 适合追求高效长文本推理的模型部署者与研究者
高效推理对长上下文语言模型至关重要，其中注意力计算和键值缓存访问占主要开销。近期工作RAT+提出一种递归增强的注意力骨干，可在推理时实现灵活的稀疏注意力。本文探究该指数衰减记忆模块是否能提升现有查询感知稀疏推理方法的效果。基于Quest、MoBA和SnapKV等代表性方法，在8个'大海捞针'任务上，RAT+在不同稀疏预算下均一致提升准确率。验证使用了RAT+论文发布的检查点以及我们在OLMo2-7B基础上继续预训练100亿tokens后加入记忆模块的版本。最后，我们提出两个解释性假设，并设计针对性实验予以支持。
原文摘要 · Abstract (English)
Efficient inference is critical for long-context language models, where attention computation and KV-cache access dominate the cost. Recent work RAT+, introduces a recurrence-augmented attention backbone that enables flexible dilated attention at inference time. In this paper, we investigate whether this exponentially decaying memory can also improve existing query-aware sparse inference methods. Using representative methods including Quest, MoBA, and SnapKV, we show that RAT+ consistently improves accuracy over standard attention across sparse budgets on eight needle-in-a-haystack tasks. We validate these gains both on the released checkpoints from the RAT+ paper and on OLMo2-7B, which we continue pretraining with the added memory module for 10B tokens. Finally, we propose two hypotheses explaining why this memory module benefits query-aware sparse inference and design targeted experiments to support them.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。