RazorAttention：通过检索头实现高效的KV缓存压缩
RazorAttention: Efficient KV Cache Compression Through Retrieval Heads
AI总结：
针对长上下文语言模型KV缓存内存大的问题，提出免训练压缩算法RazorAttention，通过为检索头保留完整缓存并引入补偿token机制，在保持性能的同时将缓存大小减少70%以上。
摘要：
Key-Value (KV) 缓存的内存和计算需求为部署长上下文语言模型带来了重大挑战。先前的方法试图通过选择性地丢弃token来缓解此问题，但这会不可逆地抹去未来查询可能需要的关键信息。在本文中，我们提出了一种保留所有token信息的全新KV缓存压缩技术。我们的调查表明：i) 大多数注意力头主要关注局部上下文；ii) 只有少数头（称为检索头）能够实质上关注所有输入token。这些关键观察促使我们对注意力头采用分离的缓存策略。因此，我们提出了RazorAttention，一种免训练的KV缓存压缩算法，它为这些关键的检索头保留完整缓存，并丢弃非检索头中的远程token。此外，我们引入了一种涉及“补偿token”的新机制，以进一步恢复被丢弃token中的信息。在多种大型语言模型（LLMs）上的广泛评估表明，RazorAttention将KV缓存大小减少了70%以上，且对性能没有明显影响。此外，RazorAttention兼容FlashAttention，使其成为一种高效且即插即用的解决方案，能够提升LLM推理效率，而无需额外开销或对原始模型进行重新训练。
英文摘要：
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.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
