NACL: A General and Effective KV Cache Eviction Framework for LLMs at Inference Time

📄 arXiv: 2408.03675v2 📥 PDF

作者: Yilong Chen, Guoxia Wang, Junyuan Shang, Shiyao Cui, Zhenyu Zhang, Tingwen Liu, Shuohuan Wang, Yu Sun, Dianhai Yu, Hua Wu

分类: cs.CL

发布日期: 2024-08-07 (更新: 2024-08-08)

备注: Accepted by ACL 2024 (main conference, long paper)

🔗 代码/项目: GITHUB

💡 一句话要点

NACL：一种通用高效的LLM推理时KV缓存淘汰框架

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: KV缓存淘汰 长上下文建模 大型语言模型 注意力机制 推理优化

📋 核心要点

  1. 现有KV缓存淘汰方法依赖有偏的局部注意力统计，且评估指标和文本长度不足，导致淘汰效果不佳。
  2. NACL框架通过单次操作实现更优的KV缓存淘汰，结合精确注意力分数和随机淘汰策略，提升鲁棒性。
  3. 实验表明，NACL在短文本和长文本任务上分别提升80%和76%的性能，同时显著减少KV缓存占用。

📝 摘要（中文）

大型语言模型（LLMs）凭借其扩展的上下文窗口，引发了AI应用创新浪潮。然而，由于长上下文建模中KV缓存的大量内存消耗，托管这些模型的成本非常高昂。尽管一些工作提出了从KV缓存中淘汰不必要的token，但它们大多依赖于累积注意力分数的有偏局部统计，并使用不具说服力的指标（如困惑度）在不足的短文本评估上报告性能。本文提出了NACL，一个通用的长上下文KV缓存淘汰框架，可在编码阶段的单次操作中实现更优和更高效的淘汰。由于NACL的效率，我们将PROXY TOKENS EVICTION中更准确的注意力分数统计与RANDOM EVICTION的多样化随机淘汰策略相结合，旨在缓解注意力偏差问题，并增强在长上下文建模任务中维护关键token的鲁棒性。值得注意的是，我们的方法在短文本和长文本任务上的性能分别显著提高了80%和76%，在保持95%以上性能的同时，KV缓存最多减少了50%。

🔬 方法详解

问题定义：论文旨在解决大型语言模型（LLMs）推理过程中，由于长上下文建模导致KV缓存占用过多内存，从而增加计算成本的问题。现有方法主要依赖于局部注意力得分，容易产生偏差，并且评估指标和评估文本长度不足以充分验证其有效性。

核心思路：NACL的核心思路是在编码阶段通过单次操作实现更优和更高效的KV缓存淘汰。它结合了基于代理token的淘汰策略（PROXY TOKENS EVICTION）和随机淘汰策略（RANDOM EVICTION），旨在缓解注意力偏差，并增强模型在长上下文建模任务中维护关键token的鲁棒性。

技术框架：NACL框架主要包含两个核心模块：PROXY TOKENS EVICTION和RANDOM EVICTION。PROXY TOKENS EVICTION利用更准确的注意力分数统计来识别和淘汰不重要的token，而RANDOM EVICTION则通过随机淘汰token来增加多样性，避免模型过度依赖某些特定的token。这两个模块协同工作，共同实现高效的KV缓存淘汰。

关键创新：NACL的关键创新在于其结合了PROXY TOKENS EVICTION和RANDOM EVICTION，从而在准确性和多样性之间取得了平衡。与现有方法相比，NACL不仅考虑了注意力分数，还引入了随机性，从而更好地应对长上下文建模中的复杂性。此外，NACL采用单次操作进行淘汰，提高了效率。

关键设计：NACL的关键设计包括：(1) PROXY TOKENS EVICTION中注意力分数的计算方式，需要仔细选择代理token以及计算注意力得分的策略。(2) RANDOM EVICTION中随机淘汰的比例，需要根据具体任务和模型进行调整，以避免过度淘汰关键token。(3) 如何有效地结合PROXY TOKENS EVICTION和RANDOM EVICTION，例如，可以先使用PROXY TOKENS EVICTION淘汰一部分token，然后再使用RANDOM EVICTION进行补充。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，NACL在短文本和长文本任务上分别实现了80%和76%的性能提升，同时将KV缓存减少了高达50%，并且保持了95%以上的性能。这些结果表明，NACL是一种高效且有效的KV缓存淘汰框架，能够显著提高LLM的推理性能和效率。

🎯 应用场景

NACL框架可广泛应用于需要处理长上下文的LLM推理场景，例如文档总结、机器翻译、对话系统等。通过减少KV缓存的内存占用，NACL可以降低推理成本，提高推理速度，并支持更大规模的模型部署。该研究对于推动LLM在资源受限环境下的应用具有重要意义。

📄 摘要（原文）

Large Language Models (LLMs) have ignited an innovative surge of AI applications, marking a new era of exciting possibilities equipped with extended context windows. However, hosting these models is cost-prohibitive mainly due to the extensive memory consumption of KV Cache involving long-context modeling. Despite several works proposing to evict unnecessary tokens from the KV Cache, most of them rely on the biased local statistics of accumulated attention scores and report performance using unconvincing metric like perplexity on inadequate short-text evaluation. In this paper, we propose NACL, a general framework for long-context KV cache eviction that achieves more optimal and efficient eviction in a single operation during the encoding phase. Due to NACL's efficiency, we combine more accurate attention score statistics in PROXY TOKENS EVICTION with the diversified random eviction strategy of RANDOM EVICTION, aiming to alleviate the issue of attention bias and enhance the robustness in maintaining pivotal tokens for long-context modeling tasks. Notably, our method significantly improves the performance on short- and long-text tasks by 80% and 76% respectively, reducing KV Cache by up to 50% with over 95% performance maintenance. The code is available at https://github.com/PaddlePaddle/Research/tree/master/NLP/ACL2024-NACL.