| title | Augmenting Attention with Exponentially Decaying Memory Improves Query-Aware KV Sparsity |
|---|---|
| authors | Xiuying Wei, Caglar Gulcehre |
| date | 2026-05-27 |
| venue | arXiv (preprint, not peer reviewed) |
| id | arXiv:2605.28640 |
| links | abs pdf hf code |
| abstract | 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. |
cite
@misc{wei2026augmenting,
title={Augmenting Attention with Exponentially Decaying Memory Improves Query-Aware KV Sparsity},
author={Xiuying Wei and Caglar Gulcehre},
year={2026},
eprint={2605.28640},
archivePrefix={arXiv},
url={https://arxiv.org/abs/2605.28640}
}