Get More with LESS: Synthesizing Recurrence with KV Cache Compression for Efficient LLM Inference

📄 arXiv: 2402.09398v2 📥 PDF

作者: Harry Dong, Xinyu Yang, Zhenyu Zhang, Zhangyang Wang, Yuejie Chi, Beidi Chen

分类: cs.LG, cs.AI

发布日期: 2024-02-14 (更新: 2024-06-12)

🔗 代码/项目: GITHUB

💡 一句话要点

提出LESS以解决大语言模型推理中的KV缓存瓶颈问题

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

关键词: 大语言模型 KV缓存 内存优化 推理效率 自然语言处理 信息保留 驱逐机制

📋 核心要点

  1. 现有KV缓存方法在需要回忆大量先前标记的任务中效果有限，导致内存瓶颈问题。
  2. 提出LESS，通过将恒定大小的缓存与驱逐方法集成，允许在后续解码步骤中查询所有标记。
  3. LESS在多种任务中表现出色，能够减少与缓存所有内容的性能差距，甚至在某些情况下匹配其性能。

📝 摘要（中文）

许多计算因素限制了大型语言模型的广泛部署。本文关注于解码过程中由键值（KV）缓存引起的内存瓶颈。现有的KV缓存方法通过修剪或驱逐相对不重要的KV对来减少缓存的内存占用，但在需要回忆大多数先前标记的任务中效果有限。为了解决这一问题，我们提出了LESS，这是一种将（几乎免费的）恒定大小缓存与基于驱逐的缓存方法简单集成的方案，使得所有标记可以在后续解码步骤中查询。LESS在多种任务中表现出保留信息的能力，能够减少与缓存所有内容之间的性能差距，有时甚至能够匹配其性能，同时保持高效。相关代码可在https://github.com/hdong920/LESS找到。

🔬 方法详解

问题定义：本文要解决的具体问题是大型语言模型推理中的KV缓存内存瓶颈。现有方法通过修剪或驱逐KV对来减小内存占用，但在需要回忆大量先前标记的任务中效果有限。

核心思路：论文的核心解决思路是提出LESS，通过将一个恒定大小的缓存与现有的驱逐方法结合，使得所有标记可以在后续的解码步骤中被查询。这种设计旨在提高信息的保留能力，减少性能损失。

技术框架：整体架构包括一个恒定大小的KV缓存模块和一个驱逐机制。解码过程中，LESS会在需要时从恒定缓存中提取信息，同时利用驱逐机制管理内存使用。

关键创新：最重要的技术创新点在于LESS的设计，使得在保持高效内存使用的同时，能够查询所有先前的标记。这与现有方法的本质区别在于，LESS不再依赖于仅保留部分KV对。

关键设计：关键设计包括恒定大小缓存的参数设置，以及驱逐机制的具体实现。损失函数和网络结构的选择也经过精心设计，以确保在多种任务中都能有效运行。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

在多种任务中，LESS能够显著减少与缓存所有内容之间的性能差距，某些情况下甚至能够匹配其性能。实验结果表明，LESS在内存使用和推理效率上均有显著提升，展示了其在实际应用中的潜力。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、对话系统和文本生成等。通过提高大型语言模型的推理效率，LESS可以在资源受限的环境中实现更广泛的应用，具有重要的实际价值和未来影响。

📄 摘要（原文）

Many computational factors limit broader deployment of large language models. In this paper, we focus on a memory bottleneck imposed by the key-value (KV) cache, a computational shortcut that requires storing previous KV pairs during decoding. While existing KV cache methods approach this problem by pruning or evicting large swaths of relatively less important KV pairs to dramatically reduce the memory footprint of the cache, they can have limited success in tasks that require recollecting a majority of previous tokens. To alleviate this issue, we propose LESS, a simple integration of a (nearly free) constant sized cache with eviction-based cache methods, such that all tokens can be queried at later decoding steps. Its ability to retain information throughout time shows merit on a variety of tasks where we demonstrate LESS can help reduce the performance gap from caching everything, sometimes even matching it, all while being efficient. Relevant code can be found at https://github.com/hdong920/LESS.