Identify Critical KV Cache in LLM Inference from an Output Perturbation Perspective

Authors: Yuan Feng, Junlin Lv, Yukun Cao, Xike Xie, S Kevin Zhou
Primary area: applications to computer vision, audio, language, and other modalities
Keywords: KV Cache Compression, Efficient LLM Inference

Abstract:
Large language models have driven numerous paradigm shifts in the field of natural language processing, achieving remarkable success in various real-world applications through scaling model size and leveraging long-sequence context reasoning.
However, the transformer architecture, which relies on self-attention, incurs substantial storage and runtime costs when handling long-sequence inference, particularly due to the generation of extensive Key-Value (KV) cache.
Recent studies aim to mitigate storage and latency issues while maintaining output quality by reducing the KV cache size, through the elimination of less critical entries, yet they rely on a basic empirical intuition of identifying critical cache entries based solely on top attention weights.
In this paper, we present the first formal investigation into the problem of identifying critical KV cache entries from the perspective of attention output perturbation.
By analyzing the output perturbation caused when only critical KV cache entries are used instead of the entire cache, we reveal that, in addition to the commonly used attention weights, the value states within KV entries and the pretrained parameters matrix are also important. Based on this finding, we propose a novel perturbation-constrained selection algorithm to identify critical cache entries by optimizing the worst-case output perturbation.
Extensive evaluations on 16 datasets from Longbench, along with detailed empirical analysis, have comprehensively confirmed the effectiveness of constraining output perturbation perspective in identifying critical KV cache. When combined with state-of-the-art cache eviction methods, it can achieve up to an additional 34\% cache memory savings while maintaining the same generation quality.

Decision: reject
Reviewer scores: 5, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1729462723171

thread_type:
official_review

<thread_object>
creator: reviewer_nkes

<content>
summary:
The paper proposes a novel method to identify critical KVs in the KVCache for effective KV eviction. Unlike previous methods that solely rely on the attention score, the paper focuses on the attention output perturbation instead. It comes up with a theoretical upper bound for the attention output $L_1$ perturbation to optimize upon. More specifically, the paper proposes to include a small faction of KVs in the critical KVCache to capture the attention score while the rest to capture the derived upper bound term $\sum_{i=1}^n N_i A_i \Vert V_{i, :}\Vert_1$. Empirically, the paper shows that combined with this critical KV identification algorithm, they can consistently outperform the base version of SnapKV, Pyramid, and AdaKV on LongBench and over LLaMA-3.1-8B and Mistral-7B.

strengths:
1. The KVCache compression problem is indeed well-motivated under the long-context scenario. Conventional approaches usually identify critical KVs based on the heuristic that tokens with a higher attention score are more important. In contrast, this work proposes a novel term derived from the attention output perturbation to capture the importance of different KVs with theoretical guarantees. 

2. The proposed method can be easily integrated into scenarios requiring critical KV identification, such as KV Compression or sparse-attention serving systems. 

3. The additional ablation studies further complement its empirical evaluations and the conclusions that the proposed method can reduce attention output perturbations compared to the heuristics-based approaches.

weaknesses:
1. The notations and writing for the mathematical derivations are a bit confusing (see the question section for details.), which is the key contribution of this work and thus needs to be improved to give a clearer and more rigorous proof.

2. Empirically, the experiments focus mostly on small models around 7B. Some results on a larger model (e.g., ~70B) would be more convincing. Additionally, the paper can add some needle-in-the-haystack-type tasks to further show the effectiveness of the proposed method in terms of critical KV identification. One of my concerns for the general eviction-based sparse attention approach is that in a multi-turn or complex reasoning task, the evicted KVs can be important to a later decoding step.

questions:
1. In Eq. (6), $L$ is defined to be the $L_1$ norm of the attention output deviation. However, in Eq. (12) - Eq. (15), $L$ directly equals the derived upper bound for the attention output deviation. This is a bit confusing and seems to be a typo to me as the author should either use $L \leq$ or $\theta =$ for the derived upper bounds instead of $L=$ for Eq. (12) - Eq. (15).

2. If I understand Theorem 3 correctly, the author is actually trying to claim $\arg\max_{N}\gamma = \arg\min_{N}\theta$ instead of $\max_{N}\gamma = \min_{N}\theta$, a simple counter-example is when $\sum_i^{n}N_i A_i = 1$ (the selected tokens can capture all the mass), in which case, $\min_{N}\theta=0$ and $\max_{N}\gamma>0$. 

3. The proof of Theorem 3 is not quite straightforward for me. Firstly, the proposed inequality $\sum_i^n N_i A_i > \sum \text{Top}_k (A, b') > 0.5$ is not guaranteed for an arbitrary set of $N_i$ even if we have assumption 1. As we can selectively choose $N_i$ to not include tokens from $\text{Top}_k (A, b')$, in which case this is not guaranteed. Secondly, even if $\sum_i^n N_i A_i > 0.5$ is guaranteed, the derivation of $\arg\max_N \gamma = \arg\min_N \theta$ is not quite straightforward. More specifically, the assumption that 
$(2-1/\sum_i^n N_i A_i)>0$  is not a sufficient condition to derive that $\arg\max_N\gamma = \arg\min_N\theta$. For instance, I can have two sets of $N_i$, the first set has $\sum_i N_i A_i = 0.51$ and the second set has $\sum_i N_i A_i = 0.99$, then I can adversarially set $\Vert V_i\Vert_1$ so that the $\gamma$ corresponds to the first set of $N_i$ is slightly higher than that of the second set of $N_i$. According to your claim, you would choose the first set of $N_i$, which gives you a slightly higher $\gamma$ but way smaller $2-1/\sum_i^n N_i A_i$ in which case $\arg\max_N \gamma \neq \arg\min_N \theta$.  Please correct me if I am wrong on this. 

Suggestions on 3: You might want to move your algorithm description before stating the theorem as I am guessing the proposed inequality $\sum_i^n N_i A_i > \sum \text{Top}_k (A, b') > 0.5$ is based on your algorithm rather than for arbitrary $N_i$. And for theorem 3, if you want to prove the same bound, I would suggest you modify your assumption 1 to be $>0.51$ rather than $>0.5$ then you can lower bound the $2-1/\sum_i^n N_i A_i$ term with $2-\frac{1}{0.51}$ and can further upper bound the current $\theta$ with $\hat{\theta}=C-(2-\frac{1}{0.51}) \gamma$ and reach to the conclusion that $\arg\max_N\gamma = \arg\min_N\hat{\theta}$. 

4. Empirical-wise, if time allows I would be happy to see some results for a larger LLM (e.g., 70B) as well as for some needle-in-the-haystack type of task.

In general, the problem is well-motivated, and the proposed idea with theoretical support is novel. I would be happy to raise the score if the authors can address my concerns during the rebuttal period.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730235766696

thread_type:
official_review

<thread_object>
creator: reviewer_19en

<content>
summary:
This paper proposes a novel method to identify essential KV cache values based on a perturbation-constrained selection algorithm. Empirical results highlight this approach's performance in small KV budgets against other KV-cache eviction methods.

strengths:
This paper introduces a novel perturbation-constrained selection algorithm. The criteria underlying the approach are well-motivated and rooted in mathematical principles with proof and derivation. The algorithm’s design is intuitive and straightforward, facilitating easy integration into existing frameworks.

weaknesses:
While the overall paper is relatively sound, there are some writing and formatting issues I would like to see addressed:

* Section 3.2: In my opinion, this section needs to be reformatted and partially rewritten to be more concise with clearer structure. As the most essential part of the paper, it is currently challenging to read, with many notations introduced ambiguously and references placed somewhat haphazardly.
* Equation Alias: The author should create an alias for $\frac{N \cdot A}{\sum_n N_i A_i}$ to simplify and streamline the equations that reference this expression.
* Assumption 1 and Theorem 3: I believe these should form their own subsection, as they represent the core approach, while the other sections primarily provide proofs and background information.

Additionally, there are several points I still find unclear, and I would appreciate clarification from the authors:

* Assumption 1 Interpretation: I am having difficulty reconciling the assumptions made in Assumption 1 with several statements in the paper:
 
> Lines 17–19: "Yet they rely... attention weights"  

> Lines 60–63: "Although the term... higher attention weights"
 
> Lines 128–129: "However... critical entries"

These statements critique the empirical reliance on attention weights as indicators of importance. Could the authors clarify how they reconcile these criticisms with their own assumption about using attention scores to represent the importance of key-value entries?

* Arbitrary Choice of $\alpha$: I also find that Appendix A does not sufficiently address the arbitrary selection of $\alpha$. Does the model’s performance change with varying $\alpha$, and would such variation impact the conclusions drawn?

* Performance and Budget Constraints: Finally, I feel that the average improvement achieved is often marginal. Moreover, the performance degradation at budget $b$ appears too severe to be practical in real-world settings. Could the authors provide insight on the specific budget value 1024 < $b < 6711$ perhaps (2048, 4096) at which both their approach and other key-value eviction strategies yield comparable performance to the baseline? Additionally, is there a discernible advantage to the proposed method under these conditions?

questions:
Please see weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730674763754

thread_type:
official_review

<thread_object>
creator: reviewer_ye2p

<content>
summary:
This research examines KV cache compression by analyzing its impact on output perturbation. The study identifies critical KV cache elements by evaluating their maximum potential effect on output disruption, taking into account both value vector norms and attention score matrices. Through extensive testing across 16 datasets, the proposed methodology demonstrates consistent performance improvements.

strengths:
1. KV cache compression represents a critical and increasingly relevant challenge in recent literature.
2. The study evaluates three distinct KV cache eviction strategies, each yielding comparable performance improvements.
3. Both the analytical framework and the implemented algorithm are reasonable.

weaknesses:
This study can be further enhaned by enriching its experiments and empirical analyses (see suggestions & questions for more details).

questions:
# Writing
On the proposed kv cache eviction algorithm, it would be better to discuss the connection and difference with FastGen [1] and MInference [2] in more details, given the similarity of the resulting algorithm. Also, as to the output perturbation, I feel it would be helpful to discuss the  existing success of this concept's adaptation in other deep learning algorithms (e.g., [3] and [4])

# Analyses
Additional theoretical analysis could strengthen the output perturbation perspective. A key area for investigation is the correlation between worst-case output perturbation predictions and actual perturbation measurements in real-world applications. 

# Experiments
Intuitively, the proposed algorithm applies not only to long-context tasks, but to all generation tasks. However, in the evaluation, only long-context tasks are used for evaluation. I highly suggest the author to evaluate the proposed algorithm in other popular LLM benchmarks (e.g.,  AlpacaEval). This is necessary to understand the performance of the proposed algorithm on a wider spectrum of scenarios. 

1. Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs
2. MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention
3. Understanding the Difficulty of Training Transformers
4. Catformer: Designing Stable Transformers via Sensitivity Analysis
</content>
</thread_object>

</thread>
<thread>
creation date:
1730685625208

thread_type:
official_review

<thread_object>
creator: reviewer_hvtb

<content>
summary:
This paper optimizes the critical KV selection for existing KV eviction methods during LLM inference. To identify critical KV, this paper uses L1 distance to measure the perturbation of attention outputs w./w.o. KV eviction under certain budgets, and draws the conclusion that both attention weights and value states projected through parameter matrix are important for reducing the output perturbation. It then proposes a perturbation-constrained selection algorithm which adopts top-k sampling to select two sets of KV: one for high attention weights and the other for both the norm of the projected value states and attention weights, respectively. Experimental results show that when integrating this selection algorithm into existing KV eviction methods like SnapKV, AdaKV, the output quality can be improved.

strengths:
1.	Well-motivated: this paper points out that current KV eviction methods relying solely on attention weights, which is not a proper metric for identifying critical KV. 
2.	Sufficient theoretical analysis: the conclusion aligns with the motivation. 
3.	Good writing: easy to follow.

weaknesses:
In general, the contributions are incremental, many claims lack supports and the experiments are not sufficient enough. 

1.	Many claims lack further explanation or theoretical supports. For example, why a subset of KV is sufficient for LLM inference? Instead of claiming that it is a commonly accepted assumption from empirical studies, more insights from theoretical perspective are expected since this paper emphasizes its theoretical contributions. 
(1) Why L1 distance is a good metric for measuring the output perturbation? L1 distance is well-known for its sparsity-sensitivity, but the attention outputs are dense. Any further explanation or theoretical supports for this metric? 
(2) What’s the relationship between output perturbation and the generation quality? Perturbation measures the deviation from the original output, but does this deviation have a significant impact on determining the output token? 
(3) Why set alpha to 0.25? The attention weights obey the power-law distribution, but why setting a fixed hyper-parameter is reasonable for any models and input distributions? Deriving this parameter directly from the attention weight distribution seems to be more reasonable.  

2.	The comparison is far from sufficiency, omitting many important baselines. For example, StreamingLLM [1] should be considered since it is an important empirical work. This paper needs this comparison to verify the gains from theorical analysis. Besides, the naïve attention mechanism (without any KV dropping) should also be included. The authors should demonstrate the generation quality gap between eviction methods and the ground-truth.  

[1] Xiao, Guangxuan, et al. "Efficient streaming language models with attention sinks." arXiv preprint arXiv:2309.17453 (2023).

questions:
Justify claims. Add more experiments comparing with baselines.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732288157246

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global Response

comment:
We deeply appreciate all reviewers' careful evaluation and insightful feedback on our work. In response, we have revised our manuscript during the rebuttal phase, especially enhancing the clarity of Section 3. Furthermore, we included additional evaluation results and discussions in the appendix.
The revised version has been uploaded. If there are any additional concerns, please let us know, and we would be happy to engage in further discussions.
The key revisions are summarized as follows:

1. Presentation and Theoretical Improvements
    * **Section 3**: Reorganized for clearer presentation and more rigorous theoretical proofs. 
    *  **Appendix F**: Added a theoretical analysis of previous methods from the perspective of output perturbation. 
2. Experimental Additions
   * **Appendix E**: Introduced results from the **Needle-in-a-Haystack** Test.
   *  **Appendix K**: Added a new instruction-following benchmark, **AlpacaEval**, to demonstrate the broader applicability of our method.
   *  **Appendix H**:  Conducted a detailed comparison with an earlier empirical study (**StreamingLLM**).
3. Broader Discussion
   * **Appendix I**: Discussed the relationship between output perturbations and final generation quality.
   *  **Appendix J**: Expanded the related work discussion and future research directions.
   * **Appendix G**:  Discussed the choice of distance metric.

We hope these revisions effectively address the reviewers' concerns and contribute to enhancing the overall quality of the paper. If there are any remaining concerns, we would be happy to engage in further discussions. We look forward to hearing from you.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732845313805

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response to Reviewers and ACs

comment:
We would like to express our sincere gratitude to all the reviewers **[Reviewer 1 (hVtB), Reviewer 2 (YE2P), Reviewer 3 (19en), and Reviewer 4 (nKeS)]** for their thoughtful and supportive feedback. We are pleased that the reviewers recognized the following key aspects of our work:

* **Well-motivated (Reasonable)** work for identifying critical KV cache from a perturbation perspective. **[Reviewers 1, 2, 3, 4]**
* **Novel algorithm** for constraining worst-case perturbation through a two-stage selection process. **[Reviewers 3, 4]**
* **Solid theoretical foundation:** Our algorithm aligns well with the theoretical analysis. **[Reviewers 1, 2, 3, 4]**
* **Broad applicability** in enhancing existing methods. **[Reviewers 1, 2, 3, 4]**

In response to the reviewers' suggestions, we have refined the presentation of our paper and incorporated additional experiments. If there are any remaining concerns, we would be happy to discuss them further and look forward to your response.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper investigates the identification of critical key-value (KV) cache entries in large language model (LLM) inference through a novel perturbation-constrained selection algorithm. The problem of efficient KV cache compression is well-motivated, and the authors propose an innovative approach that combines attention weights and value norms to constrain output perturbation. The paper presents solid theoretical foundations and demonstrates the method’s ability to integrate with existing state-of-the-art (SOTA) cache eviction algorithms, achieving additional memory savings while maintaining output quality.

However, the paper has some limitations that prevent its acceptance at this stage. While the theoretical analysis is sound, several assumptions, such as the use of L1 distance and specific hyperparameter choices, lack comprehensive justification. Empirical evaluations, although extensive, are primarily focused on smaller models and long-context tasks, leaving broader applicability underexplored. Additionally, key baselines like larger LLMs and StreamingLLM are not fully addressed, which limits the understanding of the proposed method’s comparative advantages. Some parts of the paper, particularly the mathematical notations and key derivations, remain challenging to follow, which detracts from the clarity of the contributions.

It is unfortunate that I have to recommend rejecting this paper. The topic and approach are promising, but further refinement in theoretical justification, broader experimental evaluations, and improved clarity are necessary for the work to reach its full potential.