PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Xiuying Wei, Caglar Gulcehre

arXiv:2605.28640v1cs.LG

TL;DR

Long-context inference is costly because attention computation and KV-cache access scale poorly, while query-aware sparse methods can lose accuracy at high sparsity. This paper tests whether RAT+’s exponentially decaying memory improves Quest, MoBA, and SnapKV when applied to memory-augmented states. Across eight needle-in-a-haystack tasks and sparse budgets, RAT+ consistently outperforms standard attention, including on released RAT+ checkpoints and a 10B-token-continued OLMo2-7B counterpart.

  • Problem

    Long-context attention has quadratic computation and memory costs, and existing sparse methods show accuracy degradation under challenging sparsity settings.

  • Method

    The paper applies Quest, MoBA, and SnapKV to RAT+’s exponentially decaying memory states and tests hypotheses about selection and information preservation.

  • Results

    RAT+ consistently improves sparse inference accuracy over standard attention across three methods, inference budgets, and eight needle-in-a-haystack tasks.

  • Takeaways & Limitations

    Efficient inference can be improved by designing upstream architectures that are inherently more capable with sparse inference.

  • Takeaways & Limitations

    The evaluation is limited mainly to 7B checkpoints with maximum sequence length 4K, RULER needle-in-a-haystack tasks, and Quest, MoBA, and SnapKV.

Abstract

from arXiv · show

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.

1 Introduction

The paper asks whether RAT+’s exponentially decaying memory can strengthen query-aware sparse inference, which otherwise loses accuracy under demanding sparsity settings. Across Quest, MoBA, and SnapKV, RAT+ improves sparse accuracy across budgets and motivates memory-aware upstream architectures.

  • Motivation: Query-aware sparse methods reduce attention computation or KV storage by selecting critical tokens or blocks, but accuracy degrades under challenging sparsity settings.Examples include 98% sparsity for Quest and over 75% KV-cache storage reduction for SnapKV.
  • Approach: RAT+ adds exponentially decaying memory over KV states to enable flexible-budget dilated attention at inference time.
  • Results: RAT+ substantially improves Quest, MoBA, and SnapKV accuracy over standard attention across eight needle-in-a-haystack tasks and sparse budgets.The methods target decoding-time FLOPs, prefilling-time FLOPs, and KV-cache storage, respectively.
  • Results: 34.11 and 40.03 points: SnapKV’s average improvement across eight tasks under 1/4 and 1/8 budgets, respectively.On OLMo2-7B, Quest improves from 68.0 to 98.6 on MK-2 under 1/16, while MoBA improves from 53.6 to 94.8 on MK-3.
  • Analysis: The proposed explanations are that memory improves critical-token selection and provides selected candidates with an additional path for preserving answer-related information.
  • Implication: The results suggest improving efficient inference through upstream architectures that are inherently more capable with sparse inference.

2 Method

The method augments attention with exponentially decaying memory, applies sparse selection and computation to the resulting states, and analyzes both candidate selection and information preservation.

  • 2.1 Exponentially decaying memory: RAT+ updates key and value states with an input-dependent gate, yielding exponentially decaying memory with effective length 64.The memory module adds O(1) computation and storage per token.
  • 2.2 Sparse Inference Instantiations: The study compares each sparse method applied to standard states S(K, V) against the same method applied to RAT+ states S(˜K, ˜V).
  • 2.2 Sparse Inference Instantiations: Query-aware sparse inference first selects critical KV candidates using a scoring function and then computes attention only over the selected states.
  • 2.2 Sparse Inference Instantiations: Quest and MoBA select KV blocks, whereas SnapKV selects individual KV tokens using an observation window and discards unselected entries.SnapKV thereby reduces KV-cache size and subsequent decoding FLOPs.
  • 2.3 Analyses: The analyses test whether memory improves critical-token selection by measuring strict head-level hit rate for containing every required golden answer position.
  • 2.3 Analyses: A second hypothesis tests whether memory provides an additional information path by replacing the designed selector with a random selector.Randomized evaluations use five seeds and average the resulting performance.
  • Evaluation: Figure 1 compares two checkpoint types across eight NIAH tasks, with curves indexed by per-query KV budget; Quest and MoBA retain full KV caches.

3 Experiments

The experiments evaluate RAT+ with representative query-aware sparse inference methods and analyze both accuracy and the mechanisms behind its gains. RAT+ improves performance over standard attention, with evidence pointing to better candidate selection and more informative selected states.

  • Experimental setup: The study evaluates Quest, MoBA, and SnapKV across released 7B checkpoints and an OLMo2-7B RAT+ counterpart continued for 10B tokens.Quest targets decoding-time FLOPs, MoBA prefilling-time FLOPs, and SnapKV KV-cache storage.
  • Main results: RAT+ consistently outperforms standard attention across sparse inference methods and KV budgets.The trend holds on both checkpoint types.
  • Main results: In the 1/16 budget, RAT+ improves SnapKV on S-1 from 39.2 to 84.2.
  • Main results: For Quest on OLMo checkpoints, RAT+ improves MK-2 from 76.4 to 99.4 and MK-3 from 68.0 to 98.6.The larger gains occur on harder multi-key settings, while standard attention is already near-perfect on simple NIAH-Single tasks.
  • H1 results: RAT+ yields higher hit rates for Quest and SnapKV, especially where accuracy gains are larger.The analysis ranks heads by hit rate; SnapKV measures retained answer-relevant tokens directly, while Quest reflects step-wise dynamic selection.
  • H2 results: With a random selector and a 1/4 budget, RAT+ improves Quest from 53.4 to 84.9 on MK-1 and from 39.6 to 74.4 on MV.Because candidates are random, the result supports an additional information path through memory-augmented selected states rather than better selection alone.

4 Related work

Sparse-attention research reduces inference cost by selecting limited tokens or blocks instead of attending to the full context. Query-aware methods dynamically adapt this selection to each query or sample.

  • Sparse attention is motivated by attention-map sparsity and the quadratic computation and memory cost of attention with context length.
  • Static methods use fixed patterns such as local windows, whereas query-aware methods dynamically select important tokens or blocks.
  • Quest uses min/max block statistics, MoBA uses mean-pooled block representations, and Twilight selects top-p rather than top-k elements.

5 Conclusion

The paper concludes that exponentially decaying memory improves query-aware KV sparsity across representative methods and inference budgets. It argues that efficient inference can also benefit from upstream architectures designed to work well with sparse inference.

  • The analyses attribute the gains to improved critical-token selection and an additional information path for selected candidates.
  • RAT+ substantially improves sparse inference accuracy over standard attention across three representative methods and inference budgets.
  • The results support improving efficient inference through more capable upstream architectures, not only through downstream inference methods.

Limitations

The evaluation is bounded to 7B checkpoints with maximum sequence length 4K, primarily synthetic RULER NIAH tasks, and three sparse inference methods. Longer contexts, other downstream tasks, and other sparse methods are not evaluated.

  • The experiments mainly use 7B checkpoints with a maximum sequence length of 4K, so longer context lengths are not evaluated.
  • The study focuses mainly on synthetic needle-in-a-haystack tasks from the RULER benchmark and does not evaluate other downstream tasks.Their synthetic structure supports controlled continued-pretraining experiments and tuning of both original and added parameters.
  • The evaluation covers Quest, MoBA, and SnapKV but does not evaluate other sparse inference methods.

A.1 Implementation details

The experiments evaluate three query-aware sparse inference methods across models, tasks, and KV-access budgets, using specified training and inference settings. Implementation choices cover decoding, prefilling, KV-cache reduction, and supervised fine-tuning.

  • Models and training: The evaluation uses released 7B checkpoints and an OLMo2-7B RAT+ counterpart continued-pretrained for 10B tokens with added memory parameters.Supervised fine-tuning uses a different-seed synthetic split of approximately 7M training tokens; compute budgets are also reported for pretraining, fine-tuning, and evaluation.
  • Sparse inference settings: Quest selects top-k KV blocks during decoding using block representatives built from dimension-wise minimum and maximum key values.The main configuration uses block size 64 and varies selected blocks from 16 to 4; an extreme block-size-16 setting with four blocks is also evaluated.
  • Sparse inference settings: MoBA routes each query to relevant KV blocks using mean-pooled block representations and is evaluated mainly with block size 64, plus block size 128.The methods retain the first and local blocks according to their original designs.
  • Sparse inference settings: SnapKV estimates prior KV-position importance from attention scores in a 64-token observation window and retains 1024, 512, or 256 entries.These settings correspond to 1/4, 1/8, and 1/16 KV budgets and reduce storage and subsequent decoding attention cost.
  • Tasks: The study evaluates Quest, MoBA, and SnapKV on eight RULER needle-in-a-haystack tasks with 500 examples per sub-task.The tasks retrieve numbers or long UUID-like strings from long contexts containing noise or similar key-value pairs.

A.2 Supplementary experiments

Supplementary experiments provide additional NIAH results and analyze the two proposed hypotheses. Under random selection, SnapKV outperforms Quest, unlike under the original selectors.

  • Additional results: Additional NIAH results are reported in Tables 3 and 4, with further H1 analysis in Figures 3 and 4 and H2 results in Table 5.These materials extend the main evaluation and hypothesis analyses.
  • H2 analysis: Under random selection, SnapKV performs better than Quest, whereas Quest performs better under the original selectors.The authors conjecture that SnapKV selects once after prefilling, while Quest resamples at every decoding step and experiences more dynamic noise.

A.3 License information

The section lists licenses and links for the FineWeb-Edu dataset, RULER benchmark, OLMo2-7B model, and RAT+ model, alongside performance tables and head-level figures.

  • License information: FineWeb-Edu is listed under the Open Data Commons License Attribution family, with a Hugging Face dataset link.The license information identifies FineWeb-Edu as a dataset.
  • License information: The RULER benchmark is listed under the Apache 2.0 License.RULER is identified as a dataset.
  • License information: OLMo2-7B is listed under the Apache 2.0 License, with a Hugging Face model link.The listed model is identified as OLMo2-7B.
  • License information: RAT+ is listed under the Apache 2.0 License, with a Hugging Face model link.The listed model is identified as RAT+.
  • Performance materials: Table 3 reports eight-task performance for 7B checkpoints pretrained from scratch on 100B tokens, with budgets denoting the fraction of KV entries accessed per token.The table includes MoBA block size 128 and an extreme Quest setting with block size 16 and four blocks.
  • Performance materials: Table 4 reports eight-task performance for OLMo2-7B, whose RAT+ counterpart uses 10B-token continued pretraining with added exponentially decaying memory.The table also specifies budget interpretation and alternative block-size settings for MoBA and Quest.
  • Performance materials: Table 5 reports random-selector performance under a 1/4 budget as mean ± standard deviation across five seeds.Quest and MoBA select 16 blocks of size 64, while SnapKV selects 1024 tokens.
  • Performance materials: Figures 3 and 4 show head-level hit-rate distributions for the top-50 Quest and SnapKV heads ranked by hit rate.The supplied captions identify the methods and ranking criterion but do not state a comparison outcome.
Loading 2605.28640v1…