Recycled Attention: Efficient inference for long-context language models

Authors: Fangyuan Xu, Tanya Goyal, Eunsol Choi
Primary area: applications to computer vision, audio, language, and other modalities
Keywords: long-context language model, efficiency, inference-time method

Abstract:
Processing long-context input imposes a heavy computational burden when deploying large language models. Recently proposed inference-time methods accelerate generation by attending only to local context. Despite its efficiency gains, this approach fails to capture all relevant information in the input, showing substantial performance drop in long-context benchmarks. We propose recycled attention, an efficient and effective method which alternates between full context attention and attention over a subset of input tokens. When performing partial attention, we leverage the attention pattern of a nearby token that has performed full attention and attend only to the top K most attended tokens. We evaluate our methods on RULER, a suite of tasks designed to comprehensively evaluate long-context abilities, and long-context language modeling tasks. Applying our inference method to off-the-shelf LLMs achieves comparable speedup to baselines which only consider local context while improving the performance by 2x. We further experiment with continued pre-training the model with recycled attention to improve the performance-efficiency trade-off.

Decision: reject
Reviewer scores: 6, 5, 5, 5, 6

Reviews and discussion:
<thread>
creation date:
1729409524152

thread_type:
official_review

<thread_object>
creator: reviewer_ic8w

<content>
summary:
This paper focuses on accelerating generation speed when dealing with long-context inputs using Large Language Models. Though prior methods have achieved inference speedup by different KV cache eviction policies, the eliminated tokens cannot be recovered in future generation process, thereby leading to performance decline on tasks requiring aggregating long-context information. To address this issue, the authors propose Recycled Attention, which keeps full kv cache in memory and realizes speedup by alternating between a sparse full-attention mode and a consecutive recycled attention mode. The latter only performs attention over a small subset of kv pairs identified using the attention weights produced by the full attention mode. Experiments are conducted on language modeling(Arxiv, Book, and PG19) and synthetic long-context tasks(RULER). Results demonstrate that Recycled Attention deliver higher task accuracy compared to baselines while achieving similar speedup gain.

strengths:
1. The idea of utilizing attention pattern similarity between adjacent tokens to speed up inference makes intuitive sense.
2. The attention mass analysis and empirical results on downstream tasks demonstrate the effectiveness of Recylced Attention.
3. The experiment section is comprehensive, which includes two LLMs, two types of downstream tasks and different long-context configurations.

weaknesses:
1. Details about the methods: Are there any discussion on the choice of using the first query head's results for each kv heads in a group? Why not the average of all query heads?
2. Benchmarks: The paper perform experiments on RULER, which is mostly a synthetic long-context benchmark and the reported results exhibit a notable margin compared to vanilla full attention. It would be more convincing to incorporate more realistic long-context benchmarks, e.e.,  LongBench.
3. Baselines: For KV cache eviction baselines, the mainly compared methods in this paper are H2O and StreamLLM, which are both query-agnostic KV cache eviction methods. The experiment lacks comparisons with more accurate query-aware KV cache eviction methods such as SnapKV, NACL, PyramidInfer, and etc, for further validation.
4. Experimental setting: The paper only focus on the decoding phase of LLM inference on short answer long-context tasks. From my understanding, a large fraction of major latency comes from the prefilling phase(Time-To-First-Token). The authors should include latency for prefilling stage to justify that the decoding speedup worth the sacrificed task accuracy. Moreover, for S=50 on RULER, it is equivalent  to prefilling-time eviction methods(such as SnapKV, NACL), of which the comparison is absent in the paper.

Reference:

[1] SnapKV: LLM Knows What You are Looking for Before Generation. 

[2] NACL: A General and Effective KV Cache Eviction Framework for LLMs at Inference Time. 

[3] PyramidInfer: Pyramid KV Cache Compression for High-throughput LLM Inference.

questions:
See weakness above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730022113376

thread_type:
official_review

<thread_object>
creator: reviewer_njen

<content>
summary:
This work aims to improve the inference speed of long-context large language models. The motivation is clear: restricting certain tokens to attend only to a subset of tokens during decoding, reduces computation and thus accelerates the decoding speed. The method is verified on a popular synthetic benchmark (RULER) and several language modeling datasets to evaluate long-context large language models.

strengths:
1. The motivation and paper writing are clear. Exploring inference acceleration for long-context large language models is highly meaningful. The writing in this paper is very clear, and I can easily understand the work.
2. The method is somewhat innovative. Compared to previous similar works ([1, 2, 3]), this work considers accelerating the decoding stage.

[1] Li, Yuhong, et al. "Snapkv: Llm knows what you are looking for before generation." arXiv preprint arXiv:2404.14469 (2024).

[2] Zhang, Yichi, et al. "PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling." arXiv preprint arXiv:2406.02069 (2024).

[3] Jiang, Huiqiang, et al. "Minference 1.0: Accelerating pre-filling for long-context llms via dynamic sparse attention." arXiv preprint arXiv:2407.02490 (2024).

weaknesses:
1. The benchmarks are limited. The authors validated the method on RULER and several datasets for language modeling evaluation. However, there are two vital problems:
 * RULER is just a synthetic benchmark; more evaluations on real-world tasks, such as the commonly used LongBench [1], are necessary. In addition, the widely used "needle in a haystack" [2] task for evaluating long-context LLMs is necessary.
 * The authors have validated the language modeling capability of their model through extensive experiments. However, recent works ([3, 4]) have pointed out that this metric (perplexity) is not an indicative measure.
2. The baselines are limited. The authors considered two important baselines: StreamingLLM and H2O. This is reasonable because both models apply attention only to tokens within a subset of the context during decoding. However, there are many other similar works that this paper does not address, such as SnapKV [5], PyramidKV [6], MInference 1.0 [7], etc.
3. The performance loss caused by this method is too severe. From Table 2, we can see that RecycledAttention shows a significant performance drop compared to the standard model, with a decrease of 33 points on Llama 3.1 and 32 points on QWEN-2. This level of performance loss is unacceptable in practical applications. When accelerating inference speed, we should prioritize maintaining the model's performance; simply speeding up while sacrificing performance is meaningless.

[1] Bai, Yushi, et al. "Longbench: A bilingual, multitask benchmark for long context understanding." arXiv preprint arXiv:2308.14508 (2023).

[2] https://github.com/gkamradt/LLMTest_NeedleInAHaystack

[3] Gao, Tianyu, et al. "How to Train Long-Context Language Models (Effectively)." arXiv preprint arXiv:2410.02660 (2024).

[4] Hu, Yutong, et al. "Can Perplexity Reflect Large Language Model's Ability in Long Text Understanding?." arXiv preprint arXiv:2405.06105 (2024).

[5] Li, Yuhong, et al. "Snapkv: Llm knows what you are looking for before generation." arXiv preprint arXiv:2404.14469 (2024).

[6] Zhang, Yichi, et al. "PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling." arXiv preprint arXiv:2406.02069 (2024).

[7] Jiang, Huiqiang, et al. "Minference 1.0: Accelerating pre-filling for long-context llms via dynamic sparse attention." arXiv preprint arXiv:2407.02490 (2024).

questions:
The paper is written very clearly, and I have no questions.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730378391603

thread_type:
official_review

<thread_object>
creator: reviewer_kj4i

<content>
summary:
Recycle attention combines the vanilla full attention and H2O. For every generation steps, it performs one full attention step and otherwise uses H2O. The novelty is limited.

The authors did not discuss how they decide the value of $S$ for each task. As a result, the results are questionable. Further discussion is necessary during the rebuttal period.

strengths:
The paper is well-written. Many tasks and baselines are tested.

weaknesses:
1. The novelty is very limited. Recycle Attention is essentially a straightforward interpolation between vanilla full attention and H2O. For every $S$ generation steps, it performs one full attention step and otherwise uses H2O. Therefore, it is not surprising that Recycle Attention performs well on the NIAH task, benefiting from occasional full attention. However, as a trade-off for combining the strengths of both approaches, Recycle Attention inherits their respective drawbacks: it is not as efficient as H2O and not as effective as full attention.

While combining vanilla full attention and H2O with an adaptive $S$ could have been novel, Recycle Attention leaves this to future study.

2. The emprically set $S$ makes experiments questionable. If I missed something, please correct me, but I did not find a detailed discussion on how $S$ was set for each task. In the ablation of S (Table 5 and Line 396), you conclude that there is "a different trend for different tasks." It appears that you set a specific $S$ for each task, making Recycle have optimal performance. If this is the case, it is a test data leakage, making the experiments unreliable.

questions:
Please provide a detailed discussion on how you choose S for each task; this is crucial for evaluating the soundness of your paper. Without this information, I can only give the lowest soundness score, but I am open to revising it once you provide clear guidelines for decising the value of S. 

I am also open to raising the overall rating if the authors can demonstrate that their experiments are reliable and fair.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730559924913

thread_type:
official_review

<thread_object>
creator: reviewer_pfmm

<content>
summary:
This paper proposes a new efficient LLM inference method called Recycled Attention. Observing that the top-k tokens based on attention scores at current step still hold a significant portion of the attention mass over the following steps, authors interleave full attention within sparse attention to more accurately select important tokens, balancing the advantages and disadvantages of full and sparse attention.

strengths:
Previous efficient inference methods based on KV Cache compression typically do not reuse a token once it has been evicted. Recycled Attention addresses this issue by interleaving Full Attention.
Compared to Vanilla Attention, Recycled Attention significantly reduces inference latency.
Compared to StreamingLLM and H2O, Recycled Attention achieves substantial gains on the Ruler Benchmark with nearly comparable inference latency.

weaknesses:
Recycled Attention increases the memory burden. Compared to StreamingLLM and H2O, Recycled Attention requires additional maintenance of a full KV Cache.
The benchmark is limited. Ruler uses synthetic examples for testing. Performance on real sample benchmarks, such as LongBench[1], should also be reported.
After continued pretraining, only PPL is tested. For long context text modeling, PPL is not an intuitive metric. Additional experiments on Ruler and LongBench are needed.
Stronger comparison methods are missing. For example, Quest[2], which has a similar motivation to this paper, should also be included for comparison.
[1] LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding (Bai et al., ACL 2024)
[2] Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference (Tang et al., ICML 2024)

questions:
I am curious about how H2O would perform in the attention mass overlap test in Section 2.2.
Additionally, most tokens in the Full KV Cache also come from the Recycle Steps. As the generation length increases, will there be an accumulation of errors compared to Vanilla Attention, resulting in an increasing performance gap between Recycled Attention and Vanilla Attention?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730581413153

thread_type:
official_review

<thread_object>
creator: reviewer_h8fr

<content>
summary:
This paper tries to tackle the challenges in long context LLMs. Long context LLMs create lots of KV cache during inference, therefore requiring large memory space and high bandwidth requirements. A representative line of work to address this problem is aiming at reducing the number of KV cache stored during inference. However, the authors reported performance problems for these approaches. Therefore, they proposed recycled attention. Instead of completely getting rid of some KV caches, a full attention is performed periodically generation process and a partial attention (i.e. using fewer important KV cache) is performed for the rest of the time. Evidences show that recycled attention can effectively identify the important KV cache and achieves higher performance without sacrificing too much efficiency.

strengths:
1. The authors have used recovery rate to justify the idea of recycled attention. I do think the idea itself presents a thinking of hierarchical attention subject to generation phase.
2. The presentation of this paper is clear and convincing.

weaknesses:
1. It would be good to have theoretical reasoning on the effectiveness of the recycled attention idea.
2. As mentioned in the Limitation section, it would improve the paper if the experiments can go deeper into more settings, like the different stride for different layers. These questions are likely to raised after reading the paper. It is worth to have these results and will be helpful to draw more comprehensive insights. Otherwise, the contribution of the idea seems pretty limited.

questions:
1. Why QWEN appears to get a similar performance for streaming LLM and recycled attention? This pattern is true for both Figure 4 and FIgure5. 
2.  Why the gain is minimal when the output length is small? I feel like if the output length is small, it is possible that only the first step attention is in full form, others should be fast in terms of generating the tokens.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732155307662

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response 1/2

comment:
We thank all reviewers for their reviews and helpful comments. We are delighted to see that they found our work to present an intuitive and novel idea (Reviewer h8fr
, Reviewer PFMm, Reviewer NjEN, Reviewer iC8W), with clear motivation and presentation (Reviewer kJ4i, Reviewer NjEN, Reviewer iC8W) and achieving substantial gains compared to previously proposed methods (Reviewer PFMm). Here we include new experiment results and clarifications regarding common requests from the reviewers. 

We have uploaded a new manuscript with new experiment results / discussion (highlighted in blue text or yellow background) and updated results (highlighted in red).

**Summary of changes made to the manuscript**:
* New baselines added (SnapKV). See Section 3.2.
* Results added on LongBench datasets, in Appendix section A.2. 
* Dynamic stride selection based on the similarity of queries of the current time step with that of the last full attention step, in Section 6.  
* Discussion of the suggested baseline in the related work section (Section 7).
* Discussion of aggregation methods for GQA models, in Section 3. We have also updated the experiment results for Recycled Attention with max aggregation (instead of first in the group) in all the tables.
* RULER results for continued pre-training, see Table 5 in Section 5.

We elaborate on these below.


**More baselines**:
We thank the reviewers for suggesting other relevant baselines (Reviewer PFMm, Reviewer NjEN, Reviewer iC8W) from recent work. Here we provide a discussion on them. We have also added a paragraph in the related work section of the updated PDF and updated result table to include suggested baseline (SnapKV). 

* (1) Query-aware permanent KV cache eviction method:
SnapKV[1]  (Reviewer iC8W, NjEN), NACL[2] (Reviewer iC8W)  and PyramidKV[3] (Reviewer NjEN) are query-aware KV-cache eviction methods. Among these, SnapKV is the most relevant to our method, as it uses attention scores of the last few tokens in the prompt to select tokens to keep. We focus on comparing against SnapKV and include the new experimental results in Table 2,3,4,9. Our method outperforms / performs on-par with SnapKV for both the language modeling and downstream tasks (RULER and LongBench), with a slightly faster decoding speed. 
 
* (2) Quest [5] (suggested by Reviewer PFMm): Unlike most prior work, our method dynamically selects tokens that are likely to be relevant at the current generation step. Quest [5] is the only other method that maintains the full KV cache and dynamically constructs a smaller KV cache for attention computation. While we leverage previous tokens’ attention pattern, they use the minimal and maximal key values to estimate import tokens for the current input token. This method incorporates PageAttention and Top-K cuda filtering, making inference speed comparison a bit challenging. We are working on adding this baseline as a comparison.

* (3) MInference [6] (suggested by Reviewer iC8W). This approach accelerates the pre-filling stage while we focus on accelerating the decoding stage, so not very applicable in our setting as a baseline. 

* (4) PyramidInfer [4] (suggested by Reviewer iC8W): This is a query-agnostic method and leverages accumulated attention scores to evict tokens during both the pre-filling and generation stage, similar to the H2O baselines that we have included. As demonstrated by our experiments of H2O, such query-agnostic KV cache eviction methods can prematurely evict tokens. 


[1] SnapKV: LLM Knows What You are Looking for Before Generation. NeurIPS 2024.
[2] NACL: A General and Effective KV Cache Eviction Framework for LLMs at Inference Time. ACL 2024
[3] PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling. Arxiv, 2024.
[4] PyramidInfer: Pyramid KV Cache Compression for High-throughput LLM Inference. ACL 2024.
[5] Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference. ICML 2024.
[6] MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention. NeuRIPS 2024.

**Evaluate on more benchmark datasets**:
Reviewers (PFMm, NjEN, iC8W) asked for results on another, “more realistic”, benchmark, specifically LongBench [1] . 

First, we would like to clarify that our initial evaluation dataset, RULER contains 14 tasks. Two of which are QA tasks (reported under “QA” in Table 3) that contain realistic question and answers, in addition to synthetic tasks such as NIAH. 

We include results for two datasets from LongBench (NarrativeQA and Musique) that have input length longer than 10K in Table 9 of Section A.2 in the updated PDF. Overall, we found that our method performs on-par or better with the best baseline (SnapKV) with faster decoding speed, repeating the trend we see in other benchmarks.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732760459066

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response (Further comparison with SnapKV) [1/2]

comment:
We thank the reviewers for their reply to our previous response. We address commonly raised questions among reviewers here and provide experiment results to compare our methods against SnapKV, the newly added baseline.

A few reviewers (Reviewer kJ4i, Reviewer NjEN and Reviewer iC8W) mentioned that our method is mostly on-par with SnapKV in terms of performance on current experiments. Our performance is on-par with SnapKV for the two LongBench datasets (NarrativeQA and Musique)  we added previously, and is better than SnapKV and language modelling, all with a faster decoding speed. Here, we further clarify our differences with SnapKV, and present further evidence for scenarios where our methods are better than SnapKV in terms of performance. 

One major distinction between our method as SnapKV is that we maintain the full KV cache and occasionally perform full attention to refresh the recycle cache, while SnapKV permanently evicts tokens that might be useful in later generation. Conceptually, our method will outperform SnapKV on task settings where LLM has to leverage different tokens in the context based on the tokens that it has generated. This will apply to many real-world scenarios, such as (1) LLM is tasked with generating longer text or (2) chain-of-thought reasoning that requires looking up information from the in-context tokens as LLM’s generation continues.

To provide empirical support, we present two new sets of experiments. 

**Results for two summarization tasks from LongBench**:
We have added results for two summarization datasets from LongBench (GovReport and QMSum; Table 10 in Section A.2 in the updated manuscript). For these two tasks, our method consistently outperforms SnapKV (the best baseline) for the two models we tested, especially for a smaller stride.  We have also added 11 other tasks from LongBench in Table 10 and Table 11.

**Experiment setting for a synthetic task that requires longer generation**
We further design a new synthetic task which requires the model to leverage various information in the context as the generation continues. We have added it to Section A.7 in the updated manuscript and describe the setting briefly in the reply below.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733265288855

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response to AC and reviewers (rebuttal summary)

comment:
Dear AC and reviewers,

Thank you for your service and for organizing and contributing to the review of our works. The review process has helped us greatly improve our manuscript.

We summarize the new results added to the manuscript in response to reviewers’ concerns and suggestions. 

**A new section on dynamic stride:** Multiple reviewers (Reviewer h8fr, kJ4i) suggested dynamic scheduling with an adaptive stride S, instead of the fixed stride we experimented with in our initial submission. We have added a new section (Section 6) on dynamic stride based on query similarity. Our experiment shows that compared to a fixed stride (our implementation in the submitted draft), dynamic stride achieves similar performance with less decoding time (Table 6).

**More baselines:** We have added a new baseline (SnapKV), based on the suggestions from reviewers (Reviewer PFMm, NjEN, iC8W) to all experiment settings in our initial submission (RULER: Table 2 and 3; language modelling: Table 4), and the two sets of new experiments (LongBench: Table 10 and Table 11; Chain-of-keys: Table 12) for the two models we experimented with. Our method outperforms this new baseline, with competitive/better performance and a faster decoding speed. We have also added a discussion on the other baselines suggested by the reviewers in Section 7.

**More benchmarks:** As multiple reviewers suggested us to experiment on more datasets (Reviewer PFMm, NjEN, iC8W), we have added two sets of new experiments, including 11 datasets from LongBench (Table 10 and 11 in Section A.6) and a new synthetic dataset (chain-of-key, in Section A.7). Our method performs better than KV cache eviction methods (such as SnapKV) for longer generation which requires leveraging different information in the context based on what has been generated (two summarization datasets from GovReport and QMSum from LongBench, reported in Table 10; and the new “chain-of-key” task, reported in Table 12).

**Other experiments, analysis and manuscript improvements:** We have also added more results, analysis and discussion for other parts of the paper based on reviewers’ feedbacks:
* Ablation study of attention score aggregation method for GQA models (Section 3.3, Table 7; Reviewer iC8W)
* Ablation study of the choice of S (Table 8 and 9 in Section A.1; Reviewer kJ4i)
* Results on RULER for continued pre-training with Recycled Attention (Table 5 in Section 5; Reviewer PFMm)
* H2O's performance for attention overlap analysis in Figure 2 (Reviewer PFMm)

Together, we present a simple approach (Recycled Attention) to speed up LLM inference. Our method differs substantially from previously proposed KV cache eviction methods, which focus on memory efficiency. Instead, our method achieves speed-up by dynamically constructing a small KV cache for generation based on previous tokens’ attention patterns. We conduct comprehensive experiments (three tasks, 30+ datasets on two long-context models) and show that our method performs competitively or better than baseline methods. We include analysis ablating hyperparameters. We experiment with two extensions of Recycled Attention through (1) dynamic striding and (2) continued pre-training, both directions further improve the performance-efficiency tradeoff.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduce a method for accelerating generation speed for long context lengths of LLMs by alternating between full context attentions and attention for an input token subet.

Although this paper addresses an important task, is well-written, and includes diverse experimental results, reviewers raised ciritical concerns such as lack of novelty and heuristic-based methods due to the existence of similar ideas in kv cache research.

AC also agrees with the concerns and think this paper is not sufficient for ICLR quality.

So, AC recommends rejecting this paper.