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

Authors: Yu Fu, Zefan Cai, Abedelkadir Asi, Wayne Xiong, Yue Dong, Wen Xiao
Primary area: applications to computer vision, audio, language, and other modalities
Keywords: Key-Value cache, Contextual reasoning, Efficiency inference, Large-Lauguage Model

Abstract:
Key-Value (KV) caching is a common technique to enhance the computational efficiency of Large Language Models (LLMs), but its memory overhead grows rapidly with input length. Prior work has shown that not all tokens are equally important for text generation, proposing layer-level KV cache compression to selectively retain key information. Recognizing the distinct roles of attention heads in generation, we propose HeadKV, a head-level KV cache compression method, and HeadKV-R2, which leverages a novel contextual reasoning ability estimation for compression. Our approach operates at the level of individual heads, estimating their importance for contextual QA tasks that require both retrieval and reasoning capabilities. Extensive experiments across diverse benchmarks (LongBench, LooGLE), model architectures (e.g., Llama-3-8B-Instruct, Mistral-7B-Instruct), and long-context abilities tests demonstrate that our head-level KV cache compression significantly outperforms strong baselines, particularly in low-resource settings (KV size = 64 & 128). Notably, our method retains just 1.5% of the KV cache while achieving 97% of the performance of the full KV cache on the contextual question answering benchmark.

Decision: accept (poster)
Reviewer scores: 8, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1729749951902

thread_type:
official_review

<thread_object>
creator: reviewer_vgf6

<content>
summary:
This paper introduces a method for KV-cache compression for efficient language modeling. While previous work compresses the KV-cache at the global or layer level, this paper leverages the properties of multi-head attention and operates at individual heads, achieving a finer-grained computation allocation. The method weighs the importance of each head by considering their contribution to both the retrieval and reasoning processes, resulting in a better selection of past KV cache. With all strengths combined, the method outperforms the baseline methods in multiple tasks that require long-context retrieval or reasoning.

**post review**: I raised my score from 5 to 6.

strengths:
1. This paper refines the method for head importance scoring proposed by Wu et al. (2024) and the authors justify their refinement with ablation studies.
2. This paper proposes a head-level allocation schema that can dynamically allocate memory across attention heads and layers. 
3. The performance of this method is better than baselines under the same memory budget.

weaknesses:
1) The contribution of this paper is relatively marginal in terms of methodology. The modification to the method of Wu et al. (2024) is a shift from retrieval to reasoning and a relaxation of the weighting process, while the overall design of the scoring remains the same. From the KV-cache compression perspective, it improves the granularity of prior work from layer to head without changing the overall compression logic. 

2) The application of the KV cache compression is limited to scenarios where retrieval is required. The compression/allocation is already done after reading the instruction part of the prompt, thus the attention of the generation will be constrained to the selection caches. For tasks such as Needle-in-the-Haystack or retrieval-based QA, this strategy might work, but for the rest of the tasks such as summarization or creative generation, it might fail. 

This might be a general issue for all the work of KV cache selection, but this paper does not break this limitation.

3) This work might not be practically useful in terms of efficiency. Though the authors claim that their method can be as efficient as other methods and better than the full-KV baseline in section 4.5, the analysis is based on a 7B model. For larger models, the overhead on the KV cache and attention matrix is negligible. 

Even for a 7B model, the analysis conducted in this paper shows that the decoding speed is improved only when the decoding sequence length reaches 1000 tokens. For memory usage, the difference is observable when the context length reaches 32k tokens. This might not be a typical case even for the datasets tested in this paper.

questions:
1. What is the experiment condition for the memory usage diagram in Fig 6? The x-axis here is named "context length", so what is the generation length? Also, given the pre-filling phase has not changed, can the proposed method improve the memory usage for context encoding?
2. Finer-grained memory allocation usually means more fragmented memories. From my understanding, the method prunes caches from the token, layer, and head dimensions, will this result in fragmented KV caches? Will this affect the efficiency of parallel computation?
3. It would be great if the authors could discuss and compare this work to RAG methods.

It would be great if the authors discuss other efficiency-concerned papers. A few past work discusses the general compression of KV caches without considering the retrieval process:

- In-context Autoencoder for Context Compression in a Large Language Model. Ge et al, 2024. In ICLR.
- Dodo: Dynamic Contextual Compression for Decoder-only LMs. Qin et al, 2024. In ACL.
- VCC: Scaling Transformers to 128K Tokens or More by Prioritizing Important Tokens. Zeng et al., 2023. In NeurIPS.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730427493587

thread_type:
official_review

<thread_object>
creator: reviewer_njwn

<content>
summary:
The paper proposes a new KV cache compression technique, HeadKV and HeadKV-R2, by selectively discarding or retaining the important KV cache based on different types of head. The proposed method consists of two main stages: retrieval-reasoning head identification and KV cache budget allocation based on the head's type. The paper proposes a new type of NIAH test and claim it can be used to identify heads that can both retrieve and reason about the tokens from the context based on the proposed retrieval-reasoning score estimation equation. After successfully identifying the R2-heads, the method dynamically allocate the KV cache budget to retain for each head based on its score from the previous step. The paper claims to maintain 97% performance, compared with baseline with full KV, even though only retaining 1.5% of total KV.

strengths:
- The study presents a new way to compress KV cache based on different types of attention head is novel, even though there are concurrent works that also work on the same idea.
- The performance of the proposed method surpasses other baselines considered in the paper by a considerable margin at the extreme cases when retained KV size is small.

weaknesses:
1. Even though the paper claims to successfully identify heads that can do both retrieval and reasoning for **long-context** task, which is better than the retrieval-only heads initially proposed in Wu et al. (2024), the NIAH experiment setting in the paper is not long enough (longest prompt = 8k). I believe 8k is considered to be not long enough nowadays. Can the author try longer NIAH test such as 64k or 128k to show the effectiveness of the identified R2-heads?
2. I believe there is a work called "Razorattention" [1] that is released in July which follows the same trajectory as this study, i.e. kv cache compression based on head type. Even though the paper addresses this work in their writing, I don't see any comparison in term of performance between the proposed method and that work, especially two works follow the same directory and Razorattention was released a few moths ago. It is unclear to notice the major contribution of the R2-heads from retrieval head only. Can the author benchmark and compare their performance in your experiment?
3. The estimation equation used to determine R2-head seems to be vague (or even incorrect).
- What is the first sigma, or the t-sigma, sum used for?
- The claim that this proposed estimation method can identify which head is responsible for reasoning is not convincing. Firstly, the study modifies the needle to include reasoning logics & incorrect answer, but do not consider them at all in the estimation equation. The estimation equation only considers the correct answer, c2, as the ground truth, which imo, the same as the original retrieval-identification test. What is the usage of the add-on logics and incorrect answer here if you don't consider it in the estimation?
4. I believe the paper would benefit more from ablation study showing and discussing the effect of different values of hyper-parameter alpha & beta on the performance of the methods.

[1] Hanlin Tang, Yang Lin, Jing Lin, Qingsen Han, Shikuan Hong, Yiwu Yao, and Gongyi Wang. Razorattention: Efficient kv cache compression through retrieval heads, 2024. URL https: //arxiv.org/abs/2407.15891.

questions:
Please see weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730684053082

thread_type:
official_review

<thread_object>
creator: reviewer_npoe

<content>
summary:
This paper proposes a head-level KV cache eviction method, and the authors use the retrieval capability of heads to score their importance, then less important heads evict more tokens. Based on previous retrieval-head criteria, the authors put forward two improvements: one is a more challenging retrieve-reasoning dataset, and the other is using attention weights to refine the score, making the scoring more accurate. By setting different KV cache budges for different heads, this method outperforms other baselines in various evaluation metrics while maintaining the same total KV cache size and inference latency.

strengths:
1. The KV cache budget allocation strategy maintains the total amount of KV cache constant and keeps inference time unchanged.

2. Using top-k attention weights to refine the score can enhance the retrieval-head evaluation. 

3. The proposed retrieve-reasoning dataset may benefit future works.

weaknesses:
1. How the S_h is normalized is not mentioned in this paper. In Equation (4), it seems that S_h should sum to one across all heads and all layers.

2. How the retrieval-reasoning dataset is generated is not mentioned in the paper.

3. The left subfigure in Figure 6 saids decoding times but line 512 mentioned that the decoding time includes prefilling time. This is quite confusing. In the figure, the prefill time of each method (when generating a length of 0) is squeezed into the same point on the image, making it impossible to discern their merits and demerits. Please rename the axis label, move the relevant explanation to the caption, or separate the prefill and decode into two figures.

questions:
1. The a_h in Equation (1) and (2) may need a superscript t.

2. Typo in Figure 1: The red text Prefilling Phrase should be Prefilling Phase.

3. Typo in Figure 5 and Figure 8: 'Neele-in-a-Haystack' should be 'Needle'.

4. Are those gaps between FullKV and HeadKV in the right subfigure of Figure 6 solely caused by KV cache?

5. Giving a pseudo code would be better.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730816496805

thread_type:
official_review

<thread_object>
creator: reviewer_iyge

<content>
summary:
This paper proposed a head-level Key-Value cache compression algorithm, different from Ada-KV (also head-level), they don’t perform  allocation within a single layer during the budget allocation process but for all heads. 

They conduct experiments on LongBench and LooGLE, the performance is consistently better or comparable than the full KV baseline with a reasonable amount of KV cache size.

Also, they introduced the the retrieval-reasoning head to assign higher importance score for those heads with higher attentions on the correct answer. Such head seems useful based on the experiments, though not all of them, still makes reasonable improvements on certain datasets.

They also conduct thorough analysis on long-context retrieval, latency , memory usage, etc.

strengths:
1. Reasoning head-level kV cache allocation and importance score estimation. 
2. Performance can be consistently better or comparable with the full KV setting.

weaknesses:
1. We do not see much improvement on latency and memory as compared against Ada-KV, as I believe this work is based on Ada-KV.

questions:
N.A.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731687084161

thread_type:
public_comment

<thread_object>
creator: ~sophia_fulton1

<content>
title:
Confusion Regarding Hyper-parameter $\beta$ Selection

comment:
Hi, authors. While reading this paper, I found the statement in the Experiment Settings section—"The hyper-parameter $\beta$, which controls the size of the shared budget pool, was chosen from {1.005, 1.01, 1.1, 1.2, 1.5, 2, 5, 10}, and we report the best performance"—somewhat unclear. Does this mean that HeadKV used different $\beta$ values for various budgets or datasets to achieve the reported results? 

It would be helpful if the paper could specify the exact $\beta$ values corresponding to the results presented, as this would improve the transparency of the methodology. Additionally, providing guidelines on how to select $\beta$ without extensive searching would be valuable, as this approach seems impractical in real-world applications.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732168581718

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of our revisions

comment:
**We sincerely thank the reviewers for their thorough reading and valuable feedback.**

**Contribution of the paper:**  

Our paper aims to enhance computational efficiency in long-context scenarios. By designing a fully head-level KV cache compression method during the pre-filling stage, we can retain only 1.5% of the KV cache while maintaining 97% of the performance of the full KV cache. To achieve the fully head-level KV cache compression, we first obtain the Retrieval-Reasoning distribution by refining the importance score estimation and constructing the retrieval-reasoning examples to serve as the dataset for estimation. Secondly, We design a head-level KV cache allocation strategy by creating a shared budget pool guided by the importance score distribution to determine the KV cache retention for each head. We believe that head-level KV cache compression is more meaningful and promising, as the significance of different heads can vary considerably, even within the same layer.

**Overview of Revision:**
- We have made revisions to the expressions and formulas based on the reviewers’ feedback.
    
- We used different colors for the added sections to indicate which reviewer’s comments they correspond to.
    
- We added a section of how to construct the retrieval-reasoning dataset, along with corresponding explanations and final dataset. (Appendix E)
    
- We included additional pseudo-code to clarify the settings and overhead of our proposed method. (Appendix F)
    
- We added an analysis of the hyperparameter $\beta$. (Appendix G)
    
- We added a discussion section to show the differences between context compression and KV cache compression. (Appendix H)
</content>
</thread_object>

</thread>
<thread>
creation date:
1733209306016

thread_type:
public_comment

<thread_object>
creator: ~xiang_liu10

<content>
title:
Missing Reference to Previous KV Cache Compression Work Based on Attention Head Level

comment:
Hi, authors. While reading this paper, I found you should consider referencing the paper "Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs" (ICLR 2024, Oral), which presents a related approach to KV cache compression based on attention head-level granularity for large language models (LLMs). These two papers share similar motivations, making it relevant to highlight this prior work.

Engineering-Related Question:

If each attention head retains a different number of tokens, how is attention computed? Typically, attention mechanisms assume that each head operates over the same number of tokens. I would appreciate clarification on how this issue is addressed in the current framework.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces a new key-value (KV) cache compression method that allocates KV cache across attention heads based on their importance. It also presents a variant that refines the importance score considering retrieval and reasoning capabilities, i.e., identifying heads with higher attention to correct answers using a retrieval-reasoning dataset. Results on long-context LM benchmarks and long-context QA demonstrate consistent improvements under varying KV cache budgets.

Strengths
- A new KV cache allocation and importance score estimation (iyGe, NPoe, NjWN)
- Performance is consistently better (iyGe, NjWN)
- The new retrieve-reasoning database may be beneficial for future work (NPoe)

Weaknesses
- Improvement is marginal compared to Ada-SnapKV (iyGe) - In AC’s opinion, improvements are significant with KV size 128 but less with KV size 1024, partially because when KV size = 1024 all KV cache methods achieve performance that is close to full KV.
- Limited novelty (vGF6)
- Limited application scenario – this would be more useful for understanding tasks, summarization tasks, or tasks with CoT where substantial retrieval of the given input text is required, but less useful for other tasks such as long-term text generation (vGF6).