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

Summary

The paper proposes a method for reducing the memory requirement of inference for large autoregressive language models. The 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. This 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.

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. To 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. The 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: -- 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. -- The exposition of Algorithm 2 should be improved, adding further explanation and clarification in the text. Other comments: -- 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. -- 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$). -- 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$. Minor comments: -- Lines 67 and 69: "between who they ..." should be "between which tokens they ...". -- Caption of Table 2: "This table summarises the maximum batch size ..." should be simply "Maximum batch size ...". -- Line 102: "parallel strategy" should be "parallelism strategy". -- Line 109: "Table 2.1" should be "Table 2". -- 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 ...". -- Line 132: "drawn from C4 [22]" may be rephrased as "drawn from the Colossal Clean Crawled Corpus (C4) [22]". -- Algorithm 2: $[t-W,t]$ should be $[t-w,t]$ (the History Window Size is defined as lower-case $w$). -- Line 263: "tokes" should be "tokens". -- Caption of Figure 3: The text "This figure shows the accuracy trend of ..." should be changed to "Accuracy trend of ...".

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? -- 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? -- 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?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

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.

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.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

Limitations are discussed.

Summary

This work aims to increase the throughput of inference system by reducing the size of the KV cache. - The KV cache comes from saving forward pass computations and can take up up to 5x the model size. The 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. - 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 - Theoretical analyses (Theorem 4.1) show that such process doesn't affect the attention output much.

Strengths

- The problem is well motivated and the solution makes intuitive sense. - The empirical results are promising: - The proposed method achieves up to 5x compression and doesn't sacrifice much performance. - The proposed method is compatible with 4-bit quantization.

Weaknesses

Some design decisions are not well justified, and there are insufficient baseline comparisons.

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? - The current scheme of selecting pivotal tokens is attention-based. Are there other methods / baseline to compare to? - e.g. Keep the recent window as the current method, and change the selection of distance tokens based on word frequency? - 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? - 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). - 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}$? - Equation 1: what is there a scaling of 1/t within the softmax? Writing: - The paragraph starting at line 202 could be made clearer. - Line 238: $w,r$ have been defined in Alg 2; have they been defined in the main text? - Minor point: Fig 3: might be better to use different line style / color for models / tasks. - Please fix typos, e.g. line 202, 263, 295, 308.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

There is no direct societal impact.

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.

Strengths

Ideas and math are well written. The direction of using pivotal tokens could be interesting.

Weaknesses

The authors need to do a better job in evaluation to justify the contributions. 1. There is no benchmarks on the inference efficiency (latency and throughput) with the new approach. Was inference efficiency affected by the compression ratio? 2. 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. 3. KV state is large only when sequence length is very long. However, the tested dataset have a very short sequence length. 4. Please show the table with numbers instead of figures when claiming no loss in accuracy. In LLM evaluation, every single digit matters.

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. https://arxiv.org/abs/1911.02150

Rating

3: Reject: For instance, a paper with technical flaws, weak evaluation, inadequate reproducibility and incompletely addressed ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

1 poor

Presentation

2 fair

Contribution

1 poor

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.

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.

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? ---- EDIT: Based on discussion with authors, conditioned on the inclusion of an experiment/discussion showing the limits of this approach, I have increased my score.

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? + 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?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

3 good

Contribution

3 good

Limitations

See above in weaknesses and questions.

Follow-up questions

Thank you for the clarifications and the additional evaluation on the throughput! My concern remains that even though the idea of sparsely sub-selecting tokens is interesting, some design choices seem a bit arbitrary and that there need to be more comparison / ablation. - Regarding Q2: Thank you for the experiments, though I'm not sure using a Local Window alone would be a competitive baseline. - I didn't mean to use local window only since it clearly will lose information in the distant past; I was thinking about using local window in combination with other ways to select distant tokens; e.g. non-context-based selection, or selection based on the similarity to the attention weight matrices as Deja Vu did. - I'm a bit concerned about attention-score-based selection is not ideal, since attention scores don't translate directly to "importance" due to the presence of other components e.g. $W_V$. For example, could you explain why you chose to use attention score (i.e. softmax output) instead of e.g. the norm of $a_t$ (i.e. taking consideration $W_V, W_O$, as defined in eq (1))? It seems that in Theorem 3.1, the assumptions also take into account $W_V, W_O$. - Regarding Q3: choosing the threshold to be $1/t$ seems arbitrary and need more comparison, e.g. comparing with adaptive threshold such as selecting the top-$k$ candidates. The results in Table 7 are interesting. The current example seems to suggest that a higher compression rate on the KV cache results in shorter generation. Is this true in general?

Response

Thank you for providing these additional experiments. While I appreciate them, they do not quite address my concerns. The additional datasets (RTE, OpenBookQA, COPA) also don't really require long contexts. I was hoping to see something a bit more adversarial to give readers a picture of where this method might begin to falter. For example, what about an experiment with some sort of pseudocode like this: ``` import uuid for i in range(1000): secret_key = uuid.uuid4() test_string = f"Here is a private key: {secret_key}. {2000_filler_tokens_of_c4}. What was the private key?" output = llm(test_string) # Does the output from the llm match the secret key? ``` Such a task would require ignoring the 2000 tokens of filler content and recalling the private key from the beginning of the sequence. If the cache is sufficiently small, wouldn't the cache eviction method in this paper remove the uuid from the cache since it's not needed for the 2000 filler tokens (and therefore fail the task)? Of course, OPT isn't instruction tuned so I don't expect this to work out of the box. But I would expect some sort of experiment along these lines that pushes the heuristic cache eviction algorithm into a situation where we suspect it might be suboptimal so that we have a thorough characterization of the method. Of course this might be an unusual task that isn't typically benchmarked, but it nonetheless might be something that users encounter. I think it's totally okay for the method to fail such a test and still be a valuable contribution, but it's important to give the readers a thorough picture of the method.