RazorAttention: Efficient KV Cache Compression Through Retrieval Heads

本文提出了一种新的Key-Value（KV）缓存压缩技术，用于部署长上下文语言模型时，KV缓存的内存和计算需求提出了重大挑战。先前的方法尝试通过有选择地丢弃标记来缓解此问题，但这会不可逆地删除可能需要用于未来查询的关键信息。我们的研究发现：i）大多数注意力头主要关注局部上下文；ii）只有少数头，被称为检索头，可以基本上关注所有输入标记。这些关键观察结果促使我们为注意力头使用单独的缓存策略。因此，我们提出了RazorAttention，这是一种无需训练的KV缓存压缩算法，它为这些关键的检索头保留了完整的缓存，并丢弃非检索头中的远程标记。此外，我们引入了一种新颖的机制，涉及“补偿标记”，以进一步恢复删除标记中的信息。对各种大型语言模型（LLM）进行的广泛评估表明，RazorAttention实现了超过70％的KV缓存大小减小，而不会对性能产生明显影响。此外，RazorAttention与FlashAttention兼容，使其成为一种高效且即插即用的解决方案，可以提高LLM推理效率，而无需重复训练原始模型。

本文旨在解决KV缓存对于长序列语言模型的内存和计算需求过高的问题，同时避免之前的方法会丢失重要信息的问题。

本文提出了一种新的KV缓存压缩技术，称为RazorAttention，可以在保留所有标记信息的同时减小KV缓存的大小。该技术将注意力头分为检索头和非检索头，只保留检索头所需的所有信息，同时使用“补偿标记”机制来恢复丢失的信息。

本文的实验结果表明，RazorAttention可以将KV缓存大小减小70%以上，同时不会对性能产生明显影响。此外，该技术与FlashAttention兼容，可以提高LLM推理效率而不需要重训练模型。

在这个领域中，最近的相关研究包括《FlashAttention: A Scalable Attention Mechanism for Long Sequence》和《Reformer: The Efficient Transformer》等。