Source-linked AI summary
RazorAttention: Efficient KV Cache Compression Through Retrieval Heads
Hanlin Tang, Yang Lin, Jing Lin, Qingsen Han, Shikuan Hong, Yiwu Yao, Gongyi Wang
TL;DR
Growing KV caches challenge long-context LLM deployment, and importance-based token dropping can remove information needed by future queries. RazorAttention keeps full caches for retrieval heads and compresses non-retrieval heads with compensation tokens. The authors report 70% cache compression without noticeable performance degradation.
Problem
KV-cache growth is a deployment bottleneck for long-context LLMs, while importance-based token dropping can remove information needed by future queries.
Source: §1 · pp. 1–2Method
RazorAttention retains full caches for retrieval heads and compresses non-retrieval heads using recent tokens, attention sinks, and compensation tokens.
Source: §1 · p. 2Results
RazorAttention compresses KV cache by 70% without noticeable performance degradation.
Source: §1 · p. 2Takeaways & Limitations
Its head-wise pruning criterion is compatible with FlashAttention, with negligible compensation-token computation overhead.
Source: §1 · p. 2The authors leave open why retrieval heads behave differently, how to achieve higher compression, and whether optimal configurations vary across models.
Source: §6 · pp. 8–9
Abstract
from arXiv · showhide
The memory and computational demands of Key-Value (KV) cache present significant challenges for deploying long-context language models. Previous approaches attempt to mitigate this issue by selectively dropping tokens, which irreversibly erases critical information that might be needed for future queries. In this paper, we propose a novel compression technique for KV cache that preserves all token information. Our investigation reveals that: i) Most attention heads primarily focus on the local context; ii) Only a few heads, denoted as retrieval heads, can essentially pay attention to all input tokens. These key observations motivate us to use separate caching strategy for attention heads. Therefore, we propose RazorAttention, a training-free KV cache compression algorithm, which maintains a full cache for these crucial retrieval heads and discards the remote tokens in non-retrieval heads. Furthermore, we introduce a novel mechanism involving a "compensation token" to further recover the information in the dropped tokens. Extensive evaluations across a diverse set of large language models (LLMs) demonstrate that RazorAttention achieves a reduction in KV cache size by over 70% without noticeable impacts on performance. Additionally, RazorAttention is compatible with FlashAttention, rendering it an efficient and plug-and-play solution that enhances LLM inference efficiency without overhead or retraining of the original model.
1 Introduction
RazorAttention addresses the risk that token-dropping methods discard information needed by future queries. It uses different cache strategies for retrieval and non-retrieval heads, then adds compensation tokens to reduce information loss.
1 Introduction
- Importance-based token dropping can discard information needed by future queries that differ from the main theme or arise in later conversation turns.The authors identify this risk when tokens deemed unimportant are removed from the cache.Source: §1 · p. 2
- Only a few retrieval heads can recall information across the input, while other heads primarily focus on local context or attention sinks.This motivates retaining the full cache for retrieval heads and caching only recent tokens and attention sinks for other heads.Source: §1 · pp. 2–3
- Training-free RazorAttention compresses KV cache by 70% (3X) without noticeable performance degradation.The paper reports results for contexts from 8K to 100K tokens.Source: §1 · pp. 2–3
- Compensation tokens compress dropped cache information for non-retrieval heads and further improve accuracy over direct remote-token removal.The authors report that this reduces the accuracy gap left by truncating the cache.Source: §1 · p. 2
- RazorAttention uses a head-wise pruning criterion compatible with FlashAttention and incurs negligible compensation-token computation overhead.The authors contrast this with prior importance-based methods that rely on attention weights.Source: §1 · p. 2
2 Related Work
Related work covers cache compression through quantization, token dropping, and cache sharing across attention heads. RazorAttention instead motivates head-wise pruning from differences in how attention heads use long contexts.
2 Related Work
- For typical Transformers, KV-cache compression addresses memory growth as sequence length increases.The paper focuses on typical Transformer models rather than non-Transformer architectures.Source: §2 · p. 3
- Quantization approaches compress KV caches through designed inference pipelines, error-reduction techniques, low-bit formats, or data-free distillation.The passage cites FlexGen, Atom, and QServe [1], KVQuant, KIVI, and LLM-QAT [27].Source: §2 · p. 3
- Token-dropping methods save memory by removing selected KV pairs, using strategies such as sliding windows, attention-score eviction, or layer-wise cache allocation.The passage also notes approaches based on special tokens, punctuation, key-embedding clusters, and correlations among nearby tokens.Source: §2 · p. 3
- Other methods reduce cache through attention-head sharing, ranging from one shared KV head to grouped heads or a shared latent space.The passage identifies MQA, GQA, and MLA.Source: §2 · p. 3
- RazorAttention draws on findings that some heads recall queried information, proposing head-wise pruning based on attention-mechanism interpretability.The passage connects this motivation to induction heads and a study of the property under extended inputs.Source: §2 · p. 3
3 Methodology
RazorAttention uses different cache policies for retrieval and non-retrieval heads. For ALiBi models it derives a head’s scope from attention decay; for RoPE models it identifies retrieval heads and uses compensation tokens when truncating other heads.
3.1 RazorAttention for ALiBi models
- For ALiBi models, attention to tokens farther than a head-specific scope L_h is bounded by ε, so the cache keeps recent tokens within that scope and drops more distant ones.L_h is derived from head-specific query/key matrices and LayerNorm parameters.Source: §3.1 · p. 4
3.2 RazorAttention for RoPE models
- For RoPE models, reducing retrieval-head KV-cache size causes a 16% accuracy decrease, versus 1.5% degradation when dropping remote tokens in non-retrieval heads; protecting retrieval-head caches while retaining only recent 4K tokens elsewhere preserves most performance.About 15% of heads are retrieval heads. Protecting random heads brings no performance gain.Source: §3.2 · p. 4, Table 1 · p. 5
- A compensation token represents dropped KV information; with it, RazorAttention is almost lossless when dropping 70% of the KV cache in non-retrieval heads.Remote tokens in those heads are discarded and the compensation token is added to the remaining KV cache.Source: §3.2 · pp. 4–5
3.3 Identification of Retrieval Heads
- For RoPE models, retrieval heads include echo heads that attend to a previous identical token and induction heads that attend to a preceding token followed by the current token.The method probes repeated random-token sequences, then selects the top 14% of heads by induction score and top 1% by echo score.Source: §3.3 · pp. 5–6
- RazorAttention’s RoPE procedure prunes only non-retrieval heads, keeping recent tokens and sink tokens while compressing the remaining tokens into a compensation token.Retrieval heads continue to use the original attention computation.Source: §3.3 · p. 6
4 Experiments
Experiments evaluate RazorAttention on LongBench and Needle In A Haystack, then ablate echo heads, induction heads, and the compensation token. Results report comparable LongBench performance at 3X compression and accurate recall under long inputs.
4 Experiments
- The experiments test Qwen, Llama2, Llama3, and Baichuan on LongBench and Needle In A Haystack, using H2O and StreamingLLM as comparison methods.The listed experiments run on NVIDIA GeForce RTX 4090 (24GB).Source: §4 · pp. 6–7
4.1 LongBench Evaluation
- RazorAttention achieves comparable performance to the uncompressed baseline and outperforms StreamingLLM and H2O across the tested LongBench models, even at 3X compression.The evaluation covers Qwen1.5-7B and -72B, Llama3-8B, and Baichuan2-13B.Source: §4.1 · p. 7
4.2 Needle In A Haystack Evaluation
- On Llama2-7B-80K Needle In A Haystack, RazorAttention still accurately recalls queried information under long inputs, unlike H2O, whose performance is severely degraded.The Figure 4 caption notes that H2O is incompatible with FlashAttention and produces OOM errors on longer sequences.Source: §4.2 · p. 7, Figure 4 · p. 7
4.3 Ablation Studies
- The ablation studies examine the algorithm components and configurations, including echo-head inclusion, induction-head count, and compensation tokens.The passage frames the studies as evaluating compression ratio alongside acceptable performance degradation.Source: §4.3 · p. 7
4.3.1 Importance of Echo Heads
- Adding the 1% echo heads improves retrieval performance; the paper suggests induction heads may depend on echo heads.The dependency is presented as one possible explanation, citing.Source: Figure 5 · p. 7, §4.3.1 · p. 8
4.3.2 Number of Induction Heads
- Accuracy improves as more induction heads are included, and the authors select 14% to balance compression ratio and model performance.Table 4 reports accuracy for different numbers of protected heads on Qwen1.5-7B-Chat Needle In A Haystack.Source: §4.3.2 · p. 8, Table 4 · p. 8
4.3.3 Importance of the Compensation Token
- The compensation token is critical to RazorAttention's performance and compresses much of the dropped-token information, helping maintain high accuracy despite substantial cache reduction.Figure 6 illustrates the compensation-token ablation.Source: §4.3.3 · p. 8, Figure 6 · p. 8
5 Conclusion
RazorAttention compresses the KV cache while preserving semantic information in retrieval heads and compressing remote tokens in non-retrieval heads into compensation tokens.
5 Conclusion
- RazorAttention achieves a 3X compression ratio for models using RoPE or ALiBi embeddings, with performance comparable to the original model and better accuracy and efficiency than previous methods.Its head-wise pruning criterion is compatible with FlashAttention, making it a plug-and-play compression method for extended-context inference.Source: §5 · p. 8
6 Limitation
The authors identify open questions about retrieval-head behavior, higher compression, and model-specific configuration.
6 Limitation
- The reasons attention heads behave differently and how retrieval heads operate under lengthy inputs remain open questions.The authors also see potential to increase cache reduction beyond 70%.Source: §6 · p. 8
- The authors believe cache reduction could be improved beyond the 70% already achieved.The passage presents higher compression as a future direction.Source: §6 · pp. 8–9
- The optimal RazorAttention configuration may differ across models, potentially requiring more or fewer retrieval heads.The authors say they have tested several models and plan to continue investigating configuration differences.Source: §6 · p. 9
A Appendix: Proof of Theorem 1
The proof establishes an upper bound for query-key products and uses it to show attention weights decay when positional bias becomes sufficiently large.
A Appendix: Proof of Theorem 1
- The proof aims to bound the query-key product and show that attention weights decay to zero when positional bias is significantly larger than that bound.Source: A Appendix: Proof of Theorem 1 · p. 12
- The derivation begins from the attention query and key projections of the input, then invokes the input's LayerNorm normalization.The prose identifies x as the attention-block input and says it is attained after LayerNorm.Source: A Appendix: Proof of Theorem 1 · p. 12
- The appendix combines intermediate bounds to upper-bound attention weight and derives a condition ensuring that weight is at most epsilon.It relates the condition to the logarithm of the attention weight and log(epsilon).Source: A Appendix: Proof of Theorem 1 · p. 12