Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Time
Authors: Zichang Liu
Aditya Desai
Fangshuo Liao
Weitao Wang
Victor Xie
Zhaozhuo Xu
Anastasios Kyrillidis
Anshumali Shrivastava
Venue: NeurIPS 2023 poster (2023)
Keywords: Large language model; KV Cache Compression

Abstract: Large language models(LLMs) have sparked a new wave of exciting AI applications. Hosting these models at scale requires significant memory resources. One crucial memory bottleneck for the deployment stems from the context window. It is commonly recognized that model weights are memory hungry; however, the size of key-value embedding stored during the generation process (KV cache) can easily surpass the model size. The enormous size of the KV cache puts constraints on the inference batch size, which is crucial for high throughput inference workload. Inspired by an interesting observation of the attention scores, we hypothesize the persistence of importance: only pivotal tokens, which had a substantial influence at one step, will significantly influence future generations. Based on our empirical verification and theoretical analysis around this hypothesis, we propose scissorhands, a system that maintains the memory usage of the KV cache at a fixed budget without finetuning the model. In essence, Scissorhands manages the KV cache by storing the pivotal tokens with a higher probability. We validate that scissorhands reduces the inference memory usage of the KV cache by up to 5$\times$ without compromising model quality. We further demonstrate that scissorhands can be combined with 4-bit quantization, traditionally used to compress model weights, to achieve up to 20$\times$ compression.

Seissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Time
Zichang Liu Department of Computer Science
Aditya Desai Department of Computer Science
Rice University zichangliu@rice. edu
Rice University Aditya. P. Desaiirica
Fangshuo Liao
Weitao Wang
Department of Computer Science
Department of Computer Science
Rice University Lingshuo.Liaorriceee edu
Rice University wtwangerice.eds edu
Fangshuo Liao@rice.........
Victor Xie
Zhaozhuo Xu
Department of Computer Science Rice University vyx2@rice.edi
Department of Computer Science Stevens Institute of Technology zxu790stevens. edu
Anastasios Kyrillidis
Anshumali Shrivastava Department of Computer Science Rice University & ThirdAl Corp. anshumali@ra i@ricerre..e. edu
Department of Computer Science
Rice University anastasios@ricee. edu
Abstract
Large language models(LLMs) have sparked a new wave of exciting AI applica- tions. Hosting these models scale requires significant memory resources. One crucial memory bottleneck for the deployment stems from the context window. It is commonly recognized that model weights memory hungry; however, the size of key-value embedding stored during the generation process (KV cache) can easily surpass the model size. The enormous size of the KV cache puts constraints the inference batch size, which is crucial for high throughput inference workload. Inspired by an interesting observation of the attention scores, we hypothesize the persistence of importance: only pivotal tokens, which had substantial influence one step, will significantly influence future generations. Based on our empirical verification and theoretical analysis around this hypothesis. we propose SCIS- SORHANDS. system that maintains the memory usage of KV cache under a fixed budget without finetuning the model. We validate that SCISSORHANDS reduces the inference memory usage of the KV cache by up 5 without compromising model quality. We further demonstrate that SCISSORHANDS can be combined with 4-bit quantization for further compression
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Introduction
Large language models(LLMs), trained on immense amounts of text data, have demonstrated an incredible ability to generate text that is both logically connected and contextually relevant [1-5). LLM inference follows an autoregressive fashion, generating one token at each step conditioned on the previous steps. At each step, the key-value embedding in attention is stored in memory to avoid repetitive key-value projection computation at future steps. Unfortunately, the memory of the key-value cache( KV cache), including prompts and previously generated tokens, can be surprisingly large. Using OPT-175B as an example, the impressive 175 billion parameters consume around 325 GB of memory. At the same time, at batch size 128 and sequence length 2048, the KV cache requires around 950 GB of memory, three times larger than the model weights. Considering that 8 Nvidia A100-80GB offers 640GB GPU memory, the memory usage of the KV cache is truly concerning.
LLMs are typically deployed on fixed memory hardware, and the size of model weights is also fixed once deployed. Apart from small memory buffer typically reserved for communication and computation, the rest of the available memory is for the KV cache. The size of the KV cache depends on batch size, sequence length, and model dimension. Thus, given inference sequence length, compression in the KV cache memory translates almost linearly into an increase in the batch size. And any increase in batch size is significant for high-throughput inference systems 6
Quantization and sparsity ap- proaches [4]4 have been stud- ied in LLMs to reduce the model sizes. However, compressing the KV cache remains an open but challenging problem. First, training models at the scale of hundreds of billions of parame-
Aetoooo Hevo
Sequancee
(a) Attention map at position 178
Mttenton Head
ters on a large amount of data is prohibitively expensive. Thus, an ideal compression algorithm should be applicable without training. Second, emerging ap- plications such as dialogue sys- tems require an extremely long context window. The maximum sequence length of LLMs is grow- ing to over 32K 15. The size of the KV cache also grows linearly with sequence length. For scala- bility, an ideal compression algo- rithm should reduce the memory from the sequence length dimen- sion. At last, compression should preserve LLMs' quality and in- context learning ability.
Sequuence PouB
(b) Attention map at position 228
(c) Attention map at position 278
Figure 1: Repetitive Attention Pattern. We plot the attention map at three token positions in a sentence. Only five attention heads plotted for clearer presentation. We discretize the attention score such that the high score is dark green, and the low score light green. In Figure [(a) the token at position 178 pays heavy attention to positions 27, 63, 98, etc. This pattern is also present in the attention maps of position 228 and position 278.
We go beyond the traditional model compression techniques to achieve such demanding requirements. We envision that not all tokens must be stored memory for LLM to understand the context. Just like humans can skim through an article and grasp the main idea, LLMs may also be able to skim and comprehend. It is commonly observed that the attention score from one token follows a strong power law distribution 20),20) meaning that one token will only heavily attend to small number of tokens. More importantly, we observe Repetitive Attention Pattern from different tokens in the sequence in a trained LLM( Figure D.. Certain tokens are more important throughout the paragraph. Specifically, for two different tokens, there are similarities between which tokens they are heavily attending to and similarities between which tokens they are ignoring.
Inspired by the above observation, we articulate the Persistence of Importance Hypothesis: Only pivotal tokens, which had a substantial influence at one previous step, will have significant influence at a future step. This hypothesis, true, suggests that it is possible to foresee which token is likely to be important for future generations. Fortunately, we empirically verify that later tokens in the

---

sentence mostly only attend to tokens that were heavily attended from the early tokens in sentence. And the overlapping ratio surprisingly high, over 90% in most of the transformer layers (Figure 2
Based on the above two findings, we present SCISSORHANDS that exploits the persistence of impor- tance hypothesis realize LLM inference with compressed KV cache. In Section 4 we present an efficient algorithm such that the size of KV cache is always less than predetermined budget. A theoretical guarantee justifies that such a compressed KV cache can approximate the attention output. we empirically evaluate SCISSORIIANDS and show that SCISSORIIANDS reduces the memory usage of KV cache 2 5x without compromising model quality. Reduction in the KV cache can directly result in a larger batch size. Further, we adopt quantization and show its compatibility with SCISSORHANDS.
In Section S
2 Problem Description and Related Work
This paper considers the L.LM inference workflow, specifically focusing on the memory usage for storing the keys and values in attention. Let d be the hidden dimension of the model, b be the batch size, and p be the length of prompt sentences. We are given the trained model weights, Wi E RRd Wi E Radd for the key and value projection matrix the 7th transformer layer.
The standard LLM inference consists of two stages: prompting and token generation. In the prompting stage, the model takes the prompt sentences as the input, and the key/value embedding in attention is stored as cache to reduce repetitive computation. Denote x prompe = [x1 ... xppl Zprompt E Rxxxxx as the input to attention the ht transformer layer. Denote the key cache and value cache at layer i as KC,, V* E pxx Kg = x prompe Wix V/8 -" promp Wi.
In the generation stage, the model starts with the stored KV cache in the prompting stage and generates one token at each step. At each step, the KV cache gets updated. Given the input to attention step in the 7th transformer layer æt E 1xxx1xd, K7 = [C] ot W kb Vi+1
2.1 LLM Inference Memory Breakdown
In this section, we provide the memory consumption breakdown of LLMs. The memory footprint consists of three parts: model weights, KV cache, and activation buffer. The size of model weights depends on model configuration, such as the number of transformer layers and hidden size. The size of the KV cache depends on model configurations, sequence length, and batch size. The size of the activation buffer depends on parallelism strategy, model configurations, and implementation. The size of the activation buffer considerably smaller than the previous two. As shown in Table the size of the KV cache, 2.5 x-5x x larger than model weights, can quickly become the bottleneck in memory consumption. At the same time, much research has been spent on extending the length of the context window. GPT-4-32K can process up to 32, 26, tokens [1]. Longer sequence length would make the KV cache memory problem even more severe.
Assuming L.M generates until its maximum sequence length, we summarize the maximum batch size before going out of GPU memory on a box of 8 A100 80GB GPU in Table At the GPT-3 scale with a maximum sequence length of 2048, batch size cannot exceed 35 without offloading. Small batch size limits the model inference throughput.
2.2 Efficient Attention
Computing the attention matrix necessitates time complexity of O(n)), where is the sequence length. As result, a line of work has been proposed to mitigate the computation burden of the attention mechanism [16 20). These approaches exploit low-rank sparsification to approximate the attention output. Besides, 21 realized exact efficient attention with wall-clock speed by optimizing
Table 1: The memory consumption of model weights and KV cache for three different LLMs at batch size 128 and sequence length 2048 shows that the KV cache dominates the memory consumption.
Model
# of Layer
Hidden Size
Weights (GB)
KV cache (GB)
OPT-175B
96
12288
325
1152
LLaMA-65B
8192
130
640
BLOOM
70
14336
352
950
3

---

Table 2:
Maximum batch size before hitting out of memory on a box of A100 80GB GPU when models are deployed with its maximum sequence length.
Model
OPT-175B
LLaMA-65B
BLOOM
Maximum Batch Size
34
102
the number of memory reads and writes. However, these approaches were evaluated mostly for training, focused on computation complexity, and did not address the KV-Cache memory usage introduced by auto-regressive language models.
Recently, there is active research attempting to apply quantization or pruning in LLM 8.14 However, they mostly focus reducing the size of model weights. Flexgen (7] applies quantization and sparsification to the KV cache; however, the memory of the KV cache is not reduced regarding sequence lengths. It stores the quantized KV cache for all tokens in CPU memory and loads all attention keys from CPU memory to compute attention scores. At the same time, methods such as Multi- Multi-Query- Seery--Attentio(MMM across all attention heads. MQA requires training the model from scratch, while our works focus entirely on the inference stage.
[22] change the attention design such that keys and values shared
3 The Persistence of Importance Hypothesis
We first present one interesting observation upon which the persistence of importance hypothesis is derived in Section 3.1 In Section 3.2, we discuss the hypothesis in detail with empirical verification.
Then, in Section 3.3 we provide theoretical intuition on the reason behind such model behaviors.
3.1 Repetitive Attention Pattern.
Observation. We are interested in the attention score from the position over all the words that come before in the sentence. In Figure IT we provide three attention maps sentence randomly drawn from the Colossal Clean Crawled Corpus (C4) (23 using OPT-6B. Each attention map is discretized attention score calculated at random position. We consider a score larger than as significant as indicates averaging mixing score. High attention scores are marked with dark green.
Result. High attention scores are observed at the same set of tokens from various positions in the sentence. In all three plots, we see dark green at sequence positions 27, 63, 98, 121, 152, and 177, suggesting that these tokens received high attention all three positions. We observe similar model behavior at different transformer layers with different text inputs. More plots are in Appendix A
Implication. Even though small differences exist, repetitive attention patterns evident in the attention maps. There exist specific tokens that keep receiving high attention. Meanwhile, these attention maps show sparsity: only few tokens have high attention scores.
3.2
The Persistence of Importance Hypothesis
The repetitive attention pattern suggests that specific tokens are influential throughout the sequence. stricter claim is that these tokens are the only ones that could be significant for future step. Thus, we articulate the persistence of imporiance hypothesis.
The Persistence of Importance Hypothesis. With trained autoregressive language model, only pivotal tokens, which had a substantial influence at one previous step, will have significant influence future step.
true. this hypothesis indicates the possibility foreseeing what information the previous sequences could be vital for future steps. This hypothesis is trivial when pivotal tokens include all tokens in the entire sentences. However, a much more interesting case is when pivotal tokens are a subset of previous words. This would enable to reduce the size of the KV cache by throwing away the embedding of non-oooonnnn non-impottant tokens.
Pivotal Token. One natural indication of a token's influence is the attention score. We consider token pivotal for position if this token receives an attention score larger than threshold from the token at position Let S, denote the set of pivotal tokens for position t. Sa-b denote the set of

---

1.0
1.0
OPT-6B
0.8
OPT-13B
OPT-30B
0.8
0.6
OPT-66B
OPT-6B
OPT-138
0.6
OPT-30B
0.2
OPT-668
0.5
0.0
10
20
30
40
50
60
10
20
30
40
50
60
Transtomeer Layer
Transtormer Layer
(a) Persistence Ratio
(b) Size Sp
Figure 2: Persistence ratio and the corresponding size of the pivotal token set. The persistence ratio is over 95% in most layers, with decreases at the later layers. Meanwhile, the number pivotal tokens considerably smaller than the sequence length. This suggests that the pivotal tokens of later half sentences are almost all included in the set of first halves.
pivotal tokens for every position from b. Sa-b
U
Verification. We measure persistence ratio as an empirical test the hypothesis. Persistence ratio measures how many tokens the pivotal token sets of the later part of the sentence are also in the pivotal token sets of the initial part of the sentence. Let I denote the length of the sentence. We record Si E { tokens (Ii, } who received high attention from every position until Then, we record S1+1-1 E {21:....4)) tokens in {#1...., .t
who received high attention from position after f. The persistence ratio is the intersection divided by the size of S,+1-1 Formally,
|S 0 So-sel
Persistence Ratio
|x| E S141-1, E (x11...,2 21}1
At the same time, we measure [So-vel |So-el indicates that every token substantially impacted 4 at least one position, which is the trivial case of persistence of importance hypothesis. Our test performed with OPT models [24] with different datasets such as OpenBookQA 25 and Wiki- Text [26). In our verification, set t 1, which measures the overlapping between the first and later half of the sentences. Same in Section 3.1
we a L. which suggests an average score.
Result, We present our main results in Figure First, given the current criterion of pivotal token and t value, the size of So is considerably smaller than half of the sentence length. This verifies that we are not considering the trivial case of our hypothesis. Second, the persistence ratio is generally over 95%, with dips in the later transformer layers. The pivotal token set of the later half sentences mostly included in the set of the first half sentences. Combining these two pieces of empirical evidence, see positive evidence for hypothesis test.
Implication. The hypothesis provides insights for understanding the behavior of LLMs and opens up new opportunities for reducing the KV cache memory. The hypothesis suggests the possibility of predicting the potentially influential tokens for future steps. The non-influential tokens are unnecessary to store in the memory, as they are unlikely to have high attention scores. This reduces the number of tokens stored the KV cache and the computation required at attention.
3.3 Attention Weights Decides the Pivotal Tokens
In the previous section, we verified that the significant tokens would continue to be significant. In this section, we try understand the reasons for such phenomena. We consider the token generation process of a simplified model: a single-layer transformer model with single-head attention.
#t+1 F (a:), where a, = sof tmax (1/- x. W.W (XX,) X.
(1)
Dt E IR. Ixd is row vector. X,-1 E ((t-1)xd denotes the aggregation of æ1.... FT where the jthh row is Ij. We. WK. Wy E Rdxd and Wo E [dd are the attention weights. Lastly, F I - Rixd denotes the MLP block following attention block, two-layer MLP with skip connections, given by
F(x) = x W2relu(W1z)
(2)
5

---

Algorithm Inference with Budget KV cache
Input:
Memory Budget B, Maximum Sequence Length T Key Cache K E Rnxd, Value Cache V E Raxd, where n = 0
while f < Tiax do
Model update K,, V such that n t- n n+1 if n > B then:
Compress KV cache using Algorithm Shsrrthm@ssuchh such that n S B.
end if
t + t ++1 end while
Algorithm 2 Compress KV Cache
Input: Key Cache K Re Value Cache V E Raxd, History Window Size w, Recent Window Size r, Drop Amount m, Generation Step t,
Importance Record I 0 RR
for i E [1 w. 4 do
> Consider tokens within history window Increment the counter for low score token
T + (; <
end for [:: -r] -
Keep cache within the recent window
Keep set St + Argsort (1) [m]]
Keep everything in St in K E Raxd, V E Raxd such that t- m
We are interested in the attention scores az = softmax(Y/- Wo Wo XTI)). Notice that or scales with z, Wo W xx The following theorem characterizes the behavior of &
Theorem 3.1. Let A = Wy WoWoowWk and let λk, S2. Ao denote the largest singular values of WK. Wq, Wv, Wo. respectively. Consider the transformer in (7) with normalized inputs ||
for all t. Let > 0 be constants. Assume that a63741 2 (1 - 6) ||ael22 with 8 S Tg
Then for all re satisfying xzAx]] 2 c and Axx 2 2c1 max 2 AxT it holds that
AxA (au, 3c) Wok < |0el12
xtA
|atl2 (ar. 3c)
(3)
The proof is provided in Appendix B Theorem 3.1 shows that under assumption the MLP
in (2). for all zz such that Ar, is large enough, T++1 WK,, satisfies Equation (3) The assumption on the MLP ax 2 (1 5) ||عel2 essentially requires a large cosine similarity between the input and output of F. This behavior can be empirically verified Appendix A Essentially, skip connection dominates the output because ||=|||2 >> Warelu(W(x)) resulting in a cosine
similarity close to one between input and output. Equation (3) shows that despite a factor of T|
Az,
I WT n almost scales with a,,e. Since r441 Wo W K directly affects 01+1,6, this property shows that larger ai5g will potentially imply large 0++1,6-
Our theorem shows that the property in Equation (3) property only holds for ae such that x An; large. A are trained attention weights. This condition may suggest that the trained weights A selects Fe as pivotal token. Each attention is learned identify some subspace. Only those tokens embedded inside these regions are pivotal for this attention. This would explain why only some specific tokens are always relevant.
Sequential Token Generation Under budget
In this section, we present SCISSORHANDS, which reduces the KV cache memory from the sequence
length dimension without fine-tuning the model. In Section 4.1 we describe how SCISSORHANDS
maintains the KV cache under given budget. Section 4.2 provides theoretical analysis of the algorithm and the approximation error.

---

4.1 Budget KV Cache for Single Attention Head
In this section, for the sake of the discussion, we drop the layer number notation i and batch size dimension. K2, V4 E Rexd denote for the KV cache until step t. Iy E R is a row vector that denotes the input attention at step t. The output of attention head at step can be written as,
exp((1 Ka(ll))
Ln an Vidle where at
d:
Σ expl ( (,))))))
Intuition. As shown in Section 3 the attention scores Ctt, follow a strong power-law distribution. For the autoregressive generation process, if there exists oracle such that we can identify the heavy score tokens before the future generation step, then the memory of the KV cache can be significantly reduced by only storing the heavy score tokens. Fortunately, the persistence of importance hypothesis provides us with such an oracle. It states that only historical tokens with significant contributions toward previous generated tokens will have significant contributions toward future tokens.
Challenges. I.LMs are deployed on hardware with a fixed memory. The algorithm should maintain the cache under fixed memory to meet the hard requirement. Further, LLMs are already computationally intensive. The algorithm should avoid introducing much extra burden on computation.
A fixed memory budget for one attention head is B tokens. In other words, we can store key and value embedding for B previous tokens. We describe the problem as follows,
Definition 4.1 (Sequential generation at attention head under budget B). Given stream of token embedding, including prompt and previously generated tokens, denotes their input to the head as St,....,tt,.-1 The problem of sequential generation an attention head under budget B is maintaining key cache Kr, and value cache V, such that Kr, V, E Ruxdd and n < B.
Approach, Inspired by the textbook solution of reservoir sampling and the Least Recent Usage cache replacement algorithm. SCISSORHANDS reserves fixed memory buffer for the KV cache. When the buffer is full, SCISSORHANDS drops stored but non-influential tokens from the cache. We present the main algorithm in Algorithm I and Algorithm 2
When the KV cache size exceeds the budget, SCISSORHANDS drops tokens from the KV cache according to Algorithm [ The importance record is a counter that indicates how many times a token deemed non- non-important. We choose attention scores as the importance indicators, following our methodology in Section 3.2 The importance record is collected over a history window to reduce variance. A higher counter suggests dropping from the cache. Recent tokens are always kept because of the lack of information on their importance by setting the counter for all tokens in the recent
window r to O.
With a sampled KV cache, attention output can be computed by the following estimator
n
ex
à,- di, Vrli], where de, E [xp
Overhead Tradeoff At the compression step, an extra attention computation is introduced to collect the importance measurements over a history window. However, such compression is not required at every generation step. controls the frequency, and steps after the compression have reduced attention computation because of the reduction in the KV cache. On the other hand, one can trade tiny amount of memory to avoid the overhead by maintaining the importance record during generation steps in Algorithm !
= 0.5B in experiment. Further,
Allocating Budgets Across Attention Heads. An LLM typically consists of L transformer layers where each layer has H heads. total memory budget has to be distributed over layers and heads. Within each transformer layer, the budget is distributed evenly across heads. Within the entire model, we distributed the budget according to Figure [ The rule of thumb is to allocate more budget to later layers to compensate for the lower persistence ratio.

---

4.2 Theoretical Analysis.
We study how much the tokens generated by the compressed KV cache deviate from the tokens generated by the original transformer using our simplified model in (1) Let { denote the tokens generated by the transformer with budget KV cache in Algorithm 2 with
_ = F (āz), where ā softmax (1e- W V.T Wo
Notice that when m 1, i.e., in each iteration, we drop one token with the lowest score, the cache will always maintain tokens. the ranking of the attention scores does not change in each iteration, Algorithm 2 will always drop tokens with the smallest attention scores.
For reference purposes, let (2re)=_ denote the tokens generated by vanilla transformer defined in We will bound the difference || ăzll2-
Theorem 4.1. Let X1, X2 denote the largest singular values of W and W2 in 2 Let
exp 2, W W.
Bes
E exp (1/- E, W.WWWRk
and assume that each Bi. CoL. where Vt, are sampled from power-law distribution with pdf f(x) c(x + b) Suppose that λv Ao + x1 dadK) I 3. Let Tmin and Tmax denote
the starting and maximum sequence lengths, respectively, and let B < Tmax denote the budget as in Algorithmm 2 If all E [Tmiu: mmax]. St contains only tokens with at most the largest B values of Bt.5- that is, [ = B and minjes, Beg 2 max maxjes, Be,,. then all e E (0, 1), with probability (Tmin -1)
T max exp _2(Tm 11)(11B)))mm) =/Tmae))) (1-1)²
at least Tiax exp
the following error
[k (R-2)[[ab)) b)²²
bound must hold for all 1. E [Tmin Tmax]
2.1(1 13 mm/m)
1
1/4 - 1)
E [1| 5z 3e112]
k k- (k (k)
(1 - e)²
B/TTuee -
The definition of Br, means the attention scores computed on the tokens generated by the compressed approach. Our theorem assumes that dropping the tokens depends on the attention score of the current iteration. (4) provided a bound on the expected difference between the tokens generated in the budget and the original approach. The upper bound scales with B/Tma.. When B = Tmax meaning
1/0 1
that we are keeping all of the tokens, the error becomes zero The term (k
B-C
depends the distribution that the attention scores fitted to and is always less than one. With strong power-law distribution, this term provides a further decrease to the error bound in
Empirical Evaluation
In this section, present the results that demonstrate SCISSORHANDS achieves up 5x reduction in the KV cache memory compared to the standard model with no accuracy loss. We also show that SCISSORHANDS is compatible with 4-bit quantization.
Experiment Setting. We compare the accuracy of SCISSORHANDS- OPT against the original OPT on one language model datasets C4 2 and number of few-shot downstream tasks: Hellaswag [77] [31 to evaluate few-shot tasks.
MathQA 28 PIQA 29], Winogrande (30) We use Im-eval-harness
Our experiments are conducted on NVIDIA 4 A100 40GB GPU servers.
No Accuracy Drop untill 5x. In Figure we present SCISSORHANDS' accuracy trend where 1x denotes the original OPT. In the language modeling setting, perplexity is the lower the better. For OPT-6B, perplexity is maintained until 50% of the original KV cache size for OPT-13B, For OPT-66B, perplexity is maintained until 75% of the original KV cache. We observe a flatter accuracy trend as the model size grows, which exceptionally encouraging. This suggests that SCISSORHANDS can scale with the model size. Downstream tasks are usually less sensitive to perturbation and bear more variance terms of accuracy. We evaluate the 5-shot setting and x denotes the original OPT model. For Winogrande and MathQA, accuracy is maintained even after 5x compression for OPT-66B. Similar to the language modeling setting, SCISSORHANDS performs better at larger models. Generally, accuracy is maintained with 15% 30% of the original KV cache size.

---

10.0
OPT-138
OPT-308
OFT-668
Perlexity
9.0
Hellaswag
MathQQ
8.0
PIDA
Winegrande
1X
2X
3X
4X
5X
1X
2X
3X
4X
5X
Compression
Compression
(a) Language Modeling
(b) OPT-6B Five shot
1.0
0.8
Accuraty
Accuracy
0.6
0.4
Hellaswag MathQA
laswag
MathQA
PIQA
PIQA
Winegrande
Winogrande
1x
2X
3X
4X
5X
1x
2X
3X Compression
5X
Compression
(c) OPT- 13B Five shot
(d) OPT-30B Five shot
Figure 3: Accuracy trend of SCISSORHANDS on language modeling dataset and downstream tasks with different KV cache compression. In general, SCISSORHANDS incurs accuracy drop until 5x compression on OPT-66B.
Ablation on the Importance of Pivotal Tokens We divide C4 into three subsets depending on the sequence length. C4-[256-512] contains data sequences that are longer than 256 tokens but less than 512 tokens. C4-[512- 1024] contains data sequences longer than 512 tokens but less than 1024 tokens. C4-(1024-2048) contains data sequences that are longer than 1024 tokens but less than 2048. Results are summarized in Table Local Windows refers to only keeping tokens in the recent window, while SCissORHANDSKeeps both recent tokens and pivotal tokens. We observe the perplexity of the full model degrades slightly with the growing sequence length. At all sequence lengths, SCISSORHANDS'S performance is comparable against the full cache model, while Local Window incurs a significant quality loss. This demonstrates that keeping the pivotal tokens is important to reserve model performance. It is also interesting note that at longer sequence lengths, the local window has higher accuracy. This also shows at longer sequence length, the attention mechanism in current architecture tends to focus recent context.
Table 3: Perplexity on C4 with different sequence lengths.
512)
[512- 1024)
[1024- 2048]
OPT-13B
8.7968
9.1017
9,3005
OPT-13B + Local Window
81,8297
29.3823
15,5883
OPT-T3B + SCISSORHANDS 8.7972
9.1011
9.3009
Compatible with 4-bit Quantization We test the compatibility of quantization and SCIS- SORHANDS at 2x compression. We adopt 4-bit quantization following [7]. Even Hellaswag is most sensitive based on Figure 3 adding quan- tization doesn' 't introduce compounded errors.
Table 4: Applying 4-bit quantization on top of SCISSORHANDS on Hellaswag.
OPT-6B
Original 0.702
SCISSORHANDS
SCISSORHANDS+ 4-bit 0.704
0.706
OPT- 13B
Ablation on Attention Score Error. We present the change ratio in attention score be- tween original OPT-13B and SCISSORHANDS OPT-13B at 3x compression C4 in Figure
Original 0.720
SCISSORHANDS 0.720
SCISSORHANDS+ 4-bit
0.720

---

We observe the attention score generated from SCISSORHANDS is almost the same as the original KV cache, which also echoes Theorem 4.1 The change ratio calculated as au where ag is the SCISSORHANDS attention score and ag the original score. From Figure 4 we observe that the change ratio is centered around 0. -I indicating that o, is significantly smaller compared to the original, suggesting that small portion of the important tokens are dropped in the cache. To explain the above observation of SCISSORHANDS we denote the number of tokens with the highest score as { 1o J=-- Then, for any other sets of tokens {21 that has no greater than n tokens, we can easily prove that similarity (5t Lop xt) z (2% 2.). Thus, SCISSORHANDS gives the most similar output as the original model all layers.
6 Discussion, Limitation, and Future Work
We discover repetitive attention patterns given trained language models. One interesting question that needs to be answered is whether such behavior is model architecture bias or an un- expected training outcome. For such purpose, we perform the same experiment with randomly initialized OPT, and compare it against the results presented in Section 3.1 As shown in Figure 5 the repetitive attention pattern does not exist in randomly initial- ized models. Apart from an efficiency deployment perspective, could such repetitive attention patterns contribute some known problems in language generation, such as repetitions? It may be worth investigating the relationship between repetitive attention patterns and undesired generations.
-1.0
-0.8
-0.6
-0.4
-0.2
0.0
Change Ratie
Figure 4: and SCISSORHANDS.
Score between OPT
Due to the limitation of the server in academics, the largest model we can fit is OPT-66B. We try to understand the behavior and ver- ify the generality across the differ- ent models and datasets. However, we cannot access the training pro-
(a) Attention map of the token at position 178
cess and fail to know exactly how
an LLM is trained to exhibit such behavior. Experiments with the large model create carbon dioxide emissions. However, work im- proves the efficiency of LLM, and we foresee no negative impacts.
(b) Attention map the token at position 228 Figure 5: We plot the attention map corresponding to Section 3.1 but with randomly initialized OPT. We observe repetitive attention for a randomly initialized model.
Conclusion
Inspired by our intriguing findings
that pivotal tokens exert a lasting influence on future steps, we developed SCISSORHANDS to leverage this observation to reduce the memory usage of KV cache. Our method achieves memory reductions of 5x the KV cache without compromising the performance of LLMs. Furthermore, we demonstrate the compatibility of SCISSORIIANDI with quantization techniques, opening up the possibility of reducing memory usage both the representation and sequence length dimensions.
Acknowledgement
We would like thank the anonymous reviewers for their helpful discussions and feedback. This work is supported by NSF-CCS-221 ONR-DURIP and NST-BIGDAATA-1888117
10

---

References
[1] Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportunities and risks of foundation models. arXiv preprint arXiv:2108.077 07258, 2021.
[2] Percy Liang, Rishi Bommasani, Tony Lee, Dimitris Tsipras, Dilara Soylu, Michihiro Yasunaga, Yian Zhang, Deepak Narayanan, Yuhuai Wu, Ananya Kumar, et al. Holistic evaluation of language models. arXiv preprint arXiv:2211.09110 2022.
[3] 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, 13:1877-1901, 2020.
[4] Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer. Rethinking the role of demonstrations: What makes in-context learning work? arXiv preprint arXiv:2202. /2000.22880 /2837, 2022.
[5] Stephanie CY Chan, Adam Santoro, Andrew Kyle Lampinen, Jane Wang, Aaditya Singh, Pierre Harvey Richemond, James McClelland, and Felix Hill. Data distributional properties drive emergent in-context learning in transformers. In Advances in Neural Information Processing Systems, 2022.
[6] Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury. Anselm Levskaya, Jonathan Heek, Kefan Xiao, Shivani Agrawal, and Jeff Dean. Efficiently scaling transformer inference, 2022.
[7] Ying Sheng, Lianmin Zheng, Binhang Yuan, Zhuohan Li, Max Ryabinin, Daniel Y. Fu, Zhiqiang Xie. Beidi Chen, Clark Barrett, Joseph E. Gonzalez, Percy Liang, Christopher Ré, Ion Stoica, and Ce Zhang. High-throughput generative inference of large language models with a single gpu, 2023.
[8] Zhewei Yao, Reza Yazdani Aminabadi, Minjia Zhang, Xiaoxia Wu, Conglong Li, and Yuxiong He. Zeroquant: Efficient and affordable post-training quantization for large-scale transformers. arXiv preprint arXiv:2206.012 01661, 2022.
[9] Gunho Park, Baeseong Park. Se Jung Kwon, Byeongwook Kim, Youngjoo Lee, and Dongsoo Lee. nuqmm: Quantized matmul for efficient inference of large-scale generative language models. arXiv preprint arXiv:2206. Xiv:2206.009577, 2022.
[10] Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Llm. int8 O: 8-bit matrix multiplication for transformers at scale. arXiv preprint arXiv:2208.07333 2022.
[11] Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. Gptq: Accurate post-training quantization for generative pre-trained transformers. arXiv preprinl arXiv:2210.1777 173233, 2022.
[12] Elias Frantar and Dan Alistarh. Massive language models can be accurately pruned in one-shot. arXiv preprini arXiv:2301.00774 2023.
[13] Hritik Bansal, Karthik Gopalakrishnan, Saket Dingliwal, Sravan Bodapati, Katrin Kirchhoff, and Dan Roth. Rethinking the role of scale for in-context learning: An therpretability-basee case study at 66 billion scale. arXiv preprint 1rXiv:2212.09092 2022.
[14] Guangxuan Xiao, Ji Lin, Mickael Seznec, Julien Demouth, and Song Han. Smoothquant: Accurate and efficient post-training quantization for large language models. arXiv preprint arXiv:22 1. 10438, 2022.
[15] OpenAl. Gpt-4 technical report, 2023.
[16] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In 8th International Conference Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
[17] Sinong Wang, Belinda Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint 2rXiv:2006..4777 2020.

---

[18] Beidi Chen, Zichang Liu, Binghui Peng. Zhaozhuo Xu, Jonathan Lingjie Li, Tri Dao, Zhao Song, Anshumali Shrivastava, and Christopher Re. Mongoose: learnable Ish framework for efficient neural network training. In International Conference Learning Representations, 2021.
[19] Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, and Christopher Ré. Scatterbrain: Unifying sparse and low-rank attention. Advances in Neural Information Processing Systems, 34:17413-17426, 2021.
[20] Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreca Gane, Tamás Sarlós, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J. Colwell, and Adrian Weller. Rethinking attention with performers. In 9th International Conference Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
[21] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness, 2022.
[22] Noam Shazeer. Fast transformer decoding: One write-head all you need, 2019.
[23] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with unified text-to-text transformer. arXiv e-prints, 2019.
[24] Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, Todor Mihaylov, Myle Ott, Sam Shleifer, Kurt Shuster, Daniel Simig, Punit Singh Koura, Anjali Sridhar, Tianlu Wang, and Luke Zettlemoyer. Opt: Open pre-trained transformer language models, 2022.
[25] Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can suit of armor conduct electricity? new dataset for open book question answering. In EMNLP, 2018.
[26] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016.
[27] Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can machine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, 2019.
[28] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog. 1(8):9, 2019.
[29] Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language. In Thirty-Fourth AAAI Conference on Artificial Intelligence, 2020.
[30] Sakaguchi Keisuke, Le Bras Ronan, Bhagavatula Chandra, and Choi Yejin. Winogrande: An adversarial winograd schema challenge at scale. In Communications of the ACM, 2019.
[31] Leo Gao, Jonathan Tow, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Kyle McDonell, Niklas Muennighoff, Jason Phang, Laria Reynolds, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. framework for few-shot language model evaluation. In Version vO. 0.1. Sept. Zenodo, September 2021.
12

---

Appendix
A More Results
A.1 Repetitive Attention Pattern
We provide the attention map similar Figure but from different transformer layer on the same text in Figure 6 Figure 7 Figure and and Figure A repetitive pattern and attention sparsity can be observed across layers.
AStenOAn Head
Sequence Posilion co
Attention map position 178
Adtintionn Head
(b) Attention map at position 228
Attendoon
Sequuarce ca PosiBon
(c)
Attention map position 278
Figure 6: Attention Map Layer
Attandenn Head Head
Sequence Position
Attention map position 178
Manurd Road 1
(b) Attention map at position 228
Attention Head
Sequercc Positio
(c)
Attention map position 278
Figure 7: Attention Map Layer 10
13

---

Sequencc
(a) Attention map position 178
Aduntion Hovd
Position
(b) Attention map at position 228
Artwerkin Head
Seqqueica Posidon
(c) Attention map at position 278
Figure 8: Attention Map at Layer 15
Atterttion Head
(a) Attention map at position 178
Atunndinn Head
228266 010 1315
Sequenne Resttion
(b) Attention map at position 228
Attanton
NOR SURRRU SNNN 0. 20000 Sequeence
(c) Attention map at position 278
Figure 9: Attention Map at Layer 20
A.2 Cross Layer Cosine Similarity
In Section 3.3, our analysis assumes a large cosine similarity between the input and output F. Here, provide empirical evidence to support such assumption in Figure I0 Because of the residual connection in and the domination æ, the cosine similarity between x and F(x) is extremely high.
A.3 Generated examples with SCISSORHANDS
14

---

Residual At MLP
1.00
Cosine Similarity
0.98
0.96
0.94
Cos(x, 7(x))
0.92
0.90
0.88
0.86
5
10
15
20
Transformer Layer
(a) Cosine Similarity
Residual MLP
120
100
Norm
||||
60
[[W,Relu(Wx3]]
40
20
0
5
10
15
20
Transformer Layer
(b) Norm Comparision
Figure 10: x and F(x) is high in cosine similarity
Table 5: Generated examples using OPT-13B with full cache and SCISSORHANDSal different compression ratio.
Does Asus ever plan to update the intel and nvidia display drivers on the Ul UTV laptop? The only downloadahle drivers from the Asus website are at least years behind and new features for both graphics chips have been released since then. I sent message to Asus but no response was given. Intel and Nvidia will not allow the user to download the new drivers directly. problem is the drivers won't install. states something about your notebook needs custom drivers from the manufacturer due to customized this that or the other.
Prompt
have the same problem. have a laptop with an Intel HD Graphics 4000 and an Nvidia GeForce GT 740M. can't 't install the drivers from the Asus website. have to use the drivers from the Nvidia website.
OPT-13B
have the same problem. have laptop with Intel HD Graphics 4000 and an Nvidia
3x Compression
GeForce GT 740M. can't 't install the drivers from the Asus website. have use the drivers from the Nvidia website.
Compression
have the same problem. have custom laptop and can't install the drivers.
B Proofs
B.1 Proof of Theorem 3.1
We consider the token generation process of a simplified model: single-layer transformer model with single-head attention.
Et+1 F (az). where a; = sof tmax 2/ X
(5)
St E 1xd is a row vector. X,-1 E (t(1))xd denotes the aggregation of .., -1 where
the jth is Tz. Wo, WK, Wy E Rdxp and Wo E Rpxd are the attention weights. Lastly, F IR 1xd R lxd denotes the MLP block following attention block, two-layer MLP with skip connections, given by
F(e) x + Worelu(W1z)
We are interested in the attention scores = sof tmax 2 WXT,). Notice that ot, scales with x, W WoW.WWWT. x We first re-state the Theorem 3.1 below.
15

---

Theorem B.1. Let = WyWoWaWk K and let Xs. Xq. Xv, Ao denote the largest singular values
WK Wo Wv, Wo. respectively. Consider the transformer 5 with normalized inputs ||zell2
for all t. Let c, 0 be constants. Assume that a 2 (1 6) ||aell2 with ô S (AQAAAAO
Then for all me satisfying xtA 2 c. and Axz c maxxE(() Ax it holds that
x (ar,F 3t) S Fe4 WoW S JeAu (ur, + 3e)
10el2
||ell2
As preparation of the proof, we first show two lemmas. Lemma B.1. Let T2, E Rxxm satisfies || l12 ||
1 and xtaa 2 1 - 8 for some
ó E (0, 1). Then for all y E R2 we have
byy x2y T
22 11y12
Proof. Let Ø2 x22 + Z2 where
22 x132 21;
22
22 22
Then it is easy to see that x .
0. By the Pythagorean Theorem, we have
| 1 = || 2
!1 5(2 - 6)
Therefore, we have
(1) - x.!)) 2t 2
|| x2 1
[| - x ) x1 - x3 1|
(1 - x 2 + . ?
2
25
Thus, the Cauchy-Schwarz Schwarz inequality implies |awn xay S || x2l12 |1u112 22 ||
Lemma B.2. Let E [6] be given. Suppose that x4 Au, - -=, for all j # l. Then we have (8(1), e) xf axe S ot Wx Woa: S (S(1)e + +e) x axe
Proof. Notice that
a: a= xx=
Σ aijoy WrWo
Thus, we have
Σ aisTy
WoWoWoWKet
Σ A
Therefore
Σ a
A
j=lift
b-1
Σ
Ot |.
#=1
$ coeANT
Σ ar
1=1
S CoAART
16

---

where in the second inequality we use e-1 |x S Az and in the third inequality we use Ej= otj S Ej=1 onj = This implies that
(on,e - Axy S aaWo S (ou, + ae
Now we proceed to the main body of the proof. Assume that ||ze||2 1 for all t. Using Lemma B.1 if 4,2/41 2 (1 6) ||ail12. then we have
||aella* at aoWo Jeti W & 25 WWW R ll
Recall that λg, Ax are the maximum singular values of Wi and WK. respectively. Then it holds that [ WT 2 1l2 S A ||æelll Using |æell22 1, we have
||alll" a.WoW Ie+i WetiWoWKEY
S V20dodx
Notice that
1 Cnjzit
1at112
Wr
$ voov
I λoyv
On |
Then since 8 < (oodkAndo)) )²
we have
200
206.
|aaelle" aeWoW 2441 WotWoo S Avdo
||ae|12
Since by Lemma B.2 we have
a loWo Au. exf axı
must hold that
200
Fet1 Wo Wo W xz |at+1l1
S l|al11
I2 ax +
||a ll
Since xj aze 2 c, holds that
20E
2e
<
x aæ
|1ol|12
|
which implies that
36
F4+1 W
Wo - 1laell
ax. <
ot ar
||a:|12
Therefore
AxA |0tll2
(ar,e 3c) S x S
s
(au, + 3c)
|alll
17

---

B.2 Proof of Theorem 4.1
Let {#r] ) denote the tokens generated by the transformer with budget KV cache in Algorithm with m = 1:
- = F (āz), where å tmax (v/- (V--,, W. V.T Wo
Notice that when m i.e. in each iteration, we drop one token with the lowest score, the cache will always maintain B tokens. If the ranking of the attention scores does not change in each iteration, Algorithm will always drop tokens with the smallest attention scores.
For reference purposes, let {2x} ) denote the tokens generated by a vanilla transformer defined in
5
We re-state Theorem 4. below, which bounds the difference ||,, -
Theorem B.2. Let A1, 22 denote the largest singular values of W and W2 in 1 Let
exp (1/1 2 W
Brs E exp (1/- (.
and assume that each = where Vt, are sampled from a power-law distribution with pdf f(x) = c(x + b) Suppose that Av (1 \Φ < 3.2 Let Tmin and Tmax denote the starting and maximum sequence lengths, respectively, and let B < Tmax denote the budget as in Algorithm If for all E [Tmin: Tmax). St contains only tokes with at most the largest B values of Bt. that St| B and minyes, Peg 2 maxjes,
Be,,. then for all e (0, 1), with probability
at least T max exp (-ik- ( -1) Timx exp _.2Tmin-- 1)(1 1/1-B//mmas) the following error
(2-22)((-
(1-1)²
bound must hold for all 1. (Tmin; Tiaa]]
2.1(1
B/ B/Tmax)
1-
1/1k
E [1|T - 84 1211
k- (ki 1)
1--r c)2
B
B/TTns
Define may = Iije Se). With the definition of my.j. āg can be written as
my.j exp (1/2 3e W.
āz Σ ور
WvWo; ā E ma, exp (1/- R. WoW
(8)
Our first lemma shows the Lipschitzness of the attention module.
Lemma B.3, Consider two sequences of tokens {z,Y.-1 and {vi Vi where | ll2 llvell2 = 1 for
all i E [t]. Define Y R(t-1)xd as the matrices whose ith row are x; and ys. respectively.
Let A = |xxe - yell2- Then we have
(!uWoWWrZI))
52 -
softmas (Z=WWWWXXT.. -softmax
Proof. We can decompose the difference
(ZuVomErz
softmax
softmat toft
softmax (HaWWOWAXI.))
(HoMOWaYZ.))
softmax
+ |sottmxx (Ho.WoHYY..)) softmax
By the Lipschitzness of softmax, we have
(frWWWWZII) softmax (ooWaWArZ.))
softmax
Wo K (Xi-1 Y=-1)
< Nade || |X,,1 Yr-1l12
18

---

Y (Σ ||, -vll2) S vl A. we have
Since ||xx|12 = 1 and -
Vt
tmax (z,
softmax (2.WoW
Aa
t
Similarly,
softmax (fo-oWWxZI)) - softmax (HuwOWaXI
1
[ (0,, 1 (lal-y)aaWkk
1
AQAKK |||--l1e 1|=, yell2
Since ||x; - yell2 = D. and ||c-1ll22 = -vt-1 we have
sof tmax ( (
softmax (:uWoWWrY..) )
Vt-1
S AadRa.
Combining the two bounds gives
softmax
Vt Io
1
softmax
(Hun
)/
22 22 AadxAr
t
Our second lemma shows the difference between the output of the sampled and vanilla transformer when the input is the same.
Lemma B.4. Let āg be defined as in Define by as
exp (1/2 8z WoW
b, Σ Bi WyWo: Bus E exp (2/r 3, W WoW
(9)
Assume that ||ا,ع|2 for all j [0] Then we have
l|ā - be S Avño Σ Bus
j45.
Proof. A direction computation yields
ā, by Σ (ār. Bes)
WvWo
Thus, ||ū be bell2 can be bounded
lā= - 1bll2 S \v Σ (ār. - BBG)) || 112 \v Σ (āe. Be,s)
1=1
since || 1 for all j E [t]. Now we analyze ãz,i Br Let S = St \ {t). Then maj 1
and only if j E S., For convenience, let rt,j = 1/c T, W W
W J... Thus, 3 can be written as
exp (rt,))
Bis EiEs, exp (Tr,)) Eips, exp (rr,))
Furthermore, for all j & S, we have ätj 0. For all j 8r have
exp (r
āi, EiEs, exp (re.)
19

---

Therefore, for all j E St, we have
Eigs, (T(,)
Big k,, exp (xp()())
(Eies, (xp(r)))
(EiEs, exp (rea) + Eigs, exp (xp(re.))
exp (meg)
Ligs, exp (Fr,s)
Eies. exp (rea) Eies, exp (re,) + Dies. exp (r1,)
Cn, Σ Brg
igS.
Therefore, the bound of l|ār
bell2 can be written
[-1
lā baell S \v
Σ āi. Σ
Bls
Σ
Bus
=2
E
Bas
SES,
$4
J4S.
jES.
where the last equality follows from Ejes, (tt,j =.
Our last lemma shows the Lipschitzness of the MLP in (6).
Lemma B.5. Let d1, 22 denote the largest singular values of W1, W2 (0. we have
For all æ 22 IR..
- - s (1 A1 - =22ll2
Proof. Direct computation yields
[[F(x1)) - F(x2)) = [[( + W2relu (W1z1)) (m2 + W2relu (W1x2))
< | x2ll2 |Werelu (W1z)) W2relu (Waæ2)]
< | z2ll2 + λ2 ||rel (Wiz1) relu (W1æ))
S || x2l1, A2 |[] (x) 1 x2)]l2
< || - 22ll2 + Åιλ ||||| x2l12
(1 + \1 |||| - x2ll2
where in the third inequality use the fact that relu(-) is I-Lipschitz.
Now we turn to the proof of our main theorem. Combining all of the results, we have
t-1
ag āN
Σ aj
WyWo- -
Σ
WvWo
Dj
WyWo-
Σ و
WvWo
T.
atzz
WrWo
Σ
Bijez
WvWo
1-1 Σ
+
Σ
Wy -
WrWo
Ta
Therefore, by triangle inequality, have
||a. - S ||T + ||Talla + ||Tall2
(10)
20

---

To start, the magnitude of T1 can be bounded as
7
O1y Bes))
Wv wo
On Ong(are Fee)
S
S Avvo Σ
|1 and
Tesll2
S Av
Σ
Oth
j=1
where in the third inequality we use || - fe A. and in the last equality we use
Ej=1 Dn. = 1. To bound the magnitude of Tz, we apply Lemma B.3 which shows that
|00 Bell $ 2 - A to get that
1-
12ll2
Σ (ar.) - Beg)aaz
WyWo
1-1
$ Avdo
Σ
(ar. - Be
j=0
t-1
S \v Σ
Jan.
Br,l ||
j=0
Avdo ||, Bell
s vl T\v 1| - Bell2
$ 1 1 1) Aadxxvdo
Lastly, to bound the magnitude of Ta, we use Lemma B.4 to get that
Tall S 2λv Ao
Σ Bes
j4S.
Putting things together for 100) we have
||a arll2 S Avdo
2 Σ
Bes (2λrAK +1
jES
By Lemma B.5 we can further show that
| iet1ll2 S A1de)AAoo
2 Σ Brg + (2λrAx + 1)
jæs,
(-
(I
By Theorem
B.3 we have that with probability at least - Tmax exp
(k (k-2)
Trax exp
2 Trin-
D (11--/m B ///ma)))) it holds for all t E [Tmins Tmax] that
(1-))²
(1 B B/Tmax)
F
E
Σ
Beg
5
k: (k: 1)
max
0.98(1 c)²
B B/Tmux
jes,
21

---

Given that E 1| - Telll S 24 max. we have
E [1 - Bet1llel S (1 + Au oo (2A + 2 222 +1) Amax)
S 4Av + A1 )
Thus, as long as Ao + (1 AQAx) S 2.2 we can guarantee that
E [1 - B2t1ll21 - 24 max
Thus, for all E Timin Tmax], we have that
2.1(1 B/Tmax)
1- k-1
E [1a Txl12]
k: k- (k- (k(--1)
(1 6)2
B B/Tax
B.3 Budgeted Cache
Theorem B.3. Let Bi, be sampled from some power-law distribution f(x) cx+ with support
on [0, u b) for some k and 5b. Let St be defined in Theorem B.2 and define St = {1).
Then with probability al least Tmax exp (- (k- 2) (u-b)- b) holds for all E T that
2bb (Tain -1)
T max exp
2[T 1)(1 13)
(1-c)²
Σ Se
BL,
S
B/Tama) 0.98(1 - e)²
k: h- (k- (6l--)
B/Tnnre -
B/1
-
T--
(11)
We consider the case of maintaining a budget of B by dropping the smallest Br,j 's. Assume that has pdf f(x) = c(x + b) -k with support on [0, u b). To make things precise, we first compute c
(xr + b) =kdx
g1-k- k u
To start, we notice that
f c0l b)-
(x + b) 1))->>((( ((k: 1)x + b)
= g(x)
(k- - 1) 2)
Let C Σj=1 Uy. then the expectation C is
k 1
E [[[] (L [m (1 11-b u1 1-
x(x b) kdx
k: -
(t 1)
f - w (g(u) - g(0))
k:
b²
ul ul-x((6 1)u (k 2)b)
= (t 1) 11 - wl-k (k - 1) -
(A: - 1)(k 2)
t b2-k k (k: 1) 1)u2 2- k + (I: 2)bul1k
k
b1ーk w
62-k -2- (K-
2- 2- (kk-
Let A
By Hoeffding's inequality, we have that
b -t
2c (C)²
PCC < (1 -E exp
(T ( 1)(u b)2
2r 20- have -Tk=2)((x-03>
This implies that with probability at least 1 eeepp
exp
t 1-1
C
(1 c)
k
22

---

Now, we proceed to bound j, Pes where S, = {j E [t - 1]: Brs 2 7). Equivalently, we can
bound C-1 Σ {uz 1 } ej. Its expectation is given by
[-1 I Ho 13
k - 2 E
Σ
Hes S
(t (t- 1) -)
j=1
j=1
A(1 - c) b1-k w - ['' x(x b) &
k - 2
k- -
(k- 1) 2)
A c) (b1-k -1- ul k) (0(7) g(0))
We pause here and study how small can we choose 7. Notice that
bl-k (7 (7+ (u, S 7) = (t - ()) 1).
E
bi-k - u
By Hoeffding's 's inequality again, we have that
1-1
2 (1- (1-e)(t- 10) 1)- bl-k - (8 + b)1
Σ Hu
62- w w
2(1- - 1)6² (b²-k - (7 + b))
Sexp
S
exp
(b1-k u¹ x²
Enforcing T < 2 Tmax
B gives
(7 b)1-k - & bl 1-B/rmas (bi-k x--k),
I-c
(/mama I= -1) b. Therefore
which can be satisfied as long as y 2
9(7) bi-kk - BBB B/Tmoe (bi-k - 2 )
b+ (k- - 1)7
(k 1) 2)
We further notice that
1 - B B/Tmer (bi-k u x²-k)² B/Imer - E [bl- - u ut-
b 1-k
1
This gives
I b(1 B/Tmna)) A(1 c)²
(k: TT/m -)
E
C
Tfu S 13
A e1-
6(1 B 3/Tmax)
1 1- B B/Tmaar
h: (k (k-1) 1)
A1 +)²
Notice that if u 2 5b, we have
11 1-k A = b- (Tm
b-
J u wl-k 5 0.986
Therefore
(1 B /mmex)
1 - e B B/Traxx
E c
Σ
(op S 0
k - (k: 10-11 1)
0.98(1 - c)2
226 (t- 1) (- Tk- (k-22)(a-6))
2( 2(t-1)(1 11 M/Tma)) )2) (1-c)²
holds with probability at least 1 exp union bound gives the desired result.
-epp
Taking
23


Reviews:
[{"review_id": "i2fI3bW1z0", "reviewer": "Reviewer_h9xt", "summary": "This paper presents a new KV cache mechanism that uses an attention importance score to drop out tokens from the cache in transformer models. This allows up to 5x compression of the cache without degradation on assessed tasks.", "questions": "+ Are there any other baselines you could compare against to give readers a reference point for determining whether to use this approach or another similarly situated approach?\n+ The evaluation appears to focus on 4 datasets: HellaSwag, MathQA, PIQA, and Winograde. But I wonder whether the lack of performance degradation is a function of these particular datasets. For example, wouldn't a more adversarial dataset where there is an  fact at the beginning of the text that is unimportant until the very final token generations? By the time you need to re-use the fact, it will have dropped out of the cache. The paper would be improved with a more adversarial evaluation which shows the limitations of this approach to give users a more realistic picture of its limits. Or is it as the paper hints at in the conclusion, it doesn't matter because the architectural bias would prevent the attention mechanism from attending to those tokens anyways even if they aren't dropped out?", "limitations": "See above in weaknesses and questions.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "This is likely to be useful to practictioners and researchers alike by making large model and long context inference more accessible on commodity GPUs. With appropriate discussion of limitations and caveating of claims, this should be a useful addition to the increasing toolkit for efficient use of large language models up to the 66B parameter scale.", "weaknesses": "+ The evaluation appears to focus on 4 datasets: HellaSwag, MathQA, PIQA, and Winograde. But I wonder whether the lack of performance degradation is a function of these particular datasets. For example, wouldn't a more adversarial dataset where there is an  fact at the beginning of the text that is unimportant until the very final token generations? By the time you need to re-use the fact, it will have dropped out of the cache. The paper would be improved with a more adversarial evaluation which shows the limitations of this approach to give users a more realistic picture of its limits. Or is it as the paper hints at in the conclusion, it doesn't matter because the architectural bias would prevent the attention mechanism from attending to those tokens anyways even if they aren't dropped out?\n----\nEDIT: Based on discussion with authors, conditioned on the inclusion of an experiment/discussion showing the limits of this approach, I have increased my score.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank the reviewer for the supportive comments on the significance of our work. We have carefully considered the reviewer’s constructive suggestions and updated the manuscript accordingly.\n  \n\n**Q1: The evaluation appears to focus on 4 datasets: HellaSwag, MathQA, PIQA, and Winograd. But I wonder whether the lack of performance degradation is a function of these particular datasets. Are there any other baselines you could compare against to give readers a reference point for determining whether to use this approach or another similarly situated approach?**\n\n  \n\nTo better understand  Scissorhands’s performance and limitations, we performed three sets of additional experiments. We first investigate the performance with different sequence lengths. Second, we provide generated texts at different compression ratios for more intuitive illustration. Third, we provide results on three additional datasets.\n\n - We take C4 and divide it into three subsets depending on the sequence length. C4-[256-512] contains data sequences that are longer than 256 tokens but less than 512 tokens. C4-[512-1024] contains data sequences that are longer than 512 tokens but less than 1024 tokens. C4-[1024-2048] contains data sequences that are longer than 1024 tokens but less than 2048 tokens. In addition to the full cache and Scissorhands, we implement Local Window. Local Window only keeps the neighbor tokens’ cache. All numbers are measured at 3X compression. Results are summarized in Table 4.\n    \n\n\tWe observe the perplexity of the full cache model degrades slightly with the growing sequence length. At all sequence lengths, Scissorhand's performance is comparable against the full cache model, while Local Window incurs a significant quality loss. This demonstrates that keeping the pivotal tokens is important to reserve model performance.\n\n\tIt is also interesting to note that at longer sequence lengths, the local window has higher accuracy. This also shows at longer sequence lengths, the attention mechanism in current architecture tends to focus on recent context.\n\n- We provide generation examples with OPT-13B. We compare the generation between full cache and Scissorhands at different compression ratios. The results are provided in Table 7.\n    \n\n\tWe observe exact same generation at 3X compression. At 6X compression, the generation starts to differ from the original model in the second sentence. Even though there is a degradation in generation quality at 6X compression, the generation is still coherent.\n\n- Generally, language modeling perplexity is a good indicator for downstream tasks. Following a recent paper on LLM efficiency [1][2][3], we provide results on three additional datasets in Table 6. All accuracy is measured at 3X compression. Table 5 and Table 6 show that Sccissorhand provides comparable accuracy to the full cache model.\n    \n\n[1] Frantar, Elias, et al. \"Gptq: Accurate post-training quantization for generative pre-trained transformers.\" arXiv preprint arXiv:2210.17323 (2022).\n\n[2] Xiao, Guangxuan, et al. \"Smoothquant: Accurate and efficient post-training quantization for large language models.\" International Conference on Machine Learning. PMLR, 2023.\n\n[3] Liu, Zichang, et al. \"Deja vu: Contextual sparsity for efficient llms at inference time.\" International Conference on Machine Learning. PMLR, 2023."}, {"review_id": "Bunz9pa3us", "reviewer": "Reviewer_UGwp", "summary": "The paper introduces a new way to compress KV cache for LLM serving. It showed the proposed approach can reduce the memory usage for inference up to 5 times. However, there is no evaluation on the performance on the generative datasets to justify the claim that there is no loss in the serving accuracy. ", "questions": "Multi-query attention (K V states with only one head) already becomes commonly used for production LLMs. The KV state with MQA is  (1/num_heads) smaller than the one used for the old multi-head attention used in OPT and LLAMA. I suggests authors use MQA for the baseline.\n\nhttps://arxiv.org/abs/1911.02150", "limitations": "The motivation of this approach is for efficiency deployment. However, the authors need to do a better job in showing that in the evaluation and experiments session. There was just a mismatch in the experimental setup (like choice of dataset and benchmarks) and the claims.\n\n", "rating": 3, "confidence": 5, "soundness": 1, "presentation": 2, "contribution": 1, "strengths": "Ideas and math are well written. The direction of using pivotal tokens could be interesting.\n", "weaknesses": "The authors need to do a better job in evaluation to justify the contributions. \n\n\n1. There is no benchmarks on the inference efficiency (latency and throughput) with the new approach. Was inference efficiency affected by the compression ratio?\n\n2. The motivation of the method is to improve serving, especially allowing larger batch size for online serving of LLMs. However, the tested datasets are classification tasks, which only required one single forward pass without involving the decoding loop. We don't know how this method affects the performance of generative datasets (which involved decoding loops), like those question & answers ones as well as the instruction tuning ones. \n\n3. KV state is large only when sequence length is very long. However, the tested dataset have a very short sequence length. \n\n4. Please show the table with numbers instead of figures when claiming no loss in accuracy. In LLM evaluation, every single digit matters.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are glad that the reviewer finds our hypothesis and method interesting. We have tried to address your questions carefully. We hope the reviewer can consider raising your score in light of our response.\n\n  \n\n**Q1: There are no benchmarks on the inference efficiency (latency and throughput) with the new approach. Was inference efficiency affected by the compression ratio?**\n\nThe main focus of our work is on memory usage. We agree that compressing the KV cache also influences inference efficiency, which is worthy of investigation. Our method saves the attention computation because dropping tokens from the KV cache effectively reduces the sequence length. On the other hand, the overhead we introduce stems from identifying which token to drop. The overhead is rather limited compared to the saving.\n\n  \n\nWe provide throughput measured in tokens per second in Table 8. We based the throughput measurement based on FlexGen[1]. We use a prompt length of 512 and a generation length of 1024. For all other settings, we keep the default from FlexGen. As shown in Table 8, we observe a slight improvement (1.2X ) in inference throughput. Due to the rebuttal time limitation, we could only implement a basic version. Implementation can be further improved to realize a significant gain in both memory usage and throughput.\n\n  \n\n[1] Sheng, Ying, et al. \"High-throughput generative inference of large language models with a single gpu.\" arXiv preprint arXiv:2303.06865 (2023).\n\n  \n\n**Q2: We don't know how this method affects the performance of generative datasets like those questions & answers ones as well as the instruction tuning ones. KV state is large only when the sequence length is very long. However, the tested dataset has a very short sequence length.**\n\n  \n\nTo better understand  Scissorhands’s performance and limitations, we performed three additional experiments based on the reviewer's suggestion. We first investigate sequence length’s influence on our method. Second, we provide generated texts at different compression ratios for more intuitive illustration. Third, we provide raw numbers on the tested dataset and results on three additional datasets.\n\n- We take C4 and divide it into three subsets depending on the sequence length. C4-[256-512] contains data sequences that are longer than 256 tokens but less than 512 tokens. C4-[512-1024] contains data sequences longer than 512 tokens but less than 1024 tokens. C4-[1024-2048] contains data sequences that are longer than 1024 tokens but less than 2048. We summarize the perplexity in Table 4. In addition to the full cache and Scissorhands, we implement Local Window. Local Window only keeps the neighbor tokens’ cache. All numbers are measured at 3X compression.\n    \n\n\tWe observe the perplexity of the full cache model degrades slightly with the growing sequence length. Scissorhands' performance is comparable to the full cache model at different sequence lengths. This is because Scissorhands could identify the pivotal tokens. On the contrary, there is a significant quality loss if one only keeps recent tokens.\n\n- We provide generation examples with OPT-13B. We compare the generation between full cache and Scissorhands at different compression ratios. The results are provided in Table 7.\n    \n\n\tWe observe exact same generation at 3X compression. At 6X compression, the generation starts to differ from the original model in the second sentence. Even though there is a degradation in generation quality at 6X compression, the generation is still coherent.\n\n- We set up our experiment datasets following the paper focusing on efficient inference [1][2][3]. We also test language modeling perplexity, as it is generally a good indicator of downstream task performance. We would also like to clarify that the computation during the prompting stage is calculated as if the key/value cache is dropped when obtaining accuracy.\n    \n\n\tIn Table 5, we provide accuracy numbers, as requested by the reviewer, for the datasets we used in the experiment section. In Table 6, we also provide results on three additional datasets. All accuracy is measured at 3X compression. Table 5 and Table 6 show that Sccissorhand provides comparable accuracy as the full cache model.\n\n[1] Frantar, Elias, et al. \"Gptq: Accurate post-training quantization for generative pre-trained transformers.\" arXiv preprint arXiv:2210.17323 (2022).\n\n[2] Xiao, Guangxuan, et al. \"Smoothquant: Accurate and efficient post-training quantization for large language models.\" International Conference on Machine Learning. PMLR, 2023.\n\n[3] Liu, Zichang, et al. \"Deja vu: Contextual sparsity for efficient llms at inference time.\" International Conference on Machine Learning. PMLR, 2023.\n\n  \n\nQ3: Multi-query attention is (1/num_heads) smaller than the one used for the old multi-head attention used in OPT and LLAMA. I suggest authors use MQA for the baseline.\n\nThanks for the suggestion. MQA has indeed become a popular choice for production LLMs. However, adopting MQA requires training or finetuning the LLMs. Our work focuses on inference, where we are given trained LLMs."}, {"review_id": "OJxOPDBRBF", "reviewer": "Reviewer_g2hZ", "summary": "This work aims to increase the throughput of inference system by reducing the size of the KV cache.\n- The KV cache comes from saving forward pass computations and can take up up to 5x the model size.\n\nThe paper proposes to discard part of the KV cache, based on the observation that the set of tokens that get high attention are largely shared across positions, which the paper calls the Persistence of Importance Hypothesis.\n- The proposed method only keeps the most recent tokens and \"pivotal tokens\", which are tokens whose attention scores are higher than a threshold. Empirical results using OPT show - T\n- Theoretical analyses (Theorem 4.1) show that such process doesn't affect the attention output much.", "questions": "- The main motivation of the paper is to improve throughput, and the paper has mostly discussed the improvement in terms of memory. What's the improvement on wall clock time?\n- The current scheme of selecting pivotal tokens is attention-based. Are there other methods / baseline to compare to?\n  - e.g. Keep the recent window as the current method, and change the selection of distance tokens based on word frequency?\n  - e.g. Instead of thresholding the token importance, another way could be to select tokens that contributes to e.g. 90% attention weights, similar to Nucleus sampling; how do you think this would compare?\n- Results in Fig 2: I'm not sure $|S_{0\\rightarrow t}|$ can be considered as \"considerably smaller than $t$\" based on Fig 2 (b).\n- Clarification on $\\alpha = \\frac{1}{t}$: for $S_{0 \\rightarrow t}$, does then mean for $\\tau \\in [t]$, $S_{\\tau}$ is calculated with $\\frac{1}{\\tau}$ or $\\frac{1}{t}$?\n- Equation 1: what is there a scaling of 1/t within the softmax?\n\nWriting:\n- The paragraph starting at line 202 could be made clearer.\n- Line 238: $w,r$ have been defined in Alg 2; have they been defined in the main text?\n- Minor point: Fig 3: might be better to use different line style / color for models / tasks.\n- Please fix typos, e.g. line 202, 263, 295, 308.\n", "limitations": "There is no direct societal impact.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The problem is well motivated and the solution makes intuitive sense.\n- The empirical results are promising:\n  - The proposed method achieves up to 5x compression and doesn't sacrifice much performance.  \n  - The proposed method is compatible with 4-bit quantization.\n", "weaknesses": "Some design decisions are not well justified, and there are insufficient baseline comparisons.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nWe appreciate the reviewer’s thoughtful suggestions and supportive comments. We have updated our manuscript to provide further clarification and improve writing details.\n\n**Q1: The paper has mostly discussed the improvement in terms of memory. What's the improvement on wall clock time?**\n\nThe main focus of our work is on memory usage. We agree that compressing the KV cache also influences inference efficiency, which is worthy of investigation. Our method saves the attention computation because dropping tokens from the KV cache effectively reduces the sequence length. On the other hand, the overhead we introduce stems from identifying which token to drop. The overhead is rather limited compared to the saving.\n\nWe provide throughput measured in tokens per second in Table 8. We based the throughput measurement based on FlexGen[1]. We use a prompt length of 512 and a generation length of 1024. For all other settings, we keep the default from FlexGen. As shown in Table 8, we observe a slight improvement (1.2X ) in inference throughput. Due to the rebuttal time limitation, we could only implement a basic version. Implementation can be further improved to realize a significant gain in both memory usage and throughput.\n\n[1] Sheng, Ying, et al. \"High-throughput generative inference of large language models with a single gpu.\" arXiv preprint arXiv:2303.06865 (2023).\n\n  \n**Q2: The current scheme of selecting pivotal tokens is attention-based. Are there other methods/baselines to compare to?**\n\nThe memory problem from the KV cache is rather a recent constraint. The most established direction is applying quantization to reduce the memory footprint. We believe quantization are orthogonal direction rather than a competing direction. We show the possibility of combining quantization and Scissorhands in Section 5.\n\nWe understand that adding baselines would help understand Scisscorhands’s performance and limitations. Thus, we implement a straightforward position-based method: Local Window. Local Window only keeps the neighbor tokens’ cache. We take C4 and divide it into three subsets depending on the sequence length. C4-[256-512] contains data sequences that are longer than 256 tokens but less than 512 tokens. C4-[512-1024] contains data sequences that are longer than 512 tokens but less than 1024 tokens. C4-[1024-2048] contains data sequences that are longer than 1024 tokens but less than 2048 tokens. Results are summarized in Table 4.\n\nWe observe the perplexity of the full cache model degrades slightly with the growing sequence length. At all sequence lengths, Scissorhands performance is comparable against the full cache model, while Local Window incurs a significant quality loss. This demonstrates that keeping the pivotal tokens is important to reserve model performance.\n\nIt is also interesting to note that at longer sequence lengths, the local window has higher accuracy. This also shows at longer sequence length, the attention mechanism in current architecture tends to focus on recent context.\n\n**Q3: Clarification regarding to $|S_{0\\rightarrow t}|$ and $\\alpha = 1/t$.**\n\n$|S_{0\\rightarrow t}| / t$ varies across layers. Using OPT-66B as an example, the early and later layers have a much lower ratio (around 0.2 on average), while the middle layer has a higher ratio, around 0.5 on average. For $S_{0 \\rightarrow t}$, $\\tau \\in [t]$, $S_{\\tau}$ is calculated with $\\frac{1}{\\tau}$. The main message of this observation is to show that only a subset of tokens will carry its influence further. We have updated the manuscript to make sure the descriptions are accurate and more in detail.\n\n**Q4: Equation 1 what is there a scaling of 1/t within the softmax?**\n\nWe added the $\\frac{1}{t}$ scaling in the softmax to facilitate the theoretical analysis in Section 4.2. In particular, since the size of the not subsample W_Q and W_K grows along the token generation process, we need this scaling to control the error propagation through the softmax. Alternatively, we can also use the $\\frac{1}{\\sqrt{d}}$ scaling and add an assumption to enforce that the maximum sequence generation length satisfies $T \\leq \\sqrt{d}$."}, {"review_id": "YkjUxlBnKp", "reviewer": "Reviewer_AWSm", "summary": "This work studies attention patterns of LLMs and find that high-valued token-token patterns repeat for particular pivotal tokens. Since these patterns repeat, they can be compressed by storing only the pivotal tokens thus reducing the KV-cache size during inference. \n\nRecommendation:\n\nI recommend acceptance. The paper has a novel important analysis, which is already highly valuable. The proposed KV-compression is robustly evaluated. I  have no concerns and recommend acceptance. However, since the impact of this work is in a limited sub-area (attention-specific research mostly), I would not highlight it at the conference.", "questions": "The paper is very clear, and I have no questions.", "limitations": "Limitations are discussed.", "rating": 7, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- The analysis is thorough and insightful and will guide further development of KV-cache compression\n- The experimental validation is convincing using both different model scales and different zero-shot tasks\n- Results are excellent, making this approach practically useful", "weaknesses": "OPT has some unusual patterns, for example later layers have larger variances for hidden states. Other models do not have this property and are more stable. As such, I would imagine that the results would change slightly when considering a different model (likely the results would improve in later layer compared to the OPT results reported in the paper). It would be good to have these additional results for a different model (LLaMA would be best). But even without these results the paper is experimentally robust.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank the reviewer for the supportive comments, recognizing the novelty of our analysis, and thorough evaluations. We understand that another model family would further demonstrate the generality of our hypothesis and methods. However, we couldn't finish this within the time frame due to the time constraints of rebuttal and additional experiment efforts. We will try to update the manuscript with results on LLaMA later. "}, {"review_id": "uQdYpMtLA5", "reviewer": "Reviewer_H8Ve", "summary": "The paper proposes a method for reducing the memory requirement of inference for large autoregressive language models. \nThe proposed techniques is based on the empirical observation that the high attention score tokens in a transformer self-attention layer maintain a remarkable persistence for different token positions in the input sequence. \nThis observation is used to consider a reduction of the memory required in the attention KV cache, by retaining only the information corresponding to 'pivotal tokens', which are defined as the token positions that have a high attention scores for some of the previous steps in the self-attention autoregressive generation.", "questions": "-- Regarding the persistence of the attention pattern, it seems that its manifestation critically depends on the cosine similarity between input and output of the MLP of equation (2) - could the authors elaborate on this point? \n\n-- Appendix A.2 gives empirical evidence of large cosine similarity between input and output of the MLP block. In which setup have these results been obtained? How generally can this be assumed for different downstream tasks?\n\n-- Section 6 provide evidence that a persistent attention pattern does not exist in randomly initialised models. Is my understanding correct that the current study was limited to inference, and that the authors were not in a position to explore the evolution of the attention map during training, even for smaller models?", "limitations": "The authors have acknowledged that running large language models can correspond to significant power consumption and CO2 emissions, but their work is aimed at improving the power efficiency of inference for large language models.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper addresses the important issue of the high memory cost of inference for autoregressive generative language models, associated with the size of the KV cache, which has directly affects the inference execution time for large language models.\nTo the best of my knowledge, the observation of the persistence of the self-attention pattern is original, and gives a new angle in the understanding of this very important compute block. \nThe proposed method makes use of this observation to try to reduce the memory requirements associated with the KV cache of the self-attention, without requiring fine-tuning of the language model.", "weaknesses": "The exposition of the paper should be improved. In particular:\n\n-- The authors make the claim that their technique \"reduces the inference memory usage of the KV cache by up to 5x without compromising model quality\", while from the results of Section 5 this is not achieved in all cases, and depending on the downstream task the performance is not maintained for compression larger than 2x or 3x.\n\n-- The exposition of Algorithm 2 should be improved, adding further explanation and clarification in the text. \n\nOther comments:\n\n-- Line 91 and line 93: the sequence length is denoted by $p$, while the length of the prompt sequence was defined as $s$ at line 87.\n\n-- Line 185: in the case of a single attention head, the matrices $W_Q$, $W_K$, $W_V$ should be of size $d \\times d$ (and not $d \\times p$), and the matrix $W_O$ of size $d \\times d$ (and not $p \\times d$).\n\n-- In the equation after line 215 and in the equation after line 238, on both numerator and denominator of the softmax formula, $x_t W_K$ should be changed to $x_t W_Q$.\n\nMinor comments:\n\n-- Lines 67 and 69: \"between who they ...\" should be \"between which tokens they ...\".\n\n-- Caption of Table 2: \"This table summarises the maximum batch size ...\" should be simply \"Maximum batch size ...\".\n\n-- Line 102: \"parallel strategy\" should be \"parallelism strategy\".\n\n-- Line 109: \"Table 2.1\" should be \"Table 2\".\n\n-- Caption of Figure 2: The text \"In this figure, we plot the persistence ratio and the corresponding size ...\" should be changed to simply \"Persistence ratio and corresponding size ...\".\n\n-- Line 132: \"drawn from C4 [22]\" may be rephrased as \"drawn from the Colossal Clean Crawled Corpus (C4) [22]\".\n\n-- Algorithm 2: $[t-W,t]$ should be $[t-w,t]$ (the History Window Size is defined as lower-case $w$).\n\n-- Line 263: \"tokes\" should be \"tokens\".\n\n-- Caption of Figure 3: The text \"This figure shows the accuracy trend of ...\" should be changed to \"Accuracy trend of ...\".", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer’s supportive comments on novelty and constructive suggestions on the expositions. We have updated our manuscript and provided further details below. We hope our detailed clarification will clear the doubt about the significance and position of our work.\n\n**Q1: The exposition of the paper should be improved.**\n\nWe appreciate your suggestions on the narratives. We have updated our manuscript accordingly. For example, (i) “reduces the inference memory usage of the KV cache 2-5x without compromising model quality”. (ii) a detailed explanation of Algorithm 2 is added in Section 4.2. “When the KV cache buffer is full, Scissorhands drops tokens from the memory according to Algorithms 2. The importance record is a counter that indicates how many times a token is deemed non-important. The counter for all tokens in the recent window size is set to 0. A higher counter indicates dropping from the cache. ”\n\n**Q2: Regarding the persistence of the attention pattern, is my understanding correct that its manifestation critically depends on the cosine similarity between the input and output of the MLP of equation (2)?**\n\nOur theoretical justification in Theorem 3.1 requires the input-output alignment of the MLP, but this is likely due to the internal difficulty of theoretically analyzing the property of trained weights in LLMs.  We didn’t explicitly design experiments to verify the causal relationship between the input-output alignment of the MLP and the persistence of the attention pattern due to the difficulties in controlling input-output alignment without access to training. We believe that it remains an open question to study the dependence of the persistence of the attention on the input-output alignment of the MLP, both theoretically and empirically.\n\n**Q3:  Appendix A.2 gives empirical evidence of large cosine similarity between the input and output of the MLP block. In which setup have these results been obtained? How generally can this be assumed for different downstream tasks?**\n\nWe run OPT family model on 1000 sequences randomly drawn from the C4 validation dataset. We record the Cosine similarity between the input of the MLP block and the output of the MLP block for every transformer layer. We observe a similar trend across the entire OPT family. Further, recent work also observes similar phenomena. [1] shows in Figure 5 that the input and output of both the MLP block and attention block have high cosine similarity.\n\n[1] Liu, Zichang, et al. \"Deja vu: Contextual sparsity for efficient llms at inference time.\" International Conference on Machine Learning. PMLR, 2023. \n\n**Q4: Section 6 provides evidence that a persistent attention pattern does not exist in randomly initialized models. Is my understanding correct that the current study was limited to inference and that the authors were not in a position to explore the evolution of the attention map during training, even for smaller models?**\n\nThe current work focuses on the inference stage. Unfortunately, we don’t have access to pre-train a large language model. Actually, the intention of comparing with a randomly initialized model is to spark the interest of the community to investigate the training evolution of attention maps. Investigating such phenomena during training may lead to understanding model behaviors and intuition for designing sparse attention.\n"}]

Author rebuttal:
**Response to All Reviewers:**

We thank reviewers [R1(h9xt), R2(UGwp), R3(g2hZ), R4(AWSm), R5(H8Ve)] for their thoughtful and highly supportive feedback! We were glad that the reviewers found the problem significant and interesting [R1, R3, R5], the observations and theoretical analysis insightful and valuable [R2, R4, R5], the methods novel and clear [R2, R3, R4], and the experimental results convincing and promising [R3, R4, R5].

  

We have updated the paper to incorporate constructive suggestions, which will show in the revision. We summarize the attached additional experiments:

[ R1, R2, R5] Additional experiments and baselines (Table 4, 5, 6, 7):
    

-   We investigate the performance at different sequence lengths to understand the limitation.
    
-   We provide generated texts at different compression ratios for more intuitive illustration.
    
-   We provide results on three additional datasets.
     
[R2, R3] Inference efficiency (Table 8):
    

-   Scissorhands save attention computation because of the decreased sequence length from dropped KV cache. We show a 1.2X improvement in inference throughput.

Meta-review: This paper proposes a method to reduce memory costs incurred from storing past keys and values during inference with large language models. It leverages the observation that only a few important tokens influence later predictions, and presents a method for identifying such tokens and pruning the rest from the cache. Experiments on generation tasks validate the memory savings with a limited performance decrease. Overall, all reviewers except for one agreed that the paper should be accepted. The rejecting reviewer was concerned about the short generation length of the tasks included in the initial submission, the lack of mention of speedups (instead of just memory savings), and the lack of comparison to multi-query attention (MQA). The authors provided additional experiments on longer-form generation, measured speedups (which were modest, but nontrivial), and pointed out that MQA involves additional training while their method does not. The rejecting reviewer did not provide a response to the rebuttal; I think an appropriate response would have been to raise their score given that all concerns were addressed. However, I do agree that the paper should at least mention MQA as a possible orthogonal way to obtain speedups.
Decision: Accept (poster)