In-context KV-Cache Eviction for LLMs via Attention-Gate

Thank you for your detailed and thoughtful comments. Below, we provide our responses to each of your points.

W1: Lack of analysis on data sources and impact of post-training data volume.

A: Thank you for your feedback. The data used for post-training was randomly sampled from a subset of RedPajama, containing fewer than 5,000 examples. The focus of our experiments is not on the specific content of the data but on demonstrating the efficiency of training our AG module. In our CPT approach, only the AG module is trainable while the original model remains frozen, ensuring that the training data does not affect the base model.

Additionally, we have included a supplementary experiment where all baselines are trained on the same dataset for comparison.

For baselines, ($W_q, W_k, W_v, W_o$) are made trainable, while in our method, the AG module is also trainable. The eviction ratio is fixed at 50% for StreamingLLM and H2O. Higher values indicate better performance for all metrics.

Please refer to Appendix B.1 for more results.

W2: Effectiveness limited to llama-2 and lack of analysis on other models and scales.

A: Thank you for your comment. For Mistral-7B, it cannot evict as aggressively as Llama2-7B, possibly because Mistral uses GQA, which already performs a level of compression that Llama2-7B does not. We are currently conducting tests on different model architectures, as well as models of varying scales beyond 7B. The results will be updated as soon as they are available.

W3: Missing analysis on inference efficiency and Attention-Gate overhead.

A: Thank you for pointing this out. Our primary focus is on memory efficiency, as indicated by the KV-Cache eviction ratio (%Eviction). Additionally, we have provided a detailed discussion of how the incorporation of the AG module impacts FLOPs, which can be found in Appendix A (L726–L746) and Appendix B.1 (L756–L765). The actual time overhead is influenced by the kernel design used for implementing AG. We plan to further optimize this design to obtain a more accurate estimate of the time cost.

Thanks for your detailed and insightful comments. Below, we provide our responses to your feedback.

W1: Bias in attention-gate and presence of attention sinks.

A: Thank you for raising this point. The core motivation of our method is to develop in-context eviction policies, with addressing attention bias being one of its benefits. The attention bias in prior methods, such as H2O, arises from their reliance on local information to decide which tokens to retain. In contrast, our approach leverages global information by processing the entire sequence through the AG module to generate the eviction policy, effectively avoiding attention bias. The attention sinks observed in Figure 1 are a common phenomenon first identified by StreamingLLM. Our method utilizes this phenomenon to improve the eviction process.

W2: Limitation to prefilling stage; Consideration of decoding stage.

A: Thank you for pointing this out. In fact, our method can also be applied during the decoding phase by storing only the KV-Cache of the AG module. Since the AG module has a limited number of heads, the additional storage requirements would be minimal. We plan to explore this approach further.

W3: Clarification on the use of STE.

A: Thank you for your question. As mentioned in Section 3.4 (L311–L315), the straight-through estimator (STE) is used to address the non-differentiability issue during training. In the AG module, after applying the sigmoid function, a threshold is used to determine whether to evict tokens, which introduces non-differentiability.

The STE applied in our method can be expressed as follows:

$\text{STE}(x, \tau) = (1_{x \ge \tau} - x).\verb|detach()| + x$

Here, $\tau$ is the threshold, and $\verb|detach()|$ ensures that the first term only affects the forward pass while not contributing to the backward propagation. This allows the training process to proceed smoothly while preserving the intended threshold-based behavior.

W4: Understanding why eviction improves performance.

A: Thank you for this interesting question. As mentioned in Section 4.2.2 (L421–L422), "effectively evicting redundant tokens allows the model to focus on relevant information and improves performance." A well-designed eviction policy can mask redundant tokens, enabling the model to concentrate more effectively on useful information, which may lead to better outcomes.

Additionally, we agree with your observation that eviction might act as a form of regularization by implicitly reducing noise in the input, which could be an area for further investigation.

W5: Lack of in-depth comparison on efficiency.

A: Thank you for highlighting this important aspect. We apologize for the omission. Our primary focus is on memory efficiency, as reflected in the KV-Cache eviction ratio (%Eviction). We have also included a detailed discussion of how incorporating the AG module affects FLOPs, available in Appendix A (L726–L746) and Appendix B.1 (L756–L765). The actual time overhead depends on the kernel design for implementing AG. We plan to further optimize the kernel design to obtain a more accurate estimate of the time cost.

Thank you for your thoughtful comment. We appreciate your concern that theoretical metrics like eviction ratio and FLOPs might not fully reflect practical efficiency.

To provide more clarity, we conducted tests on a single 40GB A100 GPU, comparing inference speed and peak memory usage during the prefiling phase between our implementation (with AG) and the original LLaMA2-7B model (without AG) at various prompt lengths. Below are the results for AG with a 50% eviction ratio:

Note that the model without AG uses the regular self-attention implementation whereas our model with AG is currently implemented with a for loop over the attention heads. Despite this, our method is substantially superior over the baseline, especially for long context cases. We will further develop a kernel-level optimization with triton to avoid the for loop for enhanced efficiency recently.

We hope this practical evaluation addresses your concerns and highlights the real-world benefits of our approach. Thank you again for your valuable feedback.

I acknowledge that the evicted keys or values will not be involved in subsequent computations in prefill stage. However, this is achieved with sparse attention mask if I understand correctly. And I am wondering whether this irregular sparse attention pattern be accelerated through flash attention or something like that?

Thank you for your detailed comments. Below, we provide our responses to each of your points.

W1-Q1: Lack of explanation for attention-like structure in Attention-Gate.

A: Thank you for your suggestion. The attention-like structure is a modified version of the standard multi-head attention (MHA) but uses fewer heads. We have added a more detailed formulaic description of the Attention-Gate’s network design and computation flow in Appendix A. For a more in-depth explanation of this attention-like structure, please refer to L680–L687.

W1-Q2: Unclear and potential errors in Equation (1).

A: We apologize for the typographical errors in Equation (1). These have been corrected in the revised version. We assure you that the implementation in the code is accurate and aligns with the intended design.

W2-Q1: Assessment of inference efficiency.

A: Thank you for your feedback. Our primary focus is on memory efficiency, as reflected in the KV-Cache eviction ratio (%Eviction). We have also included a detailed discussion of how incorporating the AG module affects FLOPs, available in Appendix A (L726–L746) and Appendix B.1 (L756–L765). Notably, adding AG does not increase the model’s FLOPs.

W2-Q2: Inclusion of comparisons with other eviction methods.

A: Thank you for your suggestion. Reproducing the baselines presents challenges due to architectural differences; however, we are actively working on it and will update the results as soon as they are available. As highlighted in Fig. 9 of Cold Compress’s Blog, H2O outperforms other training-free methods in certain scenarios, reinforcing the adequacy of our current baselines. Moreover, the focus of our trainable KV-Cache eviction method differs from that of training-free approaches. Our method focuses on the training phase and developing eviction policies that can automatically adapt to individual downstream tasks based on the context. In contrast, training-free methods are designed to provide an on-average eviction strategies.

W2-Q3: Analysis across varying eviction ratios.

A: Thank you for your insightful suggestion. Our current method controls the KV-Cache eviction ratio by adjusting hyperparameters, making it challenging to measure performance at a specific eviction ratio. However, as shown in Table 2 of the paper, reducing the eviction ratio generally leads to improved model performance. As a next step, we plan to integrate a controllable component into the Attention-Gate mechanism to enable multiple compression ratios from a single training process.

W2-Q4 Evaluation on long sequences.

A: Thank you for your suggestion. We have added evaluation results on the Longbench dataset to analyze the performance of our method on long sequences.

For baselines, ($W_q, W_k, W_v, W_o$) are made trainable, while in our method, the AG module is also trainable. The eviction ratio is fixed at 50% for StreamingLLM and H2O. Higher values indicate better performance for all metrics.

Please refer to Appendix B.1 for more results.

Thank you for your detailed comments. Below are our responses to your feedback.

W1: Additional relevant baselines.

A: Thank you for your suggestion. The core motivation of our method is to develop in-context eviction policies, with addressing attention bias being one of its benefits. Reproducing the baselines is challenging due to architectural differences, but we are actively working on it and will update results soon. As shown in Fig. 9 of Cold Compress’s Blog, H2O outperforms other training-free methods in some scenarios, supporting the sufficiency of our current baselines. Moreover, the focus of our trainable KV-Cache eviction method differs from that of training-free approaches. Our method focuses on the training phase and developing eviction policies that can automatically adapt to individual downstream tasks based on the context. In contrast, training-free methods are designed to provide an on-average eviction strategies.

W2: Impact of additional training phase on llama2-7b performance.

A: Thank you for your suggestion. We have added results where all baselines, including LLaMA2-7B without the AG mechanism, are continual pre-trained on the same dataset.

For baselines, ($W_q, W_k, W_v, W_o$) are made trainable, while in our method, the AG module is also trainable. The eviction ratio is fixed at 50% for StreamingLLM and H2O. Higher values indicate better performance for all metrics.

Please refer to Appendix B.1 for more results.

However, we believe our original comparisons remain fair, as our CPT approach only trains the AG module while keeping the original model parameters frozen.

Q1: Evaluation on long-text datasets for practicality.

A: Thank you for your suggestion. We have added results on the LongBench dataset to evaluate the performance of AG in long-text scenarios.

For baselines, ($W_q, W_k, W_v, W_o$) are made trainable, while in our method, the AG module is also trainable. The eviction ratio is fixed at 50% for StreamingLLM and H2O. Higher values indicate better performance for all metrics.

Please refer to Appendix B.1 for more results.

Q2: Clarification on the definition of $\beta$ in Equation (1).

A: Thank you for pointing this out. We apologize for the typo, which has been corrected in the revised version. To clarify, $\beta$ is just a hyperparameter for the auxiliary loss during training and is not required during inference.

Thank you very much for your constructive and positive comments. Here are our responses to your comments:

W1: Questionable choices of the evaluation benchmarks.

A: We apologize for any confusion caused by our choice of evaluation benchmarks. To address this concern, we have added tests on Longbench and MMLU to provide a more comprehensive evaluation.

For baselines, ($W_q, W_k, W_v, W_o$) are made trainable, while in our method, the AG module is also trainable. The eviction ratio is fixed at 50% for StreamingLLM and H2O. Higher values indicate better performance for all metrics.

Please refer to Appendix B.1 for more results.

W2: Unconventional experimental settings.

A: Thank you for your valuable feedback. Training on a general dataset followed by evaluation on various datasets closely resembles the scenario of continual pre-training, which we have already explored. For SFT, if a specific downstream task dataset is available, it can be used to train a tailored eviction policy optimized for that particular task. This scenario is also meaningful and worth investigation.