Attention Score is not All You Need for Token Importance Indicator in KV Cache Reduction: Value Also Matters

TL;DR

The paper tackles the KV-cache memory bottleneck in long-context LLMs by challenging the assumption that attention scores alone determine token importance. It introduces Value-Aware Token Pruning (VATP), which multiplies the attention score by the $\ell_{1}$ norm of the corresponding value vector to assess token importance, and preserves the initial attention-sink tokens to stabilize attention. Extensive experiments on LLaMA2-7B-chat and Vicuna-v1.5-7B across 16 LongBench tasks show VATP consistently outperforming attention-score–only baselines with minimal overhead, demonstrating the practical value of incorporating value-vector norms in KV-cache reduction. The work reveals a previously underappreciated factor in long-context inference efficiency and lays groundwork for more robust, efficient KV-cache pruning strategies in real-world deployments.

Abstract

Scaling the context size of large language models (LLMs) enables them to perform various new tasks, e.g., book summarization. However, the memory cost of the Key and Value (KV) cache in attention significantly limits the practical applications of LLMs. Recent works have explored token pruning for KV cache reduction in LLMs, relying solely on attention scores as a token importance indicator. However, our investigation into value vector norms revealed a notably non-uniform pattern questioning their reliance only on attention scores. Inspired by this, we propose a new method: Value-Aware Token Pruning (VATP) which uses both attention scores and the $ \ell_{1} $ norm of value vectors to evaluate token importance. Extensive experiments on LLaMA2-7B-chat and Vicuna-v1.5-7B across 16 LongBench tasks demonstrate that VATP outperforms attention-score-only baselines in over 12 tasks, confirming the effectiveness of incorporating value vector norms into token importance evaluation of LLMs.

Figures (3)

• Figure 1: Typical attention map (logarithmic) and value vector norm patterns in LLaMA2-7B-chat. Key observations include: (1) The $\ell_{1}$ norms are non-uniformly distributed across tokens in all layers and heads. (2) In figure (b), for most heads in layers 3-31, regardless of the input text, there are two attention sink xiao2024efficient tokens at the beginning of the text. Contrary to their massive attention scores, their $\ell_{1}$ norms are close to 0 (highlighted in red). (3) In some heads of the last layer, the second attention sink token in figure (b) has a smaller attention score than other tokens, while its $\ell_{1}$ norm is significantly larger than those of other tokens.
• Figure 2: An illustration example of Value-Aware Token Pruning (VATP) method. VATP uses the product of value vector norm and attention score to evaluate the importance of each token's KV cache. The KV cache of the token with the lowest importance score is removed. Here we use the accumulated attention score proposed by H$_2$O zhang2023ho, resulting in H$_2$O w/ VATP variant.
• Figure 3: Performance on 2WikiMultihopQA of the LLaMA2-7B-chat with varying KV Cache Ratios.