ThinK: Thinner Key Cache by Query-Driven Pruning
Authors: Yuhui Xu
Zhanming Jie
Hanze Dong
Lei Wang
Xudong Lu
Aojun Zhou
Amrita Saha
Caiming Xiong
Doyen Sahoo
Venue: ICLR 2025 Spotlight (2025)
Keywords: Large Language Models; KV Cache Compression; KV Cache Pruning
Primary area: foundation or frontier models, including LLMs
TL;DR: We propose a novel query-dependent KV cache channel pruning method to reduce the memory usage of LLMs during inference.

Abstract: Large Language Models (LLMs) have revolutionized the field of natural language processing, achieving unprecedented performance across a variety of applications. 
However, their increased computational and memory demands present significant challenges, especially when handling long sequences.
This paper focuses on the long-context scenario, addressing the inefficiencies in KV cache memory consumption during inference. 
Unlike existing approaches that optimize the memory based on the sequence length, we identify substantial redundancy in the channel dimension of the KV cache, as indicated by an uneven magnitude distribution and a low-rank structure in the attention weights.
In response, we propose ThinK, a novel query-dependent KV cache pruning method designed to minimize attention weight loss while selectively pruning the least significant channels. Our approach not only maintains or enhances model accuracy but also achieves a reduction in KV cache memory costs by over 20\% compared with vanilla KV cache eviction and quantization methods. For instance, ThinK integrated with KIVI can achieve $2.8\times$ peak memory reduction while maintaining nearly the same quality, enabling a batch size increase from 4$\times$ (with KIVI alone) to 5$\times$ when using a single GPU. Extensive evaluations on the LLaMA and Mistral models across various long-sequence datasets verified the efficiency of ThinK.  Our code has been made available at https://github.com/SalesforceAIResearch/ThinK.

Published as a conference paper at ICLR 2025
THINK: THINNER KEY CACHE BY QUERY-DRIVEN PRUNING
Yuhui Xu'
Zhanming Jie' Amrita Saha'
Hanze Dong'
Lei Wang'
Xudong Lu² Doyen Sahoo!
Aojun Zhou²
Caiming Xiong'
Salesforce AI Research
2 The Chinese University of Hong Kong
ABSTRACT
Large Language Models (LLMs) have revolutionized the field natural language processing, achieving unprecedented performance across variety of applications. However, their increased computational and memory demands present significant challenges, especially when handling long sequences. This paper focuses the long-context scenario, addressing the inefficiencies in KV cache memory con- sumption during inference. Unlike existing approaches that optimize the memory based on the sequence length, we identify substantial redundancy in the channel dimension of the KV cache, as indicated by an uneven magnitude distribution and a low-rank structure in the attention weights. In response, we propose THINK,, a novel query-dependent KV cache pruning method designed to minimize attention weight loss while selectively pruning the least significant channels. Our approach not only maintains or enhances model accuracy but also achieves a reduction in KV cache memory costs by over 20% compared with vanilla KV cache evie- tion and quantization methods. For instance, THINK integrated with KIVI can achieve 2.8x peak memory reduction while maintaining nearly the same quality, enabling a batch size increase from 4x (with KIVI alone) to 5x when using a sin- gle GPU. Extensive evaluations on the LLaMA and Mistral models across various long-sequence datasets verified the efficiency of THINK. Our code has been made available phttpps: h http:///.....eeeeeeeeeeeeeeaaaaa
INTRODUCTION
Large language models (LLMs) (Hadi et al. 2023 Brown et al. 2020 OpenAl 2023 Touvron al.] [023a bj Scao et al. 2022 Reid et al. 2024) have emerged as a dominant paradigm in natural language processing, achieving state-of-the-art performance across various tasks. A key principle, the Scaling Law (Kaplan et al. 2020), suggests that LLMs exhibit emergent abilities model size increases, improving their capacity to understand complex context and handle long sequences (Xiong et al./ (023). This growth in capacity enables LLMs to generate coherent, con- textually accurate responses and supports variety of downstream applications, such document summarization (Zhang et al. 2019 (024a)), code generation (Chen et al.. 2021b), solving mathemat-
ical problems (Hendrycks et al. 2021 ZZou al. 2023 Wang et al. 2023 Lightman al. 2023)
and conversational AI (OpenAl 2022 2023).
Despite their success in various applications, generating outputs with LLMs incurs significant com- putational and financial costs, which rise with increasing model size and sequence length. Both the training (Strubell et al. [2020] Hoffmann al. 2022 Dong et al. 2024a) and inference (Ainslic et al., 2023) stages involve frequent generation, further contributing to these costs. Consequently, efficient
L.MMs have gained traction recent years (Hu et al.
2021 Wan et al. 2023). To address these
2024 Dettmers et al. 2024 Xu et al.) (023) challenges, quantization (Frantar et al. 2022 Lin et al.
and pruning methods (Sun et al., [2023] Frantar & Alistarh 2023 Lu et al. 2024b) are employed to
reduce model size. Additionally, the key-value (KV) cache, stored in GPU memory alongside model parameters, scales linearly with both sequence length and batch size, creating a substantial memory burden when handling long sequences. Consequently, effective management of extended contexts is essential for the practical deployment of LLMs. this paper, we focus on the long- long-context scenario, aiming reduce memory consumption associated with processing lengthy sequences.

---

Published as a conference paper at ICLR 2025
Head N-1
Head 1
Head
obs
Scoreg
5.2
HEO Snapky
0.1
4.5
0.1
0.2
7.8
Key
Query
Top(Score; 2)
10011 Channel Mask
Score;bs 19:1-5 chs (1
Cache
Figure I: An illustration of the pruning procedure of THINK. Within each attention head, scores are computed for each channel, and only the top T channels out of D are selected for retention. A binary channel mask, along with the pruned keys, is then stored in the cache memory.
Specifically, the number of KV cache parameters is the product of batch size B, sequence length S, number of layers L, number of heads N, channel size per head D, i.e. K, V E B which need to be stored in the GPU memory during inference. To reduce memory and compu- tational costs during inference, efficiency can only be achieved by pruning the dimensions across S, I., NN D or applying quantization the caches. It is well-acknowledged that token importance tends to be sparse. Consequently, KV eviction algorithms have been proposed to reduce the memory footprint by addressing the sequence length dimension S (Xiao et al. 2023b Li al., 2024 Zhang
et al. 2024c Leviathan et al., 2023). Additionally, inter-layer redundancy has been explored (Liu al. 2024a Wu & Tu 2024; Brandon al., (024) to address the layer dimension L. Despite these advances, existing methods have largely overlooked the channel dimension D. In this paper, we highlight that the magnitudes across key cache channel dimensions are significantly imbalanced, and we observe a low-rank structure in attention weights. Based on these findings, we hypothesize that the channel dimension of the key cache exhibits redundancy. Consequently, we focus on explor- ing the redundancy in the KV cache along dimension D, aiming develop strategies that reduce memory costs without compromising performance.
In this paper, introduce THINK, simple yet effective method for KV cache pruning. To pinpoint the least significant channels, we formulate the problem as an optimization task, aiming to minimize the loss in attention weights caused by pruning. To effectively address this problem, we propose a novel query-dependent criterion that assesses the importance of each channel. Using this crite- rion, we then select the most critical channels in a greedy fashion. We evaluate THINK using the LLaMA (Meta) 2024) and Mistral (Jiang et al. (2023) models, and validate its effectiveness across various long-se sequence datasets. The results indicate that, when paired with token eviction and KV cache quantization methods, THINK not only maintains comparable or superior accuracy but also reduces KV cache memory costs by over 20%.
This work pioneers the investigation into the sparsity within the channels of the KV cache. Specifically, we uncover that the activated key cache is sparse for given query. This insight allows us to prune the key cache channels using query-induced norm. Building on this insight, we introduce THINK, the first channel pruning method specifically designed for KV cache. THINK reduces the dimensionality of the cache channels, leading to linear savings in memory usage. As a plug-and-play technique, THINK is orthogonal to other KV cache compression schemes (e.g. KV cache eviction, guantization). Our extensive experiments demonstrate THINK's remarkable efficiency on the LLaMA and Mistral models. Moreover, we explore the potential extension of THINK to value cache pruning (THINKV), showcasing the broad applicability of our method.
Contributions.
OBSERVATIONS
We identify several key observations that motivate approach to pruning the channels of the KV cache. Specifically, we visualize the magnitude of the KV cache and perform singular value decomposition (SVD) on the attention mechanism of the L.LaMA model.
2

---

Published as a conference paper at ICLR 2025
Magnitudes of KV cache channels. Figure 4 (in Appendix A visualizes the absolute values of the
KV cache across tokens in each channel | Consistent with previous findings (Lin et al. 2024 Xiao
et al.] 2023a; Liu et al. 2024b), in the key cache, whereas the value cache lacks obvious patterns. For instance, in layer 14 (Figure 4
observe that only certain channels have significant magnitudes
(a)), the magnitudes in the key cache are substantially higher around the 50th channel across all tokens. A similar pattern is observed in the 50th and 150th channels of the first head in layer 20 (Figure (c)) Given such observation, [Liu et al. (2024b) proposed to perform quantization over the channels of the key cache. Beyond quantization, our findings suggest that certain key cache channels with smaller contributions to the attention mechanism can be pruned. Moreover, channel quantization and pruning are orthogonal techniques, meaning they can be applied concurrently to further improve model efficiency.
Singular value analysis. We conducted singular value decomposition (SVD) (Demmel 1997) on the attention weights of the specified layers to investigate their principal components. The singular values derived from SVD capture the effective rank of the attention matrix, indicating how the information distributed across different components.
Figure (a) (in Appendix A illustrates the energy distribution of the singular values, plotted on a logarithmic scale to enhance visibility of the differences. Notably, only a few singular values exhibit high energy levels exceeding 0.01 across all heads and layers, highlighting their relative significance. This observation aligns with previous findings (Bhojanapalli et al. (02T), where a small subset of singular values often captures most of the information in attention mechanisms. In addition, the rapid decay of the energy suggests that a low-rank approximation can effectively capture the essential information in the key cache.
Figure 5 (b) (in Appendix A the normalized cumulative energy sum reveals that the top 50 singular values account for over 90% of the total energy. These findings suggest that the attention matrix is inherently low-rank (Wang et al. 22200] Chen et al., [222 Dong et al. (2024b), indicating that the key cache can be approximated using low-dimensional vectors (Singhamia al. (2024).
THINK
Notations. We use uppercase letters (e.g., X, Y) to denote scalar values and boldface uppercase letters (e.g., Q K) to denote matrices and tensors. The notation || | Ilp denotes the nn-rrrm for vectors. Unless otherwise specified, || denotes the (2-norm. The Frobenius norm is denoted by I|. The floor function is denoted by [-]. and the ceiling function is denoted by [.
3.1 PRELIMINARY STUDY OF KV CACHE OPTIMIZATION
In scenarios with extended contexts or batch processing, the main limitations in terms of mem- ory and speed are due to the handling of the KV cache size. Considering batch of requests to a Large Language Model (LLM) service that provides a long input prompt consisting of tokens [uBI: (BB the total KV cache size can be computed follows: 2 B S x L N D, where is the number of layers, N is the number of heads, D is the head dimension. The KV cache size grows linearly as the batch size B and sequence length For a model with multihead attention (MHA) Vaswani et al. 2017), such I.LAMA--BB (Touvron et al.) 2023b), context length of 2048 and a batch size of 13 require storing a 13 GB KV cache, which is equivalent the size of the model parameters. The KV cache must be transferred from off-chip memory (HBM) (Jia al. 2018) to on-chip memory (cache) for each token generated, leading to a memory bottleneck. This substantial memory demand highlights the challenges in managing large-scale models and the need for efficient memory utilization strategies. Current methods optimize the KV cache based on the sequence length S (Xiao et al.) 2023b Zhang et al., 2024c] Li et al., 2024) and precision (Hooper et al. 2024 Liu et al. 2024b). We will introduce a new method, THINK, K to optimize it from the perspective of the number of head dimensions D.
Magnitude based Pruning: Based on the observations in Figure 4 which depicts the significant variation in the magnitudes across different channels, one straightforward criterion is use the
We use the visualization code from
3

---

subject to Published as a conference paper at ICLR 2025
Table 1: Performance comparison of pruning key cache by 4p norm on LongBench.
Single Document QA
Document QA
Summarization
Few-shot Learning
Syuthetk
Mothod
NrtrQA
Quaper
MF MF-mn
HolpatQA
2WIKMMA
Musique
GovRopurt
QMSam
MulliNews
TREC
TrivintA
SAMSum
PComnt
FBc
L.ce
RB-P
Avg.
H20
0.0
23.52
17.93
34.68
42.11
33.52
19.92
22.11
22.56
23.82
41.00
90.46
40.20
5.87
69.50
56.71
51.09
37.23
411
0.3
23. 1.
17.15
34.99
40.55
31.49
19:90
21:37
22.15
23.44
40.50
90.10
40.65
5.41
69.00
58. 64
54.90
37.11
+h
0.4
23.51
15.40
34.37
40.71
31.28
20.24
21.25
22.29
22.54
38.30
69.22
39.27
5.87
66.33
58.47
54.33
36,60
tto
0.3
23:98
17.04
35.19
39.27
31.29
20.40
21.62
22.46
23.34
40.50
89.75
40.71
5.54
68.67
60.
58.52
31.40
0.4
23.76
16.25
32.19
40.33
32,13
20,60
21.30
22.25
23.20
39.50
89.61
40.24
5,66
69.00
60.09
99:45
37.22
SmipXV0.0
24:84
13.96
18.77
42.75
34.55
20.57
22.36
22.61
23.97
70.00
90.52
40.29
5.81
69.50
59
51.81
40.10
40.04
+11
41s
41,86
21.59
23.77
40.31
5.70
+6
23.75
38.35
42.43
37.96
20.30
23.41
60.00
90.30
30. 30.59
5.75
68.29
39.74
norm of the magnitude to measure the importance of different channels in key cache.
Mn.d ;.
Given pruning ratio X, We only keep T [(1 - A)DJ most important channels among the D channels of each head: I = Top (M, T)) where ||- lIo the lp norm of each channel. n E [V] and d E [1, DJ are indicators of heads and channels in key cache. (2+) NxT stores the indicators of the top I values in tensor M per head.
In Table I we present the results of key cache pruning with various pruning ratios applied to the LLaMA-3-8B model. We utilize the I and 12 norms as criteria for evaluation, and validate per- formance using the LongBench benchmark (Bai et al. 2023). Compared to the baseline methods, H2O (Zhang et al., (024c) and SnapKV V (Li et al. 2024), both with a KV length of 512, we further prune the channels of the key cache. A 30% pruning ratio maintain accuracy; however, increas- ing it to 40% results in significant performance degradation, especially for 4 norm based pruning. The results of magnitude- the channel dimension. These results also indicate the need for a better pruning metrics to achieve higher pruning ratios effectively.
pruning support our assumption that the key cache is redundant in
3.2 QUERY-DRIVED PRUNING
For each head, the attention scores are computed using the queries and keys, and then applied to the values. The formula for the attention for head i is: Attention(Q; K, V,) softmax Q.K." where Q, K, V. E RSSD When one channel of K; is pruned, the corresponding channel Q; will also be removed. We aim to find the optimal subset of channels to prune, denoted by the selection matrix S {0, 1] DXDDD, where S is a diagonal matrix with binary entries (1 for keeping channel, 0 for pruning it). To better maintain the performance after pruning the channels, we minimize the Frobenius norm of the difference between the original and pruned attention weights: mins ||,K?" S,S(K Given pruning ratio X, it can further expanded as:
VD Vi,
min 8
||Q, - Q,SK? ||
(2)
trace(S) = 1 - (1-A)D
S = ddiaag(1,15) ,2. where sj E {0,1)
To achieve the pruning goal, we define criterion for evaluating the importance of each channel and greedily select the channels with largest scores: Score,[y] = Q.[: K ||=
For simplicity, we use greedy algorithm to optimize S.
Top, 7). Here's detailed explanation of why optimizes the selection matrix. The [c] measures the magnitude of the interaction between the query and key vectors for channel j in each head i. By selecting channels with the highest interaction magnitudes, we aim to retain the most significant contributions to the attention mechanism. This criterion ensures that the selected channels preserve the primary information flow in the attention computation, thereby minimizing the loss of important information.
Observation Window, Following SnapKV (Li et al. 2024), to reduce the computation cost, only use the last Sohs window to calculate the score as the last window of input sequence recognizes highly similar attention pattern with generation: ||Q,[-Sa.. :. K [..]]] ||..
4

---

Published as a conference paper at ICLR 2025
Table 2: Performance comparison of key cache pruning on LLMMA-3-8B770B on Long- Bench. THINK (A) indicates we prune the key cache channels with a pruning ratio of X.
Single-Document
Multi-Dacument
Summarization
Fou-shoo Learning
Synthetic
Method
NrtvQA
Qusper
AlF-on
HorporA
SWIKMOA
Musigee
GovReport
QMSua
MultNNerrs
TREC
Triving
SAMSum
PCount
PRe
Loc
RE-P
Avg
LLaMAA In-tssrt KV-size Full
ALL
25.56
32.27
39.71
43.56
35.09
28.71
23.26
26.64
73.50
90.48
42.33
4,80
69.25
59.29
54.05
41,86
1aaMA---RR-Teomm KV-iine
1120
22.12
13.20
51.61
37.79
32.71
18.45
20.32
22.02
21.10
38.50
87.75
39.14
5.83
69.50
55 06.06
50.97
35.38
+THISK (0.4)
22:85
14.55
29.49
38.63
30.84
18,90
20.12
21.96
20.68
53.50
86.38
38.40
5.50
69.17
57.03
56.12
35.63
TTHSK (0.5)
23.47
14,06
25.67
38.35
30.21
17.87
19.69
21.94
19:95
38.50
87.14
38.07
4,92
69.50
57.
50.60
35.44
StapKV
21:19
13.55
32.64
18.75
25.64
18.73
18:58
21.62
20.26
45.00
88.36
37.64
5.15
68.85
58.
5132
35.90
+THINK (0.4)
22.11
14.67
32.49
36.25
28.63
18.80
18.93
21.49
20.14
44.50
38.32
5.75
69.17
58.21
55.89
35.84
-THINK (0:5)
21.79
32.03
37.52
43.50
LLaMA-3-HBB--N Iist-nccc KV-siae
512
H20
23.52
17.93
64
42.11
53.52
19.92
22.
22.56
23.82
41.00
90:46
40.20
5.87
69
69.50
56
51
37.23
+THISK (0:4)
23.76
17.80
33.80
40.39
30.70
19.09
21.82
22.51
23.78
41.00
90.16
40.67
5.15
69:25
60.77
57.58
$7.39
+THINK (0.5)
24.17
16.96
35.
19.67
30.29
18.67
21.39
22.59
23.06
41.00
89.81
40.35
5.23
69.35
60
58.34
37.29
.THIKK (0.6)
23.40)
14,83
38.47
20,
40.00
79
69.50
58.
36.44
SYapKV
24.84
36.77
40.75
30.87
22.
23.97
70.00
+THISK (0.4)
24.58
25.44
57.03
41.87
21:77
16
70.00
90.33
40.29
6.06
69:50
49.55
+THINK (0.5)
24.85
25. 10
57.06
41.58
30.34
20.80
21.61
22.44
23.66
59.50
90.39
39.70
5.84
69
61.57
59.42
60.34
+THINK (0.6) 25.98
22.77
38.37
40.44
33.19
19,90
20.84
22.21
22.55 22.55
59,00
90.32
39.12
0.39
69.50
59
19.14
58.40
39.20
LLaMA
KV-4ie
1024
HD0
25.62
22.16
36.81
41.01
33.53
25.41
46.50
90.82
41:78
69.25
99
55.50
38.70
+THINK (0:4)
25.52
21.93
57.17
41.56
22.89
47.00
41.34
+THISK (0.5)
25:41
22.19
37.64
40.92
33:27
18.66
22.17
22.22
24.84
46.50
90.34
40.59
5.20
69.50
63
38.57
+THINK (0.6)
24.06
17,80
37.85
38,63
2988
19.40
21.41
22.32
23.42
44.50
90.16
29.43
5.84
69.50
58
58.73
37.58
SnajoV 24
25..9
37 62
45.84
--109
7000000
28.15
2:5.0
23.63
723
96.56
4454T
5.36
16925
60.57
16
40.88
+THINK (0:4) 24.88
27.72
38.60 60
43.16
32.44
20.87
24.21
22.79
25. 56
71.50
90:45
40.94
5.93
89.
00.5000
62 77
59.15
41.29
+THINK (0.5)
24
27.26
30
42.83
32.09
19.56
23.55
22:46
25.
71.50
90:43
40.74
5.20
69:
62.46
59.75
41.07
+THINK (0.6) 24.46
27.55
41.95
21.89
23.08
70,00
90.19
40.30
1..MA-3-88---em KV-sive
2048
1120
25.56
26.85
39.54
44,30
32,92
21.09
24.68
23.01
26. 16
53.00
90.65
41.84
4,91
69
58.43
51.31
39.59
+THINK (0:4)
25.56
26.31
39.33 20
42.96
31.81
20.53
24.23
23.35
25. 90
53.50
90.56
41.03
5.52
69 25
62
59.00
40.05
+THINK (0.5)
25.01
25.37
38.82
42.32
51.27
20.50
23.77
23.21
26.03
$3,00
90.37
40.86
5.13
69.50
61.91
58.95
39.75
.THINK (0.6)
24.37
22.14
37.77
40.13
70,50
30.26
12.09
22:76
24 78
49.50
90
39.69
5.56
69.50 59.04
68.78
38.51
ScapKV 25.6
"25.55
4T.
44435
21.31
23.68
22544
56.46
73.30
900.50
"45.6
16525
58.67
31.32
4T.58
(0:4)
25.41
+THISK (0:5)
25:00
30.25
39:27
43.23
21.24
25.16
23.01
20.5
73.00
90.37
41.26
69.50
41.77
+THINK (0.6) 24.89
28.88
60.44
41.30
29.99
21.34
23.48
22.9
24.99
72.50
90.35
28.5
5.71
69.50
59.77
59.50
40.88
LLaMA-3-7UBBTTddie Trosieee
KV-size 128
SnapKV 25.91
39.41
43.83
49.80
51.23
27,76
22.14
21.91
23.16
58.00 39.00
91.55
43.54
12.50
72.00
48.41
63:49
41.09
.THINK (0.4)
25.64
39.20
43.601
50.22
21.70
68.00
43.24
73.00
48.01
44.00
+THINK (0:5)
26:31
38.76
44.86
48.54
28.97
21.46
22.01
22.91
67.00
91.52
43.15
12.50
72.50
47.21
60.82
43.63
11.MA-3-700--rese KV-siae $12
StapKV 27.95
45.19
48.58 50
50.97
54.53
29.78
25 .4
22.76
26. 03.05
73,50
92.65
45.07
12.50
72,50
45.21
68.22
46.27
+THINK (0:4)
27:47
45.31
48.5 57
51.22
54.32
30,05
25
22:72
26
73.50
92.13
45.53
12.50
73,
48.32
66.99
46.45
+THINK (0.5)
26.97
44.55
50.84
25.29
73,00
42.15
73.00
46.12
11MAA.1-000 Incerett KV-sine 1024
StupKV 26.80
46.21
29.93
$1.70
54.71
29.86
27.61
22.45
27.15
73,50
92.38
46.18
12.50
72.50
42.84
69.89
46.64
+THINK (0:4)
27.04
46.01
50.13
51.96
54.36
29.87
27.74
22.76 76
27.07
73.50
91. 91.88
46.35
12.50
73.00
45
05
67.87
46.69
+THINK (0.5)
27.62
46.22
48.97
51.79
$3.39
30.47
27.45
23.05
26.57
73.50
91.
43.99
12.50
72.50
47.41
66.84
46.51
1.LMAA3-700:
KV-sine
2046
StapKV
27
46.51
60
54.77
31.05
20 67
27.43
73.50
92. 03.38
45.98
12.50
72.50
46.76
+THINK (0:4) 27.13
46.25
50.04
51.72
55.03
31.19
29.75
22:47
27.28
73.50
91 18
46.37
12.50
72.50
42.66
67.77
46.75
.THINK (0.5) 27.44
46.86
29.18
51.97
53.58
31.44
29.41
22,99
27.33
73.50
91.88
43.60
12.50
72.50
44.78
66.65
46.62
3.3 IMPLEMENTATION OF THINK
We opt not prune the most recent tokens and newly generated keys. Consequently, our key-value (KV) cache has two categories: one subset consists of pruned keys with a reduced channel size, while the
Pruned Query
Pruned Key
Key
Quuery
other at their original size. Additionally, we maintain a binary mask whose memory overhead is negligible to indicate which channels have been pruned. Fig- ure [ illustrates one implementation of THINK dur- ing the decoding stage. We first prune the query using the stored mask, ensuring that the query dimensional- ity and the pruned key cache remain consistent, The pruned query is then multiplied by the pruned key, while the unpruned query is applied to the unpruned key. Subsequently, the two outputs are concate- nated. Our method removes unimportant Key cache channels to get real memory reduction.
Hit [06
Kky Cuadeo
Eosilig
Figure 2: Implementation during decoding.
5

---

Published as a conference paper at ICLR 2025
Table 3: Performance comparison of key cache pruning on disttal-7B-Instruut-obl on LongBench. THINK (A) indicates we prune the key cache channels with a pruning ratio of X.
Singls-Document
Multi-Ducument
Summarization
Feu Learning
Synthetic
Method
NrtvQA
Qasper
AlF-cn
HorpoQA
SWIGMOA
Musigee
GovReport
QMSua
MultNeerr
TREC
Triving
SAMSum
PCount
PRe
Lot
RE-1
Avg
Histral-8-Insttua---00 KV-size Full
ALL
26.63
32.99
49.34
42.77
27.35
32.87
24.24
27.10
71,00
N6.25
42.96
86.98
56.95
54.49
Misenl-7
(m- KV-tine
1120
21:21
21.81
38.87
30.42
20.36
12.30
20.58
22.61
22:10
39,00
K2.37
40.44
2.90
79.56
51
48.38
34.63
+THIKK (0.4)
21:17
21.90
39.29
29.92
20.99
12.50
20184
22.91
21.92
30.00
K2.70
40.35
2.97
79.21
51
48.32
34.69
TTHIKK (0.5)
21.67
21.80
39:48
28.74
20.65
13.14
20.57
22.83
21.78
39,00
K2.54
40.12
3.61
78.39
50.27
45.4
54.56
-THnsK (0.6)
21.04
21.30
99 56
18.68
21.20
13:97
20.13
22.52
21.81
39. 50
82.05
50.14
4.16
74.25
49
47.67
54.18
SnapKV 19.17
22..40
<2.49
36.75
"2Z.45
T5.36
19.08
2,185
21.55
87.50
B1.T5
"40224
2:3
16526
52.51
36.30
35.29
(0:4)
20.52
+THISK (0.5)
20.67
20.60
43.37
37.27
21.58
15.66
19,00
21.74
21:02
47,00
39.77
50.80
+THnNK (0.6) 21:25
20.82
36.21
21.68
16.47
19.05
21.99
20.73
45.00
83.81
38.79
4.19
66
$4.92
Mistral 7B Instruct v0.2. KV
sizze 512
1120
21.83
26.00
44. 69
32.66
23.05
14.89
23.53
23.06
24.59
42,00
85.
41.49
3.60
86.20
54.76
51.09
37.36
.THINK (0.4)
21.58
25.15
32.73
23.56
24.45
42,00
55.7
37.49
-THINK (0.5)
22.76
25.74
31.74
23.31
34
.THINK (0.6)
23:91
25.57
20.48
22,88
23:31
85
41.15
2.98
36.58
SnapKV 25:43
27.8T
68.95 9R
19:06
25:25
T6.95
23.70
22.96
24.37
57.00
85. 88
-4125
2.75
B6.35
46.46
33
40.46
+THINK (0.4) 24:27
28.46
26
38.13
24.22
16.92
23.59
23.70
24.46
67.50
85.9
42.51
2.92
K5.32
55
53.35
40.40
+THISK (0:5)
24.56
29.22
48.59
37.70
24.27
17.39
23.68
23.65
24.58
67.50
96.05
42.01
3,07
86.30
56.49
53.29
49.52
+THINK (0.6)
24.07
26.27
49.10 10
38.65
24.31
17.52
23.16
23.51
24.25
67.00
86.35
40.78
3.69
83.74
54.35 94
52.23
40.10
H2O
23.67
36.99
24.2
25:21
23.04
25.77
46,00
41.98
35.90
+THINK (0.4)
23:97
28.91
45.84
25.78
24.88
14,55
25.11
23.35
25:83
43.50
86.11
42.44
3.23
84.82
56.21
53.02
39.72
+THIKK (0.5)
23.89
25.40
46 60
35.57
24.26
14.78
24.08
23.11
25.68
44.50
86 16
42.72
3.38
83.20
88 88
52.63
38.50
-THINK (0.6)
23.87
27.76
66.32 25
15.28
24.38
14.24
24.35
23.35
25. 55.50
44.50
65 85.38
41.37
3.34
61.42
55
51.89
38.00
SnapKV 25.47
52
23 f3
2555
T9771
23,59
2,5345
58.21
89
66.16
135
1656
57
3570
4T.67
+THINK (0.4) 25
30.48
48.38
41.11
25.28
25.91
24.00
26.15
70.00
66.64
43.35
56.3
56.71
54.19
41.62
-THnkK (0.5)
25.63
40.99
25.47
60.5
86.67
-THINK (0.6) 24.69
29.3
40:64
25.23
23.6
68.00
85.85
40.96
Mistril 76
Imorrcttx0.2, KV-size 2048
H2U
25:76
31.10
29.06
40.35
26.43
16.75
27.37
23.64
26.69
55,00
86.35
42.48
2.72
86.64
50,
53.91
49.69
+THnNK (0:4)
25:40
50.80
48.45 45
19.64
26.08
16.82
27.12
23.79
26.65
$3.50
86.39
43.03
3.29
86.30
56.61
53.60
40.47
.THINK (0.5)
25.68
39.65
26,60
52,00
74.74
40.29
(0,6)
25.83
26.51
42.05
39.76
SnppEV
"5I.5
23.50
55.63
70.T0
+THINK (0.4)
25:77
32.67
48.77 70
41.06
27.07
19.14
28.91
24.37
26.88
70,00
85.37
42.75
3.61
87.38
57.21
54.03
42.27
.TRIKK (0.5)
26.44
32.94
29.02
40.86
26.84
19.49
28.46
24.51
26.72
70.00
86.50
41.75
2.78
84.70
56.47
54.15
41.98
-THINK (0.6)
26
52.53
48.73
40.95
26.77
18.92
27.40
23.97
26. 36.37
70.00
86.45
41.12
3.31
62.24
56.01
53.53
41.52
4 EXPERIMENT RESULTS
In this section, conduct comprehensive experiments to evaluate the effectiveness of THINK on performance and memory reduction. the experiments, we prune Key cache channels by default, the value cache results are in Table 8 in Appendix D
4.1 SETTINGS
Benchmark Datasets. We evaluate our proposed method against state-of-the-art KV cache com- pression methods on two widely recognized benchmarks: LongBench and feedle-in-a-Haystte LongBench (Bai et al., 2023) is designed to comprehensively assess the long context understanding capabilities of LLMs. It includes 17 datasets covering six different tasks: single-document QA, multi-document QA, summarization, few-shot learning, synthetic tasks, and code completion. The average input length of LongBench is 6,711 words, which necessitates reducing the KV cache to lower memory usage for inference. Bredle-in-a-Haysse benchmark that tests a model's 's ability to accurately locate small but crucial piece of information (the "needle") embedded within a lengthy document (the "haystack"). The random positioning of the needle in this challenge serves as critical test to determine whether KV cache compression methods can retain essential information without loss of accuracy.
(Kamradi) 2023) is recently developed
Baseline Approaches. The baseline methods in our evaluations include Heavy Hitter Oracle (H20), SnapKV and KIVI, all of which are the tate-off-hhe-rt KV cache compression methods but use different strategies. H20 (Zhang et al., 2024c) is designed to reduce memory usage by dynami- cally managing the balance between recent tokens and Heavy Hitter (H2) tokens. H2 tokens rep- resent a small set of tokens that contribute most of the value when computing attention scores. SnapKV (Li et al. 2024) introduces an automated compression mechanism that selects clustered,

---

Published as a conference paper at ICLR 2025
Table 4: Performance evaluation of combining THINK with KIVI et al. (2024b) on LongBench. THINK (0.4) indicates we prune the key cache channels with pruning ratio of X = 0.4.
Single Document
Multi Document
Summarization
Few shot Learning
Synthelic
Code
Mithod
NevQA
Qusper
MF-4
HotpetQ
DWBIMQA
Masque
GovReport
QMSum
MuttNens
TREC
TritagA
SAMSum
PCausi
PRE
Lec
RB-1
Arg.
KIVE
19.67
18,62
30.28
29.02
25.00
10.30
21.34
20.51
25, 10
63.00
85.04
40.16
4,00
8.00
58.04
3192 56.51
52.48
-THINK (0.4) 2/2
19.45
19.01
30.52
28.79
25.78
9.53
22:11
20.66
25.75
63.00
$4.62
41.54
3.50
7.00
48.92
31.77
important KV positions for each attention head, optimizing the KV cache without sacrificing per- formance. KIVI (Liu et al.) (022b) reduces memory overhead by quantizing the KV cache into lower-precision formats, significantly lowering the memory cost while preserving model accuracy.
Implementation Details. In this paper, we use LLaMA-2-7B-chh LLaMA-3-70B ------7----nstr (Meta) 2024) and bone LLMs, both accessible via HuggingFace (Wolf et al. 2020). Our THINK aims to prune chan- nels of the key cache, which is agnostic to KV cache compression methods. there is no other statement, prune the key cache by default. All the experiments A A100 GPUs. To ensure a fair comparison between KV cache compression methods and their inte- gration with THINK, we applied consistent hyperparameters across both settings. For instance, when comparing SnapKV and SnapKV integrated with THINK, we used a maximum pooling kernel size of 7 and an observation window size of 32, maintaining the same KV-size for both configurations. We compress the Key cache starting from the prefilling stage.
SaaMA-3--B-Insireee (Jiang et al. 2023) as the back-
conducted using NVIDIA
4.2 RESULTS ON LONGBENCH
Tables 2 and present the results of KV compression methods and their integration with our pro- posed channel pruning technique for the key cache (THINK) across three different base LLMs, evaluated at various KV-sizes on the LongBench benchmark. The pruning ratio of X = 0.4 indicate that 40% of key cache channels are removed ,rsulting in a 20% reduction in the total KV cache memory footprint. The following observations can be drawn: (1) Our method successfully prunes the channels of the key cache after the KV cache has been compressed using H2O and SnapKV. V. base model, our approach reduces memory usage while slightly im-
For the LLaMA-3-8B-Inst proving performance for both H2O and SnapKV. For the Dossaal-7B--nstuuttt-.. base model, our method similarly reduces memory usage, with only minor performance drop in some cases. For the larger LLaMA-3-P 3-BB base model, our method achieves comparable or superior performance after pruning 40% of the key cache channels, compared to the SnapKV baselines. (2) When Comparing SnapKV or H2O integrated with THINK in Ta SnapKV or H2O integrated with 4 or le norm [ our query-drive -diven channel pruning approach demonstrates superior performance when the pruning ratio of A = 0.4. (3) Lower pruning ratios generally result in better performance compared higher pruning ratios. (4) As the KV-size increases from 128 2048, the performance of our channel pruning method improves. Notably, with a KV-size of 2048 and pruning ratio of 0.4, our method even surpasses the performance LLaMA-3- 3--st------nnttrettt with a full KV cache. These findings suggest that our method is agnostic to the underlying KV cache compression techniques and can further enhance both performance and memory efficiency. Moreover, query-driven channel pruning proves more effective than 4 and l2 norm-based methods for channel pruming in LLMs. Ex- periment results of applying THINK directly to vanilla models are in Table which further validate the effectiveness of our method.
We further validate the efficacy of our method by applying to the KV cache quantization technique KIVI (Liu et al. (224bb) shown in Table 4 First, prune 40% of the key cache channels, followed by quantization of the remaining channels into 2-bit (implementation details provided in Appendix C.Ip. Compared the standard KIVI method, approach reduces KV cache memory by 20%, with minimal performance degradation.
4.3 RESULTS ON NEEDLE-IN-A-
Table presents the results of the Baelle-in-a-Haystad test, using the SnapKV (Li et al. 2024) approach with varying K KV-sizes, ranging from 128 2048. With a modest pruning ratio of

---

Published as a conference paper at ICLR 2025
Table 5: N Test Results
KV-size
Model
Method
A
128
512
1024
2048
SnapKV
0.0
79.6
90.2
91.2
91.7
LLaMA3-8B Instruct
Snapk V
0.4
79.6
90.3
91.2
91.7
SnapKV+THin
91.7
SnapkV
0.0
77.8
89.5
90.4
90.8
Coattal--B-Isttuutttt
Snapk K
0.4
78.6
90.1
90.6
90.9
SnTIKV+THINK
78.1
91.1
SnapKV+THINK
0.6
75.9
89.2
90.6
91.1
Table 6; Performance comparison of key cache pruning with varying recent-sizes.
Single-Decument Docament QA
Mutti-Decumont QA
Summarization
Few-shot Learning
Syuthetic
Code
Recent Sizee
NrmaQa
Qusperr
MF MF-c0
HorpottA
zWiMMQA
Mosique
GevRegort
QMSum
MlttiNeuss
TREC
TrivinQA
SAMSim
PCouni
PRe
Lec
RB-P
Arg
120 + THINK (A 0..4
24.40
27.50
45.42
35.17
24.45
13,02
27.65
23.53
26,86
$3,50
86.06
41.73
3,01
83.42
55.12
51.32 38.91
30:00
48.45
23.79
3.29
96.39
56.1
40.47
30.98
23.92
SnipKV+ + TursK (A = 0.4)
0
24.94
28.58
45.78
39.59
25.40
15:92
29,50
24.05
26.72
70.00
85.60
41.38
2,97
84.00
55.27
$2.39
40.76
32
25.77
32.67
48.70
41,06
27.07
19.14
28.91
24.37
26.88
70.00
86.37
42.75
3.61
67.38
57.21
54,64
42.27
A = 0.4, THINK consistently outperforms or matches the accuracy of the original SnapKV across both LLaMA-3 and Mistral models, regardless of KV-size. These comparisons demonstrate that the proposed query-driven channel pruning method effectively retains informative channels while discarding noisy ones. However, when the pruning ratio increases to 2 0.5, we observe a drop in accuracy with smaller KV-sizes, particularly for 128 and 512. across both LLaMA-3 and Mistral models. Despite this, THINK achieves comparable performance with SnapKV when the KV-size is larger(i.e. 1024 and 2048). Intuitively, larger pruning ratio with smaller KV-size may lead to the loss of more critical information compared to scenarios with a larger KV-size. In addition, the performance on larger KV-sizes suggests that THINK is robust for long-context tasks.
Figure 6 (a)-(d) (in Appendix B) visualize the Needle-in-a-Hage Haysstack test accuracy across varying token lengths and depths. The KV-sizes and A = 0.5, respectively. THINK preserves the retrieval capabilities of SnapKV, although there are minor numerical differences in overall accuracy (e.g., 77.8 vs. 78.6 and 90.4 vs. 90.6). THINK matches SnapKV in accuracy for the majority of token limits and depths, demonstrating consistency in performance. Furthermore, THINK successfully retrieves certain "needles" that SnapKV fails to capture, resulting in improved overall accuracy. These visualizations highlight the robustness of THINK from a fine-grained perspective, illustrating its capacity to enhance the original approach.
128 and 1024, with pruning ratios of X = 0.4
4.4 ABLATION STUDIES
Impact of Different Recent Sizes. Preserving the most recent KV embeddings (Zhang al. 2024c Li al. 2024) is important for maintaining the performance of LLMs after KV cache compression. However, tradeoff exists: increasing the recent-size allows more information be retained, but also increases the cache size. To assess its impact, we evaluate the performance of three recent-size configurations, namely 0, 32 and 128, on LongBench, using Mistral-7B. Bestral-7B-Instruee In-trructt--0.2 as the baseline model. The results are summarized in Table 6 As observed, a recent-size of 32 yields superior performance compared
0, as indicated by the averaged score on LongBench, demonstrating the importance of retaining the most recent KVs. On the other hand, the performance difference between recent-sizes of 32 and 128 is negligible, suggesting that retaining the most recent 32 KV's is sufficient to maintain optimal performance.
Performance Comparison Under the Same Memory Usage. To ensure a fair comparison, we adjust the KV-size of H2O SnapKV to match the memory usage of H2O with THINK SnapKV with THINK on Mistral- 1-1177--nntttttO... For example, the KV-size of H2O with THINK set 128. Due to channel pruning applied to the key cache, the memory consumption of H2O with
8

---

Published as a conference paper at ICLR 2025
Table Performance comparison of key cache pruning with the same memory consumption.
Sigle--ecmmen Q4
Etrrmment Q4
Summartcttion
Peweshut Learning
Synthetie
Cinde
Methods MemorycM)
NrtsQA
Queper
HhipaagA
SWRIMQQA
Musipic
GovReport
QM
MultNen
TREC
Trilsg4
SAMSuu
PCann
PRo
Iec
RB-P
Arg.
H20
Vanilla
54.5
21.29
20.54
37.56
23.65
21005
1401
20.20
22.11
21.35
38,50
K2.55
39.87
3.66
78.14
50.32
48.54
54,29
ThInK
54.4
21.17
21.90
19.29
29.92
20.99
12.30
20.64
22.91
11.92
39.00
K2,70
40.35
2.97
79.21
51.19
48.32
$4.69
Yauilla
208.0
22.13
23.53
43.24
30.92
23.36
14.56
22.92
22.77
24.23
41.50
K5.04
41.25
3.02
56.03
54.91
50.50
36,89
ThanK
208.0
21.58
26.15
44.49
32.73
23.99
15.09
23.56
23.28
24.45
42.00
53.58
42.58
3.18
85.7
$4.39
51.15
37.49
Yanill
413.0
22.90
28.45
46.16
$5.57
23.86
13.34
14.90
23.1 19
15.37
44.3
K5.54
41.97
1.23
85.82
55 96
$2.53
58.37
Tunn
412.8
23.97
28.91
45.54
35.78
34.88
14.55
25.11
23.35
25.83
45.50
86.11
42.44
3.23
54.52
56.21
53,02
38.72
Yauill
822.5
25.51
50.23
48.23
39.72
25.36
16.75
36.08
23.83
26.47
50.50
86.43
42,09
2.78
55.57
57.4
53.42
40.00
THoNK
8224
25.40
30.80
45.45
19.64
16.32
26.65
53.50
86.
43.03
86.
56.61
53,00
00.47
51.39
18.74
96.14
41.23
41/62
24.84
16
41.97
19.14
70.00
37
70
LongBBecch Performance
36
27
Thask
Vasiss
KIVI 4/4
StapKV ThicK (0.4)
4/4, Thick
SnapkV
20
KIVI 2/2, Dank (0
100
200
201
400
200
(00)
700
800
Cache Memory (M)
Batch
(a)
(b)
Figure 3: (a) presents the performance comparison with token eviction methods under identical
memory usage for Bestral-7B-Inssruct-ee with the KV cache quantization method KIVI across different batch sizes for LLaMA-2-7B-o8
while (b) illustrates the memory usage comparison
THINK (0.4) indicates we prune the key cache channels with pruning ratio of X 0.4.
THINK at a KV-size 128 is lower than that of H2O at the same KV-size. Consequently, the KV-size of H20 is adjusted from 128 to 109 to equalize memory usage. Table 7 and Figure 3a present the results of these comparisons on the LongBench benchmark. The results demonstrate that H2O or SnapKV combined with THINK consistently outperforms their counterparts without THINK while maintaining the same memory footprint. This highlights the effectiveness of inte- grating query-driven channel pruning with KV cache compression methods, enabling more efficient memory utilization and improved compression of the KV cache.
Memory Usage Comparison. To evaluate the efficiency of THINK. we follow the methodology used in KIVI (Liu et al. 2024b). We generate synthetic workloads with an input prompt length of 160 and an output length of 338. The peak memory usage is reported for the vanilla FP16 baseline, KIVI, and KIVI combined with THINK (0.4) for LL-MA-2-7B-hhe
As in Figure 3b] the memory savings from our method become increasingly evident the batch size grows, both the KIVI 2/2 and KIVI 4/4 configurations. Compared to the baseline model, our approach achieves over (from 4x with KIVI alone) increase batch size while maintaining the same memory footprint when integrated with KIVI. Model weights and KV cache the primary memory components accessed during generation. By effectively reducing the memory footprint the KV cache, our method alleviates the memory bottleneck, enabling faster generation speeds as shown in Table TO
Pruning Channels of Both Key and Value Cache In this part, we explore the impact of prun- ing channels in the value cache (Appendix D). Specifically, for KV cache compression methods, we apply different pruning ratios to the channels of the key and value eaches. Table g presents the results with LLLMA-3-8FE L-MM--3-8B-Instruu and - on the Long benchmark.

---

Published as a conference paper at ICLR 2025
When evaluating the base model LLaMA-3-88 --I-nttruct,, H2O or SnapKV with both key and value channel pruning perform comparably to H2O or SnapKV without pruning. In certain instances, the models with key and value channel pruning even outperform their non-pruned counterparts. For the base model drop. This aligns with the observations in Figure + where the Key cache shows highly unbalanced magnitudes along the channel dimension, while the Value cache exhibits more uniform magnitudes. This suggests that the Value cache has less channel sparsity, making it harder identify redun- dant channels for pruning. Nevertheless, pruning Value cache channels still contributes to further memory reduction in the KV cache.
pruning the value cache channels leads to a slight performance
RELATED WORK
In scenarios involving long contexts, the key-value (KV) cache poses the most significant computa- tional and memory burden within the attention mechanism of large language models. Reducing the KV cache is therefore key priority for optimizing deployment efficiency. To address this challenge, system-level optimizations, such as FlashAttention (Dao) 2023) and PagedAttention (Kwon et al.) ).23), have been developed to tackle this challenge. In parallel, algorithm-level optimizations are also being explored to further enhance efficiency.
KV Cache Eviction. StreamingLLMM (Xiao et al. 2023b) retains few initial tokens along with recent tokens based on the observation of attention sinks, which can leading to the loss of critical in- formation earried by the dropped tokens. H2O (Zhang et al., (222c) selectively retains small subset of tokens by greedily dropping those with lower contributions to cumulative attention. SnapKV (Li al. 2024) selects clustered important KV positions for each attention head from an "observation" window located at the end of the prompts. FastGen (Ge et al. 2023) adaptively evicts tokens from attention heads that focus on local contexts, discarding non-special tokens that surround key tokens, while standard KV cache is applied to attention heads that attend more broadly. PyramidKV (Zhang et al. 2024b) and PyramidInfer (Yang et al. (0244) take a hierarchical approach, adjusting KV cache sizes across different layers by allocating more cache to lower layers and less to higher ones.
KV Cache Quantization. SmoothQuant (Xiao et al. 2023a) enables the quantization of the KV cache to 8-bit with minimal performance degradation. Q-Hitter (Zhang al. 2024c) leverages accumulated attention scores and "Quantization Friendliness" metrics to identify tokens that are crucial for preserving the generalization capabilities of LLMs, making them suitable for KV cache quantization. Furthermore, recent studies suggest that the key and value caches should be quantized differently (Liu al.] 2024b; Hooper et al., 2024): the key cache should be quantized per-channel, while the value cache should be quantized per-token.
Structured Pruning of LLMs. Traditional structured pruning (Ma al. 2023 Ding et 2023 Lu et al., 2024a) of LLMs typically focuses on removing unimportant layers. heads, or hidden di- mensions, often leading to significant performance degradation. In contrast, our approach preserves the original architecture of the LLM and specifically targets the channel dimension within each head's s key cache. By dynamically identifying unimportant channels using data dependant criterion, our method greatly reduce the key cache size with minimal performance loss.
CONCLUSION AND LIMITATIONS
Motivated by the observation that certain channels exhibit significantly larger magnitudes and the low-rank nature of the key cache (as shown by singular value analysis), we propose THINK as a query-dependent pruning method targeting key cache channels. Optimized based on attention scores, our approach preserves essential information for each input query. THINK seamlessly in- tegrates with existing token-level KV cache pruning (Li et al. 2024 Zhang et al., 2024c) and KV cache quantization (Liu et al. 2024b), further improving inference efficiency. Extensive experiments on LongBench and Needle-in-a-Haystter benchmarks demonstrate its effectiveness, achieving com- parable or superior performance to baselines while reducing key cache size by 40%. However, our method slightly increases TTFT (Time-to-First Tokee) due to channel importance computation (see Appendix FF). Additionally, while effective for key cache pruning, value cache pruning remains an open direction for future work.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Joshua Ainslie, James Lee-Thorp Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head check- points. arXiv preprint arXiv: 13245, 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. arXiv preprint arXiv:2308. 14508, 2023.
Srinadh Bhojanapalli, Ayan Chakrabarti, Himanshu Jain, Sanjiv Kumar, Michal Lukasik, and An- dreas Veit. Eigen analysis of self-attention and its reconstruction from partial computation. arXiv preprint arXiv:2106.08885 2021.
William Brandon, Mayank Mishra, Aniruddha Nrusimha, Rameswar Panda, and Jonathan Ragan Kelly. Reducing transformer key-value cache size with cross-layer attention. arXiv preprint arXiv:2405. 12981, 2024.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Chi-Chih Chang, Wei-Cheng Lin, Chien-Yo Yu Lin, Chong-Yan Yan Chen, Yu-Fang Hu, Pei-Shuo Wang, Ning-Chi Huang, Luis Ceze, Mohamed Abdelfattah, and Kai- Kai-Chiang Wu. Palu: Compressing kv-cache with low-rank projection. arXiv preprint arXiv:2407.2111 2024.
Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, and Christopher Ré. Scatterbrain: Uni- fying sparse and low-rank attention. Advances in Neural Information Processing Systems, 34: 17413 17426, 2021a.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde De Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, al. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374 2021b.
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint (=XXi:23007.86991, 2023.
James W Demmel. Applied numerical linear algebra. SIAM, 1997.
Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. Qlora: Efficient finetuning of quantized lims. Advances in Neural Information Processing Systems, 36, 2024.
Ning Ding, Yulin Chen, Bokai Xu, Yujia Qin, Zhi Zheng, Shengding Hu, Zhiyuan Liu, Maosong Sun, and Bowen Zhou. Enhancing chat language models by scaling high-quality instructional conversations. arXiv preprint arXiv: 14233, 2023.
Hanze Dong, Wei Xiong, Bo Pang. Haoxiang Wang, Han Zhao, Yingbo Zhou, Nan Jiang, Doyen Sahoo, Caiming Xiong, and Tong Zhang. Rlhf workflow: From reward modeling to online rlhf. arXiv preprint arXiv:2405. 007863, 2024a.
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 11m inference. arXiv preprint arXiv:2402.099 09398, 2024b.
Elias Frantar and Dan Alistarh. Sparsegpt: Massive language models can be accurately pruned in one-shot. In International Conference on Machine Learning, pp. 10323-10337. PMLR, 2023.
Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. Gptq: Accurate post-training quantization for generative pre-trained transformers. arXiv preprint arXiv:2210.17322 2022.
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. arXiv preprint arXiv:2310.0180, 2023.
11

---

Published as a conference paper at ICLR 2025
Muhammad Usman Hadi, Rizwan Qureshi, Abbas Shah, Muhammad Irfan, Anas Zafar, Muham- mad Bilal Shaikh, Naveed Akhtar, Jia Wu, Seyedali Mirjalili, et al. Large language models: A comprchensive survey of its applications, challenges, limitations, and future prospects. TechRxiv, 2023.
Dan Hendrycks. Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. arXiv preprint arXiv:2
2021.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Train- ing compute-optimal large language models. arXiv preprint arXiv:2203. 15556, 2022.
Coleman Hooper, Sehoon Kim, Hiva Mohammadzadeh, Michael Mahoney, Yakun Sophia Shao, Kurt Keutzer, and Amir Gholami. Kvquant: Towards 10 million context length Ilm inference with kv cache quantization, arXiv preprint arXiv:2401. 18079, 2024.
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Alue Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.09685, 2021.
Zhe Jia, Marco Maggioni, Benjamin Staiger, and Daniele P Scarpazza. Dissecting the nvidia volta gpu architecture via microbenchmarking. arXiv preprint arXiv: 06826, 2018
Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, al. Mistral 7b. arXiv preprint arXiv:2310.06825 2023.
Gregory Kamradt. Needle In A Haystack pressure testing LLMs. Github, 2023. URL https: NetMMeesteeeeeeeenlllaiaaaaaaa Wneetdetiitttssttcckttreeeee main
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08366 2020.
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Symposium on Operating Systems Prin- ciples, pp. 61 311-626, 2023.
Yaniv Leviathan, Matan Kalman. and Yossi Matias. Fast inference from transformers via speculative decoding. In International Conference on Machine Learning, pp. 19274-19286. PMLR, 2023.
Yuhong Li, Yinghing Huang, Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. Snapkv: Lim knows what you are looking for before generation. arXiv preprint arXiv:2404.14440 14469, 2024.
Hunter Lightman, Vineet Kosaraju, Yura Burda, Harri Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let's verify step by step. arXiv preprint arXiv:2305.20050. 2023.
Ji Lin, Jiaming Tang, Haotian Tang, Shang Yang, Wei-Ming Chen, Wei-Chen Wang, Guangxuan Xiao, Xingyu Dang, Chuang Gan, and Song Han. Awq: Activation-aware weight quantization for on-device ilm compression and acceleration. Proceedings of Machine Learning and Systems, 6: 87-100, 2024.
Akide Liu, Jing Liu, Zizheng Pan, Yefei He, Gholamreza Haffari, and Bohan Zhuang. Mini- cache: Kv cache compression in depth dimension for large language models. arXiv preprint arXiv:2405. 14366, 2024a.
Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen Zhong, Zhaozhuo Xu, Vladimir Braverman, Beidi Chen, and Xia Hu, Kivi: A tuning-free asymmetric 2bit quantization for kv cache. arXiv preprint arXiv:2402.02750 2024b.
12

---

Published as a conference paper at ICLR 2025
Xudong Lu, Qi Liu, Yuhui Xu, Aojun Zhou, Siyuan Huang, Bo Zhang, Junchi Yan, and Hongsheng Li. Not all experts are equal: Efficient expert pruning and skipping for mixture-of -eeeertstteeessseeeeeeeeeee large language models. arXiv preprint 24X.:22402.14800, 2024a.
Xudong Lu, Aojun Zhou, Yuhui Xu, Renrui Zhang, Peng Gao, and Hongsheng Li. Spp: Sparsity parameter-effiieet fine-tuning for large language models. arXiv preprini arXiv:2405. 16057, 2024b.
Xinyin Ma, Gongfan Fang, and Xinchao Wang. LIm-pruner: On the structural pruning of large language models. Advances in neural information processing systems, 36:21702-21720, 2023.
AI Meta. Introducing meta llama 3: The most capable openly available I1m date. Meta AI, 2024.
OpenAl. OpenAl: Introducing ChatGPT, 2022.
URL https::/opeeai.coo www.pppaai.....oooooooo/oooooao.
ght
OpenAl. Gpt-4 technical report. arXiv preprint arXiv:2303.08777, 2023.
Machel Reid, Nikolay Savinov, Denis Teplyashin, Dmitry Lepikhin, Timothy Lillicrap, Jean- baptiste Alayrac, Radu Soricut, Angeliki Lazaridou, Orhan Firat, Julian Schrittwieser, et al. Gem- ini 1.5: Unlocking multimodal understanding across millions of tokens of context. arXiv preprint arXiv:2403.05530 2024.
Utkarsh Saxena, Gobinda Saha, Sakshi Choudhary, and Kaushik Roy. Figen attention: Attention in low-rank space for kv cache compression. arXiv preprint arXiv:2408.050: 2024.
Teven Le Scao, Angela Fan, Christopher Akiki, Ellie Pavlick, Suzana Ilié, Daniel Hesslow, Roman Castagné, Alexandra Sasha Luccioni, François Yvon, Matthias Gallé, et al. Bloom: A 176b- parameter open-access multilingual language model. arXiv preprint arXiv:2211.05100 2022.
Prajwal Singhania, Siddharth Singh, Shwai He, Soheil Feizi, and Abhinav Bhatele. Loki: Low-rank keys for efficient sparse attention. arXiv preprint arXiv:2406.0255 025442 2024.
Emma Strubell, Ananya Ganesh, and Andrew McCallum. Energy and policy considerations for modern deep learning research. In Proceedings of the AAAI conference on artificial intelligence volume 34, pp. 13693-13696, 2020.
Mingjie Sun, Zhuang Liu, Anna Bair, and Zico Kolter. simple and effective pruning approach for large language models. arXiv preprint arXiv:2306.11695, 2023.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 13971 2023a.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint arXiv:2307.009: 092888, 2023b.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention you need. Advances in neural informa- lion processing systems, 30, 2017.
Zhongwei Wan, Xin Wang, Che Liu, Samiul Alam, Yu Zheng, Zhongnan Qu, Shen Yan, Yi Zhu, Quaniu Zhang, Mosharaf Chowdhury, et al. Efficient large language models: A survey, arXiv preprint arXiv:2312.03863, 1, 2023.
Ke Wang, Houxing Ren, Aojun Zhou, Zimu Lu, Sichun Luo, Weikang Shi, Renrui Zhang, Lingi Song, Mingjic Zhan, and Hongsheng Li. Mathcoder: Seamless code integration in llms for en- hanced mathematical reasoning. arXiv preprint arXiv:2310.033311, 2023.
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang. and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04700 2020.
13

---

Published as a conference paper at ICLR 2025
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Transformers: State-of-the-aad art natural language processing. In Proceedings of the 2020 conference empirical methods in natural language processing: system demonstrations, pp. 38 38-45, 2020.
Haoyi Wu and Kewei Tu. Layer-condensed kv cache for efficient inference of large language models. arXiv preprint arXiv:2405. 00024055.03377 10637, 2024.
Guangxuan Xiao, Ji Lin, Mickael Seznec, Hao Wu, Julien Demouth, and Song Han. Smoothquant: Accurate and efficient post-training quantization for large language models. In International Conference Machine Learning, pp. 38087-38899. PMLR, 2023a.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309. /7453, 2023b.
Wenhan Xiong, Jingyu Liu, Igor Molybog, Hejia Zhang, Prajjwal Bhargava, Rui Hou, Louis Martin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, et al. Effective long-context scaling of foundation models. arXiv preprint arXiv:2309.1600 2023.
Yuhui Xu, Lingxi Xie, Xiaotao Gu, Xin Chen, Heng Chang, Hengheng Zhang, Zhensu Chen, Xi- aopeng Zhang, and Qi Tian. Qa-lora: Quantization-awwat low-rank adaptation of large language models. arXiv preprint 1341123309.1777, 2023.
Dongjie Yang, XiaoDong Han, Yan Gao, Yao Hu, Shilin Zhang, and Hai Zhao, Pyramidinfer: Pyra- mid kv cache compression for high- tighthhroughpt Ilm inference. arXiv preprint 1.Xx::2405.12322. 2024.
Zhihang Yuan, Yuzhang Shang, Yue Song, Qiang Wu, Yan Yan, and Guangyu Sun. Asvd: Activation-awar aware singular value decomposition for compressing large language models. arXiv preprint arXiv:2312.0582 2023.
Tianyi Zhang, Faisal Ladhak, Esin Durmus, Percy Liang, Kathleen McKeown, and Tatsunori B Hashimoto. Benchmarking large language models for news summarization. Transactions of the Association for Computational Linguistics, 12:39-57,20244
Xingxing Zhang, Furu Wei, and Ming Zhou. Hibert: Document level pre-training of hierarchical bidirectional transformers for document summarization. arXiv preprint arXiv: 1900:06605006 2019.
Yichi Zhang, Bofei Gao, Tianyu Liu, Keming Lu, Wayne Xiong, Yue Dong, Baobao Chang, Junjie Hu, Wen Xiao, al. Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling. arXiv preprint arXiv:2406. 02069, 2024b.
Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, al. H2o: Heavy-hitter oracle for efficient gen- erative inference of large language models. Advances in Neural Information P Processing Systems, 36, 2024c.
Aojun Zhou, Ke Wang, Zimu Lu, Weikang Shi, Sichun Luo, Zipeng Qin, Shaoqing Lu, Anya Jia, Linqi Song. Mingjic Zhan, et al. Solving challenging math word problems using gpt-4 code interpreter with code-based self-verification. arXiv preprint 27X1v.2308.0792, 2023.
A OBSERVATIONS
Figure and Figure illustrates the observations which motivates approach THINK to prune unimportant key cache channels. We conducted singular value decomposition (SVD) (Demmel) 1997) on the attention weights of the specified layers investigate their principal components.
Note that
QK" VD (K)
))
Energy; of Σ,
U, E V SVD softmax
14

---

Published as a conference paper at ICLR 2025
2.0
1.5
5
1.0
0
200
50
Channels
1ons
0.00
200
10g Channe
200 Token
50
100 Channell
200 Token
50
Channel
Taken
400
Tokef A00
400
400
6000
6000
6000
6000
(a) Key Cache (14th layer)
(b) Value Cache th layer)
(c) Key Cache (20 layer)
(d) Value Cache (20t layer)
Figure 4: Magnitude of key and value cache for LLaMA-2-7B. The first head of layer 14 and layer 20 of LLaMA-2-7B is selected visualize the magnitude of the key and value caches. We observe that the magnitudes of the key cache channels vary differently, whereas the channels of the value cache do exhibit such variation.
10.1
Layer
1.0
Head
0.9
10
.... 30
10-5
0.7
Energy
10-7
(umulated
0.5
10-11
10-11
19
0.2
300
Indix
Index
(a) Energy of singular values
(b) Cumulative energy of singular values
Figure 5: The energy and cumulative energy the singular values.
15

---

Published as a conference paper at ICLR 2025
B N TEST PERFORMANCE COMPARISON
Figure 6 visualizes the test performance comparison on Needle-in-a-Haysttac on Mistral-7B-
Instruct-vO.2.
METAATANNND D.S. the weteeaateee CorBoc Theroo Antteracc
04
11.
440
no
(a) KV-size 128, SnapKV, Acc. 77.8
Peatee udin Homan nemu Dre. INc Inclorieet sntyptachh
20
11.
100
(b) KV-size 128, SnapKV THINK (0.4) Acc. 78.6
Pegtee laiig Miwetaaeempoe Bect wnd Imterace
20
11.
(c) KV-size 1024, SnapKV, VV, Acc. 90.4
Peatir beiing Mamak taree 2i1 Bed bus mnd
40
11.
150
(d) KV-size 1024, SnapKV + THINK (0.5), Acc. 90.8
Figure 6: Needle-in-a-Haysttel test performance comparison with Mistrrl-77B-nnstrutt--00 -n02.
THINK (A) indicates prune the key cache channels with pruning ratio of X
16

---

Published as a conference paper at ICLR 2025
C IMPLEMENTATIONS
C.I IMPLEMENTATION WITH QUANTIZATION
Figure illustrates the implementation of our method when integrated with the KV cache quanti- zation method KIVI (Liu et al. 2024b). During the prefill phase, first prune the unimportant channels of X before applying quantization. In the decoding phase, each newly arrived key cache Ik is added to Xk,. Once Xx, reaches G tokens, the residual length hyperparameter in KIVI, prune and quantize the data, then concatenate it with the previously quantized ((P(X,,))
Prefill Phase
Decoding Phase
Prociees Tevece
Parning Charnelo-..
X
Xx
P(Xx)
Guaritaatoo on by Channel
A
tx
A
KVCache
Best ged
X
Q(P(Xx,)) Xx,
Xo
Q(Xv,)
X1.
Ag
((((xx,))
Xv
Ovenkoation Takon
x
Figure 7: Implementations of THINK when incorporated with KIVI.
D
VALUE CACHE PRUNING
Similar to the approach used for the Key cache, the pruning of channels in the Value cache can be guided by two primary criteria: magnitude-base based pruning and query-driven pruning. We find that driven pruning is still better than magnitude based pruning.
Qal-Saue : K?
Score,, (Q. Ki, Vall = ||softmax
Vel.
(3)
VD
J = (Score,.; T)
where K, V, RSDD We define criterion Score,,; to indicate the importance of each channel in the head of value cache. Then, only top channels are retained. Table B reported the results of pruning both key and value channels, showing that pruning the Value cache channels is harder than pruning the Key cache channels. However, pruning 30% of both the Key cache and Value cache LLLaMA---B---ma Instruct still results in acceptable performance. This demonstrates that the Value cache also has potential for pruning in the channel dimension. As depicted in Figure 4 the magniude of the Key cache along the channel dimension is highly unbalanced, whereas the magnitude of the Value cache along the channel dimension is more uniform. This indicates that the channel sparsity in the Value cache is not as significant as in the Key cache, making more challenging to identify redundant channels for pruning. We will investigate the value cache pruning strategy as part of our future work.
E PRUNING KEY CACHE ON VANILLA MODELS
To furthur demonstrate the effectiveness of our proposed THINK, we conducted additional experi- ments applying THINK directly to vanilla models, specifically LLaMA-3-8B-Instrue and Mistral- 7B- I1-- The results, as shown in Table demonstrate that our method maintains superior performance even after pruning 40% of the key cache channels. Furthermore, when the pruning ratio increased to 50%, the performance degradation remains within acceptable range. These find- ings further validate the effectiveness of THINK, not only for pruned models but also when applied directly vanilla models.
17

---

Published as a conference paper at ICLR 2025
Table 8: Performance comparison of pruning both K and V cache with different pruning ratios on LongBench. H20 + THINKV (A1+)2) indicates that the key cache channels of H2O are pruned with a pruning ratio of a and the value cache channels are pruned of a pruning ratio of X2.
Daccme-- QA
Mutti Decument Q4
Summmrrication
Few-shot Learning
Synthetic
Code
Method
Nevoa
Quoper
MF cc
Hespags
ZWIMIMQA
Mesigue
Gov Bogurrt
QM5um
Mullieen
TREC
Triviaa
SAMSiom
PCwant
PRe
Let
BR-P
Avg.
KV: size 128
H20
22.12
13.20
31.61
37.79
32.71
18.45
20.52
22,02
21.10
38,50
87.75
39.34
5,83
09.50
55.00
2057
35.
+TINKY
21.71
13.00
1308
41 M0
18.04
19.00
21.45
20.16
18.00
34.11
5.16
69.50
57.50
55
35.74
21.B
72767
70.76
83.76"
LLAMA 200 Indrect
+THINKQU 340.3)
21.86
13:79
35.26
0093
28. 19
19.23
18.1
21.30
19.76
41.3 1.300
87.00
19.95
5.78
59.50
51.84
35.81
KV-eine 512
H20
23.50
13.91
14 68
41
33.52
1992
22:11
22.55 56
23:82
41 00
90:46
40:30
5.82
69.90 90
56.71
51 AM
87.23
+THINKY (0.34..)
23.83
17:57
14 18
42.60
31,52
1999
21.17
22.25
32.85
18.5 50
00.1 11
50.0%
5.21
69.0
$9.90
56.85
37,23
SupRy
54.63
27.66
77
22.15
54.55
22.55
255T
13.97
70.00"
00.30
5BT
6035
57.03
3:31
+TiNKV (0.340.3)
24.57
24.90
38.00
44.61
34.37
20.37
21.23
21.95
23.30
56.00
90.50
39.38
5.60
69.00
61.75
58.46
40.25
KV-size 2048
H20
25.50
26,85
19.30 54
64.30
32.92
21.19
24.04
23:01
26.16
55300
90.05
41.55
4.91
09.25
58.63
51.31
39.59
+THINKV (0.3+0.3)
25.03
19 65
42.12
33.00
19:59
25.00
22.89
25.27
51.00
91.11
40:58
5:33
69.00
61:13
57.58
39.89
*THINKY V (0.440:4)
24.07
34.91
12.77
41.19
$1.42
19.60
21.67
22:70
36.52
49.00
90.61
19.78
6:00
69.00
61.81
54.08
19 19
25.88
-20.55
m0
-40.0
555.0
218
25.00
235-47
76.35
T5.305
00.56
-4746
5:17
78925
53.67
51.33-
41.58
+THINKV (0.3+0.3)
25.13
20:97
40.35
44.12
$4.64
19.94
23.62
23.03
25.30
72.55
90.78
30.46
5.35
69.00
61
57.91
41.41
+THIKVV (0.440.4) 4)
25.13
25.85
40.70
44.21
36.36
21.07
22.31
22.89
24.50
72.50
90.88
38.77
6.41
69.00
61.49
58.87
41.52
KV-siied 125
H20
21.21
21.81
38.87
30.42
20.36
12.30
20.58
20.61
22.10
39.00
82.3 22.57
40.44
2.50
70.96
51.23
48.38
34,3
THINKV (0.3+0.3)
20.71
21.40
38.01
30.66
22:28
13.87
20.13
22.45
21.07
38.50
82.33 20.20
33.69
2.54
78.56
51.55
48.28
34.46
19.17
21.40
42.93
36.76
23.44
1538
19.16
21.34
21.35
47.30
84.15
40.24
2.30
68.26
52.51
48.30
35.29
+THNKV (0.3+0.3)
19.92
20.61
42.65
37.63
23.19
15.09
15.97
21.93
20.55
45.00
84.06
39.33
2,59
66.00
51.51
47.51
34.81
KVVeine $12
H20
21.53
26,00
44.69
32,46
23,05
14.69
23.53
23.06
34.59
42.00
85.22
41.49
3,40
56.20
54.78
51,09
37.38
+THNKV (0.3+033)
22.36
24.26
44.77
30.47
22.94
14.95
22.63
22.90
25.75
41.50
85.30
40:21
3.591
30,17
54.08
50:56
36,54
3iteral Intrmae x0.2
+THINKV V (0.3+011)
22.14
25.15
45:29
31.78
23:21
14.62
25.30
22.70
24.51
41.50
85.61
41.58
2.75
34.113
54.50
51.00
37:11
SnapKy
24.44
21.81
48 95
99.80 46
25.25
16995
25.70
22.95
36.37
67.01
85.55 KM
41.26
2.70
36.55
56.46
51.41
40.66
+THINKV (0.3+0.))
24.10
23.04
47 76
58 66
25:45
17:51
23.64
22:81
23.81
56.00
95.50
19.91
1.36
$2.34
55 96
52.51
99
+THINKV (0.344.1)
1)
23.00
25.14
48.51 35
90.00
24.83
16.68
23.51
23:12
24.34
67.5 50
85.00
41
3.65
$4.34
$7.20
53.02
40.29
KV-size 1024
H20
23.67
25.55
46.4
36,99
24.82
15.02
25.21
23.04
25.77
46.00
85.93
41.98
3.24
36.57
56.40
52.75
38,90
+THNKY (0.3+0.3)
23.65
26.54
47.00
35.52
24.79
17.15
25.64
23.12
25.20
44.00
$5.38
41.67
3,46
30.14
56.53
52,86
36.23
+THNKY 11.20111
26.13
26.57
46.31
35. mg
24.92
15.34
24.58
23:35
25.55
45.50
15.0
41.97
1.57
33.64
55, 39
52.75
35,99
25.0
00.5
23.55
m
27:57
76.71
"M756
57.
4TAT
.THINKV ().340.))
25.20
20
49 17
41
25.78
19:37
24.64
25.02
25.27
49
00
95. 70
40/92
1:30
$2.05
57.15
54.55
41.00
-THINKV V (0.3+4.
25.84
10.90
49.5 56
41:44
25.70
19.02
25.21
23.75
25.72
50
00
85.00
42.55
2.44
55.76
57.55
54.10
41:45
KV-size 7048
H20
25.76
31.10
06
40.38
26.43
16.76
27.17
23.64
36.69
55.0
86.35
42.48
2.72
56.64
55.98
53.01
40,69
TTINKV (0.340.3)
25.60
25.74
47.5
35.67
26.25
17.35
24.54
23.27
36.15
51.00
87.01
43.02
2.94
$1.45
56.41
54.36
30.64
<TuNKY (0.340.1)
25.64
30.65
48.05
40.42
26.43
16,65
26.76
23.51
36.59
52.50
86.53
43.45
2.66
$3.06
56.55
53,83
40.30
Sazg&V
-3.6
R8.55
44.66
22.55
1875
53.50
"SEA7"
36.53
T0700
"LEST
$7.0
3575
48.8
+THNKV (0.3+0.3)
27.01
30.72
48.81
41.15
26.99
18.93
25.81
23.59
36.42
70.00
86.52
41.91
3.05
52.65
57.01
54.25
41.57
+THinKV (0.3+0.1)
26.22
32.09
48.96
40.53
26.70
19002
27.87
34.25
26.64
70.00
85,05
42.63
2.22
85.13
57.00
54.28
41.94
Table 9: Performance comparison of pruning key cache on vanilla models on LongBench.
Single Decument QA
Multi-Document Q4
Summarization
Few-shot Learning
Synthetic
Code
Method
NetrQA
Qasper
MF em
HotporQA
2WHRMMQA
Musiqu
GovReport
@MSaa
MottiNows
TREE
TrivtagA
SAMSion
PCaunt
PRe
Loe
RB-P
LLaMNN3 3 SD Instract, KV
KV-ize Full
Vasillaa
25.56
32:27
39.71
43.56
35.09
21.18
28.71
23.26
26.64
73.50
90.43
42.33
4.90
69.25
59.29
$4.05
41.8
K(0.))
32:26
34.77
21.10
28.63
23.13
26.38
73.50
90.58
5.21
69.50
61
94
58.37
43,27
.TmnkKi0.5)
32.80
31.81
21.5 18.58
78.35
26
90.37
69.17
42.25
+THINKG0.))
24.39
31.16
41.50
42.52
31.63
20.70
25.54
23.18
25 46
73.50
90.43
38.97
5.37
68.46
59.63
59.38
41.43
Mistral
'DD Instruct v0.2. KV
size Full
Vasilla
26.63
32.99
49.34
42.77
27.35
18.77
32.57
24.24
27.10
71.00
86.23
42.96
2.75
36.98
56,93
54.49
42.71
(THINK(0.)
27.11
33.46
48.73
41.79
28.14
18.87
32.45
26.55
27.09
71.00
86.26
43.02
3.95
36.31
56.99
54.36
42.76
26.63
33.71
49.38
42.38
26.78
18.
32.57
24/63
26.92
71.00
86.
42.83
3.11
84.65
56.75
54.04
42.53
K(0
26.43
26.37
86.28
83.11
F COMPARISONS OF GENERATION SPEED AND THROUGHPUT
In this section, we provide comparisons of TTFT (Time To First Token), TPOT (Time Per Out- put Token), memory usage, and throughput in Table and Table [1] We follow the methodology used in KIVI Liu et al. 2024b). We generate synthetic workloads with an input prompt length of 160 and output length of 338. We set batch size 300 for both KIVI and our method. As our method performs online pruning. there is a slight delay introduced during the prefilling stage due to the computation of channel importance. However, the reduction in memory usage leads to notable improvements in TPOT during inference. Specifically, after applying THINK, TPOT improves from 0.27 ms/token to 0.25 ms/token with 40% pruning, and further improves 0.24 ms/token with
18

---

Published as a conference paper at ICLR 2025
Table 10: Comparisons of TTFT (Time To First Token), TPOT (Time Per Output Token) and mem- ory usage on LLaMA-2-7B.
Method
KIVI(4/4)
KIVI(4/4) (HINK(0.4)
KIVI (4/4) + (H.NK(0.5)
Memory (GB)
61.7
53.3
$1.2
TTFT (ms)
7.0
10.4
10.2
TPOT (ms/token)
0.27
0.25
0.24
Table 11: Comparisons of throughput LLaMA-2-7B.
Method
Vanilla
KIVI(4/4)
KIVI(4/4) THINK(0.4)
KIVI + (H.5
Throughput (tokens/s)
2557
5168
5518
5676
50% pruning ratio. Regarding throughput, with the same memory usage, THINK (40% pruning) improves the throughput KIVI from 5168 tokens/s to 5518 tokens/s. At pruning ratio 50%, throughput increases further to 5676 tokens/s. Model weights and KV cache the primary mem- ory components accessed during generation. By effectively reducing the memory footprint of the KV cache, our method alleviates the memory bottleneck, enabling larger batch sizes and faster gen- eration speeds. These results highlight the potential of THINK to enhance inference performance by balancing memory efficiency with computational overhead. We are continuing to optimize the implementation to further improve performance. Memory bandwidth is major performance bottle- neck the decoding phase large language models (LLMs).
G COMPARISONS WITH SVD BASED METHODS
In this section, we compare our THINK with SVD based methods (Saxena et al. 2024 Yuan al,, 2023 Chang et al. 2024). SVD based KV cache compression method decompose the KV cache weights offline with some calibration data, relying on storing latent representations and recovering the cache during inference. Our approach is online pruning strategy that is aang-and-layy quiring no changes to the model's weights or architecture. This makes THINK lightweight and easy to integrate into existing systems. Besides, method dynamically prunes the Key cache chan- nels directly without requiring any reconstruction. We provide comparisons of our method with the base-
SVD-based approaches in Table 12 Following Palu, line model. Our results show that when the KV cache is compressed by 80%, our method preserves performance, whereas Palu experiences an accuracy drop of 1.28%. Similarly, for ASVD, we eval- uate under the same compression rate. Our method demonstrates significantly better performance, with only a 0. 12% performance drop compared to the 5.21% % drop observed for ASVD. The results demonstrate the effectiveness of THINK, with significantly less performance degradation compared to SVD-based methods under the same compression rates. Our pruning method has the potential to be integrated with SVD-based approaches. Exploring such combination could yield further advancements KV cache compression, which we aim to investigate in future work.
Mistral-7BB-Instrutt-voo
Table 12: Performance comparison of SVD based methods on LongBench.
Single Decument QA
Mtuli-Documenn UA
Sammar ization
Few-shot Learning
Syuthetic
Code
Method
NeteQA
Desper
MF wen
HolpottA
2WMMMQA
Musign
GevReport
OMSam
MuttiNwws
TREE
TrivinQA
SAMSua
PCaunt
PRe
Lac
RB-P
Avg
Mistal--B--hsttu---1222
Vasillia
26,63
32.99
49.34
42.77
27.35
18.77
32.87
24.24
27.10
71.00
86.23
42.56
2.75
86.98
56.93
54.49
42.1
PPalu
27
36.11
48.47
46.09
26.33
20.18
31.15
34.30
27.
70.00
85
41.74
2.74
73.18
51,70
54.52
41.43
28.14
18.87
27.09
96.26
56.99
42.76
26
68.81
+ASVD
16.46
13.19
28.58
21.94
22 25
8.74
17.73
20.27
22:
57.00
73.88
38.51
1.50
4.77
53.32
29.13
26.16
+THIKK10-4)
18.39
19.98
35.05
30.55
25.65
10:35
25.98
20.52
26
64.00
83.63
41.55
6,00
8.50
60.18
58 33.25
19


Reviews:
[{"review_id": "cW0lsMLQeT", "reviewer": "Reviewer_ELRT", "summary": "The paper presents a KV cache optimization method with channel wise pruning for reducing the inference cost under long context scenario. The method is based on the previously proposed efficient context compression method ( H2O, KIVI, and SNAP). Specifically, the paper utilized a query driven method to involve the query for key channel selection.", "questions": "See weakness section.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The proposed method is simple and straight forward, with direct structured pruning on the key channel.\nThe evaluation shows the pruned results achieves comparable results with the non-pruned baselines", "weaknesses": "Lack of experiments compared to directly applying to vanilla models.\nWhy the proposed method is basing on the pruned method, instead of directly on the vanilla model? The observations come from the fact of that there are large outliers alongside the channel. This not an unique attribute on the compressed model (i.e. H2O). Yet, experimenting on the compressed model introduces confounding variables which affect the analysis. The current manuscript lacks experiments directly apply the method on the Vanilla model.\n\nIn the meantime, the current manuscript lacks the details about how the K cache selected K cache is removed. This is important in order to consider the real memory benefits. Is the K cache fully removed? Or full of the KV cache are always stored in memory? If the selected channel of the K cache is removed, when new queries come, how to upgrade to the new selected K cache size? If the whole K cache are always stored in memory, how is the memory efficiency comes from?\n\nIt is also worth to indicating whether the compression is done starting at the prefilling phase, or the prefill phase is concrete. The experiments of many the previous approach (i.e. H2O) did not compress during prefill phase. Detailing it will make the understanding of the experimental results more comprehensive.\n\nTypos:\nSec 3.1 line 149 KV cach\nLine 175 k is bold?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OyJKDRokPf", "reviewer": "Reviewer_cgqU", "summary": "This paper presents ThinK, a query-dependent K cache pruning method for long-context LLM inference. ThinK takes advantage of the channel-wise sparsity property of K cache in existing KV cache tensors, and uses a greedy algorithm to remove insignificant channels and to optimize for the loss in attention weights after pruning. ThinK can be used together with many existing token eviction techniques (H2O) and KV cache compression techniques (KIVI). Experiment results show that on long-context benchmark tasks, ThinK achieves a comparable level of accuracy when used jointly with H2O or SnapKV, as compared to using H2O or SnapKV alone, but lower memory consumption.", "questions": "Please refer to \"weaknesses\".", "limitations": "", "rating": 8, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "`+` Channel-wise sparsity for K cache isn't a new thing, but this is (probably) one of the first papers that present a solid method for channel-wise K cache pruning. \n\n`+` Compatibility with existing token eviction or cache compression methods. \n\n`+` Solid and extensive evaluation. I am convinced by tables 2 and 3 that ThinK is good at detecting and pruning insignificant channels in the K cache, since the accuracy drop is very low even for a pruning ratio of 0.4 or 0.5, so this is a reasonably designed method for K cache compression.", "weaknesses": "`-` Unfortunately, V cache typically demonstrate token-wise sparsity, so this method cannot be generalized to V cache. To some people, this might seem like a smart engineering hack for a specific type of tensors with known distribution properties. **(Clear: The reviewer is clear with the presentation of section #2 after the rebuttal period)**\n\n`-` The presentation of evaluation results and the choice of evaluation metrics are very concerning to me. There is an overwhelmingly large focus on accuracy numbers (Tables 2, 3, 4, 5, and 6), but one of the most important benefits enabled by KV cache compression/pruning is increased throughput and decreased TTFT --- These two metrics do not even appear in the evaluation. Note that reduced memory usage (Figure 3a) is only an intermediate benefit you care about: Eventually, the hope is that reduced memory usage can enable larger batch size, and consequently larger throughput. Though Figure 3b has some numbers on batch size, it does not show how ThinK has the potential of boosting inference throughput. **(Clear: The authors presented updated numbers on latency / throughput during the rebuttal period)**\n\n`-` Minor issues: Typos on line 188 (\"pruning matrix\"->\"pruning metrics\"?); weird sentence on lines 214-215 (\"Based on the observations ...\"). **(Clear: The authors fixed these minor issues during the rebuttal period)**", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8ujP9NnusI", "reviewer": "Reviewer_JLGy", "summary": "This paper presents a novel pruning method named THINK to reduce KV cache memory size for long-context inputs. The authors observed a gap in optimizing the channel dimension for KV cache parameters. They show that THINK can reduce memory size by 20% with minimal accuracy loss.", "questions": "Thank you for submitting to ICLR 2025. The research problem is both interesting and timely. The motivating examples are helpful in identifying existing research gaps in channel pruning. The solution is well-reasoned, resulting in significant memory improvements.\n\nWhat are the authors' reasons for choosing the three baselines?\n\nIt would also be helpful to see the results of THINK combined with existing methods on the sequence length dimension and the layer dimension.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Timely problem\n- Novelty\n- Solid experimental results", "weaknesses": "- Comparison with related work", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5gyQ7iDgj6", "reviewer": "Reviewer_ssq7", "summary": "This paper proposes ThinK, a novel KV cache compression technique that uses query-driven channel pruning for keys and values. Specifically, channel importance is established through $|| Q.K^T ||$ magnitude, and channels with the highest scores are selected. Extensive experiments over the Llama and Mistral family of models demonstrate a 20% reduction in KV cache size when ThinK is used in conjunction with existing KV cache compression techniques.", "questions": "1. Which dataset was used to find channel importance as shown in Section 3.2? Is a pretaining dataset like WikiText a good fit for establishing channel importance for a range of downstream tasks, or task-specific data is necessary? It would be helpful to understand the generalizability of channel importance score S.\n\n2. Is query-driven pruning a better metric than SVD (as used in Eigen Attention [1]) to find channel importance? Comparison with such works can further clarify the superiority and/or limitations of ThinK.\n\n3. (Suggestion) A pruning ratio ($\\lambda$) of 0.4 implies a 40% reduction in the key cache, which means a 20% reduction for the KV cache. This should be specified in Section 4.2.\n\n4. Can the authors provide some discussion on why it's harder to compress the value cache as compared to the key cache (Table 1-7 vs Table 8 in Appendix)?", "limitations": "", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper is well-written and concise, which makes it relatively easy to follow. \n2. The authors conduct an extensive set of experiments (along with ablation studies) to show performance improvements in the form of reduced KV cache footprint as well as the possibility of enabling larger batch sizes through their technique.", "weaknesses": "1. Authors overlook relevant baselines [1-3] which also reduce KV cache size by focusing on the hidden dimension (hence effectively pruning channels). Comparison with these techniques is important to put this work in context with other such KV cache compression techniques.\n\n2. The performance benefits are not very significant since the authors show key cache compression by up to 40% and value cache up to 30%, leading to only ~20% improvement in KV cache memory footprint. \n\n[1] Saxena et al., \"Eigen Attention: Attention in Low-Rank Space for KV Cache Compression.\", ArXiv 2024.\n\n[2] Yuan et al., \"ASVD: Activation-aware Singular Value Decomposition for Compressing Large Language Models\", ArXiv 2024.\n\n[3] Chang et al., \"Palu: Compressing KV-Cache with Low-Rank Projection\", ArXiv 2024.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "YiC70Sezp8", "reviewer": "Reviewer_1rme", "summary": "This paper propose a new way to prune KV cache that focuses on reducing least significant channels. The method is query-dependent and reduces KV cache size upon previous KV compression methods while maintaining almost the same quality with previous methods. It also discovers the fact that activated key cache is sparse for a given query.", "questions": "1. Could you address the questions in weakness?\n2. I think you should be more modest about your improvement. Your method is solid and could be impactful (ex. you potentially could also help reduce computation needed in attention calculation by reducing K cache channel dimension). I suggest to make it clear about what part is your contribution and what part is the previous methods (ex. in abstract), and dig deeper into why your improvement in this paper could benefit inference system-level performance.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Very extensive benchmark scenario. Evaluated with models from 7B to 70B on 10+ datasets. The accuracy drop is indeed negligible.\n2. The method proposed is simple and works on top of previous methods. This could be a nice addition to memory-efficient transformer inference.\n3. Writing is easy to follow and method description is clear.", "weaknesses": "1. Method is dependent on query but did not show online pruning delay costs. It would be nice to show a few numbers on long the pruning takes.\n2. Did not discuss enough for value cache. This is understandable since value cache is indeed harder to compress according to my own experience. However, you should add more discussion on the implication with this. For example, you claimed that you have 2.8 peak memory saving, does this include saving with V cache? More discussion is needed on this thread to make readers better understand the use scenario of your method (this has nothing to do with the validity of your key cache pruning, it is just unclear to me how much actual benefit comes with your method and I am worried of readers may over-estimate your gain). \n3. KIVI is not your contribution. When you present the result in abstraction, please report clearly your improvement SEPARATELY. What is your gain above KIVI? Compressing K cache by 60% without accuracy drop is a very good contribution already!", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)