SnapKV: LLM Knows What You are Looking for Before Generation

大型语言模型（LLMs）在处理广泛上下文方面取得了显著进展，其中键值（KV）缓存在提高性能方面发挥了至关重要的作用。然而，随着输入长度的增加，KV缓存的增长对内存和时间效率提出了挑战。为了解决这个问题，本文介绍了SnapKV，一种创新的、无需微调的方法，可以在仍然提供实际应用中可比较的性能的情况下，有效地最小化KV缓存大小。 我们发现模型中的每个注意力头在生成过程中一致地关注特定的提示注意力特征。同时，这种稳健的模式可以从提示末尾的“观察”窗口中获得。基于这一洞察力，SnapKV通过为每个注意力头选择聚类的重要KV位置自动压缩KV缓存。我们的方法显著降低了处理长输入序列时不断增长的计算开销和内存占用。具体而言，当处理16K令牌的输入时，SnapKV相比基线模型实现了一致的解码速度，生成速度提高了3.6倍，内存效率提高了8.2倍。同时，在16个长序列数据集上，它与基线模型保持了可比较的性能。此外，SnapKV可以在单个A100-80GB GPU上处理高达380K个上下文令牌，使用HuggingFace实现并进行微小更改，仅在“草堆中的针”测试中出现可忽略的精度下降。进一步的综合研究表明，SnapKV在实际应用中具有潜力。

如何在处理长序列输入时提高大语言模型的性能和效率？

通过观察模型生成时每个注意力头的注意特征，自动压缩关键-值缓存（KV cache），选择每个注意力头的聚类重要KV位置，从而显著减少计算和内存开销。

SnapKV可以在处理16K令牌的输入时，将生成速度提高3.6倍，内存效率提高8.2倍，并在维持与基线模型相当的性能的同时，处理长序列数据集。此外，SnapKV可以在单个A100-80GB GPU上处理多达380K上下文令牌，只有微小的准确度下降，并且使用HuggingFace实现代码。

最近的相关研究包括：《Efficient Content-Based Sparse Attention with Routing Transformers》、《Reformer: The Efficient Transformer》、《Longformer: The Long-Document Transformer》等。