论文

通过指数衰减内存来增强注意力,提高查询感知的 KV 稀疏性

Augmenting Attention with Exponentially Decaying Memory Improves Query-Aware KV Sparsity

模型推理KV Cache

摘要

高效推理对于长上下文语言模型至关重要,其中注意力计算和 KV 缓存访问在成本中占主导地位。最近的工作 RAT+ 引入了循环增强注意力主干,可以在推理时实现灵活的扩张注意力。在本文中,我们研究了这种指数衰减的内存是否也可以改进现有的查询感知稀疏推理方法。使用 Quest、MoBA 和 SnapKV 等代表性方法,我们表明,在 8 个大海捞针任务的稀疏预算中,RAT+ 始终高于标准注意力的准确性。我们在 RAT+ 论文中发布的检查点和 OLMo2-7B 上验证了这些收益,我们继续使用添加的 10B 词元内存模块对其进行预训练。最后,我们提出两个假设,解释为什么该内存模块有利于查询感知的稀疏推理,并设计有针对性的实验来支持它们。