Not All Heads Matter: A Head-Level KV Cache Compression Method with Integrated Retrieval and Reasoning

First, we want to emphasize that long-context scenarios, along with KV cache compression methods designed for these scenarios, are critically important as they lay the foundation for improving computational efficiency. These approaches are essential not only for handling extensive input contexts but also for facilitating seamless transitions to long-generation tasks. By optimizing the management and compression of KV caches during the pre-filling stage, we ensure that the LLMs remain capable of processing large volumes of information while maintaining high-quality and coherent outputs, which is crucial for the next long-generation applications, such as RAG, summarization and long-form QA.

Q1: Limit Novelty

Our method mainly consists of two novelty parts: (1) refining the importance score estimation and constructing new retrieval-reasoning examples to obtain the Retrieval-Reasoning Heads distribution, and (2) performing head-level KV cache budget allocation. While Retrieval Heads already exist, effectively combining Retrieval Heads with KV cache compression remains an interesting and worthwhile problem to explore.

• Importance Score Estimation: For importance score estimation, we made modifications to the calculation method and incorporated attention scores into the process, as shown in Eq. 2. We compare the results of directly using the standard Retrieval Heads distribution with those obtained using our Retrieval-Reasoning Heads distribution (as shown in Table 1 and Figure 3). The results demonstrate that our refined importance score estimation significantly improves the model's performance across various KV size settings. Compared to the standard Retrieval Heads distribution, our modified estimating approach yields a denser and more precise distribution, enabling better guidance on the amount of KV cache retained for each head (as shown in Figure 4). The results shown in Table 2 further highlight the importance of constructing Retrieval-Reasoning examples rather than relying on standard Retrieval examples.

• Head-Level KV cache budget allocation: Next, we further explore how to better utilize the head-level importance score distribution for KV cache compression. Our method is the first to perform fully head-level KV cache budget allocation and achieves SOTA results on Longbench. Other methods, such as RazorAttention, have also recognized the existence of Retrieval Heads distributions. However, they merely use retrieval heads to decide whether to apply FullKV or StreamingLLM to retain KV cache, which leads to limited performance gains. Additionally, such methods often dramatically increase the total number of retained KV cache entries to achieve fair results. In contrast, our approach effectively balances KV cache usage and performance.

Overall, achieving a better importance score distribution and further integrating it with KV cache compression algorithms remains an interesting and under exploration problem. Our efforts on these two aspects provide a potential solution and achieve SOTA performance on the corresponding benchmarks.

Q2: Limited to retrieval-based scenarios

In this paper, we focus on KV cache compression during the pre-filling phase to deal with long-context input scenarios.  It is necessary to design different strategies during the pre-filling phase to selectively retain portions of the KV cache, thereby achieving a balance between computational efficiency and performance. The selected KV cache does not simply store information from specific token positions but retains a holistic representation of the overall input. Therefore, for summarization tasks, our proposed method also achieves better results as shown below:

This table is derived from Table 6 in our paper, where we present the results on Longbench, demonstrating that our proposed method achieves the best overall performance.

For tasks like creative generation, they rely more on the model's inherent capabilities rather than information from the input to generate the corresponding results, which does not align with the objective of this paper. In fact, generating long text is another interesting and worthwhile direction to explore, with relevant works such as [1][2].

[1] Large Language Models Still Exhibit Bias in Long Text

[2] Language Models can Self-Lengthen to Generate Long Texts

Q3: the overhead for larger models

The following code demonstrates the additional steps required for head-level KV cache budget allocation based on the obtained importance score distribution:

# Load saved importance score distribution
path = `path to the saved distribution file`
with open(path, 'r') as file:
	head_list = json.loads(file.readline())
	
# accumulate the importance score and convert to tensor
head_score_list = [np.mean(l[1]) for l in head_list.items()]
head_score_list = torch.tensor(head_score_list / sum(head_score_list))        

# accumulate the importance score and convert to tensor
head_score_list = torch.tensor(head_score_list / sum(head_score_list))

# Obtain the importance score distribution
total_attention = head_score_list.reshape(num_hidden_layers, num_attention_heads)

# Construct shared budget pool and define the minimum KV cache budget
total_pool_capacity = (base_capacity // beta) * num_hidden_layers * num_attention_heads
min_num = (base_capacity - base_capacity // self.beta)

# Head-level Allocation based on the importance score
head_capacity = torch.round(total_attention * total_pool_capacity + min_num).int()

Ideally, we only need to initialize the required head_capacity during the model's initialization phase, since our importance score distribution is static, and we do not need to adjust head_capacity during the entire dataset's execution phase. For example, in the case of the Llama-3-8B-Instruct, which consists of 32 layers with 32 heads per layer, the length of the head_score_list is 32×32=1024. For larger models, such as the Llama-3-70B-Instruct, which has 80 layers with 64 heads per layer, the corresponding length of the head_score_list is 80×64=5120. The times required to execute the above code for these models are shown as follows:

The above results indicate that although the initialization time increases with the size of the importance score distribution, the impact remains minimal because:

1. The time required for this operation is negligible compared to the decoding time (43.8s for FullKV cache to generate 512 tokens).

2. Initialization only needs to be performed once during the entire runtime, and no dynamic adjustments are required. As the dataset size increases, the initialization time can be further amortized.

Q4: computational efficiency

For computational efficiency, Figure 6 demonstrates that our KV cache method achieves comparable computational efficiency to other KV cache compression baselines under the same settings. This means our proposed method delivers significant performance improvements without introducing additional overhead.

Regarding the specific results, on Longbench, we analyzed summarization tasks across three datasets: Gov-Report, QMSum, and Multi-News. The average generation length across these datasets is 406.33, with Gov-Report having an average result length of 817.4. These statistics are based on tokenized results using the tokenizer from Llama-3-8B-Instruct. Below, we present a speed comparison between the FullKV cache method and our proposed method when the generation length is set to 512 and 1024 separately:

The KV cache compression method achieves approximately a 2x speedup compared to the original FullKV method (with FlashAttention). We believe this improvement is significant, as FlashAttention itself already provides a substantial speedup (around 4x, as demonstrated in their paper) compared to vanilla attention.

Besides the decoding latency discussed above, another point is Peak Memory Usage. While the difference in memory usage becomes more noticeable when the context length reaches 32k tokens, we believe that any reduction in GPU memory usage is meaningful, as it directly contributes to the efficiency and scalability of large language models. Furthermore, there are now numerous benchmarks exceeding 100k tokens, demonstrating that scenarios involving long contexts are highly relevant. In this work, we followed the settings of previous studies and primarily conducted tests and performance comparisons on LongBench.

Q5: experiment condition and memory usage for context encoding

For the Peak Memory Usage diagram in Figure 6, we set the generation length to 1. In our paper, we focus on the selection and eviction of KV cache during the pre-filling phase. Therefore, we set the generation length to 1 to emphasize the pre-filling stage. The experimental results show that although KV cache calculations are still required during the pre-filling phase, our proposed method, along with other KV cache compression baselines, can effectively optimize overall GPU memory usage.

Q6: Fragmented KV caches

To select KV cache at the head-level and retain different sizes for each head, this may lead to certain discontinuity issues compared to layer-level KV cache compression methods. However, the experimental results in Figure 6 demonstrate that, in sequential execution, head-level KV cache compression methods do not negatively impact overall computational efficiency.

To further analyze the impact of head-level KV cache compression in parallel computation, we conducted experiments comparing the effects of our proposed head-level KV cache method and other KV cache methods, including other head-level methods. These experiments were performed using two A6000 GPUs, with experimental settings and models consistent with those used in the decoding latency experiments in Figure 6. The results are as follows:

Compared to sequential execution, the head-level method introduces some overhead in parallel computation due to the head-level operations it requires. Therefore, there is a trade-off between performance and speed. As shown in Table 1, head-level KV cache compression methods, such as Ada-SnapKV and our proposed method, achieve better performance compared to layer-level KV cache compression methods like SnapKV and PyramidKV. Moreover, compared to Ada-SnapKV, our method achieves better performance, while maintaining the same computational efficiency.

Q7: Compare with RAG methods / other efficiency-concerned papers

Thank you for your valuable suggestions and the provided related papers. We added one discussion section in our revised version. (Appendix H)

context compression is an interesting direction closely related to KV cache compression. For example, ICAE employs an additional trained In-Context AutoEncoder to compress the input into a fixed-length memory slot, which is then used as the input to the model. From the perspective of ICAE, current KV cache compression methods can be seen as compressing input context using the model's own knowledge. For instance, the SnapKV method, which our approach builds upon, uses the last alpha tokens as the observation window and selects retained KV cache based on attention from these tokens.

Compared to context compression methods like ICAE, current KV cache compression methods are simpler, as they do not require training additional models. They are also easier to obtain higher computational efficiency since these KV cache compression methods avoid the need to rely on external models to obtain compressed inputs.

Dear Reviewer vGF6;

Thank you again for your valuable input. We believe your suggestions were incredibly insightful, and we have provided detailed responses addressing your concerns, including the motivations and innovations of our paper (Q1), its performance in other scenarios (Q2), the impact of model size on overhead (Q3), computational efficiency (Q4), and the effects of fragmented KV cache (Q5). We would like to know if our responses addressed your concerns and provided satisfactory clarification？

Additionally, we have included further discussions on related methods, particularly regarding context compression, in Appendix H. We are more than willing to engage in further discussions to refine and enhance the paper. Thanks!

Q1: Thanks for emphasizing the contribution of this paper. I understand that the two major contributions, but they might be marginal to prior works.

Q2: Thanks for clarifying the scope of this paper. From the writing of the paper, the method seems like a general approach for all kinds of text generation. It would be better if the authors could mention the limitations of their method in the paper.

Q3: Thanks for conducting these analysis. I am not concerned if the scoring or the selection process are slow but whether reducing the KV cache is necessary. This blob post about transformer flop demonstrates that the quadratic terms are only a tiny part of the overall transformer computation, given the majority of the compute is on MLPs, especially for larger models. In another word, the sequence length must be very long to make the KV cache compression to make sense.

Q4, Q5, and Q6: Thanks for providing more details. Given this paper cares about efficiency, it would be great to incorporate them in the next-version of this paper.

Additional discussion about related work: Thanks for doing that.

Overall: Given these clarifications and improvements, I am willing to increase my score to 6. Sorry for replying very late. If authors find it hard to further discuss on any of the above topics, please do them in the paper.

Q3.1: estimation equation used to determine R2-head

We refine the estimation method by focusing on the entire answers $c^2$ rather than the only token with the highest attention probability. Eq. 2 can be used to obtain the importance score for each head. For this equation:

1. (The i-sigma): We consider the tokens with the top-i highest probabilities and compute the importance score by accumulating the attention scores of the tokens that appear in the correct answer. This basic motivation here is: first, if one head is important, it can pay attention to all the tokens within the correct answer (why top-$i$). Second, tokens with a higher attention score should contribute more to the importance score if these tokens can be found inside of the correct answer (why use attention score $a_i$).
2. (The t-sigma): We follow the setup of Eq. 1 to accumulate the importance score for each head step-by-step. Ideally, The maximum score for the i-sigma (inner sum) should be 1/N. By accumulating the importance score step-by-step (the t-sigma), the maximum score for each head should also be 1, which is the same as Retrieval Heads. Importance score distributions shown in Figure 4 are normalized distributions. Compared to standard Retrieval Heads Distribution, our new distribution should be more dense, since we focus on the total correct answer rather than only focusing on the token with the highest attention score. A denser distribution plays an important role in guiding KV cache eviction, as it allows us to set the KV cache size more specifically for each head. This is something that the standard retrieval head distribution cannot achieve, as around 70% of heads receive a value of zero in the standard retrieval head distribution, making it impossible to effectively distinguish and allocate budget for these heads.

Q3.2: add-on logics and incorrect answer.

The addition of reasoning logics and incorrect answers aims to introduce complexity and context to the reasoning process, which we believe might highlight different heads depending on whether the head supports accurate reasoning patterns or not. Here’s how we envisioned their role:

• Aligning with the requirements of contextual reasoning: Based on the in-depth analysis on the contextual reasoning dataset, we know that the answer to the corresponding question will still appear in the input but with various distractors. Therefore, the model continues to rely on the retrieval-and-paste mechanism to obtain the true answer. The original retrieval heads estimation method did not account for this phenomenon, but we address it by adding logic and simulating incorrect answers to achieve a more accurate distribution.

• Introducing diverse reasoning paths: By incorporating both reasoning content and incorrect answers, we are simulating two different potential reasoning paths. The incorrect answer acts as a distractor and we hope to find heads that concentrate on the correct answer even though the incorrect answer has almost the same structure with the correct answer. We expected the correct answer $c^2$ to be treated as the ground truth in the estimation equation, similar to the original Retrieval Heads estimation method.

• Focusing on the correct reasoning path: The purpose of constructing retrieval-reasoning examples is to obtain the importance score for each head, which then guides the head-level KV cache budget allocation. Therefore, our goal is to identify the important heads rather than those focused on incorrect answers. By emphasizing the important heads, the heads that focus on incorrect answers are naturally ignored, as they share the same shared global budget pool.

• Aligning with the standard retrieval heads estimation: we followed the setup in obtaining the Retrieval Heads and determined the Retrieval-Reasoning Heads distributions based on the NIAH experiments. Since NIAH only outputs the corresponding results, we choose to focus on the correct answer to ensure that the maximum importance score each head can achieve in Eq. 2 is 1. Adding additional logic would disrupt this property, potentially affecting the final distribution.

We agree that retrieval heads play an important role in guiding head-level KV cache allocation, as shown by the Head-R results in Table 1, which also significantly outperform the other baselines. However, due to the sparsity issue, they are less effective in guiding head-level KV cache budget allocation. To address this, we: (1) incorporated retrieval-reasoning examples, and (2) refined the importance score estimation. The results in Table 2 demonstrate that using only retrieval examples as defined in retrieval heads (Head-R) or plus refining the importance score estimation method (Head-ER) does not achieve optimal performance. Therefore, both constructing retrieval-reasoning examples and refining the importance score estimation are necessary for optimal results (Head-R2).

Q1: NIAH experiment setting in the paper is not long enough (longest prompt = 8k)

In the NIAH experiments, we followed the settings from SnapKV and PyramidKV and conducted the experiments within the maximum training length supported by each model. Since the maximum supported length for the Llama-3-8B-Instruct model is 8k, the maximum length in Figure 5 is set to 8k. In Figure 8, we present the NIAH experiment results for the Mistral-7B-Instruct model, with a maximum length of 32k. The experimental results on Mistral are consistent with those on Llama-3, which demonstrates the effectiveness of our method in the NIAH task.

Q2: RazorAttention benchmark

Thank you for your suggestion. In Table 1, we also provide the results of head-level KV cache budget allocation based on the standard retrieval head distribution (Head-R). Since RazorAttention has not published their code and did not provide details on the experimental setup used to obtain their results, we implemented their approach and conducted comparative experiments on Llama-3-8B-Instruct by ourself.

Following the settings provided in the PyramidKV codebase, we use a window size of 8 and an attention sink size of 4 for reproduction. Taking six QA datasets from Longbench as examples, their average length is 8640. Therefore, to ensure a fair comparison, we need to maintain a consistent number of total KV cache in the model after performing KV cache eviction. When the KV size is 128, we can obtain a total of (128 - 8 - 4) * 32 * 32 = 118,784 KV cache entries for those retrieval heads to maintain a full KV cache. Considering the average length of 8640, the number of retrieval heads that can maintain a full KV cache is 118,784 / 8640 ≈ 14. Therefore, we set the number of retrieval heads (chosen based on the retrieval score) to 5, 10, and 20 for a fair comparison. Results are shown below:

When the KV size is 1024, we can obtain a total of (1024 - 8 - 4) * 32 * 32 = 1,036,288 free KV cache budget and the numbers of Retrieval heads with full KV cache should be: 1,036,288 / 8640 ≈ 120. Therefore, we set the number of retrieval heads (maintaining a full KV cache) to 100 and 500 for a fair comparison:

The results indicate that RazorAttention heavily relies on setting a large number of retrieval heads to maintain full KV cache in order to achieve good performance. This suggests that simply using streamingLLM on non-retrieval heads is insufficient for effectively retaining important information. Under the same KV cache settings, our method significantly outperforms RazorAttention, as we dynamically allocate KV cache size based on importance scores and use SnapKV to select the retained KV cache for each head.

Another noteworthy point when reproducing RazorAttention is that its method of compressing information from dropped tokens into a “compensation token” is highly time-consuming. When only the top k=1000 retrieval heads maintain full KV cache, the time required to run NartvQA is three times longer than with top k=5 retrieval heads (16 minutes vs. 5 minutes), which contradicts the goal of KV cache compression.

Q4:the effect of different values of hyper-parameter

We included one section about hyper-parameter $\beta$ into our revised version in Appendix G. The only hyper-parameter introduced by our method is $\beta$, which defines the size of the shared global budget pool $B$. Other hyper-parameters, such as the number of instruction tokens $\alpha$, are kept consistent with the settings provided in the PyramidKV codebase. We also ensure that all other hyper-parameters are consistent across both the baselines and our proposed method. For the hyper-parameter $\beta$, as we said in Section 4.1 Line 305, it was chosen from {1.005, 1.01, 1.1, 1.2, 1.5, 2, 5, 10} and we report the best performance.

For $\beta$, a smaller value represents a larger shared budget pool $B$, meaning that KV cache allocation relies more heavily on the importance score distribution for allocation. The results in appendix G show that Head-R2 performs better with a smaller $\beta$, indicating that our retrieval-reasoning head distribution is more effective in guiding KV cache budget allocation.

Q1: S_h normalization

Thanks for pointing out that problem. S_h​ is L1-normalized after collecting the importance score distribution. We ensure that the sum of S_h​ equals 1 to guide the subsequent head-level KV cache budget allocation.

Q2: How to construct retrieval-reasoning dataset

Thanks for pointing out, We included the details on how to construct the retrieval-reasoning dataset in the appendix E of the revised version.

For guiding head-level KV cache compression, we need to obtain the importance score for each head. To achieve this, we manually construct specific examples to ensure that the model relies on heads rather than internal knowledge to answer questions during the Needle-in-a-Haystack experiment. Therefore, we construct retrieval-reasoning examples based on retrieval examples by introducing different reasoning paths into examples to emphasize the contextual reasoning ability. One constructed retrieval-reasoning example is shown in Figure 2. In addition to the current example, we reverse the question to create a total of two examples for the Needle-in-a-Haystack experiment.

Following the setup outlined in the Retrieval Heads paper, in the Needle-in-a-Haystack experiment, we use the model's maximum training length as the maximum haystack length and evenly select 5 different length values as the actual haystack lengths. For each haystack length, the question is inserted at 10 different depths, uniformly distributed from the start to the end of the current haystack length. In total, we generate 100 examples per model to collect Retrieval-Reasoning Head distributions.

Q3: Misleading between decoding latency and decoding time

Thanks for your feedback. We fixed those errors in the revised version. We should change the y-axis in the left subfigure of Figure 6 to "time". The description in Line 512 is correct—decoding latency includes both the pre-filling time and the decoding time. Pre-filling refers to the KV cache eviction phase performed after each example is encoded by the model, while decoding refers to the generation of the output after the pre-filling phase is completed. Therefore, when the generation length is set to 1, the decoding latency reflects the time required for the model to encode the current input and perform pre-filling. The detailed decoding latency results when the generation length is set to 1 for our strong baseline method Ada-SnapKV and our method are shown below:

Based on the average results from three rounds, our method does not introduce significant additional time. In contrast, Ada-SnapKV may require extra time to compute attention and perform sorting to determine the corresponding allocation strategy. Our Retrieval-Reasoning Head distribution is a static distribution, allowing the allocation strategy to be obtained with minimal overhead. However, the pre-filling time is still negligible compared to the decoding time when generation length is relatively large.

Q4: Typo Error & pseudo code

Thanks for your careful reading, and we will fix those errors in the new version. We also add a pseudo code in Appendix F. Our code is based on the PyramidKV and Ada-KV implementations, and we will make our code publicly available. Compared to Ada-SnapKV, our main difference lies in constructing the corresponding allocation strategy for each head. Since the obtained standard Retrieval Heads distribution and Retrieval-Reasoning Heads distribution are static and do not change with modifications to the input, we can perform a one-time initialization when the model is loaded.

Q5: gaps between FullKV and HeadKV solely caused by KV cache

In the Peak Memory Usage results shown in Figure 6, we compared memory usage while maintaining consistency in the input context length. Unlike FullKV, all other KV cache compression methods limit their modifications to the selective eviction of KV cache. Therefore, we can conclude that the memory gap between FullKV and the other KV cache compression methods is entirely caused by differences in KV cache related operation. We attempted to further trace different tensors and analyze the tensors stored in GPU memory. However, we were only able to identify tensors corresponding to model weights, inputs, and RoPE, which remain constant between FullKV and all other KV cache compression methods. We could not directly trace tensors related to the KV cache. This limitation is likely due to the use of additional variables and data structures within the transformers framework to manage KV cache storage.

Q1: Computational Efficiency compared to Ada-KV

Thank you for your summary and it is very accurate and comprehensive. While maintaining a consistent KV cache, our method significantly outperforms other baselines, including the Ada-KV method. This indicates that we can achieve results comparable to Ada-SnapKV and other baselines while using a smaller KV cache. For example, in Table 4, on the Llama-3-8B-Instruct model, our proposed method HeadKV-R2 achieves a score of 32.51 with a KV size of 256, surpassing Ada-SnapKV’s score of 32.31 with a KV size of 1024. The same conclusion holds for the Mistral-7B-Instruct model, where HeadKV-R2 achieves 32.24 with a KV size of 256, compared to Ada-SnapKV’s 31.98 with a KV size of 1024.

We really appreciate your attention!

In the revised version, we have added an analysis of the hyperparameter $\beta$ in Appendix G. Conducting a grid search to identify the optimal hyperparameter is a common practice in NLP research. In our case, we maintained consistency within the benchmark datasets, meaning we did not fine-tune $\beta$ for each individual dataset to obtain the best possible results. While such dataset-specific tuning could lead to better overall results, it would deviate from the original intent of our paper. As shown in the analysis in Appendix G, our proposed method consistently outperforms the best baseline, Ada-SnapKV, under various $\beta$ settings. This demonstrates the robustness of our approach.

Regarding the selection of $\beta$, we recommend using a smaller value as an initialization and use a small calibration dataset to determine the $\beta$. This aligns with the experimental results in Appendix G, where smaller $\beta$ values yield better performance. Additionally, a smaller $\beta$ indicates that the model relies more on the provided importance score distributions to allocate budget. When $\beta$ is set to an extremely large value, causing the shared budget pool $B$ to reduce to zero, our algorithm degenerates into SnapKV. Thus, SnapKV can be considered the lowest bound of our proposed method.

Thank you for your interest and for pointing out this issue!
We will include the relevant citations and highlight this prior work in future version. Our code is based on AdaKV [1], with the distinction that we use retrieval-reasoning heads to guide the head-level KV cache budget allocation. For the implementation, the flash_attn_varlen_func function provided in flash-attn[2] handles variable-length sequences for different heads to address the issue of unequal lengths. More settings and implementation details can be found in the code published with the AdaKV [3].

[1] Ada-KV: Optimizing KV Cache Eviction by Adaptive Budget Allocation for Efficient LLM Inference

[2] https://github.com/Dao-AILab/flash-attention/blob/c4b9015d74bd9f638c6fd574482accf4bbbd4197/flash_attn/flash_attn_interface.py#L1051

[3]https://github.com/FFY0/AdaKV