Sequence can Secretly Tell You What to Discard

大型语言模型（LLMs）在广泛的任务上表现出色，但需要大量的GPU内存和消耗大量的计算资源。除了模型权重外，KV缓存占用的内存随着序列长度的增加而线性增加，成为推理的主要瓶颈。在本文中，我们介绍了一种新的优化KV缓存的方法，显著降低了其内存占用。通过全面的调查，我们发现在LLaMA2系列模型中，（i）相邻标记的查询向量之间的相似度非常高，（ii）当前查询的注意力计算可以仅依赖于前面查询的一小部分注意力信息。基于这些观察，我们提出了CORM，一种KV缓存驱逐策略，可以在不微调模型的情况下动态保留重要的键值对进行推理。我们验证了CORM在LongBench的六个任务中，将KV缓存的推理内存使用降低了高达70％，而且性能表现没有明显的下降。

本文试图解决LLMs推理中KV cache占用大量内存的问题，提出一种动态保留重要键值对的缓存淘汰策略。

本文提出的CORM策略通过利用相邻token的查询向量的相似性，只保留先前查询的一小部分的注意力信息，从而显著减少KV cache的内存占用，而不会影响模型性能。

本文通过实验验证，CORM策略可以将KV cache的内存占用减少高达70%，并且在LongBench的六项任务中没有明显的性能下降。本文的方法可以为LLMs推理的内存优化提供新思路。

最近的相关研究主要集中在LLMs的内存优化方面，例如使用分布式计算、低精度计算等方法。