PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling
Authors: Zefan Cai
Yichi Zhang
Bofei Gao
Yuliang Liu
Yucheng Li
Tianyu Liu
Keming Lu
Wayne Xiong
Yue Dong
Junjie Hu
Wen Xiao
Venue: COLM 2025 (2025)
Keywords: KV Cache Compression
Inference Acceleration
TL;DR: We developed PyramidKV, a novel and effective KV cache compression method.

Abstract: In this study, we investigate whether attention-based information flow inside large language models (LLMs) is aggregated through noticeable patterns for long context processing. Our observations reveal that LLMs aggregate information through Pyramidal Information Funneling where attention is scattering widely in lower layers, progressively consolidating within specific contexts, and ultimately focusing on critical tokens (a.k.a massive activation or attention sink) in higher layers. Motivated by these insights, we developed PyramidKV, a novel and effective KV cache compression method. This approach dynamically adjusts the KV cache size across different layers, allocating more cache in lower layers and less in higher ones, diverging from traditional methods that maintain a uniform KV cache size. Our experimental evaluations, utilizing the LongBench benchmark, show that PyramidKV matches the performance of models with a full KV cache while retaining only 12% of the KV cache, thus significantly reducing memory usage. In scenarios emphasizing memory efficiency, where only 0.7% of the KV cache is maintained, PyramidKV surpasses other KV cache compression techniques, achieving up to a 20.5 absolute accuracy improvement on TREC dataset. In the Needle-in-a-Haystack experiment, PyramidKV outperforms competing methods in maintaining long-context comprehension in LLMs; notably, retaining just 128 KV cache entries enables the LLAMA-3-70B model to achieve 100% Acc. performance, matching that of a full KV cache.

Published as a conference paper at COLM 2025
PyramidKV: Dynamic KV Cache Compression based on Pyra- midal Information Funneling
Zefan Cai', Yichi Zhang", Bofei Gao², Yuliang Liu?, Yucheng Lit, Tianyu Liu Keming Lu Wayne Xiong', Yue Dong", Junjie Hu!, Wen Xiao",
Uniiersitty of Wisconsin Madison *Peiigg University Nanjing University University of Surrey Swwen SevennUniivee University of California Riverside Microsoft zefncai@gmail.com https://github www.///eaa- htw.:/eefa--aaa/Prrmiiiik
Abstract
In this study, we investigate whether attention-based information flow inside large language models (LLMs) aggregated through noticeable patterns for long context processing. Our observations reveal that LLMs aggregate information through Pyramidal Information Funneling where attention is scattering widely in lower layers, progressively consolidating within specific contexts, and ultimately focusing on critical tokens (a.k.a massive activation or attention sink) in higher layers. Motivated by these in- sights, we developed PyramidKV, a novel and effective KV cache compres- sion method. This approach dynamically adjusts the KV cache size across different layers, allocating more cache in lower layers and less in higher ones, diverging from traditional methods that maintain a uniform KV cache size. Our experimental evaluations, utilizing the LongBench benchmark, show that PyramidKV matches the performance of models with a full KV cache while retaining only 12% of the KV cache, thus significantly reduc- ing memory usage. In scenarios emphasizing memory efficiency, where only 0.7% of the KV cache is maintained, PyramidKV surpasses other KV cache compression techniques, achieving up to 20.5 absolute accuracy experiment,
improvement on TREC dataset. In the Boodle-n----Hyssiee PY outperforms competing methods in maintaining long-context comprehension in LLMs; notably, retaining just 128 KV cache entries en- ables the LLAMA-3-70BB model to achieve 100.0 Acc. performance.
Introduction
Large language models (LLMs) (Achiam et al. 2023 Touvron et al. 2023a b: jiang et al.
2023) are integral to various natural language processing applications, including dialogue systems (Chiang et al. 2023), document summarization (Fabbri et al. 2019a), and code completion (Roziere
2023 . These models have recently been scaled up handle long contexts (Fu et al. 2024 Ding et al. 2024 Zhu et al. (023) Chen et al. 2023), with GPT-4 processing up to 128K tokens and Gemini-pro-I.5 handling 1M tokens. However, scaling LLMs to extremely long contexts naturally leads to a significant delay due to the quadratic computation of attention over long contexts. A common solution to mitigate such inference delays involves caching the key and value states (KV) of previous tokens (Waddington et al. 2013), with the trade-off of requiring extensive GPU memory storage. For instance, maintaining KV cache for 100K tokens in LLaMA-2 7B requires over 50GB of memory, while a 2K context requires less than 1GB of memory (Wu et al. 2024).
To tackle these memory constraints, recent studies have explored the optimization of KV caching, including approaches such as low-rank decomposition of the KV cache (Dong et al.)
2024) pruning non-essential KV cache (Zhang et al. 2024 Li et al., 2024 Ge et al. 2023). Notably, it has been shown that maintaining merely 20% of the KV cache can preserve substantial level of performance (Zhang et al. 2024). Moreover, extreme compression of the

---

Published as a conference paper at COLM 2025
(b) Selected KV cache (fixed position)
(c) Selected KV cache (fixed length)
(d) Selected KV cache (dynamic length)
(a) Full KV cache
Layer 31
Layer
Prompt
Query
Prompt
Query
Prumpi
Query
Prompt
Query
Dense
Streaming LLM
SnapKV /H20 PyramidKV (Ours)
Illustration of PyramidKV compared with existing KV cache compression methods. (a) Full KV has all tokens stored in the KV cache in each layer; cache size increases the input length increases. (b) StreamingL.MM (Xiao et al. (023) only keeps few initial tokens with fixed cache size in each layer. (c) SnapKV (Li et al. 2024) and H2O (Zhang et al. 2024) keep a fixed cache size across Transformer layers, and their selection is based on the attention score. (d) PyramidKV maintains pyramid-like cache sizes, allocating more cache budget to lower layers and less to higher layers. This approach to KV cache selection better aligns with the increasing attention sparsity observed in multi-layer Transformers
Figure 1:
KV cache for tasks of longer contexts (e. retrieval augmented generation or RAG for short) can drastically improve efficiency and further reduce resource use. However, questions about the universal applicability of these strategies across all layers of an LLM remain open. (1) Are these KV cache strategies applicable to all layers? (2) Is it computationally efficient to use the same KV cache size across layers as previous studies have done? These considerations suggest a need for an in-depth, more nuanced understanding of KV cache optimization in LLMs.
To examine these questions, we aim to systematically investigate the design principles of the KV cache compression across different layers, specifically tailored to the behaviors of the attention mechanism. We first investigate how information flow is aggregated via attention mechanisms across different layers multi-document question answering (QA), classic task involving long contexts. Our analysis identifies a notable transition of attention distribution from a broad coverage of global contexts to a narrow focus of local tokens over layers in LLMs. This pattern suggests an aggregated information flow where information is initially gathered broadly and subsequently narrowed down to key tokens, epitomizing the massive attention phenomenon. Our findings provide unique insights beyond the previously documented " "massive activation" " (Sun et al. 2024 that very few activations exhibit significantly larger values than others when calculating multi-head attention in LLMs and "attention sink" (Xiao et al. 2023) that keeping the KV of initial tokens will largely recover the performance of window attention.
Building on these insights on how information flows are aggregated through a pyramid pattern, we design novel and effective KV cache pruning approach that mirrors the geometric shape, named PyramidKV. As shown in Figure unlike the lixed-and-san
length KV cache pruning common in prior works (Zhang et al. 2024 Ge et al.) 2023 Li et al., 2024], PyramidKV allocates more KV cache to the lower layers where intormation is more dispersed and each KV holds less information while reducing the KV cache in higher layers where information becomes concentrated in fewer key tokens. To the best of our knowledge, PyramidKV is the first KV cache compression method with varied cache retention across layers, tailoring cache amounts to the informational needs of each layer.
We conducted comprehensive experiments on LongBench (Bai al. 2023) using 17 datasets across various tasks and domains with three backbone models
LLaMa-3-70B-Instruue and Mistral-7B [[iang et al. 2023)). The results show that P preserves performance using just 12.0% of the KV cache (KV Cache size 2048) on the LongBench benchmark and significantly outperforms other methods in extreme conditions, retaining only 0.7% of the KV cache. Moreover, PyramidKV outperforms baseline models (H20 (Zhang et al. 2024), SnapKV (Li et al.) 2024), StreamingLLM (Xiao et al., 22023) across all tested cache sizes (64, 96, 128, 256), with its advantages most pronounced at smaller
2

---

Published as a conference paper at COLM 2025
Attention Weights Hearmag Layer
Quer Parter
Naddrrs
Localized Attention
Kry Poctbott
Altantion Weights Heatmap Layer
Attention Weights Heatmap Layer
Attention Weights Heatmap Layer
Attention
Massive Attention
Quey
PosAtans
Kay Poelflon
May houllers
Figure 2: Attention patterns of retrieval-augmented generation across layers in LlaMa (ouu vron al. 2023a/b) reveal that in the lower layers, the model exhibits broad-spectrum mode of attention, distributing attention scores uniformly across all content. In the middle layers, attention becomes more localized within each document, indicating refined infor- mation aggregation (dotted red triangular shapes in layers 6 and 10). This culminates in the upper layers, where "massive attention" focuses on few key tokens (concentrated attention bars after layer 18), efficiently extracting essential information for answers.
cache sizes. In the Needle In Haystack experiment, PyramidKV effectively maintains the long-context comprehension in LLMs, outperforming than competing methods. Remark- ably, with PyramidKV, retaining only 128 KV cache entries allows the L.LaM-3-77B--nnsse model to achieve 100.0 Acc. performance, matching the performance of a full KV cache.
2 Related Work
There has been a growing interest in addressing LLMs' memory constraints on processing long context inputs. FastGen (Ge et al. 2023) introduces adaptive KV cache management strategy that optimizes memory use by tailoring retention tactics to the specific nature of attention heads. SnapKV (Li et al. 2024) improves efficiency by compressing KV caches via booecting/cuustrinngg selecting clustering significant KV positions based on their attention scores. Heavy Hitter
Oracle (H20) (Zhang et al., 2024) implements a dynamic eviction policy that effectively balances the retention of recent and historically significant tokens, optimizing memory usage while preserving essential information. StreamingLLM (Xiao al., 2022) enables LLMs trained on finite attention windows to handle infinite sequence lengths without fine-tuning, thus expanding the models' applicability broader contexts.
Pyramidal Information Funneling
To systematically understand the attention mechanism over layers in LLMs for long-context inputs, we conduct a fine-grained study focusing on the multi-document question answering (QA) task. The model is presented with multiple interrelated documents and prompted to generate an answer for the given query. The main target is to investigate how the model aggregates dispersed information within these retrieved documents for accurate responses.
3

---

Published as a conference paper at COLM 2025
In particular, we focus our analysis of the LLaMa (Tourron et al. 2023a b) and visualize the distribution and behavior attention scores over layers. lo assess the distinct behaviors of each multi-head self-attention layer, compute the average attention from all heads within each layer. Figure 2/shoss the attention patterns of one QA example over six different layers (i.e., 0, 6, 12, 18, 24, and 30).
We identify an approximately uniform distribution of attention scores from the lower layers (e.g,, the 0th layer). This suggests that the model operates in broad-spectrum mode at the lower layers, aggregating information globally from all available content without prioritizing its attention on specific input segments. Notably, distinct transition to more localized attention pattern within each document emerges, as the model progresses to encode information at the middle layers (6th to 18th layers). In this phase, atlention is predominantly directed towards tokens within the same document, suggesting more refined aggregation of information within individual contexts.
This trend continues and intensifies in the upper layers (from the 24th to the 30th layer), where we observed the emergence of 'massive attention'' phenomena. In these layers, the attention mechanism concentrates overwhelmingly on a few key tokens. This pattern of attention allocation, where extremely high attention scores are registered, signifies that the model has aggregated the essential information into these focal tokens. Such behavior underscores a sophisticated mechanism by which LLMs manage and streamline complex and voluminous information, culminating in the efficient extraction of the most pertinent data points necessary for generating accurate answers.
PyramidKV
4.1 Preliminaries and Problem Formulation
In an autoregressive transformer LLM, the generation of the i-th token requires that the attention module computes the query, key, and value vectors for all previous i- tokens. To speed up inference process and avoid duplicate computations, the key and value matrices are typically stored in the GPU memory. While the KV cache enhances inference speed and reduces redundant computations, it can consume significant memory when dealing with long input contexts. To optimize memory usage, strategy called KV cache compression is proposed (Zhang et al., 2024: Xiao et al. 2023 Li et al., 2024), which involves retaining only minimal amount of KV cache while preserving much information as possible.
In a LLM with m transformer layers, we denote the key and value matrices in the I-th attention layer respectively as K", vl E Raxd VI E [0, m - 1] when encoding a sequence of u tokens. The goal of KV cache compression to seek two sub-matrices K" V E R xd from the full matrices K' and V', given a cache budget k' < n for each layer E [0, m - 1] while maximizing performance preservation. A LLM with KV cache compression only uses K and V in the GPU memory for inference dataset D, and obtains a similar result to a full model according to evaluation scoring metric, i.e. score(K' K! v D) z score (K), v D).
4.2 Proposed Method
In this section, we introduce our method, PyramidKV, based on the pyramidal information funneling observed across different layers in $ PyramidKV consists of two steps: (1) Dynamically allocating different KV cache sizes/ / budgets across different layers (8 (44..1)
and (2) Selecting important KV vectors in each attention head for caching (.4.22)
4.2.1 KV Cache Size/Budget Allocation
Previous work on KV cache compression et often allocates fixed KV cache size across LLM layers. However, as our analysis in $ demonstrates, attention patterns are not identical across different layers. Particularly dense attention is observed in the lower layers, and sparse attention in higher layers. Therefore,
et al. 2024 Zhang et al. 2024 Xiao et al. 2023
4

---

Published as a conference paper at COLM 2025
using fixed KV cache size across layers may lead to suboptimal performance. These approaches may retain many unimportant tokens in the higher layers of sparser attentions while potentially overlooking many crucial tokens in the lower layers of denser attentions.
Thus, we propose to increase compression efficiency by dynamically allocating the cache budgets across layers to reflect the aggregated information flow based on attention patterns. Specifically, PyramidKV allocates more KV cache to the lower layers where information is more dispersed and each KV state contains less information, while reducing the KV cache in higher layers where information becomes concentrated in a few key tokens.
Following the common practice in KV cache compression (Li et al. 2024 Xiao et al. 2023), we first retain the KV cache for the last a tokens of the input across all layers, as these tokens have been shown contain the most immediate task-related information, where hyperparameter, controlling the number of last few tokens being included in the KV cache. For simplicity, we call these tokens "instruction tokens", which is also referred as "local window" in previous literature (Zhang et al. 2024 Li al., 2024 Xiao et al. 2023).
Subsequently, given the remaining total cache budget Klotal [0= Kf that can be used over all transformer layers (noted as m), we first determine the cache sizes for the top and bottom layers, and use an arithmetic sequence to compute the cache sizes for the intermediate layers to form the pyramidal shape. The key intuition is to follow the attention pattern in aggregated information flow, reflecting a monotonically decreasing pattern of important tokens for attention from lower layers to upper layers. We allocate /total/
p--1 =ttoal /B- m) for the top layer and 1o = (2 / - --n-1 for the bottom layer, where & is hyperparameter to adjust the pyramid's shape. The hyperparameter B is still required to determine the top layer. Once the top layer is identified, the budget of the bottom layer can be calculated by summing the budgets across all layers and equating this sum to the total budget. Once the cache sizes of the bottom and top layers are determined, the cache sizes for all intermediate layers are set according to an arithmetic sequence, defined as
K2
kn-1
kJ
KD
XL.
(1)
111
4.2.2 KV Cache Selection
Once the KV cache budget is determined for each layer, our method needs to select specific KV states for caching within each layer in LLMs. As described in the previous section, the KV cache of the last a tokens, referred to as instruction tokens, are retained across all layers. Following SnapKV (Li et al., 2024], , the selection of the remaining tokens is then guided by the attention scores derived from these instruction higher attention scores are deemed more relevant to the generation process and are thus their KV states are prioritized for retention in the GPU cache.
receiving
In a typical LLM, the attention mechanism in each head h is calculated using the formula:
A² = softmax @² ( /dd),,
(2)
where du denotes the dimension of the key vectors. Following (Li et al. 2024), we utilize pooling layer A to avoid the risk of being misled by some massive activation scores.
To quantify the importance of each token during the generation process, we measure the level of attention each token receives from the instruction tokens, and use this measurement to select important tokens for KV caching. Specifically, we compute the score of selecting i-th token for retention in the KV cache as si: in each attention head h by:
Σ
jE
where [n - n] is the range of the instruction tokens. In each layer and for each head h, the top k' tokens with the highest scores selected, and their respective KV caches are retained. All other KV caches are discarded and will not be utilized in any subsequent computations throughout the generation process.
5

---

Published as a conference paper at COLM 2025
5
Experiment
We conduct comprehensive experiments to evaluate the effectiveness of PyramidKV on performance preserving and memory reduction.
5.1
Experiment Setup
We maintain fixed constant KV cache size for each layer for the baseline methods. In contrast, PyramidKV employs varying KV cache sizes across different layers. To ensure fair comparison, adjusted the average KV cache size in PyramidKV to match that of the baseline models, to keep the total memory consumption of all methods the same. We set = 20 and We the same prompt for each dataset in all experiments.
5.1.1 Backbone LLMs
We compare PyramidKV against baselines using state-of-the-ae open-sourced LLMs, namely -LaMa-3-8B-Instreee
Histral-7B-Instruce (fiang et al.] (023) and LLaMa-3-70B-Instme Testing examples are evaluated in generative format, with answers generated by greedy decoding across all tasks to ensure a fair comparison.
5.1.2 Datasets
We use LongBench (Bai et al. 2023) to assess the performance of PyramidKV on tasks involving long-context inputs. LongBench is meticulously designed benchmark suite that tests the capabilities of language models in handling extended documents and complex information sequences. This benchmark was created for comprehensive multi-task evalua- tion of long context inputs. It includes 17 datasets covering tasks such as single-document 2018 Dasigi et al. 2021], multi-documeet QA YYng et al., 2018 Ho et al. 2020), summarization ( (Huang et al. |021 Zhong et al. 2021 Fabbri et al. 2019b),
QA (Kččský
few-shot learning (Li and Roth 2002) Gliwa et al.) 2019 Joshi et al. 2017), synthetic, and code generation (Guo et al. 2023 Liu et al. 2023b The datasets teature an average input length ranging from 1,235 to 18,409 tokens (detailed average lengths can found in Ta ble 1,, necessitating substantial memory for KV cache management. For all these tasks, adhered to the standard metrics recommended by LongBench (Bai et al. 2023) (i.e., F1 for
QA, Rouge-L. for summarization, Acc. for synthetic and Edit Sim. for code generation.) We refer readers to more details at Appendix
5.1.3 Baselines
We compare PyramidKV with three baselines, all of which keep the same KV cache size across different layers, with different strategies for KV cache selection.
StreamingLLM (SLM) (Xiao al. 2023 is efficient framework that enables LLMs to accept infinite input length.
Heavy Hitter Oracle (H20) (Zhang et al. 2024) is KV cache compression policy that dynamically retains a balance of recent and Heavy Hitter (H2) tokens.
SnapKV (SKV) (Li al. 2024 automatically compresses KV caches by selecting clustered important tokens for each attention head.
FullKV (FKV) caches all keys and values for each input token in each layer. All methods are compared to the FullKV simultaneously.
5.2 Main Results
The evaluation results from LongBench [Bai et al. 2023 are shown in Table and Figure In Figure 3| we report the average score across datasets for 64, 96, 128, and 256 case sizes. In Table I] we report the results for two different KV cache sizes with 64 and 2048. These two sizes represent two distinct operational scenarios the mefory-efficient scenario and the

---

Published as a conference paper at COLM 2025
Performnnc Compatison
Performance Compartoon
Performacee Comparison
02
-20
A20
Avg STORe
Store
11
1500
190
KV Cache Size
Cache Süe
Size
Figure 3: The evaluation results from LongBench (Bai et al., 2023) across 64, 96, 128 and 256 cache sizes at LLaMa----B-Inssrus (Left), Mistral-7B-Instruu (Middle) and LLaMa-3- 70B-Instruct (Right). The evaluation metrics are the average score of LongBench across datasets. PyramidKV outperforms H2O (Zhang et al. (024),, SnapKV (Li al. 2024) and StreamingLLM (Xiao et al. 2023), especially in small KV cache sizes.
Single Document
QA
Documenn QA
Summarszttion
Few shot Learning
Synthetic
Code
Method
Nr
Casper
MP MF-en
HotpotQA
2WiKMMQA
Musiqua
GovReport
GMSuu
MultNews
TBBC
TrivisOA
SAMSum
PCount
PRe
Le
BB-F Arg.
3619
4539
9151
10614
2113
5177
1235
LiaMaa--6l--retrmee KV Size Full
FKV
25.70
29.75
41.12
45.55
35,87
22.35
25.43
23.03
26.21
73.00
90.56
41.88
4.67
69.25
58.05
50.77 41.46
LlaMa-3-88--ne 1--8--n-tttcct...co KV Size 64
SKV
19.86
9.09
27.89 09.89
37.34
28.35
18.17
15,86
20.80
16.41
38.50
85.92
36.32
5.22
69.00
51.7%
48,
38 33,05
HDO
20.90
11.34
27.03
37.25
30.01
17.54
18.29
21.49
19.13
38.00
84,70
69.33
33.89
SLM
17:44
8,68
35.37
15,46
20.06
14.64
69.50
Ouns
21.13
11.18
30.26
35.12
23.76
16.33
21.65
19.23
69.50
52.61
Size
2048
SKY
25.86
29.55
41.10
40.99
35,80
21.81
25.98
23.00
26.45
73.50
90.56
11.66
5.17
69.25
56.65
49.91 11.35
SLM
21.71
25.78
38.15
10.12
32.01
16.86
23.14
22. 64
26.45
70.00
83.22
31.75
5.74
68.
50
53.50
15.56
37
H20
25.56
26.85
39.35 544
44.30
32.52
21.09
24.68
23.01
26.16
33.00
90.56
41.84
4.91
69.25
56.
49.68
39.35
Ours
25.40
29.71
40 25
44.76
35.32
21.98
26.83
23.30
26.19
75.00
90.56
42,14
5.22
69.25
58.76
51.18 41.49
Mistral Catttaall-D-nnstraa KV Size Full
FKV
33.07
20
43.02
27:33
32.91
26.99
71.00
86.98
Mistral-78-Irean KV Sizee
SKV
16.94
17:17
39.51
36.87
22.26
15:18
14.75
20.35
21.43
37.50
84.16
37.28
4.30
61.13
38.4
30.72
SLM
15.01
13.84
28.74
30.97
24.50
13:42
13.25
19:
19.17
35. 50
76.91
29.61
4.67
27.33
03
38.71
35.
25.60
H20
18.19
19:01
37. 40
30.18
22
13.77
16,00
21.52
21.95
37. 00
81.02
38.62
5.00
66.03
03.54
40.06 30,88
Ours
20.91
20.21
39.91
33.57
22.87
15.70
17.31
21.23
21.41
$1.00
61,98
36.96
3.38
60.83
$1.52
37.99
32.19
Mistral-7B-Inetee wi-Irrruct KV
2048
25.89
32.93
42.96
21.17
SLM
20.31
35.25
24.31
27.47
55.
24.51
71.95
32.35
H2O
25.76
31. 10
49.03
40.76
26.52
17:07
24,81
23. 64
26 60
55.00
86.35
42.48
5.50
88.15
49
46.57 39.95
Oun
25.55
32.21
48.97
42.26
27.50
19.36
26.60
23.99
26.73
71.00
86.25
42.94
4.30
90
53.12
47.21 41.63
LlsMa-3-7008 Instruct, KV Size Pull
FKV
27.75
46.48
29.45
52.04
54,90
30.42
32,57
22.27
27.58
73.50
92.46
45.73
12.30
72.50
40.96
63.91 46.55
LlaMa-3-70B Instruct, KV Size
SKV
18:05
72.50
63.33
SLM
22.07
23.53
27:31
43.21
51.66
23.85
16.62
19
15.20
39:50
76.89
33.06
12.00
72.50
40
50,20
35
H2O
25.45
31.61
33.23
48.25
50.30
24.58
20.03
21. 50
21.39
42.00
90.36
41.58
12.00
71.50
65.83
38.16 39.94
Ours
25.47
36.71
42.29
47.08
46.21
28.30
20.60
21.62
21.62
64.50
89.61
41.28
12.50
72.50
45.34
56.50
42.01
LlaMa-3-708
-00-------T--rsttrcce KV
Size 2015
SKV
26.73
45.18
47.91
52,00
55.24
30.48
28.76
22.35 35
27.31
72.50
92.38
45.58
12,00
72.50
41.52
69.27 46.36
SLM
41.01
35:97
46.53
$2.96
20.81
27.16
91.55
72.00
H20
27.67
46.51
27.55
72.50
Ous
27.22
46.19
48.72
51.62
54.56
29.76
73.50
91.88
12,00
72.50
41.36
Table 1: Performance comparison of PyramidKV (Ours) with SnapKV (SKV), H2O, StreamingLLM (SLM) and FullKV (FKV) on LongBench for LaMMa--8B--Isstruct, Mistral- 7B-Instruct and LaMMa-3-70B-nntruue compression methods across various KV Cache sizes and L.MSs. The performance strengths of PyramidKV V are more evident in small KV Cache sizes (i.e. KV Size = 64).
PyramidKV generally outperforms other KV Cache
performance-preserviiee scenario, respectively for trade-off between memory and model performance. In (n/ppnddix Nj we report results of KV cache sizes with 64, 96, 128 and 2048.
Overall, PyramidKV preserves the performance with only 12% of the KV cache and it consistently surpasses other method across a range of KV cache sizes and different backbone models, with its performance advantages becoming particularly pronounced in memory-
7

---

Published as a conference paper at COLM 2025
LLaMA-3-70B 8K Context Size
aa) FullKV, KV Size = Full, acc 100.0
14
54.0
(b) PyramidKV, KV Size=128, acc 100.0
350
(c) SnapKV, KV Size=128, acc 98.6
11.2
294
(d) H2O, KV Size= 128, acc 82.3
Figure 4: Results of the Fact Retrieval Across Context Lengths ("Needle In A HayStack") test in LlaMa-3-7BB-Insteee with 8k context size in 128 KV cache size. The vertical axis of the table represents the depth percentage, and the horizontal axis represents the length.
constrained environments where only about 0.8% of the KV cache from the prompt is retained. Upon examining specific tasks, PyramidKV demonstrates a notably superior performance on the TREC task, few-shot question answering challenge. This suggests that the model effectively aggregates information from the few-shot examples, highlighting the potential for further investigation into in-context learning tasks.
Notably, we initially observe the pyramidal attention patterns from the visualization analysis on the multi-document QA task (Figure 2), but the pyramid heuristic has demonstrated its effectiveness on a range of other Lonch tasks (e.g., single-document QA, In-Context Learning), suggesting its promising generalizability beyond multi-document QA.
The performance advantage of PyramidKV increases as the KV cache memory decreases. By focusing optimizing budget allocation across layers, PyramidKV accurately allocates resources in memory-constrainee scenarios, ensuring that retained information is effectively preserved to maintain model performance. Moreover, in long bench results shown in
8

---

Published as a conference paper at COLM 2025
Table 1 even in the performance-presst preeerriireeeeeigggggg scenario (i.e., KV cache size 2048), PyramidKV improves the performance over baseline methods and even outperforms FullKV.
Among the 16 datasets, the tasks where our proposed method performs slightly worse than the baseline are mostly saturated (e.g,, HotpotQA, Musique, etc under the LlaMa-3- 8B-Instruct setting with KV Size 64, as shown in Table 1). In these cases, is only marginally inferior to the baseline and remains competitive. Conversely, on tasks with greater potential for improvement (e-g., Qasper, Mi-en, TREC, TriviaQA, etc under the same setting), our method significantly outperforms the baseline. Consequently, the overall average performance of our method surpasses that of the baselines. Notably, these tasks include several In-Context Learning tasks (i.e., TREC), our method enjoys best performance gain at In-Context Learning tasks.
method
5.3 Discussion and Insights
5.3.1 PyramidKV Preserves the Long-Context Understanding Ability
We conduct the "Fact Retrieval Across Context Lengths" (Needle In A Haystack) experi- ment (Liu et al. 2023a Fu et al. 2024), which is a dataset designed to test whether a model can find key information in long input sequences, to evaluate the in-context retrieval capa- bilities of LLMs when utilizing various KV cache compression methods. For this purpose, employ LaMa---70B-Instruue as our base, with context lengths extending up to 8k. We compared several KV cache compression techniques (PyramidKV, SnapKV (Li et al. 2024), and H2O (Zhang al., 2024)) at cache sizes of 128 and full cache. The results, presented in Figure 4| The results demonstrate that with only 128 KV cache retained, PyramidKV effectively maintains the model's ability to understand short contexts, and shows only modest degradation for longer contexts. In contrast, other KV cache compression methods significantly hinder the performance of LLMs. Notably, for the larger model (LlaMa-3-70B- Instruct), PyramidKV achieves 100.0 Acc. performance, matching the results FullKV, thereby demonstrating its ability to preserve long-context comprehension with a substan- tially reduced KV cache. We adopt the haystack setting of haystack formed from a long corpus for the Needle In A Haystack task as Wu et al. (2024).
5.3.2 PyramidKV Significantly Reduces Memory with Limited Performance Drop
In this section, we study how sensitive the methods are with different sizes of KV cache. We report the KV cache memory reduction in Table We evaluate the memory consumption of LLaMaa--SB-Imstruuct Specifically, we evaluate the memory consumption of all methods with fixed batch size of 1, a sequence length of 8192, and model weights in fp16 format. We observe that PyramidKV substantially reduces the KV cache memory across different numbers of cache sizes. We also present that the allocation strategy and score-based selection add minimal complexity in the inference phase as Appendix L
cache size
Memory
Compression Ratio
QMSum
TREC
IriviaQA
PCount
PRe
Lec
512
428M
6.3%
22.80
71.50
90.61
5.91
69.50
58.16
1024
856M
12.5%
22.55
71.50
90.61
5.91
69.50
58.16
2048
1712M
25.0%
22.55
72.00
90.56
5.58
69.25
56.79
Full
6848M
100.0%
23.30
73.00
90.56
5.22
69.25
58.76
Table 2: Memory reduction effect and benchmark result by using PyramidKV. We conducted a comparison of memory consumption between the Llama-3-8B-Insttan model utilizing the Full KV cache and the Llama-3-8B-Instee model compressed with the PyramidKV.
Additional results with 64, 96 and 128 KV cache sizes with Ma-Ma-3-8B-Instruee at 8k context length, LlaMa-a-70B-Instrue at 8k context length, and Mistral-7B-Instte (Jiang et al. (023) at 32k context length available in Anpppddixx

---

Published as a conference paper at COLM 2025
6 Conclusion
In this study, we investigate Pyramidal Information Funneling, the intrinsic attention patterns of Large Language Models (LLMs) when processing long context inputs. Motivated by this discovery, we design a novel KV cache compression approach PyramidKV that utilizes this information flow pattern. Our method excels in memooy-consirained settings, preserves long-context understanding ability, and significantly reduces memory usage.
References
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, lige Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. ari preprint arXiv:2303.08774 2023.
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, et al. Longbench: A bilingual, multitask benchmark for long context understanding. ar Xiv preprint ar S.fi::338.14500 2023.
Liang Chen, Haozhe Zhao, Tianyu Liu, Shuai Bai, Junyang Lin, Chang Zhou, and Baobao Chang. An image is worth 1/2 tokens after layer 2: Plug-and-play inference acceleration for large vision-language models. arXiv preprint ar Ktf::243300677 2024a.
Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. Longlora: Efficient fine-tuning of long-context large language models. Xix preprint #rXiv:2309.12307, 2023.
Zhuoming Chen, Ranajoy Sadhukhan, Zihao Ye, Yang Zhou, Jianyu Zhang, Niklas Nolte, Yuandong Tian, Matthijs Douze, Leon Bottou, Zhihao Jia, and Beidi Chen. Magicpig: Lsh sampling for efficient IIm generation, 2024b. URL
org/ abs 2410. 16179
Wei-Lin Chiang, Zhuohan Li, Zi Lin, Ying Sheng, Zhanghao Wu, Hao Zhang, Lianmin Zheng, Siyuan Zhuang, Yonghao Zhuang, Joseph F. Gonzalez, Ion Stoica, and Eric Xing. Vicuna: An open-source chatbot impressing gpt-4 with 90%* chatgpt quality, March 2023.
Pradeep Dasigi, Kyle Lo, Iz Beltagy, Arman Cohan, Noah A Smith, and Matt Gardner. A dataset of information-ssekinn questions and answers anchored in research papers. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 4599-4610, 2021.
Yiran Ding, Li Lyna Zhang, Chengruidong Zhang, Yuanyuan Xu, Ning Shang, Jiahang Xu, Fan Yang, and Mao Yang. Longrope: Extending 1lm context window beyond 2 million tokens. arviv preprint ar arXio:2402.13755 2024.
Harry Dong, Xinyu Yang, Zhenyu Zhang, Zhangyang Wang, Yuejie Chi, and Beidi Chen. Get more with less: Synthesizing recurrence with kv cache compression for efficient Ilm inference. Xrv preprint ar Xio:2402.00939, 2024.
Alexander Fabbri, Irene Li, Tianwei She, Suyi Li, and Dragomir Radev. Multi-news: A large-scale multi-document summarization dataset and abstractive hierarchical model. In Anna Korhonen, David Traum, and Lluís Mårquez, editors, Proceedings of the 57th Annual Meeting of the Association for iomputational Linguistics, pages 1074-1084, Florence, Italy, July 2019a. Association for Computational Linguistics. doi: 1011855//1//99..01 https://aclaattoology wtwphhh/oooog..rr///111122
URL
Alexander Richard Fabbri, Irene Li, Tianwei She, Suyi Li, and Dragomir Radev. Multi-news: A large-scale multi-document summarization dataset and abstractive hierarchical model. In Proceedings of the 57th Anmual Meeting of the Association for Computational Linguistics, pages 1074-1084, 2019b.
Yao Fu, Rameswar Panda, Xinyao Niu, Xiang Yue, Hannaneh Hajishirzi, Yoon Kim, and Hao Peng. Data engineering for scaling language models to 128k context. ar Xiv preprint ar 1001122402.0017, 2024.
10

---

Published as a conference paper at COLM 2025
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. arXio preprint #rXiv:230.001801, 2023.
Bogdan Gliwa, Iwona Mochol, Maciej Biesek, and Aleksander Wawer. Samsum corpus: A human-annotated dialogue dataset for abstractive summarization. EMNLP-IJCNLPP 2019, page 70, 2019.
Daya Guo, Canwen Xu, Nan Duan, Jian Yin, and Julian McAuley. Longcoder: A long-range pre-trained language model for code completion. arXiv preprint arXio:2306.14893, 2023.
Chi Han, Qifan Wang, Wenhan Xiong, Yu Chen, Heng Ji, and Sinong Wang. Lm-infinite: Simple on-the-fly length generalization for large language models. arXiv preprint (xXiv:2308.16137, 2023.
Xanh Ho, Anh-Khoa Duong Nguyen, Saku Sugawara, and Akiko Aizawa. Constructing a multi-hop qa dataset for comprehensive evaluation of reasoning steps. In Proceedings of the 28th International Conference on Computational Linguistics, pages 6609-66255 2020.
Luyang Huang, Shuyang Cao, Nikolaus Parulian, Heng Ji, and Lu Wang. Efficient attentions long document summarization. In Proceedings of the 2021 Conference of the North Ameri- can Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 1419-1436, 2021.
Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint ar Xrix:2310.06882, 2023.
Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir H Abdi, Dongsheng Li, Chin-Yew Lin, et al. Minference 1.0: Ac- celerating pre-filing for long-context Ilms via dynamic sparse attention. arviv preprint arXiv:2407.0244 022499 2024.
Mandar Joshi, Eunsol Choi, Daniel S Weld, and Luke Zettlemoyer. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 1601-1611, 2017.
Tomáš Kočiský, Jonathan Schwarz, Phil Blunsom, Chris Dyer, Karl Moritz Hermann, Gábor Melis, and Edward Grefenstette. The narrativeqa reading comprehension challenge. Transactions of the Association for Computational Linguistics, 6:317-328, 2018.
Wonbeom Lee, Jungi Lee, Junghwan Seo, and Jaewoong Sim. Infinigen: Efficient generative inference of large language models with dynamic kv cache management, 2024. URL
https://arxiv.orga .rgg abs/ 2406 19707
Xin Li and Dan Roth. Learning question classifiers. In COLING 2002: The 19th International Conference Compulational Linguistics, 2002.
Yuhong Li, Yingbing Huang, Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. Snapkv: Llm knows what you looking for before generation. arXiv preprint ar Xiv:24444 14469, 2024.
Nelson F. Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. Lost in the middle: How language models use long contexts, 2023a.
Tianyang Liu, Canwen Xu, and Julian McAuley. Repobench: Benchmarking repository-level code auto-completion systems, 2023b.
Baptiste Roziere, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Gat, Xiaoqing Ellen Tan, Yossi Adi, Jingyu Liu, Tal Remez, Jérémy Rapin, et al. Code llama: Open foundation models for code. ar Xiv preprint 2tt.:::33...9555 2023.
11

---

Published as a conference paper at COLM 2025
Mingjie Sun, Xinlei Chen, J Zico Kolter, and Zhuang Liu. Massive activations in large language models. arx preprint ar (xXi::4402.1766, 2024.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint ar XTX::22322.13771 2023a.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arv preprint (tt::2307.00288, 2023b.
Daniel Waddington, Juan Colmenares, Jilong Kuang, and Fengguang Song. Kv-cache: A scalable pigh-peeformance web-object cache for manycore. In 2013 IEEE/ACM 6th International Conference on Utility and Cloud Computing, pages 123-130. IEEE, 2013.
Lean Wang, Lei Li, Damai Dai, Deli Chen, Hao Zhou, Fandong Meng, Jie Zhou, and Xu Sun. Label words are anchors: An information flow perspective for understanding in-context learning. Xiviv preprint ar arXit:2305.14160, 2023.
Wenhao Wu, Yizhong Wang, Guangxuan Xiao, Hao Peng, and Yao Fu. Retrieval head mechanistically explains long-context factuality. Xiv preprint ar (xXi:2444..1557, 2024.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient stream- ing language models with attention sinks. arXiv preprint ar fttt:2399.145553 2023.
Dongje Yang, XiaoDong Han, Yan Gao, Yao Hu, Shilin Zhang, and Hai Zhao. Pyramid- infer: Pyramid kv cache compression for high-throughput Ilm inference. ar preprint arXit:2405.1220 0oti22255..222,,. 2024.
Zhilin Yang, Peng Qi, Saizheng Zhang, Yoshua Bengio, William Cohen, Ruslan Salakhutdi- nov, and Christopher D Manning. Hotpotqa: A dataset for diverse, explainable multi-hop question answering. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pages 2369-226622 2380, 2018
Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, et al. H2o: Heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems, 36, 2024.
Ming Zhong, Da Yin, Tao Yu, Ahmad Zaidi, Mutethia Mutuma, Rahul Jha, Ahmed Hassan, Asli Celikyilmaz, Yang Liu, Xipeng Qiu, et al. Qmsum: new benchmark for query- based multi-domain meeting summarization. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 5905-5921 5921, 2021.
Dawei Zhu, Nan Yang, Liang Wang, Yifan Song, Wenhao Wu, Furu Wei, and Sujian Li. Pose: Efficient context window extension of lims via positional skip-wise training. ar Xiv preprint arXiv:2309.104000 2023.
12

---

Published as a conference paper at COLM 2025
A Limitations
Our experiments were limited to three base models: LLAMA-3-BB-Intttiie LLAMA-3-70B- Instruct and Mistral-7B-Inntructt While these models demonstrated consistent trends, the robustness of our findings could be enhanced by testing broader array of model families, should resources permit. Additionally, our research conducted exclusively in English, with no investigations into how these findings might be transferred to other languages. Expanding the linguistic scope of our experiments could provide more com oompeeeee understanding of the applicability of our results globally. Based on our results at LongBench and Weedle-in-a-HaySStad experiment, P generally works decently in most of the language tasks (i.e., Single-Document QA, Multi-Document QA, Summerization Few- Shot In-Context Learning, etc.). Although we observe that PyramidKV performs better in some tasks (i.e., Few-Shot In-Context Learning) compared with some other tasks (i.e., Summerization), we have not observed cases that the decoding result collapses some tasks. This remains new topic for future work to explore.
B Future Work
Our investigation on PyramidKV highlights considerable opportunities for optimizing KV cache compression by adjusting the number of KV caches retained according to the distinct attention patterns of each layer (or even for each head). For instance, the retention KV cache for each layer could be dynamically modified based on real-time analysis of the attention matrices, ensuring that the compression strategy is consistently aligned with the changing attention dynamics within LLMs. Furthermore, our experiments indicate that PyramidKV significantly surpasses other methods in few-shot learning tasks, suggesting promising applications of KV cache in in-context learning. This approach could potentially enable the use of more shots within constrained memory limits.
Alcention Weights Heatmmp Layer 0
Attention Weights Heatmap Layer
Cusry
Query
Localized
Kay Roaitions
Key Foethone
Attention Weights Heatmap Layer 18
Attention Weightt Heettmpp Layer 24
Attention Weights Heatmmp Layer 30
Attention
Quar Rotto
Massive Attention
Query
Key Pooloota
Sky
Kay Prailaro
Figure Attention patterns of retrieval-augmmened generation across layers in Mistral-7B- Instruct model (Jiang et al. 2023
13

---

"attention sink" Published as a conference paper at COLM 2025
Alcention Weigles Healmap Layer
Attantion Waights Heathop Layer
Attention
Quiry Nation
Cuap S0tam
Localized
Altention
Postions Altention Weights Heatmap
Kay Peu tore
Mey hartions
Altention Weights Layer
Weghts Heatmap Layer
Attention
Sink
Massive Attention
Qary faaters
Quary
Key Postiors
Figure 6: Attention patterns of retrieval-augmented generation across layers in Mixtral- 8x7B-Instruct Mixture-of-Experts model.
C Related Work
Interpretation of LLMs Prior research has shown that attention matrices in LLMs are typically sparse (Chen et al. 2024a Xiao et al. 2023 Zhang et al. 2024), focusing dispro- portionately few tokens. For instance, Xiao et al. (2023) identilied phenomenon, where maintaining the Key and Value (KV) states of the first few tokens can substantially restore the performance of windowed attention, despite these tokens not being semantically crucial. Similarly, Sun et al. (2024) identified "massive activations" pattern, where minority of activations show significantly larger values than others within LLMs. Interestingly, these values remain relatively constant across different inputs and act as critical bias terms in the model.
Further explorations in this field reveal distinct patterns across various attention heads and layers. et al. (2024) observed that certain atfention heads consistently target specific prompt attention teatures during decoding, Additionally, Wang et al. (2023) discovered that In-Context Learning scenarios, label words in demonstration examples serve as semantic anchors. In the lower layers of an LLM, shallow semantic information coalesces around these label words, which subsequently guide the LLMs' final output predictions by serving reference points. Recently, Wu et al. (2024) revealed that special type of attention head, the so-called retrieval head, is largely responsible for retrieving information. Inspired by these findings that the attention mechanism exhibits varying behaviors across different layers, we discovered that "Massive Activation" does not consistently manifest across all layers in long context sequences; instead, it predominantly occurs in the upper layers. Additionally, we identified a novel trend of information aggregation specific to long-context inputs, which will be further explained in 3
KV Cache Compression There has been a growing interest in addressing LLMs' memory constraints processing long context inputs. FastGen (Ge et al. 2023) introduces adaptive KV cache management strategy that optimizes memory use by tailoring retention tactics to the specific nature of attention heads. This method involves evicting rongerange contexts from heads that prioritize local interactions, discarding non-special tokens from heads focused special tokens, and maintaining standard KV cache for heads that engage
14

---

Published as a conference paper at COLM 2025
broadly across tokens. SnapKV [ii et al. 2024)
improves efficiency by compressing KV positions based on their attention scores. implements dynamic eviction policy that
caches via selecting/ chanterinn sigmificant KV
Heavy Hitter Oracle (H20) (Zhang et al. 2024 effectively balances the retention of recent and historically significant tokens, optimizing memory usage while preserving essential information. IL.e (Xiao et al. 2023 enables LLMs trained on finite attention windows to handle infinite sequence lengths without fine-tuning, thus expanding the models' applicability
broader contexts. LM-
Infinite (Han et al. (202) allows LLMs pre-trained with 2K or 4K-long segments to generalize to up to 200M length inputs while retaining perplexity without parameter updates.
While these approaches have significantly advanced the efficient management of memory for LLMs,, they generally apply fixed KV cache size across all layers. In contrast, our investigations into the attention mechanisms across different layers of LLMs reveal that the attention patterns vary from layer to layer, making a one-size-fits-all approach KV cache management suboptimal. In response to this inefficiency, we propose a novel KV cache compression method, called PyramidKV that allocates different KV cache budgets across different layers, tailored to the unique demands and operational logic of each layer's attention mechanism. This layer-specific strategy takes a significant step toward balancing both memory efficiency and model performance, addressing a key limitation in existing methodologies.
D Pyramidal Information Funneling
Figure and Figure 6 shows the attention patterns of one QA example over six different layers (l.e., 0, 6, 12, 18, 24, and 30) for Mistral-7D-Instruct model and Mixtral-&x7B-Instan Mixture-of-Experts model. Figure
Figure 6 demonstrate that the Pyramidal Informa- tion Funneling phenomenon is also evident in both the Mistral model and Mixtral model. The results reveal that, akin to Llama-like models, Mistral exhibit progressively narrowing attention focus across layers. This supports the universality of the Pyramidal Information Funneling phenomenon across diverse model families. We hope this addresses your concern and underscores the generalizability of our findings.
Our analysis uniquely examines attention metrics across all transformer layers, from 0 to 30, leading to the discovery of key phenomenon we term Pyramidal Information Funneling.
Lee et al. 2024) conducted a limited investigation into attention patterns, focusing only the lower layer (layer 0) and a single upper layer (layer 18). While Lee et al. (2024) noted that attention becomes more skewed in upper layers, it did not provide a fine-grained observation of attention patterns across all layers. In contrast, our study reveals several novel findings:
Localized Attention: We observe that attention progressively narrows its focus, targeting specific components within the input sequence.
Massive Attention Mechanism: In the upper layers, attention heavily concentrates
on small set of critical tokens. Notably, these tokens are not limited to the leading positions, as observed in Lee et al. (2024), but also appear at regular intervals across the sequence. The discrepancy arises from differences in input settings, with Lee et al. 2024) identifying massive attention only at the initial tokens.
These insights motivated us to propose token-selection method based on the highest attention scores in the upper layers, rather than solely relying on tokens from earlier
positions.
To the best of our knowledge, Chen et al. 2024b has not analyzed attention patterns across transformer layers.
Therefore, although Lee et al. 2024 and Chen et al. (2024b) are considered contemporane- ous with our work, making comparison unnecessary, the perspective of our observation is considered novel compared with [ee et al. (2024) and Chen et al. (2024b). Moreover, although Lee et al. (2024) also observed attention patterns, the method we proposed based
15

---

Published as a conference paper at COLM 2025
attention score
Selected Key and Value
Layer 31
Layer
0.0
Input ids
Local Window
Figure 7: Illustration of PyramidKV. At the lower level of the transformer, the PyramidKV selects more keys and values based on the exhibited average attention pattern. Fewer keys and values at the higher level are selected based on the massive activation pattern, where we observe that attention scores are concentrated over local regions.
on our observations is significantly different from novelty our work.
Lee et al. (2024), further highlighting the
E Details of Proposed Method
Based on the pyramidal information funneling observed across different layers, PyramidKV consists of two steps: (1) Dynamically allocating different KV cache sizes /bddgets across different layers; and (2) Selecting important KV vectors in each attention head for caching
as Figure 7
Our decision to use an arithmetic sequence driven by three key factors:
Alignment with Pyramidal Information Funneling Pattern: Empirical observa- tions reveal a pyramidal information funneling pattern, where lower layers exhibit dispersed attention while higher layers concentrate on fewer tokens. Inspired by this, we adopt the arithmetic sequence design to align with this natural progression.
Superior Empirical Performance: Through extensive experimentation across di- verse datasets, we compared various methods, including the arithmetic sequence and adaptive approaches. Results consistently showed that the arithmetic sequence method outperformed others.
The arithmetic sequence method introduces minimal computational overhead compared to adaptive approaches, which require dynami- cally computing cache budgets across layers.
Computational Efficiency:
To perform KV cache eviction, and release process of torch.gather:
torch.gather. Below, outline the memory allocation
Index Selection:
Identify the positions of the elements to extract from the input
tensor.
Memory Location Calculation: Compute the specific memory locations of the elements to be extracted using the strides of the input tensor across each dimension.
16

---

Published as a conference paper at COLM 2025
Allocate memory to create a output tensor and copy the selected elements to their corresponding positions in the output tensor.
Output Tensor Creation:
Memory Management: Since torch. torch.gather is not an in-place operation, creates a new tensor to store the results, while the memory of the original input tensor is released.
The speed-up offered by PyramidKV is complementary to that achieved through tensor parallelism and pipeline parallelism, these approaches are not mutually exclusive. Pyra- midKV can be seamlessly integrated with both tensor parallelism and pipeline parallelism.
In our current setting, we focus on scenarios with long input sequences and short output sequences. Accordingly, we reset the KV cache only after the prefill phase. As a result, during the first forward pass, each attention head may discard tokens based on attention scores computed from the full input sequence and the most recent token. While our implementation resets the cache only after the prefill, it is indeed possible to update the cache more frequently-for instance, by discarding one token per forward pass or discarding multiple tokens after set number of steps. Discarding multiple tokens after an amount of forward passes tends to be more efficient, as it reduces the overhead associated with frequent token selection and eviction.
In our approach, the term "instruction tokens" refers to the final tokens in the input sequence, which are retained in the KV cache across all layers. Our prompt format typically consists of a system message, long input document, and query; as result, the instruction tokens may include the tail portion of the query. In our experiments, we fix a = 8 across all datasets, which is usually much shorter than the full query length. These instruction tokens are used to compute attention scores over the full input sequence, guiding the selection of tokens to retain in the KV cache. This design is in line with prior work, such as SnapK V,, which uses a similar mechanism often referred to as a "local window." It's important to note that this mechanism does not allow us to "cheat" by selectively preserving only tokens that contain the correct answer- there is no guarantee that such tokens are retained. Instead, the goal to preserve tokens that are deemed important by the model's attention mechanism. Rather than undermining the integrity of the evaluation, this strategy embodies the core motivation of KV cache compression: leveraging attention signals to identify and retain a small subset of key tokens that approximate the performance of full-context inference.
To keep the experimental setup consistent with existing baselines such as SnapKV, we currently follow their practices to use simple heuristic: selecting the last few tokens as instruction tokens. We acknowledge that more principled selection strategies could improve performance and represent valuable direction for future work. Regarding the claim that fewer instruction tokens yield better results: this arises because we always retain the instruction tokens in the final KV cache. Increasing the number of instruction tokens reduces the number of remaining slots for important tokens selected via attention, thus limiting compression flexibility and often degrading performance. This is supported by our empirical finding that using 8 or 16 instruction tokens outperforms longer configurations. That said, we agree that more instruction tokens could lead more accurate estimation of token importance. To balance both concerns, we are considering separating the roles of instruction tokens: one set for computing attention, and another (potentially smaller) set for forced retention. We believe this refinement could improve overall performance.
Our method allows flexible adjustment of both the parameter of for last tokens and the retained KV cache budget to suit different tasks. For instance, in tasks that require extensive long-context reasoning, the retained KV cache budget can be increased. Similarly, for tasks with longer queries, more last tokens can be used to ensure sufficient context coverage.
Details of Evaluation
We use LongBench (Bai et al. 2023 to assess the performance of PyramidKV on tasks involving long-context inputs. LongBench is a meticulously designed benchmark suite that tests the capabilities of language models in handling extended documents and complex
17

---

Published as a conference paper at COLM 2025
Datose
NI
Qayne:
MF-mn
Hetp404
Mesore
GesRqpont
MultNess
TRDC
TArROA
SAMSum
Lst
RB-F
Fergh
4:51
1131
.P14
617
6756
11:01
0090
1255
0183:77
00076:
800306:
093070
0.77244
038--11
000009
00073>
616776
200999
000699)
889927
123
030863
00033
8 335876
E208000
0211434
00012:0
8046777
E30390
223454
0201493
30244
IUDSEES
348
0311250
(15292
8 444822
0323801
2439877
038533
0334486
0390335
0096236
828982
333750
0 338003
63000
Table 3: KV cache compression ratios for LongBench.
information sequences. This benchmark was created for multi-task evaluation of long context inputs.
We present the details of metrics, language and data for LongBench at Table 4
We run all the experiments on NVIDIA A100.
In Figure and 1 "KV-cache size" means the average number of KV cache retained for each layer and each head after pruning, For example, cache size of 128 means that each attention head in each layer retains 128 tokens. Therefore, the total retained KV tokens across the entire model is 128 x #layers X #heads. The original (unpruned) context length corresponds to the model's maximum sequence length: 32K tokens for Mistral models and 8K tokens for LLaMA-3 models. At inference time, the actual context length depends on the input sequence length, which can vary across datasets.
The reported KV-cache size refers to the average number of tokens retained per layer and per head. While PyramidKV supports adaptive token allocation across layers, we report the average size for consistency and fair comparison with existing baselines such as SnapKV and H2O, which use uniform or non-adaptive budgets. This helps isolate the benefits of our compression strategy under comparable conditions.
The compression rate is computed as the ratio the retained KV cache sizes (e.g,, e.g., 64, 96, 128, 256, and 2048) to the average input sequence length, which are listed in the third row of Table 1 (e.g., 18,409 409 for NrtvQA, 11,214 for Musique, 10,614 for QMSum, and 11,141 for PCound). We explicitly report the actual compression rates alongside the corresponding KV cache sizes both Table 1 and Figure 3 as Table
The KV cache size we refer to (e.g., 64, 96, 128, 256) denotes the number of retained key-value pairs after compression, not the model's context length. The input context length remains unchanged in our experiments and matches the maximum sequence length supported by the model (e.g,, 32K for Mistral, 8K for LLaMA-3). PyramidKV operates after full attention over the entire input, aiming to reduce the KV cache used in the decoding stage. This compression is especially relevant for congect inference, where storing all key-value pairs becomes a bottleneck.
G License
LongBench: MIT
H Handle Rotary Embedding after Tokens are Removed in PyramidKV
We keep the rotary embedding unchanged after tokens are removed, that LLMs can still capture the exact position information even if the tokens are removed. StreamingLLM (Xiao) et al. 2023) shows that rolling kv cache with the correct relative position is crucial for maintaining performance. This is because StreamingLLM is designed to mainly handle unlimited context sizes, where contexts exceed the LLM's fixed context length. Without changing the rotary embedding after token removal, LLMs would receive rotary embedding of a non-monotonic position sequence. For example, after the first KV cache compression, LLMs might receive the input position embedding as [0, 2, 3, 309,, 3097, ,0066), and the position embedding of the generated sequences could be [1005, 1006, 1007,...I. The position sequence of [0, 2, 3,0%,,......... 409,, 1005, 1006, 1007, .. is non-monotonic se- quence, which may negatively hurts the performance. In contrast, our targeting settings will not process unlimited context size. For example, given input sequence of 4012 length, after
18

---

Published as a conference paper at COLM 2025
Dataset
Source
Avg len
Metric
Language
#data
Mngle-Document QA NarrativeQA
Literature, Film
18,409
F1 Fi
English
200
Qasper
Science
3,619
English
200
MultifiedQQAeen
Multi-field
4,559
F1
English
150
QA
HotpotQA
Wikipedia
9,151
F1 F1
English
200
2WikiMultihopOA
Wikipedia
4,887
English
200
MuSiQue
Wikipedia
11,214
Fi
English
200
Sunmarization
GovReport
Government report
8,734
Rouge-L. Rouge-L Rouge-L
English English
200
QMSum
Meeting
10,614
200
MultiNews
News
2,113
English
200
Few-shof Learning
TREC
Web question 5,177
Accuracy (CLS) Fi Rouge-I.
English
200
TriviaQA
Wikipedia, Web
8,209
English
200
SAMSum
Dialogue
6,258
English
200
Synthetic Task
PassageCount
Wikipedia Wikipedia
11,141
Accuracy (EM)
English English
200
PassageRetrieval-en
9,289
Accuracy (EM)
200
Code Completion
LCC RepoBench-P
Github
1,235
Edit Sim Edit Sim
Sython/C#/Jave Python/Java
Github repository 4,206
Table 4: An overview of the dataset statistics in LongBench (Bai et al. 2023 'Source' denotes the origin of the context. 'Accuracy (CLS) refers to classification accuracy, while' 'Accuracy (EM)" refers to exact match accuracy.
KV cache compression, the position sequence would be [0,4, ,6,...,39277, 3987, 008.,987710001 4012], and the position sequence of the generated tokens would be [4013, 4014, By keeping the rotary embedding unchanged after the tokens removed, the LLM avoids monooonoo position sequences, and the LLM can capture the exact position information even if the tokens are shifted. Our preliminary results show that rolling KV cache with the correct relative position will slightly decrease the performance.
I Ablation Study
In this section, we present an ablation study for hyperparameters and allocation strategies.
Based our observations of the attention pattern, find that relatively stable, linear arithmetic decrease aligns more closely with the underlying structure of the pattern. We conduct experiments comparing various allocation strategies.
We conducted hyperparameter testing on the original development sets of 16 datasets in LongBench. The parameter 8 demonstrated remarkable stability, showing minimal sensitivity to varying hyperparameter settings, which highlights its robustness. Conversely, consistently produced superior results when set to 8 16. Consequently, these values were adopted for subsequent experiments. In Appendix H.2 and H.3, we further analyzed the impact of hyperparameter selection on KV cache budget allocation across different layers. The experiments reaffirmed that ß had negligible influence on the outcomes, underscoring its stability. Meanwhile, a continued to deliver optimal results at values of and 16.
I.1 Allocation Srategies
Based on our observations of the attention pattern, we find that a relatively stable, linear arithmetic decrease aligns more closely with the underlying structure of the pattern.
19

---

Published as a conference paper at COLM 2025
We conduct experiments comparing various pyramidal allocation strategies (i.e., linear decay strategy, geometric decay strategy and exponential decay strategy) with a cache size of 64 as Table 5 to confirm that linear strategy indeed optimal preferable.
We also propose three adaptive allocation baselines, which are based on the entropy, Gini coefficient, and sparsity of the attention values at each layer. The weight of each layer is calculated based on its corresponding metric (entropy, Gini coefficient, or sparsity), and the budget allocated accordingly. Specifically:
Layers with higher entropy receive higher weights. Each layer's s entropy calculated based the the layer's attention.
Entropy-based allocation:
Gini coefficient-based allocation: Layers with higher Gini coefficients receive higher weights. Each layer's 's Gini coefficient is calculated based on the the layer's attention
The empirical results as Table consistently showed that the linear strategy outperformed its counterparts, establishing it as the most effective approach for our use case. The experiment strengthens the rationale for choosing the specific allocation method.
Mingle-Decument QA
Mull--DocunentOO QA
Summarization
Few-shot Learning
Synthelle
Code
Stra.
NitQA
Cesper
AiF-en
HelpolQA
2WikiMOA
Musique
GovRsport
@MSuum
MulthNiww
TREC
IIivaaQA
SAMSum
PCaunt
PRR
1.45 R8-F
Geo.
20.51
15.04
29.4
31.93
26.41
16.6
18.32
21.68
18.81
52
87.51
36.15
5.18
69.17
53.11
41.1
34.36
Exp.
20.58
14.82
28.74
34.34
26.24
16.11
18.41
21.63
18.75
52.00
87.94
36.36
5.19
69.17
54.
43.21
34.23
Lin.
21.13
11.18
30.26
35.12
23.76
15.17
18:33
21.65
19.23
58.00
88.31
37,07
5.23
69.50
52.61
45.71
34,76
Eahropy:
18.32
27.22
33.21
1336
17.26
19.87
17.09
87.31
34.29
5.09
68.91
50.12
32.71
Gini.
32.67
Table 5: Ablation study of allocation strategies.
1,2 Hyper Parameter &
We present the study of a for LlaMa-3-8B-Instruue 128 KV cache size budget at Table We find that small alpha value (i.e., 8, 16) leads to better performance than larger alpha value (i.e., 24, 32, 40, 48).
Single-Docurent
Multi-Ducument
Learnngg
Synthetic
Ave
NAtOQA
Qssper
MF-m
HotpoiQA
2WILIMOA
Musique
GovRepart
OM45um
MultiNews
TREC
TrminQA
SAMSum
PCountt
PRe
RB-F
16
23.37
16.21
33.95
38.24
27.28
19.71
21.93
20.86
6000
88.75
38
5.88
69
53
37.19
24
22.85
14.51
32.26
38.38
28.36
20.33
19.55
21.72
20.72
54.50
88.71
38 46
548
69
83
53.65
36.61
32
23.01
14.54
31.88
38.66
29.90
19.16
19.20
21.83
20.52
49.50
87.01
38.01
5.75
69.50
02
84
36.25
21.70
13.06
30.14
36.78
27.34
18.8
18.72
21.37
19.79
44.00
87.74
38
69.25
53.89 35.21
Table 6: Ablation on &.
I.3 Hyper Parameter B
One topic want to analyze for our ablation study is the selection of B, which can determine the staircase. The smaller 8 is, the gentler the staircase is; the larger ß is, the steeper the staircase is. We want to investigate the effect of step size the final result. Results on 128 KV cache size and AlaMa-3-8B-Instruce are shown in Table 7 The results at Table 7 show that using relatively small value of ß yields better outcomes, and PyramidKV is generally robust to the selection of B.
J Integation with MInference
We would like to clarify that PyramidKV and MInference Jiang et al. (024) are complemen- tary approaches addressing different aspects of KV cache optimization. Specifically:
20

---

Published as a conference paper at COLM 2025
Single-Document
Mulli-DDccment
Summarization
Few-shol eerming
Synthetic
Code
Avg
NAYQA
Quaper
MF-4n
HolpoiAA
2WIkMQQ4
Musique
GovRepurt
QMSum
MultiNews
TREC
TriviaQA
SAMSumi
I'Couni
PRR
R6-F
33.79
20.95
36.
18
21.71
16.24
33.39
39.89
27.94
18.38
19.76
22.32
21.20
66.50
88.98
38
5.46
69.50
56.
49.23 37.25 50.11 37.32
16
21.74
14.86
33.64
39.18
28.17
18.77
19.57
22.25
21.48
66.50
89.69
38.87
5.82
69.50
14
22.53
16.31
33.50
40.50
28.15
19.26
19.66
22.39
21.38
65.50
90.02
38.56
5.75
69.50
57.51
49.71 37.51
Table 7: Ablation on p.
MInference focuses on accelerating the generation of KV caches during the pre- filling stage of LLM inference.
contrast, PyramidKV targets efficient KV cache management during LLM decod- ing.
To evaluate their respective strengths, we compared PyramidKV and MInference on Long- bench using KV cache size of 128. The results demonstrated the superior performance of PyramidKV.
Furthermore, we demonstrate that MInference and PyramidKV V can be seamlessly integrated to achieve highly efficient inference while maintaining performance comparable to full attention. The results of MInference combined with PyramidKV, evaluated on Longbench with KV cache size of 128, as PyramidKV + MInference hybrid approach.
Single Document QA
Colti-Decumntt QA
Summarization
Few-shot Learning
Syathetic
Code
Arg.
Stra.
Nrt
Qasper
MF-en
HolpoiQA
2WIKIMQA
Musique
GavRlpprt
QMSum
MultiNews
TRBC
TrivisOA
SAMSuum
PCaunl
PR
BB-P
PymmidKV
39.56
50.16
Minference
31.71
Table 8: Comparison between PyramidKV, MInference and Minference-PrrmiidKV hybrid method.
In summary, we demonstrate that PyramidKV outperforms MInference on Longbench. Furthermore, when integrated with Minference, PyramidKV enhances its performance even further.
K Comparison with PyramidInfer
Our work differs from PyramidInfer in two key aspects:
Decay Strategy: While PyramidInfer Yang et al. (2024) employs geometric decay strategy, our method adopts an arithmetic decay strategy. We argue that the rela- tively stable and linear nature of arithmetic decay better aligns with the behavior of the attention mechanism. This strategy is derived from empirically observed atten- tion patterns, aiming to closely match them. Notably, our approach also achieves superior results, as demonstrated in the experimental results presented in the table below.
Token Selection: PyramidInfer discards tokens in earlier layers, preventing them from being reconsidered in later layers. In contrast, our method allows previously discarded tokens to be re-evaluated in higher layers, recognizing that these tokens may still hold relevance at different stages of the model's processing.
Pyramidal Information Funneling Pattern: A key contribution of our work lies in identifying and leveraging the pyramidal information funneling phenomenon within attention mechanisms. Through in-depth analysis, we observe that attention tends to disperse in earlier layers and progressively concentrates crucial tokens in higher layers. This insight forms the foundation of our arithmetic decay strategy, ensuring that our method aligns more naturally with these intrinsic patterns.
21

---

Published as a conference paper at COLM 2025
Despite some similarities between the two approaches, these differences lead to significantly distinct outcomes. As shown in Table 9 our method consistently outperforms PyramidInfer, highlighting the effectiveness of our design choices.
Single Document Q4
Mulr--ocumnntt QA
Summaricatio
Few what Learning
Synthetic
Code
Arg.
Stra.
NrtaA
Qesper
MFennn
HerpoQQA
2WikIMOA
Musique
Gon Report
@M
MulliNeww
TREC
TriiaQA
SAMSum
PComnt
RB-F
Pyramidinfer 20-02
12.77
25.21
35.81
25,83
16.88
18.2
21.78
18.32
51.00
88.56
35.76
5.61
6922
6112 3344 69.30
5321
PyramiJV
21.13
14,18
30.26
35.12
23.76
16.17
18.33
21.65
19.23
58.00
88.31
37.07
5.23
52.1
4574 4577434476
Table 9: Comparison between PyramidKV and Pyramidinfer.
L Py will cause minimal extra inference overhead.
The allocation strategy and score-based selection add minimal complexity in the inference phase compared to the computation required for next-token predictions as Table 10 Each row shows the setting of using a specific "[Prompt length, Generation length]" combination. We show the inference speed comparison between total inference time, time for allocation strategy and time for score-based selection on LlaMa-3-SB-Instrucc. Each cell is the latency measured in seconds. Furthermore, our budget allocation can be calculated before infer- ence, requiring only a one-time computation. Thus, PyramidKV will cause minimal extra inference overhead.
Prompt Length
Generation Length
Inference Time
Allocation Time
Selection Time
512
512
18.26
0,0000003
0,0194
512
1024
34.69
0.000002
0.0133
512
2048
70.69
0.000003
0.013
512
4096
138.62
0.000005
0.013
1024
512
17.32
0.000002
0.0131
1024
1024
34.67
0,000002
0.01288
1024
2048
70.21
0.000005
0.01296
1024
4096
138.61
0.000003
0.01297
2048
512
17.48
0.000004
0.0128
2048
1024
34.78
0,000006
0.0129
2048
2048
69.50
0,000003
0.01297
2048
4096
138.59
0.000003
0.013
4096
512
17.58
0.000002
0.013
4096
1024
34.93
0.000004
0.0129
4096
2048
69.65
0,000002
0,013
4096
4096
138.87
0.000002
0.013
Table 10: Extra inference overhead of PyramidKV
M Inference Speed Comparison
PyramidKV does not require extra computation time for budget allocation at inference by de- sign. We show the inference speed comparison between PyramidKV and baselines on LlaMa- 3-8B-Instruct as Table 11 Each row shows the setting of using a specific "[Prompt length, Generation lengthJ" combination. Each cell is the latency measured in seconds. PyramidKV does not sacrifice the speed. PyramidKV provides performance improvement and memory saving while runs at a comparable speed compared with baselines (i.e. SnapKV (Li et al.) 2024), StreamingLLM (Xiao et al. 2023) and H2O (Zhang et al. 2024)). That's 's because the allocation strategy requires very limited additional complexity in the inference/ / phase compared with computation required for generation as Appendix T.
22

---

Published as a conference paper at COLM 2025
Prompt Length
Generation Length
H2O
SnapKV
StreamingLLM
PyramidKV
512
512
18.47
18.25
18.96
18.26
512
1024
35.10
34.76
36.20
34.69
512
2048
70.21
69.60
72.35
70.69
512
4096
140.80
139.42
146.37
138.62
1024
512
17.63
17.34
18.12
17.32
1024
1024
35.16
34.61
36.17
34.67
1024
2048
71.02
69.17
72.37
70.21
1024
4096
140.51
138.83
146.09
138.61
2048
512
17.64
19.54
18.22
17.48
2048
1024
35.09
34.76
36.29
34.78
2048
2048
70.84
69.56
72.46
69.50
2048
4096
140.16
139.55
145.22
138.59
4096
512
17.75
17.67
18.40
17.58
4096
1024
35.20
35.08
36.46
34.93
4096
2048
70.02
69.26
72.58
69.65
4096
4096
139.87
138.57
144.98
138.87
Table 11: Performance comparison across different configurations and methods.
N PyramidKV Excels in all KV Cache Size Limitation
The evaluation results from LongBench et al. 2023 are shown in Table 12 Table 13) Tnd 14 We report the results using LlaMa-3 - IB-Instuuct, LlaMa-3-70B-Insiii and Mistral-B5-Instructllangg et al. 2023) for different KV cache sizes.
Overall, PyramidKV consistently surpasses other method across range of KV cache sizes and different backbone models, with its performance advantages becoming particularly pronounced in memory-constrainee, environments. Upon examining specific tasks, Pyra- midKV demonstrates notably superior performance on the TREC task, a few-shot question answering challenge. This suggests that the model effectively aggregates information from the few-shot examples, highlighting the potential for further investigation into in-context learning tasks.
23

---

Published as a conference paper at COLM 2025
Single Document QA
Multi- Document Q4
Summarization
Fewrshot Learning
Synthetic
Method
NIvQA
Qeoper
Mi--nn
HolpolQA
2WikIMQA
Masiqu
GovRsport
MMuun
MultiNww
TREC
DriviaOA
SAMSum
PCount
PRe
RB- Avg.
18409
3639
4239
9151
4887
11214
8734
10614
2115
5127
8209
6258
11141
9289
1236 4206
BaMMa-3-8B-Insteae KV Size Full
IKV
25.70
29.75
11.12
15.55
35.87
23.03
26.21
90.56
58.05
50.77
1laMa-3-8P-nssrrun
Size 64
SKV
19.86
9.09
27.8 27.89
37.34
28.35
18.17
15.86
20.80
16.41
38.50
85.92
36.32
5.22
09.00
51.78
48.38
33.05
H20
20.80
11.34
27.00
37.25
30.01
17.94
18.29
21. 49
19.15
38.00
84.70
37.76
5.63
69.33
53.44
50.15 33,89
SLM
17:44
8.68
22:25
35.37
31,51
15.97
15.46
20.06
14.64
38,00
72.33
29:10
5.42
69.50
46.14
45.09 30.43
Ours
21,13
14.18
30.26
35.12
23.76
16.17
18.33
21.65
19.23
$8.00
68.31
37.07
5.23
69.50
52.61
45.74
34.76
LlaMa- Instruct, KV
Size 96
SKV
10.34
28.65
19.
H70
21.55
11:21
28.73
37.66
30.12
19.57
57
20 -44
67.63
38.47
69 00
50.16
34.57
SLM
18.67
8.43
21.98
38.35
30.59
16.37
17.53
19:84
18:41
41.00
73.92
29.38
5.50
69.50
45.61 31.58
Ours
21.67
15.10
33.50
39.73
26.48
17:47
19.64
22 26
20.49
61.50
87.38
38.18
6.00
69.
55.30
46.79
36.29
LLaMa-3-68-Irettran KV Sizee 125
SKV
21.19
13.55
32.64
38.75
29.64
18.73
18.98
21. 1262
20.26
45,00
88.36
37.64
5.13
68.85
15.84
51,82
35.50
H2O
13:20
37.79
18.45
22. 02
38. 50
87.75
5,83
50
55.06
30.97
35.37
SLM
9.65
29.50
18/03
20.11
20
Ous
21.40
16.92
33.79
39.73
28.72
19.96
20.95
66.50
89.35
5.92
69.00
56.49
07.95
1laMa-3-88--nai KV Size = 2048
SKV
25.86
29:55
41.10
44.99
35.80
21.81
25.98
23.40
26.46
73.50
90.56
41.66
5.17
69.25
56.65
49.94 41.35
SLM
21.71
25.78
38.13
00.12
32.01
16.56
23:14
22. 64
26.45
70.00
83.22
31.75
5.74
68.50
53.50
15.58 37
H20
25.56
26.85
54
44.30
32.92
21.09
24.68
23.01
26.16
35.00
90.56
41.84
4.91
69.25
166
49.68 39.35
Ouni
25.40
29.71
40
25
44.76
35.32
21.98
26.83
23.30
26.19
73.00
90.56
42.14
5.22
69.25
58.76
51,18 41.49
Table 12: Performance comparison of PyramidKV (Ours) with SnapKV (SKV), H2O, StreamingLLM (SLM) and FullKV (FKV) on LongBench for LaMMa--8B-Inssrruct Pyra- midKV generally outperforms other KV Cache compression methods across various KV Cache sizes and LLMs. The performance strengths of PyramidKV are more evident in small KV Cache sizes. Bold text represents the best performance.
Single- Divcumentttttttttt QA
Multi-Decument QA
Summarization
Few-shot Learning
Synthetic
Method
NrtVA
Qesper
Mi--nn
HolpolOA
2WIkMQAA
Musique
GovReport
QMSum
MultiNews
TREC
TrivwsA
SAMSum 8209
PCount
PRe
RB-P Avg.
18609
3619
4539
9151
4887
11214
8234
10614
2113
5177
6258
11111
9289 1236 4206
7B-ttral---Instuuctt KV Size Full
FKV
33.07
43,02
27:33
32.91
24.21
26.99
71.00
98
Miistral/BB.Instruuu KV
SKV
16.91
17.17
39.51
36.87
22:26
15.18
14.75
20.35
21.45
37.50
84.16
37,
1.30
61.13
02
38
30.72
SLM
15.01
13.84
28.74
30.97
24.50
15:42
13.25
19:46
19.17
35. 55.50
76.91
29.61
4.67
27.
33
38.71
35.20
25.60
H20
18.19
19.01
37. 40
30.18
22:22
13.77
16.00
21.52
21.95
37.00
81.02
38.62
5.00
66.03
03.54
40.66 30,88
Ours
20.91
20.21
39.94
33.57
22.87
13,70
17.31
21.23
21.41
54.00
81.98
36.96
3.38
60.83
44.52
37.99
32.19
Mistral-7B Inst--rstttttt.... KV
Size
SKV
18,80
43.29
23.08
15.94
16.65
43.50
60.10
45.53
SLM
15.48
30.03
23.93
16.76
19:15
19.19
41.50
28
H20
19.44
20:81
32.39
21.51
14:43
17.68
22-40
21.99
38.
82.51
39.94
6.06
77.48
42.43
32.67
Ouss
20.35
21.37
41.15
34.94
21.85
15.51
18.21
21. .6.6
21.43
65.00
83.60
39.60
4.30
00
45.83
39.38
34.08
Mistral-78--nsttie
KV Size = 128
SKV
19.16
21.46
43.52
38.60
23.35
16.09
17.66
21. 84
21.47
47.50
84.15
40.24
5.00
69.
31
46.96
42.97 34.96
SLM
16.57
14.68
30.19
22.64
12.34
18.05
18.96
19.19
43.50
74.22
29.02
4.30
29.
48
23
36.16
27:57
H20
33.56
28
12.93
22.61
39.00
82.37
Ouns
21.75
22.03
61.32
34.06
22.79
18.58
21.63
66.00
66.90
$6.96
11.28
35.72
Misiral-7BB-Inttrucc KV
SKY
25.89
32.93
08.56
42.96
27.12
19.02
26.56
21.17
26.69
70.00
86.27
12.57
5.50
88.90
50
46.72 01.56
SLM
20.31
64
45.72
35.25
24.31
12:30
27.47
21.57
24.51
68.
71.95
31.19
5:00
22.5 56
43.38
37.08 32.35
H20
25.76
31.10
49.03
40.76
26.52
17.07
24.81
23.64
26.60
36.00
86.35
42.49
5.50
85.15
49.9
46.57 39.95
Ours
25.33
32.21
48.97
42.26
27.50
19.36
26.60
23.97
26.73
71.00
86.25
42.94
4.30
87.90
53.12
47.21 41.63
Table 13: Performance comparison of PyramidKV (Ours) with SnapKV (SKV), H2O, StreamingLLMM (SLM) and FullKY (FKV) on LongBench for Mistral-7B-Inntructt Pyra- midKV generally outperforms other KV Cache compression methods across various KV Cache sizes and LLMs. The performance strengths of PyramidKV are more evident in small KV Cache sizes. Bold text represents the best performance.
24

---

Published as a conference paper at COLM 2025
Single Document QA
Mult-Dacument
Summarization
shw Learning
Synthetic
Method
NrtvA
Qeper
MPan
HotpolQA
2WikMMQA
Musique
GorRepor
QMSumn
MultiNowo
TREC
TrinsQA
SAMSum
PCount
PRe
RB-P
Avg.
18409
3619
4799
9151
4887
11214
8734
10614
2113
5127
8209
625g
11141
9289
1236
4206
LisMa-3-7UB Instruct, KV Size Full
FKV
48
27.58
40,96
LlaMa-3-708--nnttne KV
SKV
23.92
31.09
36.54
46.66
50.40
18,05
21.11
19.79
41.50
91.06
40.26
12.00
72.50
43.33
57.62
39
SIM
22.07
23.53
27.31
43.21
51.66
25,85
16,62
19:74
15.20
39. 00.50
76.89
33.06
12.00
72.50
40.23
50.20
35.47
H20
25.45
31.61
33:23
08.25
50.30
21.88
20.03
21.50
21.39
12.00
90.36
11.58
12,00
71.50
53.83
$8.16
39.91
Ours
25.47
36.71
42.29
47.08
46.21
28.30
20.60
21.62
21.62
64.30
89.61
41.28
12.50
72.50
45.34
36.50
42.01
1aMMa-37708 Instruct, KV
23.31
72.50
H20
25.30
47.39
50.61
20.87
21.80
22.03
41
12.00
43.84
Ouns
25.47
37.61
64.00
17.33
15.36
27.91
21.05
21 .00
22.31
66.00
91.45
12.36
12,00
72.50
05.12
36.89
42.43
LlaMa-3-70B Instrrctt, KV Size 128
SKV
26.22
37 49
45.70
50.86
52.82
28.50
20.58
72
22.55
53.00
91.61
41.43
12,00
71. 50
45.06
60.50 42.58
SLM
21.25
29.12
29.21
10.20
86
21.50
19.55
19:12
20.61
18.00
76.60
33.21
12.00
72.50
60.65
51.03
36.53
H20
25.61
33.02
37.74
47.77
51.16
20.7%
23. 35.35
42.00
40.35
50.20
72.50
1aMMa-3-708-Iretmaa KV
Sizee 2045
SKV
26.73
45.18
47.91
52.00
55.24
30.48
28.76
27:31
72.50
92.38
45.58
12.00
72.50
41.52
69.27
SLM
26.69
41.01
35.97
46.55
52.98
25.71
27.81
20:81
27. 16
69.00
91.55
44.02
12.00
72
00
41
68.76
43.96
H20
27.67
46.51
89.54
51.69
53.85
29.57
28.57
22.79
27.5
99.00
92.63
45.94
12,00
72.50
11.99
63.90 45.33
Oun
27.22
46.19
48.72
51.62
51.56
29.76
22. 50
27.27
73.50
91.88
15:17
12.00
72.50
11.36
68.12 06.55
Table 14: Performance comparison of PyramidKV (Ours) with SnapKV (SKV), H2O, StreamingLLM (SLM) and FullKV (FKV) on LongBench for LlaMa-3-70B-Instree midKV generally outperforms other KV Cache compression methods across various KV Cache sizes and LLMs. The performance strengths of PyramidKV are more evident in small KV Cache sizes. Bold text represents the best performance.
Pyra-
25

---

Published as a conference paper at COLM 2025
With small budget, our proposed method enables more effective allocation, better pre- serving useful attention information. Second, with a large budget, such allocation becomes less critical, it is sufficient to cover the necessary information. To further illustrate this phenomenon, we have included an ablation study titled "Attention Recall Rate Experiment" as Figure 8 The results show that with small budget, PyramidKV improves the attention recall rate (the percentage of attention computed using the keys retrieved by the method and the query, relative to the attention computed using all keys and the query.). However, with larger budget (i.e., 2k KV Cache Size), the improvement decreases. For 64, 128, 256, 512, 1024 and 2048 KV Cache sizes, PyramidKV's average attention recall rate improvements are 1.87%, 0.64%, 0.61% %,, 0.56%, 0.47% and 0.36%.
devertion bacdll ACHVES Codha STH
Antertine Fne Sacal hawe ACOIE LIVE Cacha 20y
atart Sera Arom Layen cycohe
Facie
Figure 8:
Attention recall rate (the percentage of attention computed using the keys retrieved by the method and the query, relative the attention computed using all keys and the query.) comparison of PyramidKV and SnapKV.
O LongBench results for 128 context length
We conducted additional experiments using sequence length of 128k as TTble 15) The results, summarized in the table below, showcase the model's performance with extended context lengths. These findings provide further validation of the scalability and robustness of our approach.
with
Single Document
Multi Document
Summaricattion
Few-shot Leariing
Synthetic
Code
Avg.
Method
NAoQA
Qasper
MF-m
HolpotQA
IWVKIMQA
Musique
GovRepurt
QMMSum
MuliNNwa
THEC
TaviaQA
SAMSum
PCmunt
RB-P
SnapKV
10.54
5.59
2022
7.49
24.13
61.50
St
744
1921
25
PyramidKV
8.42
10.75
20.37
34.50
Table 15: Comparison of PyramidKV with baselines at 128k context length.
P PyramidKV Preserves the Long-Context Understanding Ability
We perform Fact Retrieval Across Context Lengths ("Needle In A HayStack") (Liu et al., 2023a Fu et al. 2024) to test the in-context retrieval ability of LLMs after leveraging different KV cache methods. We conducted the aaadll-i----Hassted experiment using various LLMss
26

---

Published as a conference paper at COLM 2025
and LLaMA-3-70B-Lnstruucce var- ious KV cache sizes (i.e., 64, 96, and 128) and various methods (i.e., FullKV, PyramidKV, H2O and StreamingLLM), PyramidKV achieves Acc. performance closest to FullKV, while other methods show significant decreases. is worth noting that PyramidKV with 128 KV cache size achieves the same 100.0 Acc. performance compared with FullKV with 8k context size for LLaMA--370B-Instruee
(i.e., Misttral7B--nsttuut-32kkk
Figure 9 different cache size (64, 96 and 128, respectively).
Figure 10
Figure 11 show the results of Mistral-7B-Instruct (Jiang et al. 2023 with
Figure 12
Figure show the results of LlaMa-3-BB-Instruct with different cache size (64, 96 and 128, respectively).
Figure 13
Figure 15 Figure 16
Figure 17 show the results of LiaMa-3-70B-Instuue with different cache size (64, 96 and 128, respectively).
Model
Length
KV Cache
Full KV Acc.
PyramidKV Acc.
SnapKV Acc.
H2O Acc.
Mistral-7B
32k
64
100.00
80.50
43.90
48.40
Mistral-7B
32k
96
100.00
90.50
72.20
59.10
Mistral-7B
32k
128
100.00
91.60
80.10
64.90
LLaMa-3-8B
8k
100.00
92.90
62.00
31.90
I.LMa-3-8B
8k
96
100.00
95.80
80.70
44.20
LLaMa-3-8B
8k
128
100.00
97.40
87.40
49.10
LLaMa-3-70B
8k
64
100.00
99.60
76.20
47.30
LLaMa-3-70B
8k
96
100.00
98.60
94.40
69.90
LLaMa-3-70B
8k
128
100.00
100.00
98.60
82.30
Table 16: Recall Accuracy performance from Fact Retrieval Across Context Lengths ("Needle In HayStack")
27

---

Published as a conference paper at COLM 2025
Mistral-7B-vO. 32K Context Size
160
3000
FullKV, KV Size = Full, acc 100.0
(b) PyramidKV, KV Size-64, acc 80.5
(c) SnapKV, KV Size-64, acc 43.9
(d) H2O, KV Size-64, acc 48.4
Figure 9: Results of the Fact Retrieval Across Context Lengths ("Needle In A HayStack") test Mistral-7B-Instruct with 32k context size KV cache size. The vertical axis of the table represents the depth percentage, and the horizontal axis represents the token length. PyramidKV mitigates the negative impact of KV cache compression on the long-context understanding capability of LLMs.
28

---

Published as a conference paper at COLM 2025
Mistral-7B-v0.2 -32K Context Size
(a) FullKV, KV Size = Full, acc 100.0
(b) PyramidKV, V, KV Size-96, acc 90.5
114
332
340
(c) SnapKV, KV Size-96, acc 72.2
20
200
(d) H2O, KV 96, acc 59.1
Figure 10: Results of the Fact Retrieval Across Context Lengths ("Needle In A HayStack") test Mistral-7B-Instruet with 32k context size in 96 KV cache size. The vertical axis of the table represents the depth percentage, and the horizontal axis represents the token length. PyramidKV mitigates the negative impact of KV cache compression on the long-context understanding capability of LLMs.
29

---

Published as a conference paper at COLM 2025
Mistral-7B-v0.2 332K Context Size
(a) FullKV, KV Size Full, 100.0
(b) PyramidK V, KV 122e-128, acc 91.6
200
100
(c) SnapKV, KV Size= 128, acc 80.1
30
DC
(d) H2O, KV 122e-128, acc 64.9
Figure 11: Results of the Fact Retrieval Across Context Lengths ("Needle In A HayStack") test Mistral-7B-Innsruett with 32k context size in 128 KV cache size. The vertical axis of the table represents the depth percentage, and the horizontal axis represents the token length. PyramidKV mitigates the negative impact of KV cache compression on the long-context understanding capability of LLMs.
30

---

Published as a conference paper at COLM 2025
LLaMA-3-8B 8K Context Size
110
w1 (a) FullKV, KV Size Full, acc 100.0
210
1800
(b) PyramidKV, KV Size-64, acc 92.9
Aa (c) SnapKV, KV Size=64, ace 62.0
(d) H20, KV Size- 64, acc 31.9
Figure 12: Results of the Fact Retrieval Across Context Lengths ("Needle In A HayStack")
test in LlaMa-3-SB-Instruc) with 8k context size in 64 KV cache size. The vertical axis of the table represents the depth percentage, and the horizontal axis represents the token length. PyramidKV mitigates the negative impact of KV cache compression on the long-context understanding capability of LLMs.
31

---

Published as a conference paper at COLM 2025
LLaMA-3-8B 8K Context Size
(a) FullKV, KV Size = Full, acc 100.0
(b) PyramidKV, KV Size-96, acc 95.8
1900
101.2220.000000000000000. (c) SnapKV, KV Size-96, acc 80.7
(d) H2O, KV Sizc=96, acc 44.2
Figure 13: Results of the Fact Retrieval Across Context Lengths ("Needle In A HayStack") test in MaaMa-3-SB-Instreee with 8k context size in 96 KV cache size. The vertical axis of the table represents the depth percentage, and the horizontal axis represents the token length. PyramidKV mitigates the negative impact of KV cache compression the long-context understanding capability of LLMs.
32

---

Published as a conference paper at COLM 2025
LLaMA-3-8B 8K Context Size
320
130
(a) FullKV, KV Size Full, acc 100.0
(b) PyramidKV, KV Sizc=128, acc 97.4
80
1:2000/0080000000000000000000000 (c) SnapKV, KV Size=128, acc 87.4
(d) H20, KV Size=128, acc 49.1
Figure 14: Results of the Fact Retrieval Across Context Lengths ("Needle In A HayStack") test LaMa-3-8B-Instruue with 8k context size in 128 KV cache size. The vertical axis of the table represents the depth percentage, and the horizontal axis represents the token length. PyramidKV mitigates the negative impact of KV cache compression on the long-context understanding capability of LLMs.
33

---

Published as a conference paper at COLM 2025
LLaMA-3-70B 8K Context Size
(a) FullKV, KV Size = Full, acc 100.0
11.0
330
3::00000 (b) PyramidKV, KV 6ize-64, acc 99.6
(c) SnapKV, KV Sizc-64, acc 76.2
(d) H20, KV 6ize=64, acc 47.3
Figure 15: Results of the Fact Retrieval Across Context Lengths ("Needle In A HayStack") test in LlaMa-3-70B with 8k context size in 64 KV cache size. The vertical axis of the table represents the depth percentage, and the horizontal axis represents the token length. PyramidKV mitigates the negative impact of KV cache compression the long-context understanding capability of L.Ms.
34

---

Published as a conference paper at COLM 2025
LLaMA-3-708 8K Context Size
330
300
60
w18.8/8888...................... (a) FullKV, KV Size Full, acc 100.0
114.
(b) PyramidKV, KV Size=96, acc 98.6
(c) Snapk V, KV Size=96, acc 94.4
(d) H2O, KV Size=96, acc 69.9
Figure 16: Results of the Fact Retrieval Across Context Lengths ("Needle In A HayStack") test in LlaMa-3-70B with 8k context size 96 KV cache size. The vertical axis of the table represents the depth percentage, and the horizontal axis represents the token length. PyramidKV mitigates the negative impact of KV cache compression on the long-context understanding capability of LLMs.
35

---

Published as a conference paper at COLM 2025
LLaMA-3-70B - -KK Context Size
www. (a) FullKV, KV Size Full, acc 100.0
114
(b) PyramidKV, V, KV Size=128, acc 100.0
0.0
(c) SnapKV, KV Size=128, acc 98.6
(d) H20, KV 122,=128, acc 82.3
Figure 17: Results of the Fact Retrieval Across Context Lengths ("Needle In HayStack") test in LlaMa-3-70B with 8k context size in 128 KV cache size. The vertical axis of the table represents the depth percentage, and the horizontal axis represents the token length. PyramidKV mitigates the negative impact of KV cache compression on the long-context understanding capability of LLMs.
36

---

Published as a conference paper at COLM 2025
Q
Attention Patterns across heads in the Bottom Layer
Retrieval heads are predominantly located in the higher layers. Notably, no retrieval heads are observed bottom layers. further investigate, we conducted additional experiments on the bottom layer to analyze the attention patterns of the heads as Figure 18 indicate the absence of "massive attention" in any individual head.
Our findings
Figure 18: Attention patterns of retrieval-augmmented generation across heads in the bottom layer in LlaMa.
R PyramidKV's integration with future models
While previous results have shown the pyramidal pattern in LLaMA, Mistral, and Mixtral, is also important consider how PyramidKV might perform on models with significantly different layer counts or attention mechanisms, such as those with recurrent layers or alternative forms of sparse attention. PyramidKV generally performs better on larger models with more layers; for example, it achieves stronger results on LLaMA-3-70B (80 layers) compared to LLaMA-3-8B (32 layers). For models using sparse attention, we expect it to remain effective without substantial parameter retuning, as sparse attention still builds on self-attention transformers, and the trend of increasing attention sparsity from lower to higher layers is well-documented. In contrast, its performance models with recurrent layers is less predictable and warrants further investigation in future work.
37


Reviews:
[{"review_id": "MPBrDicd2v", "reviewer": "Reviewer_DCU1", "summary": "This paper introduces PyramidKV, a novel Key-Value (KV) cache compression method for Large Language Models (LLMs) designed to improve efficiency during long context processing. The core idea is based on the observation of a \"Pyramidal Information Funneling\" pattern in LLMs. This pattern shows that attention is widely scattered in the lower layers, progressively consolidates in the middle layers, and finally focuses on a few critical tokens (massive activation or attention sinks) in the higher layers. Leveraging this insight, PyramidKV dynamically adjusts the KV cache size across different transformer layers. It allocates a larger cache budget to the lower layers (where information is dispersed) and progressively smaller budgets to the higher layers (where information is concentrated). The authors conduct extensive experiments using the LongBench benchmark with LLaMa-3-8B-Instruct, LLaMa-3-70B-Instruct, and Mistral-7B-Instruct models. The results demonstrate that PyramidKV can match the performance of models using a full KV cache while retaining only a small fraction (e.g., 12%) of the cache, thereby significantly reducing memory usage.", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "1.  The \"Pyramidal Information Funneling\" observation offers a clear and intuitive motivation for the proposed dynamic KV cache allocation strategy.\n2. The proposed idea is novel and reasonable.\n3. This paper has conducted a thorough evaluation using the LongBench benchmark, covering 17 datasets across various tasks, three different LLM architectures (LLaMa-3 8B & 70B, Mistral-7B), and multiple KV cache size settings.", "reasons_to_reject": "See questions below", "questions_to_authors": "1. While you've shown the pyramidal pattern in LLaMa, Mistral, and Mixtral, how do you foresee PyramidKV performing on models with significantly different layer counts or attention mechanisms (e.g., models with recurrent layers or different types of sparse attention)? Would the arithmetic sequence parameters need significant retuning?", "extra_scores": {}}, {"review_id": "SKUPBfsxyS", "reviewer": "Reviewer_dYNL", "summary": "### Overview\n\nThis paper proposes a cute technique for dynamic KV cache quantization. Despite the method being very empirical, the paper does a good job of justifying it through careful analysis of attention patterns as well as strong empirical results. The main paper itself is a complete product, but extensive additional content is included in appendices. The appendices could do with some cleaning up, but this does not detract from the paper.\n\n### Notes\n\n**Introduction, RW, and abstract**: Good analysis and motivation of the problem. My only quibble is that the diagram is slightly misleading. The most recent tokens are always included in the forward pass.  I do understand that doing so would make the image less Pyramid shaped.\n\n**Attention Analysis**: Insightful, concise, clear, and well supported. Figure 2 is stellar. Strongly motivates the design.\n\n**Methods**: I think 4.1 and 4.2.1 are clear. 4.2.2 I found a bit confusing. How often do we reset the KV cache? Does every forward pass involve every head discarding one token, computed from the existing long term cache and the most recent token to leave the instruction cache? Better clarification or a visual would help.\n\n**Experiments**: \n1. Models: Llama3 and Mistral are fine models for comparison. I would like to see mistral replaced by slightly more recent models if possible.\n2. Datasets: I think LongBench is the appropriate benchmark for this kind of work.\n3. Baselines: I think these choices of baselines are satisfactory. \n4. Hyperparameters: I was happy to see the extensive sweeps of KV cache size and resulting accuracy.\nIn addition, I think that the results themselves are strong. I think Figure 4 could use some formatting tweaks to become a vectorized graphic.", "questions": "", "limitations": "", "rating": 9, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "Good analysis of important problems at the frontier of LLMs, a clever method to solve them, and rigorous experiments. Extensive appendices and good visualizations to boot.", "reasons_to_reject": "KV Cache compression is a massive field and not my primary area of research. The existing baselines seem fine to me if a bit out of date.", "questions_to_authors": "In 4.2.2, how often do we reset the KV cache? Does every forward pass involve every head discarding one token, computed from the existing long term cache and the most recent token to leave the instruction cache?", "extra_scores": {}}, {"review_id": "VB2NN4UVPi", "reviewer": "Reviewer_KRVU", "summary": "In long-context models, the KV-cache consumes a large amount of GPU memory, and attending over a long context requires a great deal of compute.  As a result, a number of research groups have recently focused on the task of KV-cache compression, and an obvious way to reduce the size of the cache is to prune it, by keeping only the most important tokens, rather than all tokens.\n\nThe authors make an empirical observation, which is that in pre-trained models, the number of \"important tokens\" differs from layer to layer.  In particular, attention in lower layers seems to be widely dispersed, while attention in higher layers concentrates on only a few key tokens.  Thus, an optimal pruning scheme can prune more tokens from higher layers, while retaining more tokens in lower layers.  \n\nThe authors design and test a KV-cache pruning scheme following this principle.  A block of tokens at the end of the prompt are designated as \"instruction tokens\".  These tokens attend to the rest of the prompt, and in each layer, the attention scores for all heads are summed to get an \"importance score\" for each token.  The most important tokens are then retained, following a fixed pyramid-shaped schedule which allocates more resources to the lower layers. \n\nThe authors test their technique on a number of long-context datasets using the llama and mistral pre-trained models, and show that it outperforms other pruning strategies on at least some datasets.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "Using decreasing KV-cache sizes in different layers is a compelling idea to me, because it somewhat resembles the convolution + pooling strategy that has long been used convnets for vision, and has thus far been absent from language models.  The author's analysis of attention patterns in pre-trained models is especially valuable, and they provide detailed and convincing experimental results in both the paper and the appendix.  This seems like it could be a worthwhile contribution to the field.\n\nUpdate: I have raised my score following discussion.", "reasons_to_reject": "The biggest problem with this paper is that authors do not adequately describe how to interpret their results.  \n\nIn particular, Figure 3 and Table 1 give results in terms of \"KV-cache size\".  What does \"size\" mean in this case?  Is it the full KV-cache size (i.e the total context length) before pruning, or is it the size of the KV-cache after pruning?  If it is the total context length before pruning, then those sizes are laughably small -- 256 tokens is hardly any context at all.  However, if it is the KV cache size after pruning, then what was the original (Full KV) context length?  Moreover, PyramidKV allocates a different KV-cache size to each layer, so does the number refer to the average size per layer, or the total size summed over all layers?  The authors claim that PyramidKV has been tested at different compression rates, but the actual compression rate in the result tables is not given, nor can it be calculated from the information that is given.  \n\nPyramidKV works best at small cache sizes, e.g. 64.  However, no LLM uses a context length as small as 64; realistic KV-cache sizes even for small (2B) models will be at least 2K or higher.   The authors should discuss this.  \n\nA second potential issue that needs to be discussed in more detail is the role of the instruction tokens.  For Q/A tasks, do the instruction tokens include the query?  Figure 1 indicates that the answer might be \"yes\".  But if so, then that's cheating -- you are using the query to do full attention over the entire sequence, and then keeping only those tokens which have the correct answer, so of course accuracy will be high.  \n\nIf the instruction tokens do not include the query, they still seem to play an outsized role in determining which tokens are \"important\"; if the instruction tokens are poorly chosen or irrelevant to the task, then the wrong tokens will be kept in the KV-cache.  I am particularly concerned about the author's claim that fewer instruction tokens yield more accurate results; I would expect more instruction tokens to yield a more accurate estimate of which tokens are important.   This should be discussed.\n\nI would also like to see a more detailed empirical analysis of attention patterns in pre-trained models.  It would be nice to have a chart that shows a layer-by-layer breakdown of how sharply attention is focused on just a few tokens.  For example, you could compute a number K for each layer, where K is the number of top-K tokens that needs to be retained to keep the result of attention within certain bounds.  Such a chart would be more objective than eyeballing hand-picked attention matrices.", "questions_to_authors": "See above.   My score of \"reject\" is based on the fact that I can't interpret the results.  If those questions are resolved I will increase my score.\n\nOne further question: how does the idea of \"most important tokens\" combine with \"most recent tokens / sliding window attention\"?  Intuitively, it is not surprising to me that attention in lower layers is less focused, but I would also expect it to be more local, which means that StreamingLLM sliding window attention could be employed at the lower levels, thus providing more token budget for the higher levels.  A more comprehensive analysis of attention patterns would be valuable here (see above).", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper introduces PyramidKV, a KV cache compression method based on the observation of "Pyramidal Information Funneling", where attention patterns become increasingly focused from lower to higher layers. The method dynamically allocates larger cache budgets to lower layers and smaller budgets to higher layers. Reviewers had initial concerns about presentation, which the authors addressed in the discussion phase. Reviewers were enthusiastic about the novelty of this work; experimental evaluation was exhaustive, covering two model families across 17 datasets from LongBench. This method achieves some strong results, matching full KV cache performance while only retaining ~12% of the cache.
Decision: Accept