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
Venue: ICLR 2025 Poster (2025)
Keywords: Key-Value cache
Contextual reasoning
Efficiency inference
Large-Lauguage Model
Primary area: applications to computer vision, audio, language, and other modalities

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.

Published as a conference paper at ICLR 2025
NoT ALL HEADS MATTER: A HEAD-LEVEL KV CACHE COMPRESSION METHOD WITH INTEGRATED RETRIEVAL AND REASONING
Yu Fu, 1 Zefan Cai,² Abedelkadir Asi, Wayne Xiong, Yue Dong, Wen Xiao University of California, Riverside, University of Wisconsin-Madison, "iicroooft yfu0930uer edu,, wxiao@microsoft.co
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 equally im- portant for text generation, proposing layer-level KV cache compression to selec- tively retain key information. Recognizing the distinct roles of attention heads generation, we propose VeadKV, 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 reason- ing capabilities. Extensive experiments across diverse benchmarks (LongBench, LooGLE), model architectures (e.g.. Llama-3-8B-Inse Instruct, Mistral-7B-IIsttruc)) and long-context abilities tests demonstrate that our head-level KV cache com- pression 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. All code and data are available http://giibb..oo///YY/U/U///
INTRODUCTION
Modern Large Language Models (LLMs) increasingly support extremely long inputs: GPT- (Achiam et al.) 2023), Llama-3 (Dubey et al. 2024), and Qwen-2 (Yang et al. 2024) handle
up to 128K tokens, while Claude (Anthropic 2024) supports up to I million tokens. These exlended
capacities improve performance on tasks like dialogue generation (Li et al. 20024 Yi et al. 2024)
question answering (Ho al. 2020 Xu et al. 2023), and summarization [iiao & Carenini 2019
Koh et al. 2022). As input lengths increase, memory usage and latency grow significantly due to self-attention in transformers. To improve inference speed and efficiency, most LLM inference
sists of two phases: prefilling for input processing and decoding for token generation, with key and value states from attention cached for reuse (KV cache). However, as input length increases, KV cache memory grows rapidly, posing significant challenges for storage and efficiency.
To address this, KV cache compression methods (Xiao al. 2024 Li al. 2024d Cai al. 2024
Feng et al. 2024) have been proposed, typically using token eviction to optimize retention per layer or head during prefilling, reducing memory without impacting performance. However, none have explored varying K cache size across individual heads. Inspired by prior observations (Voita et al.) 2019 Wu et al. 2024 Zheng et al. (024)) that attention heads vary importance for generation, we propose HeadK V, head-level KV cache compression method that allocates KV cache budgets based on head importance distribution using a novel retrieval and reasoning importance estimation. Specifically. heads deemed more important allocated larger KV cache budgets, while less impor- tant ones receive smaller allocations, optimizing memory usage without sacrificing performance.
In addition to allocating KV cache across attention heads rather than layers, key aspect of our approach is distributing cache budgets based on head importance measures. Prior work (Wu et al.) 2024) proposed method to identify retrieval heads, using importance estimation to assess each head's

---

Published as a conference paper at ICLR 2025
For Current Head
Answer: John's
Question: faveriie thing the younger one between John and Mary?
Top-N Attention
Importance
RJAN
+ SUM
Importance Score Estimation
Rest Attention
0086
sold
much he John is 12 years old. Mary years old. Mary's favorite thing is to... John's faverite thing is
did our best to
Reason
Wrong Copy
Correct Copy (N tokens)
Context
Context
Needle Sentences
Total Heads: Num of Layers" Nuum of Heads
Head-Level Allocation
Local Windows
Basic Budget
Dynamic Budget
Head3
Head4 Head5
Heads
Prefilling Phase
After Compression
Head3
31 41466
Head4
Layers
Head5
Importance
Our proposed head-level KV cache compression method consists of two steps: (1) Head- Level Importance Score Estimation (upper part): important heads that contribute to the contextual reasoning ability are identified using Needle-in-a-Haye -aays-- tests. (2) Head-Level KV Cache Allo- cation (lower part): KV cache budgets for each head during the prefilling phase are allocated based on the importance score distribution identified in the first step.
Figure I:
role in retrieving relevant context. We integrate their measure with our head-level KV cache com- pression as baseline, observing improved performance over layer-level cache compression.
However, we argue that allocating larger KV cache budgets solely to retrieval heads is insufficient for tasks like contextual Question Answering (QA), which requires both retrieval and reasoning to handle long input contexts effectively. To address this, we propose an importance-score estimation method that jointly evaluates each head's retrieval and reasoning capabilities for KV cache allocation. Specifically, illustrated in the Importance Score Estimation section of Figure I we construct questions that require both retrieval and reasoning abilities of LLM. For instance, in the provided example, the model must first identify "who is younger between John and Mary" (John) by referencing the context of their ages, and then retrieve John't 's favorite thing. We then estimate the importance score of each head based the attention scores generated by the model while answering the question. Using the estimated importance scores, we allocate the KV cache budget for each individual head, meaning that heads demonstrating greater importance retrieval and reasoning retain larger portion of the KV cache, as shown in the Head-Level Allocation section of Figure D Within each head, we then retain only the most relevant KV cache entries, following the strategy proposed in Li et al. (2024d).
We conduct experiments on various benchmarks requiring both retrieval and reasoning abilities, including QA tasks from LongBench (Bai et al.) (2024) and LooGLE (Li et al. (024b), using back- bone models such Llama-3-8B-Inee Instruct (Dubey al., 2024) and Mistral-7B-Insinn 2023a). Our results demonstrate that the head-level KV cache compression method outperforms previous approaches on nearly tasks. Furthermore, the allocation strategy based on our esti- mated importance scores-rede reflecting both retrieval and reasoning outteeirsoou allocation strategies based on retrieval ability alone. In challenging scenarios like needle-in-a-Hee aaa----aattcck and reasoning-in-at anassssstcckk tests, methods effectively preserve the model's retrieval and reason- ing capabilities. Finally, experiments on memory and latency reveal that our approach significantly reduces memory usage and decoding latency, comparing with the original full KV..
(Jiang al.)
2

---

Published as a conference paper at ICLR 2025
2 RELATED WORK
2.1 ATTENTION HEADS
Multi- Head Attention, a fundamental component of Transformer architectures Vaswwnn et al. 2017). has been extensively analyzed to understand the roles of individual attention heads (Voita et al. 2019 Olsson et al. 2022: Jin et al., 2024; Wu et al. 2024,Zheng et al. 2024). often with goals such as pruning redundant heads Shim al., 2201: Kwon al. 2022) or enhancing inter- pretability Olsson et al. 2022: Jin et al.) |024 et al. (024). For example, Voita et al. (2019) observed that only a small subset of heads plays key role in machine translation, typically man- aging positional information, syntactic parsing, or focusing on rare words. Similarly, Olsson et al. (2022) identified 'induction heads' that implement a copy-and- annd-pastte mechanism, enabling models to replicate previously encountered sequences. Additionally, Zheng al. (022) provided a thorough overview of recent efforts to characterize the diverse functions of attention heads, categorizing them into four types: Knowledge Recalling, In-Context Identification, Latent Reasoning, and Expression Preparation. A closely related study (y et al. (2024) discovered "retrieval heads" that play a cru- cial role in knowledge acquisition, using feedle-in-a-Hasstee functionality serve as the foundation for our head-level KV cache compression methods, designed to jointly preserve both retrieval and reasoning capabilities during the compression process.
tests. These insights into head-level
2.2 KV CACHE COMPRESSION
Improving computational efficiency for LLMMs, particularly for extremely long inputs, has attracted
considerable research interest (Shazeer, 2019 Chen et al. 2023 Dao 2024: Zhou et al. 2024 Yo
ct al. 2024), including caching key and value vectors in Multi-Head Attention to improve generation
efficiency (Ge et al. 2024at Xiao al. 2024 Zhang al. 2023 Li et al. 2024d Cai al. 2024).
One challenge with full cache approaches is that the memory usage increases linearly with input length, leading to significant memory management challenges. Various KV cache compression techniques have been proposed to reduce memory usage and improve inference efficiency. For example, Giamo et al. (2024) addressed the 'attention sink' issue with StreamingI retaining only the first k tokens KV cache. Zhang et al. (2023) applied the Heavy Hitter Oracle strategy to select key cache entries, while Li et (2024d) used the attention scores of the last tokens to identify relevant entries. Cai et al. (2024) introduced PyramidKV, assigning smaller cache budgets to higher layers based on attention matrix patterns.
While these methods have improved performance and efficiency, they largely depend on rigid layer- level allocation strategies, which may not fully optimize KV cache allocation for downstream tasks. Feng et al. (2024) proposed dynamic head-level allocation using attention scores but still relied on layer-level budgeting. Similarly, Tang et al. (2024) employed Retrieval Heads distribution for head- level allocation but retained Full KV cache for key heads, limiting true head-level compression. In contrast, our approach allocates KV cache solely based on head-level importance scores, inde- pendent of layer constraints, resulting in more effective compression and outperforming baselines without increasing decoding latency or memory usage.
METHOD
This section presents our proposed head-level KV cache compression method, which consists of three key components: (1) identifying important heads and calculating head-level importance score distributions (3.1). (2) using these distributions to efficiently allocate KV cache budgets across heads and (3) determining which Key and Value vectors to retain within each head (3.3..
3.1 HEAD-LEVEL IMPORTANCE SCORE ESTIMATION
Accurate budget allocation at the head level requires identifying which heads are most and least important for the given task. By leveraging this importance distribution, we can assign larger KV cache budgets to more critical heads and smaller budgets to less significant ones. To achieve this, we propose a novel importance-score estimation method, inspired by Wu et al. (2024), which allows us effectively estimate the importance of each attention head for optimal KV cache allocation.
3

---

Published as a conference paper at ICLR 2025
Retrieval Example
Reasoning-Retrieval Example
What is the favorite thing of the younger one between John and Mary?
q: What thing to Beijing?
Ncedle:
Needle:
John 12 years old. Mary 13 years old.
The best thing do Beijing take walk Chaoyang Park have cup Espresso the evening.
Mary's favorite thing to take walk in Chaoyang Park and have cup of Espresso the evening. John's favorite thing to play busketball at the local gym and enjoy smoothie aflerward.
22
Figure 2: Comparison of examples for head identification: Needle-in-a Hayssaaa test example from Wu et al. (2024) for identifying Retrieval Heads distribution (left), and our proposed Needle-in-a- Haystack example for identifying Retrieval- Retrieval-Reasoning Heads distribution (right).
Retrieval Heads Wu et al. 2024 use the Bosele-in-a-Haystee test and custom retrieval ex-
amples, as shown in Figure 2 to estimate the importance score for each head. these examples, question q that cannot be answered using the model's parametric knowledge is paired with an answer & (the "Needle") inserted into a haystack c at different positions The model is required retrieve the exact answer from the combined input (ks, c).
During each decoding step i, a head h earns a fraction of the importance score if (1) the token with the highest attention score arg max(a;) matches the generated token, and (2) the token part of the inserted answer k. The final importance score for each head h is calculated accordingly:
N Sn = E No,
-== { =0 0
if arg max E otherwise
where No
(1)
where N is the length of the inserted answer kt and at is the attention score on the combined input from head at I-sh decoding step. By using various settings of the inserted position ps, they obtain the head-level Retrieval Heads distribution.
Directly using this distribution poses two issues: (1) it focuses on the retrieval-and-perr and-t mechanism, lacking consideration for the contextual and reasoning skills needed for complex questions; and (2) the distribution is too sparse for effective head-level KV cache allocation, with nearly 70% of heads receiving an importance score of zero due the strict exact-match requirement (Wu et al. 2024).
Retrieval- Retrieval-Reasoning (R2) Heads To address these issues, we propose a new importance score estimation method that accounts for both the retrieval and reasoning abilities of the heads, enabling a more accurate assessment of their significance.
First, construct retrieval-reasoninn the retrieval examples from Wu et al. (2024), as shown in the Retrieval-Reasoning Example part of Figure We further modify the inserted needle into three parts: & (r, c². c²), where r is the reasoning step, and c² and c² are different answers to the refined question q. The model must reason with r to retrieve and generate the correct answer c², avoiding the wrong answer c².
examples by adding explicit contextual reasoning steps to
Secondly, we refine the estimation method by focusing on the entire correct answer c² (Correct Copy Figure [ since all tokens are relevant the question q. This approach aligns with Contextual QA, which requires both retrieval and reasoning abilities. By considering the full correct answer, the importance score each head h no longer depends solely on the token with the highest attention score. The importance score for head h is calculated as follows:
N
N
{ % top-i(af)) E c²
S₁
M
where M₁
(2)
0
otherwise
hutt:///ghhhn..c////aaaaaa//// NeedlclnAHHasstack
4

---

Published as a conference paper at ICLR 2025
where a; aj represents the i-th highest attention score from head h, and top-i(af,) is the token with the i-th highest score at f-th decoding step. We compute the importance score by considering the entire correct answer and increasing the number of tokens evaluated per head (Importance Score Estimation in Figure T Intuitively, heads with higher attention on the correct answer should receive higher importance scores. We further refine the score using attention weights, yielding more accurate distribution, as shown in Eq. 2
3.2 HEAD-LEVEL KV CACIE ALLOCATION
With the importance scores estimated for each head, we can identify the key heads and allocate the KV cache budget accordingly. In this section, we explain how to incorporate these distributions into head-level KV cache allocation.
In Multi-Head Attention, for each head h in layer 1, the embedded input X = 1axd mapped into different subspaces using the query W key W K. and value W. E [dxxd, matrices:
Preliminary (2,..2,....
2 = XW E RRndd K =XX E Rundd; Vg = XW E Raxd
To optimize memory and enhance efficiency, KV cache compression methods (Xiao et al. 2024 Li et al.) 2024d Cai et al.] 2024) are employed to discard unimportant KV cache entries while preserving performance. For each head h, the compressed KV cache is reduced to K), RRxds and V Rexdn, where < resulting in significant improvement computational efficiency.
Head-level Allocation Previous works on KV cache compression during the prefill phase Xiao al. 2024:1 Li et al. 2024d Cai et al. 2024 Feng al.) (2024) are limited to layer-level allocation, using either uniform or dynamic budgets per layer, but treating all heads within a layer equally. While While/Feng al. (2024) incorporate head-level information, their approach still depends layer- level allocation as a prerequisite.
Building on the head-level importance distributions, we propose comprehensive KV cache alloca- tion strategy. Each head h is initially assigned a fixed KV cache size with an associated importance score Sp.- To allow dynamic allocation, of the budget from each head, leaving the remainder as the basic budget. This process is illustrated in the Head-Level Allocation section of Figure [ The budget pool B is then distributed among the heads in proportion to their importance scores S₁. The importance score distribution S was L1- normalized to ensure that the sum of S₁, equals The final head-level KV cache allocation is as follows:
create shared budget pool B by extracting portion
b B - 3-3xLX1 b, -b b -
(4)
where b is the initial fix budget for each head, 3 is a hyper-paramette to control the size of the dynamic budget pool, I. and H is the numbers of layers and heads of current L.M respectively.
The last a instruct tokens are preserved before forming the dynamic budget pool B to guide the selection process, as detailed in Section 3.3 The retained KV cache for each head includes: (1) the basic budget (b - )). (2) the dynamic budget S₁, x B, proportional to its importance score, and (3) the last instruct tokens.
3.3 KV CACHE SELECTION
After determining the number of KV cache entries to retain using the above algorithm. we apply an
attention-based selection strategy from prior works (Li et al. 2024d Cai et al. 2024 Feng et al.) 2024) to keep the most relevant entries. Specifically, the last a instruction tokens (local windows) guide KV cache selection for each head. Attention scores from these local windows to the remaining tokens are aggregated through a pooling layer, with higher-scoring tokens considered more impor- tant and retained in the cache.
5

---

Published as a conference paper at ICLR 2025
4 EXPERIMENTS AND ANALYSIS
This section outlines the experimental setup, including KV cache baselines and implementation details. We also conduct additional experiments that: (1) emphasize the importance of enhancing contextual reasoning in importance score estimation (4.3): (2) use the Moodle-n--a-Hysteee and Reasoning a-Hasstack tests to demonstrate how our head-level KV cache compression improves long-context retrieval and reasoning (4,4); and (3) provide a comprehensive comparison with previ- ous methods, showing our approach delivers superior performance while maintaining computational efficiency 4.5
4.1 EXPERIMENT SETTINGS
Models and Datasets We compare our head-level KV cache compression method against strong baselines using two open-source L.Mss: Llama-3-8B-Instruet (Dubey al. 2024) and Mistral-7B- Instruct /Jiag et al. 2023a). The evaluation is based on two benchmarks for long-context un- derstanding: LongBench (Bai et al.) 2024) and LooGLE (Li et al., 2024b). For LongBench, we datasets from the Single-Doe QA and Multi-Doc QA categories to assess contextual reasoning. For LooGLE, we focus on the Long Dependency QA task, which includes four QA-related tasks.
Dataset details are in Appendix Table
Baselines and Settings We evaluate three strong KV cache compression methods baselines, ensuring all retain the same number of KV cache entries for fair comparison:
1) SnapKV et et al. 2024d) uses the last a tokens as local windows to guide KV cache selection. Attention scores from these windows to the remaining tokens are pooled to cluster information and guide the selection process.
2) PyramidKV (Cai et al. 2024) follows a pyramid attention pattern, allocating more KV cache lower layers to retain key information, while reducing the budget for higher layers where information is already aggregated.
3) Ada-KV (Feng et al., 2024 dynamically allocates budgets to heads within each layer based on their concentration degrees, and can be combined with SnapKV PyramidKV. Ada- SnapKV is used as the baseline due to its superior performance over Ada- ada-PyramidKV.
Our proposed head-level KV cache compression method also requires strategy guide KV cache selection after allocating the budget. Therefore, we use the SnapKV method to guide the selection process for each head. We set the size of the local windows method. The hyper-parameter 3, which controls the size of the shared budget pool, was chosen from (1.005, 1.01, 1.1, 1.2, 1.5, 2, 10), and report the best performance.
8 for both the baselines and our
4.2
MAIN RESULTS
Table lists the evaluation results for contextual tasks from the LongBench and LooGLE bench- marks. Our head-level KV cache compression method consistently outperforms strong baselines, especially with 64 and 128 KV cache configurations. In resource-constrained settings, precise KV cache allocation is crucial. Layer-level methods allocate a fixed cache size to all heads within a layer, making it difficult retain essential information. Ada-SnapKV improves this by allowing dynamic allocation within layers, but still relies on fixed layer-level budgets. In contrast, our head-level strat- egy allocates dynamically across individual heads, retaining critical information by adjusting the budget based on each head's 's importance.
We perform head-level allocation using both the standard Retrieval Heads distribution (HeadK Reasoning Heads distribution (HeadK V-R2) for global KV cache allocation. This
and our Retrieval-Reasoning combination leads to superior performance across benchmarks. Notably, integrating the Retrieval- Reasoning Heads distribution significantly improves results over the standard Retrieval Heads dis- tribution, highlighting the impact of our approach. Our passes the Full-KV cache average (32.90), achieving 32.95 with 1024 KV cache. Overall, both the head-leve allocation strategy and Retrieval-Reasoningg Reasoning distribution are key to these performance gains.
distribution even sur-

---

Published as a conference paper at ICLR 2025
Method
Single-Doc QA
Multi-Dcc QA
Avg.
Long dependency QA
Avg.
NartoQA
Qusper
MIFeen
HotpotQA
ZWILIMQA
Musique
Doc.A
Info. Retrieval
Timeline
Computation
Lims- 1-88-nnemee, KV Size Full
FalKy
30:71
Liuma sB Instruci, KV Size 128
Snapky
PynmidKY
22.01
17.05
31.52
39.27
28.99
18.34
26 20
8.89
9.63
0.61
6.72
6.46
Ad-SnupVV
10.3
HeadKVR
23.49
25.39
38.15
42.45
32.84
19.95
50.38
8.87
10.35
0.78
7.52
6.85
HoaIKV-R2
32.00)
7.16
Llama 3. SU Instruct, KV
Size 1024
SnapKV
25.76
27.30
38.38
43,60
34,81
20.07
31.65
11.34
7.22
7.38
PynmidKY
25.38
26.83
44.09
21.49
31.49
11.41
6.97
SnapKV
25.79
29:24
38.74
43,93
36.34
19.79
32.31
11.41
7.08
HeadKV-R
24.85
30.94
39.82
43.52
20.37
52 68
11.67
7.38
HoaiKV-R2
24.66
50.82
39.56
43.97
36.47
22.24
32.95
9.02
11.51
0.47
7.85
7.21
Mismml-7R--nniiuu.
PallKY
26.63
32.99
09.34
42.77
27.33
18.78
32.98
12:17
13.52
0.49
10.03
9.55
Mistral 7D Instruct, KV Size 128
StapKV
21.47
PyrmidyV
21.76
21.99
43.72
332.76
22.73
15.59
26.42
10.64
11.9
0.47
8.69
7.93
Ada-StupKV
21.57
24.59
28.09
HeadKVVR
23.97
29.60
48.40
39,06
26.31
18.15
51.01
11.43
13.04
0.53
10.26
8.82
Mistral 7B Insttrcct, KV Size 1024
SnapKV
31.90
10.54
9.16
PynmidKV
24.08
30.05
49.17
40.49
26.41
18.90
31.54
11:77
14.51
0.51
10.19
9.25
Md-SnopKV
31.98
HeadKVV.
25.87
31.44
49.55
41.95
27.09
19.88
32.65
12.21
14.17
0.50
10.58
9.37
HesIKV-R2
14.33
Table 1: Performance comparison on the LongBench and LooGLE benchmarks for Llama-3-88 Instruct and Mistral-7B-Instruuct Our head-level KV cache compression method outperforms all baselines, especially in low-resource settings (KV size = 128). It even exceeds the FullKV result (32.90) on Llama-3-8B-Instee (KV size 1024, 32.95), highlighting the benefits of incorporating contextual reasoning for head selection.
Llama-3 1arma---BB-Instruue
Mistral- IB-Isructt
32
32
Avg. Score
30
30
28
FullKV
28
SnapKV
26
PyramidKV
26
Ada-SnapKV
HeadKV-R
24
HeadKV-R2
200
400
600
800
1000
200
400
600
800
1000
KV size
KV size
Figure 3: Results for different KV cache sizes (64, 128, 256, 512, 1024), showing average accuracy across six datasets from the LongBench benchmark with average input length of 8,683 tokens. Notably, a KV cache size of 64 retains just 0.7% of the total tokens.
In addition, present the results for various retained KV cache sizes (64, 128, 256, 512, 1024) in FigureB. with detailed results available in Appendix Table
4.3 RETRIEVAL- HEADS
we propose to improve the standard Retrieval Heads distribution by in- corporating retrieval-reasoninng examples and refining importance score estimation to better capture
As detailed in Section 3.1

---

Published as a conference paper at ICLR 2025
Retrieval Heads
EnhanccedRRttriee Heads
Retrieva Seetriena-Ressoinng Heads
300
0.84
0.035
-304
0.12
00,00
0.08
-302
0.04
-001
0.02
Layers
Layers
Figure 4: Head visualization for Llama-3-80 1--33-8B-mmmmt In-tructt results. The Retrieval Heads distribution is sparse to effectively differentiate between heads, while our Retrieval-Reasoninn Heads has denser distribution for such differentiation. See Appendix Figure for Mistral-7B- Mistral-7B-Instruuee BB-sntuuctt results.
Single-Doc
Multi-Doc QA
Method
Avg.
NartvQA
Qusper
MF-en
HotpotQA
2WikiMQA
Musique
Llama-3-88- 8-B---3-u-- KV Size=128
HeadKV-R
23.49
25.39
38.15
42.45
32.84
19.95
30.38
HeadKV ER
23.33
25.86
40.28
43.25
33.23
20.28
31.04
HcadKV-R2
21.80
29.19
41.89
43.73
35.01
20.40
32.00
Mistral- Instrucct Size=128
HeadKV-R
23.97
29.60
48.40
39.66
26.31
18.13
31.01
HcadKV-ER
23.23
28.70
48.10
41.39
27.31
17.39
31.02
HeadKV R2
25.04
27.95
48.48
41.28
27.65
18.05
31.41
Table 2: Ablation study results the LongBench benchmarks. HeadKV-R leverages the standard Retrieval Heads distribution. HeadKV-ER uses the retrieval examples from Wu et al.] (2024) but with our proposed importance score estimation method. HeadK V-R2 leverages both proposed importance score estimation method and the retrieval-reasoning examples.
contextual reasoning and identify relevant heads. We also conduct an ablation study to evaluate the impact of these modifications.
Table presents the ablation study results, while Figure 4 provides visualizations for each distribu- tion. Alongside the standard Retrieval Heads and Retrieval- Heads distributions, we intro- duce the Enhanced-Retriee Retriieval Heads distribution, using retrieval examples with our modified impor- tance score estimation method. Comparing Retrieval Heads [HeadKV-R) and Enhanced-Retrieval Heads (HeadKV-ER) reveals that focusing on the entire needle, rather than specific tokens, improves performance. Figure shows that the Retrieval Heads distribution is sparse, while the Enhanced- Retrieval and Retrieval -eaeraii--aaninnggg distributions are much denser. The strict constraints on the Re- trieval Heads distribution result in most heads receiving score of zero, leading to worse results when incorporating Retrieval Heads distributions.
While the EnhanceedRRetreeve Heads distribution improves performance slightly, it remains rely retrieval examples and lacks full consideration contextual reasoning. In contrast, the Retrieval- Reasoning Heads distribution, reflecting both retrieval and reasoning abilities, consistently outper- forms other methods, underscoring the value of incorporating retrieval-reasoning examples.
4.4 LONG-CONTEXT RETRIEVAL AND REASONING
We conduct the Needle-in-a -aree-ii----heee Hayst-cck test to assess the long-context retrieval ability of different KV cache compression methods. As illustrated in Figure 5 we set the KV size to 128 for all meth- ods and keep the other hyperparameters consistent with previous experiments. Results from the 1amma---BB-Inttrute demonstrate that our head-level KV cache compression method effectively re- tain important information compared with other strong baselines, verifying the superior long-cod congcontext retrieval ability of our proposed method. However, these tests only retrieve the inserted needle from the haystack and paste into the generation, lacking assessment of contextual reasoning ability
8

---

Published as a conference paper at ICLR 2025
Liama B-B-Itstteet
Ads-SsapyY Avorage
BeadKV-R
HeadKV--R 70
Centeut Longth
Length
Figure 5: Meedle-in-at -nna--aystack test results on Llama-3-8B --Istruct with KV cache = 128. We build our head-level KV cache method based on SnapKV and proposed method significantly outperform all strong baselines. Moreover, our Retrieval-Reasont Heads distribution maintains and improves long context retrieval ability. Results Mistral-7B- IB-Inttuuct can be found in Appendix Figure 8 which are consistent with results on Slama-3-8B-Inseee
Method
Ok
2k
4k
8k
Avg.
0k
1k
2k
4k
8k
16k
32k
Avg.
FullKV
63.00
59.40
55.80
56.60
50.40
57.04
61.40
55.20
52.40
40.80
40.40
35.00
31.80
Llama-3 8B Instruct, KV Size=128
Mistral 7B Instruct, KV Size=128
SnapKy
60.80
57.00
54.80
52.60
45.60
54.16
56,60
50.60
46.60
34.40
35.40
33.20
29.00
40.83
PyramidVV
61.20
58.00
52.80
52.60
47.60
54.44
58.40
50.60
80
36.00
36.20
31.20
28.40
41.09
Ada-Sn Ada-SnapKV
61.80
59.20
53.80
53.60
46.60
55.00
58.00
51.20
46.20
35.40
36.40
31.60
28.80
41.09
HendKV-R
61.60
57.00
52.60
55.00
49.60
55,16
58.40
54.00
50.80
38.00
37.40
31.80
30.20
42.94
V-R2
62.60
60.40
55.00
55.80
50.40
56.84
60.20
54.20
51.20
37.20
37.40
32.60
28.88 30
43.09
Table 3: Reasoning-in-te Hn-a-Haystack test results with KV cache 128. The final results are averaged across QAI-QAS tasks for each length. Unlike the Needle-in in-a-Haystach test, this test inserts mul- tiple needles into the haystack, requiring the model to reason through them to extract the correct answer.
in long-conter scenarios. This long-context contextual reasoning ability is erucial for many tasks, such as question answering (QA), summarization (Kuratov et al., 2024 Li et al., 2024c). Therefore, conduct the Reasoning-in-a-Paaystaa test evaluate the long-context reasoning ability of each KV cache compression method across different scenarios.
We follow the setup from Kuratov et al. (2024) to conduct the Deasoning-in-a--Hyssee test. This test enhances the bAbl benchmark (Weston et al. 2015), designed for reasoning evaluation, by using text from the PG19 dataset (Rae al. (020) as the haystack. Reasoning needles from bAbl are inserted into the haystack, and the model must retrieve and reason through them to generate the correct answer. We use the dataset from Kuratov al.] (2024). averaging results across QAI-QA5 tasks for evaluation. Examples of tasks are shown in Figure 12
As shown in Table our head-level KV cache compression method significantly outperforms strong baselines, demonstrating its superior long-context reasoning. By incorporating retrieval-reasoning examples, our method achieves even better accuracy, particularly with the model. Notably, combining head-level KV cache allocation with the standard Retrieval Heads distri- bution also yields improved results over other baselines. This is due to two factors: first, as shown in Figure 4 there is overlap between Retrieval and -eaecviinn Heads, indicating heads may serve multiple roles. Second, since the bAbl benchmark contains the answer within the inserted needle (see Figure 12), emphasizing retrieval alone helps our method locate the needle.

---

Published as a conference paper at ICLR 2025
Decoding Latency
Peak Memory Usage
FulKKV
300
SnapKY
PyramidKV
Peak Memorr (GB)
250
Ada-SnapKV
Times (s)
HeadKV
200
150
100
30
1000
2000
3000
4000
4k
8k
16k
32k
64k
128k
256k
Generation Length
Context Length
Figure 6: The Decoding Latency and Peak Memory Usage results. Our proposed method maintain the comparable computational efficiency with other KV cache compression baselines.
4.5 MEMORY & LATENCY
In this section, we evaluate the computational efficiency of our head-level KV cache compression method using the Mistral-7B-Instieee IB-Instruct model with a maximum sequence length of 32K, and FlashAt- tention (Doo 2024) as the default setting.
To assess the decoding latency of each method, we use 32K-length data from the Reasoning in-a- Haystack test as input and set various generation lengths (1, 512, 1024, 2048, 4096) for comparison. As shown the Decoding Latency of Figure our proposed method achieves the same decoding latency as other KV cache compression methods while maintaining performance closest to the Full KV cache. Notably, the decoding latency includes both the pre-filling time and the decoding time. Therefore, we conclude that the pre-filling time for our method and other baselines (such as Vy Ada-SnapKVI V) is almost negligible, as shown at the starting point (generation length 1) of Decoding Latency in Figure 6
In addition to decoding latency, KV cache compression methods also aim to reduce memory usage during the decoding phase. Therefore, provide the Peak Memory Usage results, shown in the Peak Memory Usage Figure All results for Peak Memory Usage are averaged over three trials. Our proposed methods achieve performance comparable other KV cache compression baselines, significantly reducing memory usage compared to the Full KV cache.
5 CONCLUSION AND FUTURE WORK
In this paper, we introduce V-R2 a head-level KV cache compression method with two key components. We propose strategy for allocating KV cache budgets across attention heads based on their importance scores and develop novel approach to estimate these scores by constructing examples that account for both retrieval and reasoning abilities. By allocating larger KV cache bud- gets to more important heads and smaller budgets to less important ones, our approach efficiently retain important KV cache than layer-level KV cache compression methods. We thoroughly evaluate HeadKV across multiple benchmarks, models, and long-context ability tests, with comprehensive results demonstrating that our method achieves superior performance while maintaining computa- tional efficiency.
For future work, further exploration of various types attention heads could offer valuable in- sights, such those involved in in-context learning (Olsson et al.) 2022) and truthfulness (Li al. 2023). For example, placing greater emphasis on truthfulness heads could help mitigate issues like hallucination, thereby improving the factual accuracy of model outputs. Additionally, would be worthwhile to investigate the development of general task-specific score estimation algorithm. One potential direction leverage gradients from specific tasks allocate KV cache budgets more ef- fectively, enabling tailored compression to enhance model performance across diverse applications.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Iige Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, al. Gpt-4 technical report. arXiv preprint arXiv:2303 087744 2023.
Anthropic. The claude 3 model family: Opus. sonnet, haiku, March 2024. URL https://
www -ww-cdnn.. ww.. anw.aopih.ppc..ooo com. com/ dew.aabblcseeeeeeeeeeeee
3. påf
Accessed: 2024-07-09.
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li. LongBench: A bilin- gual, multitask benchmark for long context understanding. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Proceedings of the 62nd Annual Meeting of the Association for Com- putational Linguistics (Volume 1: Long Papers), pp. 11199133, Bangkok, Thailand, August 2024. Association for Computational Linguistics. doi: 10.1 1633//1/2222..cllloooo 172. URL ac
Ittps://accaante
ogy wrg.or///ooo wwtpo///02...111ooo..oe
172
Zefan Cai, Yichi Zhang, Bofei Gao, Yuliang Liu, Tianyu Liu, Keming Lu. Wayne Xiong, Yue Dong, Baobao Chang, Junjie Hu, and Wen Xiao. Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling, 2024. URL
Charlie Chen, Sebastian Borgeaud, Geoffrey Irving, Jean- Caan-Bapiiste Lespiau, Laurent Sifre, and John Jumper. Accelerating large language model decoding with speculative sampling, 2023. URL
https://arxiv. org/ alss 2302 01315
Tri Dao. FlashAttention-2:: Faster attention with better parallelism and work partitioning. In Inter- national Conference on Learning Representations (ICLR), 2024.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The Ilama 3 herd of models. arXiv preprint arXiv:2407.2 1783, 2024.
Yuan Feng, Junlin Lv, Yukun Cao, Xike Xie, and S. Kevin Zhou. Ada-kv: Optimizing kv cache eviction by adaptive budget allocation for efficient IIm inference, 2024. URI /https://arxiv. org//bs/ 2407 11550
Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. Model tells you what to discard: Adaptive KV cache compression for LLMs. In The Twelfth International Confer- ence on Learning Representations, 2024a. URL http:tt:: /ttps://opaae /ttp:////p///v/eeeeee ew. net. forum?id UNFFPDPMyO
Tao Ge, Hu Jing, Lei Wang, Xun Wang. Si-Qing Chen, and Furu Wei. In-context autoencoder for context compression in large language model. The Twelfth International Confer- on Learning Representations, 2024b. AtLp:ttp::/ http::s//ppeervvieeo net. formm?idd URE 34ZUGJE
Xanh Ho. Anh-Khou Duong Nguyen, Saku Sugawara, and Akiko Aizawa. Constructing multi- hop QA dataset for comprehensive evaluation of reasoning steps. In Donia Scott, Nuria Bel, and Chengqing Zong (eds.), Proceedings of the 28th International Conference Computational Linguistics, pp. 6625, Barcelona, Spain (Online), December 2020. International Com- mittee on Computational Linguisties. doi: 10. 1/1.:///1//222.oooooooooooo/
URL https:
7/ac httpl//gy..rgg///..oooooo. maa main. 580
Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chap- lot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. Mistral 7b, 2023a. URL https:
//axxiv.ore/aab//3310.o
Huiqiang Jiang, Qianhui Wu, Chin- Yew Lin, Yuqing Yang, and Lili Qiu. LLMLingua: Com- pressing prompts for accelerated inference of large language models. In Houda Bouamor,
11

---

Published as a conference paper at ICLR 2025
Juan Pino, and Kalika Bali (eds.), Proceedings of the 2023 Conference Empirical Meth-
ods in Natural Language Processing, pp. 13358-13376, Singapore, December 2023b. Associ-
ation for Computational Linguistics. doi: 10.18653/w1/2022 122.1666/.1/222...mm.......
wwwweemniinnnieeeeeeeeeeemm URL https:
pacclantholoso ogy. org/2023. emn p 1-pppmmp- ma main. 82.5
Huiqiang Jiang, Qianhui Wu, Xufang Luo, Dongsheng Li, Chin- Yew Lin, Yuqing Yang, and Lili Qiu. LanggLLMiinnaaa Accelerating and enhancing LLMs in long context scenarios via prompt compression. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar fe Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 1658-1677, Bangkok, Thailand, August 2024. Association for Computational Linguistics. doi: URL https://aclanntolooyy.oooo www.:r//////..... -ong.
1111866/111222..1lllll 91
Zhuoran Jin, Pengfei Cao, Hongbang Yuan, Yubo Chen, Jiexin Xu, Huaijun Li, Xiaojian Jiang, Kang Liu, and Jun Zhao. Cutting off the head ends the conflict: mechanism for interpreting and mitigating knowledge conflicts in language models, 2024. URL htt..tttst///// htttps://arxi..oott
abs/2402. 18154
Huan Yee Koh, Jiaxin Ju, Ming Liu, and Shirui Pan. An empirical survey on long document sum- marization: Datasets, models, and metrics. ACM Computing Surveys, 55(8):1-35, December 2022. ISSN 1557-7341. doi: 10.1145/3545176. URL http://dx..t. doi. wtt.::ii......111111..1
3545176
Yuri Kuratov, Aydar Bulatov, Petr Anokhin, Ivan Rodkin, Dmitry Sorokin, Artyom Sorokin, and Mikhail Burtsev. Babilong: Testing the limits of Ilms with long context reasoning-ing seatiin--i---haaee -n--aaas-i--ahcckk,
2024. Utth::s:::/ http//aavvvv...////i///../.a..va www.vv.rrg//oooo.ooooo..o abs 22222446.1110
Woosuk Kwon, Sehoon Kim, Michael W. Mahoney, Joseph Hassoun, Kurt Keutzer, and Amir Gho- lami. A fast post-training pruning framework for transformers. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Sys- lems, 2022. URL https:: /wpp/pppeeeieeeeeeeeeeee net forumrrrm?mmmmm d== D.
Jia-Nan Li, Quan Tu, Cunli Mao, Zhengtao Yu, Ji-Rong Wen, and Rui Yan. Streamingdialogue: Prolonged dialogue learning via long context compression with minimal losses, 2024a. URL
https://arxiv>
http://axwww..ggaa/////////////a wttg/bbs//43..... 08312
Jiaqi Li, Mengmeng Wang, Zilong Zheng, and Muhan Zhang. Loogle: Can long-context language
models understand long contexts?, 2024b. URL http::ttt////ett /ttp:i///aaiviee/ wttp:org///.....o..///// 04939
Kenneth Li, Oam Patel, Fernanda Viégas, Hanspeter Pfister, and Martin Wattenherg. Inference-time Eliciting truthful answers from language model. In Thirty- Conference on Neural Information Processing Systems, 2023. URLtstss: wwwp:eviwwweeeeeff...oo forum?
intervention:
ld aLLuYpn83y
Mo Li, Songyang Zhang, Yunxin Liu, and Kai Chen. Needlebench: Can Ilms do retrieval and reasoning million context window?, 2024c.
org/ abs///bs//200
11963
Yuhong Li, Yingbing Huang, Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. Snapkv: Lim knows what you are looking for before generation, 2024d. URL httpa/////t//////.....o../////// www.///////bbb.....//// /
144.9.4..4449 14469
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads, 2022. abs / mass/2299.111 11895
hattps://arxivv
Guanghui Qin, Corby Rosset, Ethan Chau, Nikhil Rao, and Benjamin Van Durme. Dodo: Dy- namic contextual compression for decoder-only LMs. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Proceedings of the 62nd Annual Meeting of the Association for Com- putational Linguisties (Volume I: Long Pupers), pp. 9961-9975, Bangkok, Thailand, August
12

---

Published as a conference paper at ICLR 2025
2024. Association for Computational Linguistics. doi: 10.18653w11/222:acllooongg URL
Jack W. Rae, Anna Potapenko, Siddhant M. Jayakumar, Chloe Hillier, and Timothy P. Lilli- crap. Compressive transformers for long-range sequence modelling. In International Confer- ence on Learning Representations, 2020. URL https: Dettps::/openreviiee net htw.se/wwee//foummeeeemm forum?id-
SylKikSYDH
Noam Shazeer. Fast transformer decoding: One write-head all you need, 2019. URL https:
Kyuhong Shim, Iksoo Choi, Wonyong Sung, and Jungwook Choi. Layer-wise pruning of trans- former attention heads for efficient language modeling, 2021. URL https://axiiv..o
abs/2110.0320
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:
/aax:iv.rg//bbb/2....0 15891
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, ukasz Kaiser, and Illia Polosukhin. Attention all you need. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Ad- vances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL http:s/s//..ccceee.........eeres cctpppprrccpppppprrppppppp Files/paper/2017. 11l1://5ee225557776220200000
Guyon, U. Von
Paper.
Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. In Anna Korhonen, David Traum, and Lluís Mårquez (eds.), Proceedings of the 57th Annual Meeting of the Associ- ation for Computational Linguistics, pp. 5797-5808, Florence, Italy, July 2019, Association for Computational Linguistics. doi: 10.18653/v1/P19 1580. URL https:/aclanthologg
orq/P19 PBq/P19-1580
Jason Weston, Antoine Bordes, Sumit Chopra, Alexander M. Rush, Bart van Merriënboer, Armand Joulin, and Tomas Mikolov. Towards ai-complete question answering: A set of prerequisite toy tasks, 2015. htttttt/////xxii wwtp:ca.///aaaattaa
/ttE::/:/aaxx...ooo///////// C5698
Wenhao Wu, Yizhong Wang, Guangxuan Xiao, Hao Peng, and Yao Fu. Retrieval head mecha- nistically explains long-context factuality, 2024. URL 15574
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. The Twelfth International Conference on Learning Rep- resentations, 2024. URL ht..stts: roperrevise Lew net./ forum? 1oa.ummmddNNG77151111
Wen Xiao and Giuseppe Carenini. Extractive summarization of long documents by combining global and local context. In Kentaro Inui, Jing Jiang, Vincent Ng. and Xiaojun Wan (eds.), Pro- ceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-JJCNLP), pp- 3011-3021 Hong Kong, China, November 2019. Association for Computational Linguistics. doi: 10.18653/v1/11--110 /tt///ttpe.//aallannooooo.oo////
URL
1298
Fangyuan Xu, Yixiao Song, Mohit lyyer, and Eunsol Choi. A critical evaluation of evaluations for long-form question answering. 2023. URL
An Yang, Baosong Yang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Zhou, Chengpeng Li, Chengyuan Li, Dayiheng Liu, Fei Huang, et al. Qwen2 technical report. arXiv preprint arXiv:2407.10671, 2024.
Lu Ye, Ze Tao, Yong Huang, and Yang Li. Chunkattention: Efficient self- attention with prefix-ave aware kv cache and two-phase partition, 2024. URL, Wtllhhttps///axxiii httpp///////wawoooo/////////////
wwtp /2402. 15220
13

---

Published as a conference paper at ICLR 2025
Zihao Yi, Jiarui Ouyang, Yuwen Liu, Tianhao Liao, Zhe Xu, and Ying Shen. A survey recent advances in Ilm-based multi-turn dialogue systems, 2024, URL httph://p//aaa. http://axxivv.ooe wrg. /bs/
2402.1801
Zhenyu Zhang, Ying Sheng. Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Re, Clark Barrett, Zhangyang Wang, and Beidi Chen. H2o: Heavy- hitter oracle for efficient generative inference of large language models. In Thirty-seventh Confer- ence on Neural Information Processing Systems, 2023. URL forum? Rerum? P
Zifan Zheng, Yezhaohui Wang. Yuxin Huang, Shichao Song, Mingchuan Yang. Tang, Feiyu Xiong, and Zhiyu Li. Attention heads of large language models: A survey, 2024. URL https:
wwxxiv.orgaab//25500 03752
Zixuan Zhou, Xuefei Ning, Ke Hong, Tianyu Fu, Jiaming Xu, Shiyao Li, Yuming Lou, Luning Wang, Zhihang Yuan, Xiuhong Li, Shengen Yan, Guohao Dai, Xiao-Ping Zhang, Yuhan Dong, and Yu Wang. A survey on efficient inference for large language models, 2024. URL https:
/arxiivccmm wlaaiiivrrgaaaaaa.eaeaaaa www.:/rggaaabssae 2404. 14294
14

---

Published as a conference paper at ICLR 2025
A
DETAILE RESULTS
In Table
wc show the detailed results of Figure 3 in the main paper.
Method
Single- Doe
Molt-Doe QA
Long dependency
Avg
NarvOA
Queper
MI co
HorposAA
ZWEIMOA
Musique
DecQA
Info. Retneval
Timelaee
Computatio
Liano- 3 Instruct KV SizcoFull
32.07
39.71
43.37
35.28
21.15
32,50
8.75
Liuit- 5. Instruct, K.V Suzce64
SKY
20.51
12.80
37.02
25.91
17.02
24.16
PyramiaKV
21.17
13.66
29.34
14.86
23.46
15.65
23.06
6:27
9.31
0.63
6.86
6.27
AdvSKV
22.26
17:30
37.56
17.85
9.08
Ours Copy
22.67
23.54
37.51
37.45
29.76
19:01
29.32
8.00
10.51
0.58
6.68
6.64
Ours. Mix
23.21
25,33
38.71
40.64
31,33
19.35
29.76
9.46
10.66
0.1
6.92
6.91
122
Liamo 3 81) Instract,
KV Size-128
SKV
22.11
15:79
31.01
41.12
29.20
19.35
76.43
8.36
9.46
0.79
6.56
6.29
Py
22.01
17:05
31.52
39.27
8,39
6.72
SKV
22.99
19:95
34.22
42.97
30.52
20.15
28.52
9007
10.3
0.54
6.59
6.63
Copy
23.49
25.39
32,84
19,95
8:87
7.52
Dons Mix
21:80
29.19
41.89
43.73
35.01
20.40
32,00
9.60
11.13
0.67
7.22
7.16
1.01
IJuns
Site-256
SKV
23.38
20:18
37.65
42,80
33.23
20.01
29.54
9.04
10.59
0.53
7.55
6.92
PyramiaKV
23.94
20.27
36.27
42.51
31.44
19,99
29,07
8.66
0.55
6.98
AdV
24.20
24:63
37.95
43.64
33:27
20.00
30.62
9.29
11.23
0.62
7.10
7.06
Ours Copy
25.83
20:04
39.90
42.56
33.58
20.57
31.54
9.05
11.15
0.52
7.22
6.99
Ours, Mix
24.68
30.49
38.59
44.32
36.41
20.54
32.51
9.47
11.56
7.65
Liamo 5 MB Instruct. KV Size-5
SKV
25.47
23.75
43.66
33.98
19.83
30,89
9.00
7:34
PyramiaKV
24.69
23.65
35.10
43.25
31.16
20.06
29.66
8.90
10.62
0.74
7.57
6.96
25.46
19.83
9.22
Ours Copy
23.84
29:21
39.79
44.41
36.09
20.59
32.52
9.13
11,61
0.56
7.12
7.11
1.2
Ours Mix
24.75
29.75
38.05
44.43
36,48
21.67
32,51
9.34
11,26
0.56
7.54
7.18
Liamo
3 Enstract, KV Size-10124
SKV
25.76
27:50
38.78
43,40
34,81
20.07
31,65
9.61
11.34
0.53
7.22
7.18
PyramiaKV
25:38
26.83
36.50
44.09
14.24
21,49
31,49
8.98
11.41
0.53
6.96
6.97
Adv
25.79
20:24
38.74
43.93
34
19.79
32.31
8.65
11.41
0.53
7.71
7.08
Copy
24.85
30.94
39.82
43.52
36.58
20.37
32.68
9.20
11.67
7:71
Dans Mix
24.66
30.82
99.56
43.97
36.47
22.24
32,95
9.02
11.51
0.47
7.88
7.21
1.2
FKV
26.63
32.99
49.34
42.77
27.35
18.78
32.08
12:17
15.53
0.49
10.03
9.55
Mistral- 770-
SKV
19.05
18:63
98.16
11.34
21 31.30
13.81
23,56
10.41
11.49
0.46
9.38
7.94
PyramiIKV
20.1
19.61
38.05
32.18
22.57
15.26
24.81
10.64
11.69
Ada SKV
22.70
21.28
42.39
24.35
22.40
14,85
26.33
10.56
11.50
0.45
8.72
7.81
Ours Copy
24.23
25:22
46.02
38.82
26.05
17.41
29.63
10.94
13.14
0.63
8:46
Durs, Mix
21:77
26.57
48.30
40.12
26.76
16.21
29.97
11.19
13.94
9.87
Mistral-BB- 7B-lallree,, KV Sizem
128
SKY
21.47
21.95
33.88
21.53
13.53
20.03
10.86
12.34
PyramiaKV
21.76
21.98
43.72
32.76
22.73
15.59
26.42
10.64
11.9
0.47
K.69
7.93
AVV
59
35.57
11.14
Ours: Copy
25.97
29.60
48.40
39.56
31
18.13
31.01
11:43
13.04
0.53
10.26
8882
1.5
Ours, Mix
25.04
27.95
41.28
27.65
18.05
11.44
10.20
Mistral 7B Instroct, KV
Size-156
SKV
22.26
24.94
36.76
25.16
14,93
28.72
11.07
12.39
PyramicKV
21.42
25:36
47.94
38.75
25.52
15,30
29.10
11.57
12.38
0.56
9.51
8.50
Ada-SKV
25.81
27.04
48.56
38.52
25.34
16.42
29.91
11.67
13.57
0.52
10.53
9007
Copy
24.98
29:31
49.01
41.36
27.16
17.34
31.53
11.94
13.30
10.95
Dun, Mix
26.94
31.02
50.76
42.11
26.14
18:47
32.24
12.37
13.88
0.48
9.86
9.15
1.005
Mistral-7B--Issmrue
Sirum512
SKV
24.18
25:87
45.74
38.84
25 48
15.04
30.19
11.96
13:47
0.52
10.50
9.11
PyramilKV
23.07
28.97
16.50
11:34
Ada-SV
24.22
20:92
48.96
401.6
25.55
17.45
31.03
12:12
14.53
0.52
10.57
9.44
Ours Copy
24.97
30.94
49.45
42.25
26.34
18.54
32.08
12.09
13.65
0.62
10.94
9.38
1.01
Ours Mix
25.59
31.33
50.26
42.66
27,20
19.37
32.74
11.62
0.50
9.97
1.005
Mintral-7B Intrrct KV Sire:
10:44
SKV
25.38
30.22
49.20
41.54
18.08
31.50
11.69
13.89
10.54
PyramiJKV
24.28
30.05
49.17
401.49
26.43
18,50
31.54
11.77
14.51
0.51
10.19
9.25
24.82
31.49
41.18
27.38
31.98
11.96
9.92
Ours Copy
25.87
31.44
49.5
41.98
27.09
19.88
32.63
12.21
14.17
0.50
16.58
9.37
1.5
Oun, Mix
25.64
32.54
50.49
41.50
27.88
18.89
32,87
11.94
14,93
10.49
Table 4: Details results for different KV cache (64, 128, 256, 512, 1024) on -- and Mistral-7B-IIntruue
B DATASET DETAILS
details of the datasets used in our experiments. To evaluate the retrieval and contex- tual reasoning abilities of different KV cache compression methods, we use six question-answericn
15

---

Published as conference paper at ICLR 2025
datasets from LongBench (Bai et al., 2024 and four additional question ansswerng datasets from LooGLE (Li et al. 2024b) as benchmarks. For the LooGLE datasets, we randomly selected 100
examples to form the datasets used in this study.
Source
Label
Task
Task Type
Eval metric
Len
Language
Nums
LongBench
NarQA
NatrativeQA
Single-Doc, QA
FI
18,409
EN
200
LongBench
Qusper
Qasper
SingleDocc Doc. QA
FI
3,619
EN
200
LongBench
MF-en
MulnFieldQQAeen
Single-Doc, QA
F1
4,589
EN
150
LongBench
HespoeQA
HotpotQA
Mulli-Doc: QA
F1
9.151
LongBench
2WikiMQQ
2WikiMulibopQA
Multi-Doe: QA
F1
4,887
LangBench
Musique
Musique
QA
LeoGLE
Doc.QA
Comprehension&&ree
Long, Dependoncy QA
FI
15,498
EN
LeoGLE
Info Retrieval
Multiple information retrieval
Long Dependency QA
[1]
14,808
EN
100
LooGLE
Timeline
Timeline reorder
Long Dependency QA
F1
15,425
EN
100
LooGLE F
Computation
Computation
1 Long Dependency QA F1
17,001
FN
100
Table 5: Details of Datasets.
C LONGBENCH RESULTS
We list the evaluation results on the remaining datasets from LongBench (Bai et al. 2024) in Table Our head-level KV cache compression method also outperforms the other baselines, particularly summarization and Code datasets across -lama-3-8B 8---Inssruct and Mistral-7B BB-nnttuct.
Summarization
Sow-Sho Learning
Synthctic
Code
Method
Avg.
GovReport
QMSum
MulaNNws
TREC
TriviaQA
SAMSum
PCoust
PRe
LCC
RB-P
Llama:
Size=Full
PallKV
28.71
23.26
26.64
73.5
90.48
42.33
4,80
69.25
$9.29
54.05
47.23
Llama-3 3-80--3-8 3-8t-Instracs, KV Sizem128
SnapKV
19.83
21.80
21.41
65.50
89.72
38.71
5.75
69.00
58.74
54.57
44.50
Ada-SmapKV
20.89
22.11
21.68
70.50
90.82
39.20
5.91
69.50
59.75
54.86
45.52
HeadKV-R
HeadKV-R2
21.76
22.16
23.94
71.50
90.19
38.88
6,60
69.50
61.05
60.21
46.58
Mistral-7B. I8-Instrcct, KV Size-Pull
FallKV
32.87
24.24
27.10
71.00
86.23
42.79
2.75
86.98
$6.93
54.49
48.54
Mistral-7B-Instnatt KV Sixem
SnapKV
20.76
Ada-SnapKV
21.13
22.76
22.25
68.50
85.60
41.05
3.33
62.54
52.68
49.25
42.93
HeadKV-R
22.19
22.86
22.57
69.50
85.46
41.16
3.56
74.49
$1.60
50.89
44.73
HeadKV-R2
24.30
23.48
24.18
70.50
85.54
40.72
4.83
72.63
55.49
51.39
45.31
Table Results on LongBench (KV cache=128),
3.05
0016
0.012
0.012
0012
-305
0.01
0.10
-3.02
0.03
0.06
6.034
0.01
0,032
3.90
0.000
Layers
Layers
Layers
Figure 7: Head visualization results on Mistral-7B-Instruee Our enhanced distribution can also cover the important retrieval heads.
16

---

Published as a conference paper at ICLR 2025
Mistrat-T-llastrctt
FulkY Ancrage soue 100.00
SuppY Average soare 50.10
PyramidKV Average wore
Ada SaagKV Averago sore 95.00
1
HeadKY V Www.VVV....eeaa... Aesrage score
20
Teaodli V-122 Acerage scorer Vr. DI
Conteet Length
Figure 8: Needle-in-a- test results on Mistral-7B-Inntrucce The results are consist with those on Llama-3-88 Bartna---B--nnae In--rutt and our Retrieval- Reasonning Heads distribution (HcadKV-R2 still outperforms other strong baselines.
D REASONING-IN-A-HAYSTII DETAILS
As shown in Figure| 12
we provide examples for each task used in the Reasoning-in-e -nne--iii--Hatactt test" We use the dataset curated by Kuratov et al. (2024) for the Deasonnin-a-Haysteee
test. Specif- ically, they use the "Input" listed in Figure 12 as the needle and split it into different sentences using dot. These sentences are then inserted into various positions within the haystack to form the final test examples.
As shown in Figure 13 provide detailed results across different context lengths and tasks. Our proposed head-level KV cache compression method achieves the best average score among all meth- ods. Notably, on the Jamma338BB-Instruce model, our method surpasses the Full KV cache on the QA3 task, the most challenging task among the five, as indicated in Kuratov et al. (024). This sug- gests that our Retrieval-Reasoninn and important content, leading to generating the correct answers.
Reasoning Heads distributions are better at capturing semantically relevant
E CONSTRUCT
EXAMPLES
For guiding head-level KV cache compression, we need to obtain the importance score for each head. To achieve this, manually construct specific examples ensure that the model relies heads rather than internal knowledge to answer questions during the Needle-in-a-Haasttnn Therefore, we construct retrieval-reasoning examples based on retrieval examples Wu al. (2024) by introducing different reasoning paths into examples to emphasize the contextual reasoning ability. One constructed Setrivval-Reaoone example shown in the right of Figure In addition that example, we reverse the question to create total of two Retrieval-Reasee RRasoning examples.
experiment.
Following the setup outlined in Wu et (2024) in the Needle-ine sondle-in-a-Haaee Haysstac experiment, we use the model's 'ss maximum training length as the maximum haystack length and evenly select 5 different length values 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 detrieval-Reasoning Head distributions.
The addition of reasoning content and incorrect answers cl aims to introduce complexity and context to the reasoning process, which we believe could highlight different heads depending on whether the head supports accurate reasoning patterns not. Here's 's how envisioned their role:
17

---

Published as a conference paper at ICLR 2025
Retrieval-Reasooing Examplel
Retrieval-Reasoning Example2
4: What favorite thing younger one between John and Mary?
4: What is the favorite thing older one John and Mary?
Needle:
Needle:
John is 12 years old. Mary is 13 years old.
John is 12 years old. Mary is years old. Mary's favorite thing is take walk in Chanyang Park and have cup of Espresso in the evening. John's favorite thing is play basketball the local gym enjoy smoothic afterward.
r
Mary's favorite thing to take walk in Chaoyang
22
Park and have cup Espresso in the evening.
John's favorite thing play basketball the local gym and enjoy a smoothic afterward.
Constructed Retrieval-Reasonicnn examples. They are used to conduct Needle-in-a- Haystack experiment to obtaín the Retrieval-Reasonie Heads distribution.
Figure 9:
Aligning with the requirements of contextual reasoning: Based on the in-depth analysis on the contextual reasoning dataset, we know that the answer corresponding question will still appear in the input but with various distractors. Therefore, the model continues rely the retrieval-and- paste mechanism to obtain the true answer. The original retrieval heads estimation method did not account for this phenomenon, but 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 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 e² 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 Retrieval-Reasoning ex- amples 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 naturally ignored, they share the same shared global budget pool.
Aligning with the standard Retrieval Heads estimation: we followed the setup in Wu et al. (2024) and determined the Retrieval-Reasoning Heads distributions based on the Needle-in-az Haystack experiment. Since Boodle-in-a-Haystee experiment only outputs the correct answer, we choose focus on the correct answer ensure that the maximum importance score each head can achieve in Eq. iis 1. Adding additional logic would disrupt this property, potentially affecting the final distribution.
F
PSEUDO CODE
Codes shown Listing T demonstrates the core steps required for head-level KV cache budget allocation based the obtained importance score distribution.
The obtain Beadi__hhad_dte budget function dynamically determines the budget for each head based on the previously obtained importance score distribution as shown in Eq. 4 Since the importance score distribution is static and does not change with the input, this function only needs to executed once during model initialization.
The head ky function allocates budgets for different heads. First, we use SnapKV as the selection strategy to guide the specific selection process. After obtaining the budget and determining the selection strategy, we perform the selection and obtain the final retained results based on the budget.
18

---

Published as a conference paper at ICLR 2025
Soad saved importance score distrihution o9en [data_path, (r") file:
head_list lson. (fi.ee.eeclieee)) accumulete the importance secre and convert to tensor
head_acare_list Ing. for in head_list.iteme heac_acore_liae torch,tenaor (head_score_l1s)
aum saadd__sooee_a__________________
importance score distribution
10
total_ate _attalttntt/on head_score__ 1ocore_1st,, 1 Inun wwwddenndnaaaaeeeeee
num_attention_
11
Construct shared budget pool and define the minimum RY cache budget
12
cotal_pool_capecii (base_capscity // betal nun_hidden_leyere min_rum (base_capecitg
hun.atzention_bedd
13
base_cspacity // self.beta)
14
Head-level lllaaation based the importance
15 head_capacity torch. round (total (total_attentiine
jotal_pool_cppaciee + min .int .int
17
def head_kV(....)
18
calcuiate atta_score from the observation windows to inpu:. (SnapXV)
19
celeul_atta_ (self, key_states, query_atates) Sorted basad on attn_score and obtain indices.
20
indices aetn_acoreesort
22
Obtain cached index for each head.
beadLidx range heads);
24
cache_index indices [head_idx]]
ceead_capecityy leyer__di [head_[6x]]
25
Expend the indices to match the head dimension for gathering. (Bane as SnepKV) cache_index pache_index.verr viewil,
-1, 1) (-1, -1, -1, head_éim)
27
# Gather the compressed past key and value states based on the selected indices. (Same SnapKV)
28
Lop_Rcache - arigin_hesdi _ey__tttteee_e_ee states [head] gather (dim=2, index=cache_ index)
tcahe origin_heads seaddd_valee__te states [head idx] gather (dimm2, indercacceecccceeeeeeeee 4 Merge with obervation windows
30
31
selected_k torch.cat (ttop_ crigin_heads. key_dd__yy_ttteee [head_idx] [t : -self.
vindow_aizer dim,d)m=2)
32
se selacted_v torch.cst ((tpp_ (ttop_Veache, ariain_headda hee_eeeaa_____ lee_ssstee hhead_ [dx] [1, -self.
windov_sizes, dilm2)
33
together
35
beade_key_stater append (selecteckk view(-1, head_dim) wtw.tetttaaeeeteeeeeee...o... viiiee head_dim)
36
# Marge together
38
weads_key_statee teren.cat (heads seads_kee key_atates, dineD)
heads_value_stster torch. (heade_vaiue seeed__vaiu_statee statess,, dim=D)
40
return srada_kay_atates, heads saadd__alle__taaee
Listing 1: Implementation of VeadKV in Pseudo PyTorcc style.
Llama-3-88-- 3-B----ruut
Mistral- 7B-Inntruct
33
33
32
Score
31
30
FulIKV SnapKV
29
29
28
PyramidKV
28
Ada-SnapKV
27
HeadKV-R
27
26
HeadKV-R2
1.005
1.01
1.1
2
10
1.005
1.01
111
12
1.5
Beta
Beta
Figure 10: Results for different 8, showing average accuracy across six datasets from the LongBench benchmark.
G HYPER-PARAMETER ANALYSIS
The only parameer introduced by our proposed method 3, which defines the size of the shared global budget pool B. Other hyper-parameters, such as the number of instruction tokens o,
19

---

Published as a conference paper at ICLR 2025
kept consistent with the settings provided in the PyramidKV codebase We also ensure that all other hyper-parameter are consistent across both the baselines and our proposed method. For the nyper-parameter B, as said in Section [4.1 it was chosen from 1.005, 1.01, 1.1, 1.2, 1.5, 2, 5, 10 and we report the best performance according the average score results rather than choosing one beta for each dataset. The details results cache=128 are shown in Figure [I
1iaaa-----Innttruee and -B-ssruct with KV
For hyper-parameter 3, a smaller value represents larger shared budget pool B according to Eq. meaning that KV cache allocation relies more heavily on the importance score distribution for alloca- tion. The results show that Head-R2 performs better with a smaller B, indicating that our Retrieval- Reasoning head distribution is more effective in guiding KV cache budget allocation. It is also worth noting that both Head-R and Head-R2 outperformed Ada-SnapK which is current SOTA method, across the different 3 values we set for (aama-3-8--Instmue and Mistral-7B- IB--truucctt This demonstrates the stability performance of our proposed head-level KV cache compression method.
H DISCUSSION
In addition to KV cache compression methods, another approach to mitigating the long-context scenarios during the pre-filling stage and enhancing the model's ability to handle long texts is context compression (Jiang et al. 2023b; 2024; Ge et al., 2024b Qin et al. 2024). Within the field of context compression, there also different directions exploring how compressing prompts. For example, Jiang et al. (2023b 2024) proposed defining an additional module to compress the original prompt while ensuring that the meaning and coherence of the prompt remain intact before and after compression. The compressed text will serve as the input for the subsequent L.Mss. On the other hand, [e al.] (2024b) and Qin et al. (2024) adopted different compression strategy by compressing original prompt into memory slots instead of text. Ge et al. (2024b) proposed ICAE method, which employs an additional trained in-context autoencoder fixed-length memory slots, which are then used as the input to LLMs. They pretrained and fine- tuned the proposed in-context autoencoder to enable its capacity and generalization. Qin et al. (2024) further utilize the model's internal hidden states as the retained content after compression, with an additional scorer employed for the corresponding selection.
compress the input into
The target of KV cache compression is the KV cache itself, whereas the goal of context compression is the original input or the corresponding representations derived from the input. From the perspec- tive of context compression, current KV cache compression methods can be viewed as compressing the input context using the model's 's own knowledge, without relying on trainable modules for selec- tion. For instance, the SnapKV Li al. (024d) method, which our approach builds upon, uses the last a tokens as the observation window and selectively retained KV cache based on attention from these tokens. Compared context compression methods like LLMLingua (Jinng et al. 2023b) and ICAE (Ge et al.] 20245), current KV cache compression methods are simpler, as they do not require defining or training additional components. They ar also easier achieve higher computational efficiency, as these KV cache compression methods avoid relying on external components to obtain compressed inputs and do not require recomputing the key and value matrices.
LIMITATION
Although head-level KV cache compression methods can maintain computational effciency com- pared to other baselines, they may introduce additional overhead due to the extra effort required to dynamically manage each head especially in the parallel execution. The synchronization and computational costs associated with head-level operations may reduce decoding speed in multi-GPU environments, leading to trade-off between performance and efficiency. Hence, an important chal- lenge remains in how to further optimize head-level KV cache compression to mitigate parallel overhead and maintain practical deployment benefits.
20

---

Published as a conference paper at ICLR 2025
Llamm-3--BB-nstruutt size=128
Llama--88-Instruct::
32
31
32.0
29
31.5
27
31.0
HeadKV validation
HeadKVites
26
1.005
1.01
1.1
1.2
1.5
5
10
1.005
1.01
1.1
1.2
1.5
10
Beta
Beta
Figure 11: Results for different We leverage scikit-learn to extract 15% of each of the six datasets validation (HeadK V-validation), while the remaining data used the test set (HeadK test). The black line represents the optimal /3 value determined based on the validation set.
Input: John travelled the hallway. Mary journeyed bathrom. Daniel went back to the hathroom. John moved to the bedroom.
qa1
Question:
Where Mary?
Target: bathroom
Daniel took the milk there. John jourmeyed the garden. Daniel went back hallway. Daniel journeyed to the bathroom. Daniel dropped the milk. Daniel took the milk there. John grabbed the apple there. Sandra journeyed hallway. Sandra went back garden.
Input:
kitehen. John the
qa2
Question: Where the apple? Target: hallway
Mary picked up apple. John went to the garden. Sandra travelled office, Sandra took milk, John weat bedroom, Sandra went to the kitchen. John journeyed Mary left the apple. Mary travelled the office. Sandra to the office. Daniel went to the hallway, Sandra discarded the milk.
Input:
office.
qa3
Question: Where the milk before the office?
Target: kitthenn
Input: hallway
qa4
Question:
What is the bathroom east of?
Target: bedroom
Input: Fred picked up the football there. Fred gave the football Jeff. Bill went back to the bathroom. Jeff grabbed the milk there. Jeff gave the footbell to Fred. Fred handed the football to Jeff. Jeff handed the football to Fred. Fred gave the football to Jeff.
qa5
Question: Target: Jeff
give football
Figure 12: Examples for each task used in Reasoning in--a- -aaaiinn-nn----astaaaea test.
21

---

Published as a conference paper at ICLR 2025
Liama-3-8B-44 struct
Mistral-7 Mistral-7Binnstrun B-instruct
Tulli Average Trom: $7.36
PullY Average fcom:: 45.24
300
30
-80
10
80
82
29
29
93
20
51
52
57
43
54
48
52
47
Le
SnepkY Score: 16
Snapity Average 41.09
82
022
40
52
24
60
48
46
16
Average
PyramistV Average
123
10
S*
46
422
15
34
33
27
22
32
923
-20
-20
1k
24
4k
4k
34k
20c
Ado KV. average Score: 55.00
Ada Snapkly Average Score: 42.54
100
110
023
10
82
20
24
82
20
20
2:
25
24
12
25
22
12
443
Aurfolltr Average Score: 55.16
Ours Reziivval. Avarage Score: 42.94
022
-40
-60
31
25
110
57
45
43
43
423
14
60
73
44
41
300
33
53
42
160
14
14
359
30:
Langth
Content Langth
Figure 13: Detail results for Reasoning-in-a-Haae test on five datasets.
22


Reviews:
[{"review_id": "3ZkFHnUeuV", "reviewer": "Reviewer_iyGe", "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. \n\nThey 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.\n\nAlso, 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.\n\nThey also conduct thorough analysis on long-context retrieval, latency , memory usage, etc.", "questions": "N.A.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. Reasoning head-level kV cache allocation and importance score estimation. \n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0YGl7Hm4we", "reviewer": "Reviewer_NPoe", "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.", "questions": "1. The a_h in Equation (1) and (2) may need a superscript t.\n\n2. Typo in Figure 1: The red text Prefilling Phrase should be Prefilling Phase.\n\n3. Typo in Figure 5 and Figure 8: 'Neele-in-a-Haystack' should be 'Needle'.\n\n4. Are those gaps between FullKV and HeadKV in the right subfigure of Figure 6 solely caused by KV cache?\n\n5. Giving a pseudo code would be better.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The KV cache budget allocation strategy maintains the total amount of KV cache constant and keeps inference time unchanged.\n\n2. Using top-k attention weights to refine the score can enhance the retrieval-head evaluation. \n\n3. 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.\n\n2. How the retrieval-reasoning dataset is generated is not mentioned in the paper.\n\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vJhJO6JIbK", "reviewer": "Reviewer_NjWN", "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.", "questions": "Please see weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "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.\n- 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?\n2. 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?\n3. The estimation equation used to determine R2-head seems to be vague (or even incorrect).\n- What is the first sigma, or the t-sigma, sum used for?\n- 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?\n4. 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.\n\n[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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "O6CoyJbEYb", "reviewer": "Reviewer_vGF6", "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.\n\n**post review**: I raised my score from 5 to 6.", "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?\n2. 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?\n3. It would be great if the authors could discuss and compare this work to RAG methods.\n\nIt 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:\n\n- In-context Autoencoder for Context Compression in a Large Language Model. Ge et al, 2024. In ICLR.\n- Dodo: Dynamic Contextual Compression for Decoder-only LMs. Qin et al, 2024. In ACL.\n- VCC: Scaling Transformers to 128K Tokens or More by Prioritizing Important Tokens. Zeng et al., 2023. In NeurIPS.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "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.\n2. This paper proposes a head-level allocation schema that can dynamically allocate memory across attention heads and layers. \n3. 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. \n\n2) 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. \n\nThis might be a general issue for all the work of KV cache selection, but this paper does not break this limitation.\n\n3) 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. \n\nEven 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)