Recycled Attention: Efficient inference for long-context language models

Thank you for your review, we are glad to see the reviewer found our work to be intuitive with comprehensive experiment settings. Please see our response below:

[W1] Attention aggregation method for GQA

We thank the reviewer for the suggestions on looking into different methods for aggregating attention scores in the same query group. We experimented with different methods to aggregate attention scores for GQA models and included a discussion in the updated manuscript (Section 3). Please see our general response.

[W2] More benchmark evaluation

We thank the reviewers for suggestions on including experiment results on LongBench. We have added results in Table 9 in Section A.2 of the updated manuscript.

[W3] More baselines

We thank the reviewers for suggesting alternative baselines of query-aware KV cache eviction strategies, which are indeed relevant to our method. We have included SnapKV as a representative approach in Table 2,3,4,9 in the updated manuscript. Please refer to our general response for detailed discussion of baselines.

[W4] Latency of pre-filling stage

The reviewer has a good point that our paper focuses on the decoding stage and does not optimize for the pre-filing stage, as we mentioned in our submission (in e.g. Section 2.1). We’d argue that this is a valid motivation, in line with previous work (e.g. SnapKV[1]), and is orthogonal to pre-filling stage optimization.

Thank you for your review and suggestions. Please see our response below:

[W1] The comprehensiveness of evaluation benchmark

We thank the reviewer for suggesting an alternative benchmark, including LongBench and an needle-in-a-haystack task [0]. We would like to clarify that we test on 14 tasks in RULER, which include 8 variants of needle-in-a-haystack task. The NIAH task [0] that the reviewer is referring to corresponds to the setting of niah_single_2 in RULER (mentioned in Appendix B of RULER[1]), which is reported in the paper.

For LongBench, we have included the results in Table 9 of the updated PDF, please refer to the general response for more details.

The reviewer also mentioned that evaluating on perplexity is “not indicative”. While solely relying on language modeling might not comprehensively measure a model's long context performance, it is still valuable to report language modeling performance. Thus, we report a combination of language modeling perplexity and downstream tasks from RULER.

[0] https://github.com/gkamradt/LLMTest_NeedleInAHaystack [1] RULER: What’s the Real Context Size of Your Long-Context Language Models?. COLM, 2024. https://arxiv.org/pdf/2404.06654

[W2] The comprehensiveness of baseline methods

We thank the reviewer for suggestions of alternative baseline methods. We have added experiment results for SnapKV in the updated manuscript, which we believe is the most relevant baseline for our work. Please refer to our general response for discussions on the other baseline methods.

[W3] Performance loss:

The reviewer mentioned that our approach incurs performance loss compared to the vanilla method. First, we would like to point out that our method incurs smaller performance loss compared to baseline methods. Our ablation study of varying K and S (reported in Section A.1 in the updated PDF) shows that increasing these two values can reduce performance loss at the cost of less speed-up, providing a performance-efficiency trade-off. Second, applying the max pooling methods (please refer details to Section 3.3 in the updated PDF) further closes the performance gap for both end tasks and language modeling.

Thank you for your review. We would like to clarify some confusions.

[W1] Is Recycled attention an interpolation of H2O and full attention?

We would like to clarify that our method is different from switching between H2O and full attention. H2O identifies tokens (“heavy-hitters”) based on cumulative attention scores during decoding. In contrast, our method identifies a subset of tokens to keep in the recycle cache based on the attention score of a single previous token (most recent token where full attention was performed). Thus, our method does not require access to attention scores for every decoding step, unlike H2O, which requires extra steps to compute with FlashAttention. In fact, the H2O only baseline already has a higher latency than our approach. Alternating between full and H2O would inevitably make the method less efficient.

The reviewer has a good point that our method provides a middle ground between full attention and KV cache eviction method (such as H2O). However, we would like to clarify that our gain does not merely come from occasionally performing full attention. One evidence demonstrated by our experiment is the performance of the StreamingLLM++ baseline, which performs full attention at the same rate of our method. Table 2 shows that its performance on RULER is significantly worse than our method and close to that of StreamingLLM. Instead, our gain comes from maintaining the full KV cache and flexibly selects the subset of tokens to attend to, based on the previous token’s attention pattern.

[Comment 1]: studying adaptive S

We thank the reviewer for suggesting to experiment with an adaptive S. We report new experiment results for adaptive S based on query similarity, please refer to our general response, and Section 6 in the updated PDF for details and results.

[W2] Selection of S

We would like to clarify that we did not tune S based on test-set performance. Intuitively, setting a smaller stride S will be more computationally expensive (as it entails performing full attention more often) but also more effective (as the recycle cache is refreshed more often). Thus, setting a different S provides a different performance-efficiency trade-off.

We report results for a fixed S which enables empirical speed-up compared to vanilla attention. We note that we ensure all baselines have the same S to ensure fair comparison. Theoretically Recycled Attention reduces attention operation and data movement, yet setting a small stride (e.g. S=2, performing full attention every other step) does not enable speed-up empirically due to compute overheads, which also applies to baseline methods. We further reported results (for both performance and efficiency) Table 8 in Section A.1 for language modeling on arxiv and two RULER tasks, varying K and S for different tasks. This indeed shows that having a smaller S can boost performance yet at the cost of less speed-up. For instance, comparing row 3 and 4 (S=32 and 16), we see that perplexity is lower with smaller stride and yet inference time is longer. Yet, the reviewer raise a good point that it will be helpful to include a comprehensive analysis of K and S on different tasks, which we include for all 14 RULER tasks below, as well as the updated PDF.

RULER (averaged across 14 tasks) for LLama-3.1-8B

We see that compared to language modeling where increasing S is beneficial for performance, increasing S is not as beneficial for RULER (comparing row [4] and [6]), compared to increasing K (comparing row [4] and row [9]). We note that this suggests different tasks might have a different set of (K,S) that will achieve the best performance-efficiency trade-off, and it is possible to choose K and S based on a small validation set.

Your response and the new results have raised further concerns about the tuning of hyperparameters based on the test data in your experiments.

Q1

For language modeling, we did a small pilot study exploring the value of $S$ (2, 5, 10) and chose 10 as that achieves efficiency gain over the vanilla baseline, while smaller stride does not enable speed-up compared to vanilla, though performs better than $S=10$.

As clarified in my previous comment, "when S=2, as you mentioned, your method does not speed up, and if the results between your method and the baselines are comparable under S=2, then your method does not demonstrate efficiency or effectiveness advantages." The key issue is whether the smaller stride values ($S=2$ and $S=5$) outperform the baseline, instead of whether smaller strides outperform larger strides; if they do not, the results could be seen as tuned to the test set, especially given that you did test with $S=2$ and $S=5$ but did not report those results.

To address this concern, please demonstrate that your method outperforms the baselines when $S=2$ and $S=5$ to convince me (even though these settings have worse efficiency, which I acknowledge).

By the way, I have acknowledged that a smaller $S$ improves performance but reduces efficiency, twice. Please avoid repeating this information to make our discussion more productive.

Q2

At any stride $S$, our approach outperforms StreamingLLM++

This seems overly broad now, especially considering the narrow range of strides ($S=5, 10, 15, 20$) used in your comparisons. To strengthen this claim, please show me the comparison with StreamingLLM++ at $S=2$. Because from Table 12, it appears that when $S < 10$, your method actually takes more time and performs worse than the vanilla baseline. Therefore, it is important to show whether your method can still outperform StreamingLLM++ at $S=2$, to justify the claim.

Q3

There is some confusion regarding the stride values ($S$) used in Table 6. Specifically, for the Dynamic method with $QC = 5$ and $s=0.8$, $S=25$ is used for the Arxiv dataset and $S=24$ for the Book dataset. In contrast, the Fixed baseline uses $S=10$ for both datasets. How were these stride values chosen for the Dynamic method in Table 6? If $S$ was tuned for the Dynamic method but not for the Fixed baseline, is this a fair comparison? Additionally, some of the stride values in Table 6 (e.g., 17, 31, 36, 38) seem too unusual and lack clear justification. Could you clarify how these stride values were selected?

Thank you for your review! We are encouraged to see that the reviewer found our work to achieve better performance compared to previous method. Please see our response below:

[W1] Limited benchmark

We thank the reviewers for suggesting alternative benchmarks such as LongBench for our experiments. While it is true that RULER primarily contains synthetic tasks such as NIAH, we note that it also contains realistic tasks such as question answering. We have included LongBench results in Table 9 in the updated PDF and please refer to our general response for more discussion.

[W2] Limited evaluation of continued pretraining (CPT)

We thank the reviewer for suggestions to evaluate the CPT model on the RULER task, which we have updated in Table 6 in the updated PDF. We observe a similar improvement as the language modeling task. We note that we continued pre-trained the model on a relatively small amount of tokens due to our limited compute resource, and it is possible that further CPT can lead to more gains.

[W3] More baselines.

We thank the reviewer for suggesting QUEST as an alternative baseline. We have added another baseline (SnapKV), which is more relevant for our method and we are working on adding a comparison to QUEST. Please refer to our general response about discussion on alternative baselines.

[Q1] H2O performance on attention mass overlap.

Thank you for suggesting to add H2O’s performance on attention mass overlap in section 2.2. We have included it in figure 2 in the updated manuscript. H2O performs better than StreamingLLM, as reflected in our language modeling experiments.

[Q2] Error pattern as the generation length increase:

We do not anticipate there will be accumulation of error, as we recycle the attention pattern from the nearest token (at the maximum, S-1 steps away in a fixed scheduling setting). This is supported by our language modeling experiments, for which we report performance on the last 256 tokens.

Thank you for your question! We conducted the attention mass overlap experiment on the arxiv dataset for Figure 2, and you are right that the performance of H2O and Recycled Attention is close. This is indeed reflected in the language modeling performance in Table 4 (note that we reported performance for S=10 in Table 4; but we did not refresh the Recycled Attention in the attention mass analysis, hence S=50).

However, if we consider the needle-in-a-haystack task from RULER and conduct the attention mass analysis for generating the output, for context length of 8192 and $K=1024$, Recycled Attention recovers over 97% of attention mass while StreamingLLM and H2O recovers less than 90%, as reflected in the results of RULER experiments. This is because H2O uses cumulative attention score to decide tokens to keep and might have evicted the target value from the KV cache. Thank you for the suggestion on analyzing attention mass overlap in different scenarios and we will add this to our updated manuscript.

Thank you for your review and we are glad to see the reviewer found our work to be intuitive and convincing. Please see our response below:

[W1] Theoretical reasoning on effectiveness of Recycled Attention

We have empirical support for our method. Our key intuition for recycled attention is that neighboring tokens are likely to place attention mass on similar sets of tokens in the context. And this is justified by the recovery rate of attention mass when inference with recycled attention (fig 2), as mentioned by the reviewer.

[W2] Setting a different stride for different layers

We thank the reviewer for suggestions on exploring dynamic stride and report new experiment results for dynamic scheduling based on query similarity in section 6 of the updated PDF, please also refer to the general response. In response to the reviewer’s suggestion of setting a different stride at different layers, we have included an analysis on per-layer effective stride in Section A.6 in the appendix, finding that earlier layer having a larger stride.

[Q1] QWEN’s recovery rate:

As the reviewer mentioned, Figure 4 shows that recycled attention’s recovery rate is closer to (although still outperforming) StreamingLLM, compared to LLaMA-3.1-8B. We included the analysis to understand differences between the model, and hypothesized that it might be why we observed better performance for LLaMA-3.1-8B with Recycled Attention compared to QWEN-2.

[Q2] Generation setting:

What we meant by “output length is very small, the efficiency gain will be minimal (line 486)” refers to a setting where the target output length is smaller than the stride (e.g., LLM replies single word response such as “yes” or “no”). We will clarify this in the revision. In most use cases, however, LLM generates long-form responses, in which case the efficiency gain depends on the stride S at which full attention is performed, instead of the length of the tokens generated.

Thanks authors for their detailed response. After reading the general response, I have some further questions listed below:

1. The authors stated that PyramidInfer is a "query-agnostic method and leverages accumulated attention scores to evict tokens during both the pre-filling and generation stage". This is indeed not true. PyramidInfer adopt a similar strategy to SnapKV: during pre-filling stage, it uses only the weighted average of attention weights of recent sequence S_r(which is equivalent to the observation window in SnapKV) to evict unimportant KV pairs. During decoding, it employ the same strategy using a sliding recent sequence window.
2. Another important distinction between RecycledAttention and compared baselines(including SnapKV) is the retention of complete KV cache. This is to say, RecycledAttention still suffer from massive memory usage when applied to large LLMs and long-context tasks. According to results on RULER and LongBench, RecycledAttention is mostly on par with SnapKV in terms of accuracy, slighter faster than SnapKV in speed. I suggest the authors also report the memory footprint of each method to comprehensively reflect the efficiency gain of each approach.
3. Follow the previous comment, my opinion is that since RecycledAttention focus on improving decoding speed(not memory), the paper would be much stronger if the authors could demonstrate the compatibility of RecycledAttention with other memory-efficiency techniques to fully support the claim of "efficient inference".

Dynamic scheduling: Multiple reviewers (Reviewer h8fr, Reviewer kJ4i) suggested dynamic scheduling with an adaptive stride S. We experiment with dynamic scheduling based on query similarity. Concretely, instead of performing full attention at a fixed stride, we decide to perform full attention for the current decoding step or not based on the similarity of query embeddings with the query embedding of the last full attention step. Our experiments show that Recycled Attention can be further improved with dynamic scheduling: it achieves better efficiency-performance trade-off compared to static stride (what we reported in the submitted version of the paper) for both the language modeling tasks and two RULER tasks. We describe the method, experiment setting as well as experiment results in Section 6 and Table 6 in the updated PDF.

GQA aggregation methods: Reviewer iC8W asked about how to aggregate attention scores for GQA models. In our submission, we used the attention score of the first query head to choose top K tokens for the entire query group. We later experimented with taking the average and the maximal score, finding that taking the max performs the best. We have added a discussion in section 3 of the updated PDF (highlighted in blue) and included an ablation study in Table 7 in the appendix. We have also updated the results table (table 2, 3, 4) for recycled attention with max attention scores, showing better performance for both RULER and language modeling evaluation.

Chain-of-key: We define the synthetic task as “chain-of-key” and describe it below.

The context consists of names of keys, each of which contains N number of words, for instance: apricot-waggish.The model is tasked to generate a sequence which consists of a list of keys from the context, such that the first word of the next key is the last word of the current key. For example: waggish-fishery, fishery-mosquito, mosquito-perfume, perfume-panda, panda-juice, juice-willow, willow-bronco, bronco-creditor, creditor-bathhouse, bathhouse-woman Please refer to Table 13 in the updated manuscript for example input.

Evaluation: We evaluate the length of the valid chain. A valid chain needs to satisfy two criteria: (a) the key must be in the context and (b) the first word of the current key must be the last word of the previous key. Please refer to Table 14 for example output and their scores.

We report performance as well as decoding time for all methods. As our experiment shows that LLaMA-3.1-8B is unable to perform this task (accuracy of 0.11 with vanilla attention setting), we conduct an experiment with LLaMA-3.1-70B base model.

Results We find that Recycled Attention consistently outperform baselines that evicts token from KV cache (SnapKV, StreamingLLM) as well as the StreamingLLM++ method, which perform full attention occasionally. SnapKV achieves an accuracy of 0.11, meaning that it is only able to generate a valid key for the first step. We also find that decreasing stride consistently improves performance for Recycled Attention. Please refer to Section A.7 in the Appendix for more details.