Title: Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Time
Venue: NeurIPS 2023
Accepted: yes
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.

Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Time

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 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 to 5 × without compromising model quality. We further demonstrate that SCISSORHANDS can be combined with 4-bit quantization for further compression

1 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 a 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, at a 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, 7].

Quantization and sparsity approaches [8-14] have been studied 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 parameters on a large amount of data is prohibitively expensive. Thus, an ideal compression algorithm should be applicable without training. Second, emerging applications such as dialogue systems require an extremely long context window. The maximum sequence length of LLMs is growing to over 32K [15]. The size of the KV cache also grows linearly with sequence length. For scalability, an ideal compression algorithm should reduce the memory from the sequence length dimension. At last, compression should preserve LLMs' quality and incontext learning ability.

0

4

8

12

16

20

24

28

32

36

40

44

48

52

56

60

64

68

72

76

80

84

88

92

96

100

104

108

112

116

120

124

128

132

136

140

144

148

152

156

160

164

168

172

176

180

184

188

192

196

200

204

208

212

216

220

224

228

232

236

240

244

248

252

256

260

264

268

272

276

Sequence Position

0

1

2

3

4

Attention Head

(a) Attention map at position 178

0

4

8

12

16

20

24

28

32

36

40

44

48

52

56

60

64

68

72

76

80

84

88

92

96

100

104

108

112

116

120

124

128

132

136

140

144

148

152

156

160

164

168

172

176

180

184

188

192

196

200

204

208

212

216

220

224

228

232

236

240

244

248

252

256

260

264

268

272

276

Sequence Position

0

1

2

3

4

Attention Head

(b) Attention map at position 228

0

4

8

12

16

20

24

28

32

36

40

44

48

52

56

60

64

68

72

76

80

84

88

92

96

100

104

108

112

116

120

124

128

132

136

140

144

148

152

156

160

164

168

172

176

180

184

188

192

196

200

204

208

212

216

220

224

228

232

236

240

244

248

252

256

260

264

268

272

276

Sequence Position

0

1

2

3

4

Attention Head

(c) Attention map at position 278

We go beyond the traditional model compression techniques to achieve such demanding requirements. We envision that not all tokens must be stored in 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 [16-20], meaning that one token will only heavily attend to a small number of tokens. More importantly, we observe Repetitive Attention Pattern from different tokens in the sequence in a trained LLM( Figure 1). 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 a significant influence at a future step. This hypothesis, if 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 a sentence. And the overlapping ratio is 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 importance hypothesis to realize LLM inference with a compressed KV cache. In Section 4, we present an efficient algorithm such that the size of KV cache is always less than a predetermined budget. A theoretical guarantee justifies that such a compressed KV cache can approximate the attention output. In Section 5, we empirically evaluate SCISSORHANDS and show that SCISSORHANDS reduces the memory usage of KV cache 2 -5 × 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.

2 Problem Description and Related Work

This paper considers the LLM 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, W i K ∈ R d × d , W i V ∈ R d × d for the key and value projection matrix at the i th 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 a cache to reduce repetitive computation. Denote x i prompt = [ x i 1 , ..., x i p ] , x i prompt ∈ R b × p × d as the input to attention at the i th transformer layer. Denote the key cache and value cache at layer i as K i , V i ∈ R b × p × d , K i 0 = x i prompt W i K , V i 0 = x i prompt W i V .

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 at step t in the i th transformer layer x i t ∈ R b × 1 × d . K i t +1 = [ K i t , x i t W i K ] , V i t +1 = [ V i t , x i t W i V ] .

2.1 LLMInference 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 is considerably smaller than the previous two. As shown in Table 1, the size of the KV cache, 2.5 × -5 × 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,768 tokens [15]. Longer sequence length would make the KV cache memory problem even more severe.

Assuming LLM 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 2. 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 a time complexity of O ( n 2 ) , where n is the sequence length. As a result, a line of work has been proposed to mitigate the computation burden of the attention mechanism [16-20]. These approaches exploit low-rank or sparsification to approximate the attention output. Besides, [21] realized exact efficient attention with wall-clock speed by optimizing

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 on 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-Query-Attention(MQA) [22] change the attention design such that keys and values are shared across all attention heads. MQA requires training the model from scratch, while our works focus entirely on the inference stage.

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 t over all the words that come before it in the sentence. In Figure 1, we provide three attention maps of a sentence randomly drawn from the Colossal Clean Crawled Corpus (C4) [23] using OPT-6B. Each attention map is a discretized attention score calculated at a random position. We consider a score larger than 1 t as significant as 1 t indicates an 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 at 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 are evident in the attention maps. There exist specific tokens that keep receiving high attention. Meanwhile, these attention maps show sparsity: only a 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. A stricter claim is that these tokens are the only ones that could be significant for a future step. Thus, we articulate the persistence of importance hypothesis .

The Persistence of Importance Hypothesis. With a trained autoregressive language model, only pivotal tokens, which had a substantial influence at one previous step, will have a significant influence at a future step.

If true, this hypothesis indicates the possibility of foreseeing what information in 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 us to reduce the size of the KV cache by throwing away the embedding of non-important tokens.

Pivotal Token. One natural indication of a token's influence is the attention score. We consider a token pivotal for position t if this token receives an attention score larger than threshold α from the token at position t . Let S t denote the set of pivotal tokens for position t . S a → b denote the set of

0

10

20

30

40

50

60

Transformer Layer

0.5

0.6

0.7

0.8

0.9

1.0

Pesistence Ratio

OPT-6B

OPT-13B

OPT-30B

OPT-66B

0

10

20

30

40

50

60

Transformer Layer

0.0

0.2

0.4

0.6

0.8

1.0

|

S

0

t

|

t

OPT-6B

OPT-13B

OPT-30B

OPT-66B

(b) Size of

S

0

→

t

pivotal tokens for every position from a to b .

Verification. We measure persistence ratio as an empirical test the hypothesis. Persistence ratio measures how many tokens in 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 l denote the length of the sentence. We record S 1 → t ∈ { x 1 , ...x t } , tokens in { x 1 , ..., x t } who received high attention from every position until t . Then, we record S t +1 → l ∈ { x 1 , ...x t } , tokens in { x 1 , ..., x t } who received high attention from position after t . The persistence ratio is the intersection divided by the size of S t +1 → l . Formally,

At the same time, we measure | S 0 → t | t . | S 0 → t | = t indicates that every token substantially impacted at least one position, which is the trivial case of persistence of importance hypothesis . Our test is performed with OPT models [24] with different datasets such as OpenBookQA [25] and WikiText [26]. In our verification, we set t = l 2 , which measures the overlapping between the first and later half of the sentences. Same as in Section 3.1, we set α = 1 t , which suggests an average score.

Result. We present our main results in Figure 2. First, given the current criterion of pivotal token and t value, the size of S 0 → t 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 is mostly included in the set of the first half sentences. Combining these two pieces of empirical evidence, we see positive evidence for our 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 in the KV cache and the computation required at the 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 to 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.

x t ∈ R 1 × d is a row vector. X t -1 ∈ R ( t -1) × d denotes the aggregation of x 1 , . . . , x t -1 , where the j th row is x j . W Q , W K , W V ∈ R d × d and W O ∈ R d × d are the attention weights. Lastly, F : R 1 × d → R 1 × d denotes the MLP block following attention block, a two-layer MLP with skip connections, given by

Algorithm 1 Inference with Budget KV cache

Algorithm 2 Compress KV Cache

Input : Key Cache ¯ K ∈ R n × d , Value Cache ¯ V ∈ R n × d , History Window Size w , Recent Window Size r , Drop Amount m , Generation Step t , Importance Record I ← ⃗ 0 ∈ R t for i ∈ [ t -w,t ] do ▷ Consider tokens within history window I ← I + α i < 1 t ▷ Increment the counter for low score token end for I [: -r ] ← 0 ▷ Keep cache within the recent window Keep set S t ← Argsort ( I ) [: -m ] Keep everything in S t in ¯ K ∈ R n × d , ¯ V ∈ R n × d such that n ← n -m

We are interested in the attention scores α t = softmax ( 1 / t · x t W Q W ⊤ K X ⊤ t -1 ) . Notice that α t,j scales with x t W Q W ⊤ K x ⊤ j . The following theorem characterizes the behavior of x t W Q W ⊤ K x ⊤ j

Theorem 3.1. Let A = W V W O W Q W ⊤ K and let λ K , λ Q , λ V , λ O denote the largest singular values of W K , W Q , W V , W O , respectively. Consider the transformer in (1) with normalized inputs ∥ x t ∥ 2 = 1

̸

Then for all x ℓ satisfying x ℓ Ax ⊤ ℓ ≥ c and x ℓ Ax ℓ ≥ ϵ -1 max j ∈ [ t ] ,j = ℓ x j Ax ⊤ ℓ , it holds that

The proof is provided in Appendix B. Theorem 3.1 shows that under an assumption on the MLP in (2), for all x ℓ such that x ℓ Ax ⊤ ℓ is large enough, x t +1 W Q W ⊤ K x ⊤ j satisfies Equation (3). The assumption on the MLP a t x ⊤ t +1 ≥ (1 -δ ) ∥ a t ∥ 2 essentially requires a large cosine similarity between the input and output of F . This behavior can be empirically verified in Appendix A. Essentially, skip connection dominates the output because ∥ x ∥ 2 ≫ ∥ W 2 relu ( W 1 x ) ∥ 2 , resulting in a cosine similarity close to one between input and output. Equation (3) shows that despite a factor of x ℓ Ax ⊤ ℓ ∥ a t ∥ 2 , x t +1 W Q W ⊤ K x ⊤ j almost scales with α t,ℓ . Since x t +1 W Q W ⊤ K x ⊤ j directly affects α t +1 ,ℓ , this property shows that a larger α t,ℓ will potentially imply a large α t +1 ,ℓ .

Our theorem shows that the property in Equation (3) property only holds for x ℓ such that x ℓ Ax ⊤ ℓ is large. A are trained attention weights. This condition may suggest that the trained weights A selects x ℓ as a pivotal token. Each attention is learned to 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.

4 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 a given budget. Section 4.2 provides a 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. K t , V t ∈ R t × d denote for the KV cache until step t . x t ∈ R 1 × d is a row vector that denotes the input to attention at step t . The output of an attention head at step t can be written as,

Intuition. As shown in Section 3, the attention scores α t,i follow a strong power-law distribution. For the autoregressive generation process, if there exists an 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. LLMsare 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 an attention head under budget B ) . Given a stream of token embedding, including prompt and previously generated tokens, denotes their input to the head as { x 1 , . . . , x t , . . . } . The problem of sequential generation at an attention head under budget B is maintaining a key cache ¯ K t and value cache ¯ V t such that ¯ K t , ¯ V t ∈ R n × d and n < B .

Approach. Inspired by the textbook solution of reservoir sampling and the Least Recent Usage cache replacement algorithm, SCISSORHANDS reserves a 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 1 and Algorithm 2.

When the KV cache size exceeds the budget, SCISSORHANDS drops tokens from the KV cache according to Algorithm 2. The importance record is a counter that indicates how many times a token is deemed 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 w 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 0.

With a sampled KV cache, attention output can be computed by the following estimator

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. m controls the frequency, and we use m = 0 . 5 B in our experiment. Further, steps after the compression have reduced attention computation because of the reduction in the KV cache. On the other hand, one can trade a tiny amount of memory to avoid the overhead by maintaining the importance record during generation steps in Algorithm 1.

Allocating Budgets Across Attention Heads. An LLM typically consists of L transformer layers where each layer has H heads. A 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 2. 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 { ˜ x t } T t =0 denote the tokens generated by the transformer with budget KV cache as in Algorithm 2 with m = 1 :

Notice that when m = 1 , 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 2 will always drop tokens with the smallest attention scores.

For reference purposes, let { x t } T t =0 denote the tokens generated by a vanilla transformer defined in (1). We will bound the difference ∥ x t -˜ x t ∥ 2 .

Theorem 4.1. Let λ 1 , λ 2 denote the largest singular values of W 1 and W 2 in (2). Let

and assume that each β t,j = cv t,j , where v t,j are sampled from a power-law distribution with pdf f ( x ) = c ( x + b ) -k . Suppose that λ V λ O (1 + λ 1 λ 2 )(1 + λ Q λ K ) ≤ 1 2 . Let T min and T max denote the starting and maximum sequence lengths, respectively, and let B ≤ T max denote the budget as in Algorithm 2. If for all t ∈ [ T min , T max ] , S t contains only tokens with at most the largest B values of β t,j , that is, | S t | = B and min j ∈ S t β t,j ≥ max j / ∈ ˆ S β t,j , then for all ϵ ∈ (0 , 1) , with probability t at least 1 -T max exp ( -ϵ 2 b 2 ( T min -1) ( k -2) 2 ( u -b ) 2 ) -T max exp ( -2( T min -1)(1 -B / T max ) 2 (1 -ϵ ) 2 ) , the following error bound must hold for all t ∈ [ T min , T max ]

The definition of β t,j 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 1 -B / T max . When B = T max , meaning

1

(

/

k

-

1)

that we are keeping all of the tokens, the error becomes zero. The term k -( k -1) ( 1 -ϵ B -ϵ ) depends on the distribution that the attention scores are fitted to and is always less than one. With a strong power-law distribution, this term provides a further decrease to the error bound in (4).

5 Empirical Evaluation

In this section, we present the results that demonstrate SCISSORHANDS achieves up to 5 × 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 [23] and a number of few-shot downstream tasks: Hellaswag [27], MathQA [28], PIQA [29], Winogrande [30]. We use lm-eval-harness [31] to evaluate few-shot tasks. Our experiments are conducted on NVIDIA 4 A100 40GB GPU servers.

No Accuracy Drop untill 5 × . In Figure 3, we present SCISSORHANDS's accuracy trend where 1 × 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 is exceptionally encouraging. This suggests that SCISSORHANDS can scale with the model size. Downstream tasks are usually less sensitive to perturbation and bear more variance in terms of accuracy. We evaluate the 5-shot setting and 1 × denotes the original OPT model. For Winogrande and MathQA, accuracy is maintained even after 5 × 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.

1X

2X

3X

4X

5X

Compression

7.5

8.0

8.5

9.0

9.5

10.0

Perlexity

OPT-13B

OPT-30B

OPT-66B

(a) Language Modeling

1X

2X

3X

4X

5X

Compression

0.0

0.2

0.4

0.6

0.8

1.0

MathQA

PIQA

Winogrande

(b) OPT-6B Five shot

1X

2X

3X

4X

5X

Compression

0.0

0.2

0.4

0.6

0.8

1.0

MathQA

PIQA

Winogrande

(c) OPT-13B Five shot

1X

2X

3X

4X

5X

Compression

0.0

0.2

0.4

0.6

0.8

1.0

MathQA

PIQA

Winogrande

(d) OPT-30B Five shot

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 3. 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 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.

Compatible with 4-bit Quantization We test the compatibility of quantization and SCISSORHANDS at 2 × compression. We adopt 4-bit quantization following [7]. Even Hellaswag is most sensitive based on Figure 3, adding quantization doesn't introduce compounded errors.

Ablation on Attention Score Error. We present the change ratio in attention score between original OPT-13B and SCISSORHANDS OPT-13B at 3 × compression on C4 in Figure 4.

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 is calculated as α s -α o α o where α s is the SCISSORHANDS attention score and α o is the original score. From Figure 4, we observe that the change ratio is centered around 0. -1 indicating that α s is significantly smaller compared to the original, suggesting that a small portion of the important tokens are dropped in the cache. To explain the above observation of SCISSORHANDS, we denote the n number of tokens with the highest score as { x top _ n t } T t =0 . Then, for any other sets of tokens { x ′ t } T t =0 that has no greater than n tokens, we can easily prove that similarity ( x topB t , x t ) ≤ ( x ′ t , x t ) . Thus, SCISSORHANDS gives the most similar output as the original model at 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 a model architecture bias or an unexpected training outcome. For such purpose, we perform the same experiment with a 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 initialized models. Apart from an efficiency deployment perspective, could such repetitive attention patterns contribute to some known problems in language generation, such as repetitions? It may be

1.0

0.8

0.6

0.4

0.2

0.0

worth investigating the relationship between repetitive attention patterns and undesired generations.

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 verify the generality across the different models and datasets. However, we cannot access the training process and fail to know exactly how an LLM is trained to exhibit such behavior. Experiments with the large model create carbon dioxide emissions. However, our work improves the efficiency of LLM, and we foresee no negative impacts.

7 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 5 × in the KV cache without compromising the performance of LLMs. Furthermore, we demonstrate the compatibility of SCISSORHANDS with quantization techniques, opening up the possibility of reducing memory usage in both the representation and sequence length dimensions.

References

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.07258 , 2021.

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.

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.

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.12837 , 2022.

Stephanie CY Chan, Adam Santoro, Andrew Kyle Lampinen, Jane X Wang, Aaditya K 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.

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.

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.

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.01861 , 2022.

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.09557 , 2022.

Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Llm. int8 (): 8-bit matrix multiplication for transformers at scale. arXiv preprint arXiv:2208.07339 , 2022.

Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. Gptq: Accurate post-training quantization for generative pre-trained transformers. arXiv preprint arXiv:2210.17323 , 2022.

Elias Frantar and Dan Alistarh. Massive language models can be accurately pruned in one-shot. arXiv preprint arXiv:2301.00774 , 2023.

Hritik Bansal, Karthik Gopalakrishnan, Saket Dingliwal, Sravan Bodapati, Katrin Kirchhoff, and Dan Roth. Rethinking the role of scale for in-context learning: An interpretability-based case study at 66 billion scale. arXiv preprint arXiv:2212.09095 , 2022.

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:2211.10438 , 2022.

OpenAI. Gpt-4 technical report, 2023.

Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020 . OpenReview.net, 2020.

Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768 , 2020.

Beidi Chen, Zichang Liu, Binghui Peng, Zhaozhuo Xu, Jonathan Lingjie Li, Tri Dao, Zhao Song, Anshumali Shrivastava, and Christopher Re. Mongoose: A learnable lsh framework for efficient neural network training. In International Conference on Learning Representations , 2021.

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.

Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea 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 on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021 . OpenReview.net, 2021.

Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness, 2022.

Noam Shazeer. Fast transformer decoding: One write-head is all you need, 2019.

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 a unified text-to-text transformer. arXiv e-prints , 2019.

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.

Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? a new dataset for open book question answering. In EMNLP , 2018.

Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016.

Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics , 2019.

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019.

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.

Sakaguchi Keisuke, Le Bras Ronan, Bhagavatula Chandra, and Choi Yejin. Winogrande: An adversarial winograd schema challenge at scale. In Communications of the ACM , 2019.

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. A framework for few-shot language model evaluation. In Version v0. 0.1. Sept . Zenodo, September 2021.

Appendix

A More Results

A.1 Repetitive Attention Pattern

We provide the attention map similar to Figure 1 but from a different transformer layer on the same text in Figure 6, Figure 7, Figure 8 and Figure 9. A repetitive pattern and attention sparsity can be observed across layers.

0

4

8

12

16

20

24

28

32

36

40

44

48

52

56

60

64

68

72

76

80

84

88

92

96

100

104

108

112

116

120

124

128

132

136

140

144

148

152

156

160

164

168

172

176

180

184

188

192

196

200

204

208

212

216

220

224

228

232

236

240

244

248

252

256

260

264

268

272

276

Sequence Position

0

1

2

3

4

Attention Head

0

4

8

12

16

20

24

28

32

36

40

44

48

52

56

60

64

68

72

76

80

84

88

92

96

100

104

108

112

116

120

124

128

132

136

140

144

148

152

156

160

164

168

172

176

180

184

188

192

196

200

204

208

212

216

220

224

228

232

236

240

244

248

252

256

260

264

268

272

276

Sequence Position

0

1

2

3

4

Attention Head

(b) Attention map at position 228

0

4

8

12

16

20

24

28

32

36

40

44

48

52

56

60

64

68

72

76

80

84

88

92

96

100

104

108

112

116

120

124

128

132

136

140

144

148

152

156

160

164

168

172

176

180

184

188

192

196

200

204

208

212

216

220

224

228

232

236

240

244

248

252

256

260

264

268

272

276

Sequence Position

0

1

2

3

4

Attention Head

(c) Attention map at position 278

0

4

8

12

16

20

24

28

32

36

40

44

48

52

56

60

64

68

72

76

80

84

88

92

96

100

104

108

112

116

120

124

128

132

136

140

144

148

152

156

160

164

168

172

176

180

184

188

192

196

200

204

208

212

216

220

224

228

232

236

240

244

248

252

256

260

264

268

272

276

Sequence Position

0

1

2

3

4

Attention Head

(a) Attention map at position 178

0

4

8

12

16

20

24

28

32

36

40

44

48

52

56

60

64

68

72

76

80

84

88

92

96

100

104

108

112

116

120

124

128

132

136

140

144

148

152

156

160

164

168

172

176

180

184

188

192

196

200

204

208

212

216

220

224

228

232

236

240

244

248

252

256

260

264

268

272

276

Sequence Position

0

1

2

3

4

Attention Head

(b) Attention map at position 228

0

4

8

12

16

20

24

28

32

36

40

44

48

52

56

60

64

68

72

76

80

84

88

92

96

100

104

108

112

116

120

124

128

132

136

140

144

148

152

156

160

164

168

172

176

180

184

188

192

196

200

204

208

212

216

220

224

228

232

236

240

244

248

252

256

260

264

268

272

276

Sequence Position

0

1

2

3

4

Attention Head

(c) Attention map at position 278

0

4

8

12

16

20

24

28

32

36

40

44

48

52

56

60

64

68

72

76

80

84

88

92

96

100

104

108

112

116

120

124

128

132

136

140

144

148

152

156

160

164

168

172

176

180

184

188

192

196

200

204

208

212

216

220

224

228

232

236

240

244

248

252

256

260

264

268

272

276

Sequence Position

0

1

2

3

4

Attention Head

(a) Attention map at position 178

0

4

8

12

16

20

24

28

32

36

40

44

48

52

56

60

64

68

72

76

80

84

88

92

96

100

104

108

112

116

120

124

128

132

136

140

144

148

152

156

160

164

168

172

176

180

184

188

192

196

200

204

208

212

216

220

224

228

232

236

240

244

248

252

256

260

264

268

272

276

Sequence Position

0

1

2

3

4

Attention Head

(b) Attention map at position 228

0

4

8

12

16

20

24

28

32

36

40

44

48

52

56

60

64

68

72

76

80

84

88

92

96

100

104

108

112

116

120

124

128

132

136

140

144

148

152

156

160

164

168

172

176

180

184

188

192

196

200

204

208

212

216

220

224

228

232

236

240

244

248

252

256

260

264

268

272

276

Sequence Position

0

1

2

3

4

Attention Head

(c) Attention map at position 278

0

4

8

12

16

20

24

28

32

36

40

44

48

52

56

60

64

68

72

76

80

84

88

92

96

100

104

108

112

116

120

124

128

132

136

140

144

148

152

156

160

164

168

172

176

180

184

188

192

196

200

204

208

212

216

220

224

228

232

236

240

244

248

252

256

260

264

268

272

276

Sequence Position

0

1

2

3

4

Attention Head

(a) Attention map at position 178

0

4

8

12

16

20

24

28

32

36

40

44

48

52

56

60

64

68

72

76

80

84

88

92

96

100

104

108

112

116

120

124

128

132

136

140

144

148

152

156

160

164

168

172

176

180

184

188

192

196

200

204

208

212

216

220

224

228

232

236

240

244

248

252

256

260

264

268

272

276

Sequence Position

0

1

2

3

4

Attention Head

(b) Attention map at position 228

0

4

8

12

16

20

24

28

32

36

40

44

48

52

56

60

64

68

72

76

80

84

88

92

96

100

104

108

112

116

120

124

128

132

136

140

144

148

152

156

160

164

168

172

176

180

184

188

192

196

200

204

208

212

216

220

224

228

232

236

240

244

248

252

256

260

264

268

272

276

Sequence Position

0

1

2

3

4

Attention Head

(c) Attention map at position 278

A.2 Cross Layer Cosine Similarity

In Section 3.3, our analysis assumes a large cosine similarity between the input and output of F . Here, we provide empirical evidence to support such an assumption in Figure 10. Because of the residual connection in F and the domination of x , the cosine similarity between x and F ( x ) is extremely high.

A.3 Generated examples with SCISSORHANDS

0

5

10

15

20

Transformer Layer

0.86

0.88

0.90

0.92

0.94

0.96

0.98

1.00

Cosine Similarity

Residual At MLP

Cos

(

x

,

(

x

))

(a) Cosine Similarity

0

5

10

15

20

Transformer Layer

0

20

40

60

80

100

120

Norm

Residual At MLP

||

x

||

||

W

2

Relu

(

W

1

x

)||

(b) Norm Comparision

B Proofs

B.1 Proof of Theorem 3.1

We consider the token generation process of a simplified model: a single-layer transformer model with single-head attention.

x t ∈ R 1 × d is a row vector. X t -1 ∈ R ( t -1) × d denotes the aggregation of x 1 , . . . , x t -1 , where the j th row is x j . W Q , W K , W V ∈ R d × p and W O ∈ R p × d are the attention weights. Lastly, F : R 1 × d → R 1 × d denotes the MLP block following attention block, a two-layer MLP with skip connections, given by

We are interested in the attention scores α t = softmax ( 1 / t · x t W Q W ⊤ K X ⊤ t -1 ) . Notice that α t,j scales with x t W Q W ⊤ K x ⊤ j . We first re-state the Theorem 3.1 below.

Therefore, we have

̸

̸

Theorem B.1. Let A = W V W O W Q W ⊤ K and let λ K , λ Q , λ V , λ O denote the largest singular values of W K , W Q , W V , W O , respectively. Consider the transformer in (5) with normalized inputs ∥ x t ∥ 2 = 1 for all t . Let c, ϵ > 0 be constants. Assume that a t x ⊤ t +1 ≥ (1 -δ ) ∥ a t ∥ 2 with δ ≤ ( cϵ λ Q λ K λ V λ O ) 2 . Then for all x ℓ satisfying x ℓ Ax ⊤ ℓ ≥ c and x ℓ Ax ℓ ≥ ϵ -1 max j ∈ [ t ] ,j = ℓ x j Ax ⊤ ℓ , it holds that

As a preparation of the proof, we first show two lemmas.

Lemma B.1. Let x 1 , x 2 ∈ R 1 × m satisfies ∥ x 1 ∥ 2 = ∥ x 2 ∥ 2 = 1 and x 1 x ⊤ 2 ≥ 1 -δ for some δ ∈ (0 , 1) . Then for all y ∈ R 1 × m we have

Proof. Let x 2 = x ∥ 2 + x ⊥ 2 where

Then it is easy to see that x ⊥ 2 x ⊤ 1 = 0 . By the Pythagorean Theorem, we have

Thus, the Cauchy-Schwarz inequality implies

Lemma B.2. Let ℓ ∈ [ t ] be given. Suppose that x ℓ Ax ⊤ ℓ > ϵ -1 ∣ ∣ x j Ax ⊤ ℓ ∣ ∣ for all j = ℓ . Then we have ( S ( t ) ℓ -ϵ ) x ⊤ ℓ ax ℓ ≤ x ⊤ ℓ W ⊤ K W Q a t ≤ ( S ( t ) ℓ + ϵ ) x ⊤ ℓ ax ℓ

Proof. Notice that

Thus, we have

Therefore

̸

̸

̸

̸

where in the second inequality we use ϵ -1 ∣ ∣ x j Ax ⊤ ℓ ∣ ∣ ≤ x ℓ Ax ⊤ ℓ and in the third inequality we use ∑ t -1 j =1 ,j = ℓ α t,j ≤ ∑ t -1 j =1 α t,j = 1 . This implies that

Now we proceed to the main body of the proof. Assume that ∥ x ℓ ∥ 2 = 1 for all ℓ . Using Lemma (B.1), if a t x ⊤ t +1 ≥ (1 -δ ) ∥ a t ∥ 2 , then we have

Recall that λ Q , λ K are the maximum singular values of W Q and W K , respectively. Then it holds that ∥ ∥ W Q W ⊤ K x ⊤ ℓ ∥ ∥ 2 ≤ λ Q λ K ∥ x ℓ ∥ 2 . Using ∥ x ℓ ∥ 2 = 1 , we have

Notice that

Then since δ ≤ ( cϵ λ Q λ K λ V λ O ) 2 , we have

Since by Lemma (B.2), we have

It must hold that

Since x ⊤ ℓ ax ℓ ≥ c , it holds that

which implies that

Therefore

B.2 Proof of Theorem 4.1

Let { ˜ x t } T t =0 denote the tokens generated by the transformer with budget KV cache as in Algorithm 2 with m = 1 :

Notice that when m = 1 , 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 2 will always drop tokens with the smallest attention scores.

For reference purposes, let { x t } T t =0 denote the tokens generated by a vanilla transformer defined in (5). We re-state Theorem 4.1 below, which bounds the difference ∥ x t -˜ x t ∥ 2 .

Theorem B.2. Let λ 1 , λ 2 denote the largest singular values of W 1 and W 2 in (6). Let

and assume that each β t,j = cv t,j , where v t,j are sampled from a power-law distribution with pdf f ( x ) = c ( x + b ) -k . Suppose that λ V λ O (1 + λ 1 λ 2 )(1 + λ Q λ K ) ≤ 1 2 . Let T min and T max denote the starting and maximum sequence lengths, respectively, and let B ≤ T max denote the budget as in Algorithm 2. If for all t ∈ [ T min , T max ] , S t contains only tokes with at most the largest B values of β t,j , that is, | S t | = B and min j ∈ S t β t,j ≥ max j / ∈ ˆ S t β t,j , then for all ϵ ∈ (0 , 1) , with probability at least 1 -T max exp ( -ϵ 2 b 2 ( T min -1) ( k -2) 2 ( u -b ) 2 ) -T max exp ( -2( T min -1)(1 -B / T max ) 2 (1 -ϵ ) 2 ) , the following error bound must hold for all t ∈ [ T min , T max ]

Define m k,j = I { j ∈ S t } . With the definition of m k,j , ˜ a t can be written as

Our first lemma shows the Lipschitzness of the attention module.

Lemma B.3. Consider two sequences of tokens { x i } t i =1 and { y i } t i =1 where ∥ x i ∥ 2 = ∥ y i ∥ 2 = 1 for all i ∈ [ t ] . Define X t -1 , Y t -1 ∈ R ( t -1) × d as the matrices whose i th row are x i and y i , respectively. Let ∆ t = ∥ x t -y t ∥ 2 . Then we have

Proof. We can decompose the difference as

By the Lipschitzness of softmax, we have

Since ∥ x t ∥ 2 = 1 and ∥ X t -1 -Y t -1 ∥ 2 = ( ∑ t -1 j =1 ∥ x j -y j ∥ 2 ) 1 2 ≤ √ t -1∆ t , we have √

∥ ∥ softmax ( x t W Q W ⊤ K X ⊤ t -1 ) -softmax ( x t W Q W ⊤ K Y ⊤ t -1 )∥ ∥ 2 ≤ t -1 t λ Q λ K ∆ Similarly,

Since ∥ x t -y t ∥ 2 = ∆ t and ∥ Y t -1 ∥ 2 = √ t -1 , we have

Combining the two bounds gives

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 ˜ a t be defined as in (8). Define b t as

Assume that ∥ x j ∥ 2 = 1 for all j ∈ [ t ] . Then we have

Proof. A direction computation yields

Thus, ∥ ˜ a t -b t ∥ 2 can be bounded as

since ∥ ˜ x j ∥ 2 = 1 for all j ∈ [ t ] . Now we analyze ˜ α t,j -β t,j . Let ˆ S t = S t \ { t } . Then m k,j = 1 if and only if j ∈ ˆ S t . For convenience, let r t,j = 1 / t · ˜ x t W Q W ⊤ K ˜ x ⊤ j . Thus, β can be written as

Furthermore, for all j / ∈ ˆ S t , we have ˜ α t,j = 0 . For all j ∈ ˆ S t , we have

Therefore, for all j ∈ ˆ S t , we have

Therefore, the bound of ∥ ˜ a t -b t ∥ 2 can be written as

Our last lemma shows the Lipschitzness of the MLP in (6).

Lemma B.5. Let λ 1 , λ 2 denote the largest singular values of W 1 , W 2 in (6). For all x 1 , x 2 ∈ R d , we have

Proof. Direct computation yields

where in the third inequality we use the fact that relu ( · ) is 1 -Lipschitz.

Now we turn to the proof of our main theorem. Combining all of the results, we have

Therefore, by triangle inequality, we have

To start, the magnitude of T 1 can be bounded as

where in the third inequality we use ∥ x t,j -˜ x t,j ∥ 2 = ∆ t and in the last equality we use ∑ t -1 j =1 α t,j = 1 . To bound the magnitude of T 2 , we apply Lemma B.3, which shows that ∥ α t -β t ∥ ≤ 2 √ t -1 t λ Q λ K ∆ t to get that

Lastly, to bound the magnitude of T 3 , we use Lemma B.4 to get that

Putting things together for (10), we have

By Lemma B.5 we can further show that

By Theorem B.3, we have that with probability at least 1 -T max exp ( -ϵ 2 b 2 ( T min -1) ( k -2) 2 ( u -b ) 2 ) -T max exp ( -2( T min -1)(1 -B / T max ) 2 (1 -ϵ ) 2 ) , it holds for all t ∈ [ T min , T max ] that

Given that E [ ∥ x t -˜ x t ∥ ] ≤ 2∆ max , we have

Thus, as long as λ V λ O (1 + λ 1 λ 2 )(1 + λ Q λ K ) ≤ 1 2 , we can guarantee that

Thus, for all t ∈ [ T min , T max ] , we have that

B.3 Budgeted Cache

Theorem B.3. Let β t,j be sampled from some power-law distribution f ( x ) = c ( x + b ) -γ with support on [0 , u -b ) for some k > 2 and u ≥ 5 b . Let S t be defined in Theorem B.2, and define ˆ S t = S t \ { t } . Then with probability at least 1 -T max exp ( -ϵ 2 b 2 ( T min -1) ( k -2) 2 ( u -b ) 2 ) -T max exp ( -2( T min -1)(1 -B ) 2 (1 -ϵ ) 2 ) it holds for all t ∈ T that

We consider the case of maintaining a budget of B by dropping the smallest β t,j 's. Assume that v j has pdf f ( x ) = c ( x + b ) -k with support on [0 , u -b ) . To make things precise, we first compute c

To start, we notice that

Let C = ∑ t -1 j =1 v j , then the expectation of C is

Let ∆ = b 2 -k -( k -1) u 2 -k +( k -2) bu 1 -k b 1 -k -u 1 -k . By Hoeffding's inequality, we have that

This implies that with probability at least 1 -exp ( -2 ϵ 2 ∆ 2 ( t -1) ( k -2) 2 ( u -b ) 2 ) we have

Now, we proceed to bound ∑ j / ∈ ˆ S t β t,j where ˆ S t = { j ∈ [ t -1] : β t,j ≥ γ C } . Equivalently, we can bound C -1 ∑ t -1 j =1 I { v j ≤ γ } v j . Its expectation is given by

We pause here and study how small can we choose γ . Notice that

By Hoeffding's inequality again, we have that

Enforcing ∑ t -1 j =1 I { v j ≤ γ } ≥ T max -B gives ( γ + b ) 1 -k ≤ b 1 -k -1 -B / T max 1 -ϵ ( b 1 -k -u 1 -k ) , which can be satisfied as long as γ ≥ ( ( B / T max -ϵ 1 -ϵ ) 1 1 -k -1 ) b . Therefore

We further notice that

This gives

Notice that if u ≥ 5 b , we have

Therefore

holds with probability at least 1 -exp ( -ϵ 2 b 2 ( t -1) ( k -2) 2 ( u -b ) 2 ) -exp ( -2( t -1)(1 -B / T max ) 2 (1 -ϵ ) 2 ) . Taking a union bound gives the desired result.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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 ...\".\n\nQUESTIONS:\n-- 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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-30"}

comments: 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. 

Recommendation:

I 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.

STRENGTHS:
- The analysis is thorough and insightful and will guide further development of KV-cache compression
- The experimental validation is convincing using both different model scales and different zero-shot tasks
- 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.

QUESTIONS:
The paper is very clear, and I have no questions.

LIMITATIONS:
Limitations are discussed.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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\nWEAKNESSES:\nSome design decisions are not well justified, and there are insufficient baseline comparisons.\n\nQUESTIONS:\n- 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\nLIMITATIONS:\nThere is no direct societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nIdeas and math are well written. The direction of using pivotal tokens could be interesting.\n\nWEAKNESSES:\nThe 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\nQUESTIONS:\nMulti-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\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 1, "DATE": "2023-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThis 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.\n\nWEAKNESSES:\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?\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.\n\nQUESTIONS:\n+ 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?\n\nLIMITATIONS:\nSee above in weaknesses and questions.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-27"}