Technical Summary: CriticalKV

Problem Statement

Large Language Models (LLMs) based on the Transformer architecture face significant deployment challenges due to the quadratic computational and storage costs associated with the Key-Value (KV) cache during autoregressive generation. As sequence lengths increase, the KV cache expands dramatically, leading to high GPU memory overhead and increased I/O latency.

While recent cache eviction methods (e.g., H2O, Scissorhands, SnapKV) attempt to mitigate this by pruning non-critical entries, they rely heavily on empirical heuristics. Specifically, these methods assume that entries with higher attention weights are the most critical. The authors argue that this assumption lacks formal grounding and raises two key questions:

1. What criteria truly determine the criticality of a KV cache entry?
2. Is reliance on attention weights alone sufficient for identifying critical entries?

The paper posits that current methods neglect the crucial role of value states and pretrained parameter matrices, potentially leading to suboptimal eviction decisions and significant quality degradation (compression loss).

Methodology

The paper introduces CriticalKV, a formal framework for identifying critical KV cache entries based on output perturbation analysis.

Theoretical Framework

The authors define the critical entry identification problem as minimizing the perturbation in the attention output when replacing the full KV cache with a subset of critical entries.

• Perturbation Definition: The output perturbation $L$ is defined as the $L_1$ distance between the attention output of the full cache ($o$) and the selected subset ($\hat{o}$): $L = \|o - \hat{o}\|_1$.
• Upper Bound Derivation: Through theoretical analysis, the authors derive an upper bound $\theta$ for this perturbation. Crucially, this bound reveals that perturbation is influenced not only by attention weights ($A$) but also by the projected value states ($V W_O$), where $W_O$ is the output projection matrix.
• Key Insight: The derived bound demonstrates that minimizing perturbation requires considering both the attention weights and the magnitude of the value states projected through the parameter matrix. Relying solely on attention weights minimizes a "relaxed" upper bound, which is theoretically suboptimal.

Algorithm: Perturbation-Constrained Selection

To operationalize this theory, the authors propose a two-stage greedy selection algorithm (Algorithm 1) designed to minimize the derived upper bound of the worst-case perturbation:

1. Stage 1 (High Attention Weights): A fraction of the cache budget ($b' = b \times \alpha$, where $\alpha=0.5$) is allocated to select entries with the highest attention weights. This stage ensures that the cumulative attention weight of the selected subset exceeds 50% of the total, satisfying a key assumption for the theoretical bound.
2. Stage 2 (Perturbation Constrained): The remaining budget ($b''$) is used to select entries that jointly maximize the product of attention weights and the $L_1$ norm of the projected value states ($A_i \|V_{i,:}\|_1$). This step directly targets the minimization of the derived perturbation upper bound.

The algorithm is designed as a plug-and-play enhancement that can be integrated into existing cache eviction pipelines (e.g., SnapKV, AdaKV, HeadKV) without altering their underlying accumulation processes.

Key Contributions

1. Formal Problem Definition: The paper is the first to formalize the critical cache entry selection problem from the perspective of output perturbation, moving beyond empirical heuristics.
2. Theoretical Insight: It identifies that attention weights alone are insufficient for identifying critical entries; the value states projected through the parameter matrix are equally essential.
3. Novel Algorithm: It proposes a perturbation-constrained selection algorithm that optimizes the worst-case output perturbation.
4. Universal Enhancement: The method serves as a universal enhancement for existing state-of-the-art (SOTA) eviction techniques, requiring negligible computational overhead.

Experimental Results

The authors evaluated CriticalKV by integrating it with three SOTA eviction methods (SnapKV, AdaKV, HeadKV) across three distinct LLMs (Llama-3.1-8B, Mistral-7B-v0.3, Qwen2.5-32B) on 29 datasets from the Ruler and LongBench benchmarks.

• Performance Gains: On average, the algorithm reduced the compression loss by more than half across all tested configurations.
  • On the Ruler benchmark (40% cache size), integrating the method with HeadKV on Qwen2.5-32B reduced the average loss from 13.7% to 3.4%.
  • On LongBench, the method improved performance in 88 out of 90 test cases across long-dependency domains.
• Robustness: The improvements were consistent across varying cache sizes (from 2.5% to 100%) and different model architectures.
• Efficiency: The algorithm introduces negligible computational overhead. The prefilling time-to-first-token (TTFT) increased by only ~0.06s for a batch size of 1 at 32K context length, and decoding latency remained identical to the baseline eviction methods.
• Perturbation Analysis: Empirical analysis confirmed that the algorithm reduces output perturbation in over 92% of attention heads for Llama models and significantly lowers perturbation in final-layer hidden states.

Significance and Claims

The paper claims to offer a new, formally grounded perspective for cache eviction. By shifting the focus from intuitive attention-weight heuristics to a theoretical minimization of output perturbation, CriticalKV addresses a fundamental limitation in current eviction strategies.

The authors emphasize that their work:

• Provides a theoretical justification for why value states matter in cache eviction.
• Offers a practical, low-overhead solution that substantially improves generation quality under memory constraints.
• Opens promising avenues for future research in optimizing critical cache selection through theoretical analysis of output perturbation.

The code is publicly available, and the authors note that the repository hosts an enhanced version of their approach in a follow-up study.