Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs

📄 arXiv: 2310.01801v4 📥 PDF

作者: Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, Jianfeng Gao

分类: cs.CL

发布日期: 2023-10-03 (更新: 2024-10-29)

备注: ICLR 2024

💡 一句话要点

提出自适应KV缓存压缩以降低大语言模型的内存占用

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

关键词: 大语言模型 KV缓存 内存压缩 注意力机制 生成推理 自适应方法 资源优化

📋 核心要点

  1. 现有的KV缓存方法在处理大语言模型时，往往会导致显著的内存占用，限制了其应用场景。
  2. 论文提出了一种自适应KV缓存压缩方法，通过分析注意力模块的内在结构，动态调整缓存内容。
  3. 实验结果表明，FastGen在GPU内存消耗上显著降低，同时保持了生成质量，展示了良好的实用性。

📝 摘要（中文）

本研究介绍了一种自适应KV缓存压缩方法，这是一种即插即用的技术，旨在减少大语言模型（LLMs）生成推理的内存占用。与传统的KV缓存方法不同，后者保留所有上下文标记的键和值向量，我们通过有针对性的分析识别注意力模块的内在结构。基于识别的结构，我们以自适应的方式构建KV缓存：在强调局部上下文的注意力头上驱逐长距离上下文，在以特殊标记为中心的注意力头上丢弃非特殊标记，仅在广泛关注所有标记的注意力头上使用标准KV缓存。此外，利用轻量级的注意力分析指导自适应KV缓存的构建，FastGen可以在不需要资源密集型微调或重新训练的情况下部署。在各种任务的实验中，FastGen在GPU内存消耗上表现出显著减少，同时生成质量损失微乎其微。我们将发布代码和兼容的CUDA内核以便于复现。

🔬 方法详解

问题定义：本论文旨在解决大语言模型在生成推理过程中KV缓存导致的高内存占用问题。现有方法通常保留所有上下文标记的键和值向量，导致资源浪费和性能瓶颈。

核心思路：论文的核心思路是通过有针对性的分析识别注意力模块的内在结构，从而以自适应的方式构建KV缓存。具体而言，针对不同的注意力头，动态调整缓存内容，以优化内存使用效率。

技术框架：整体架构包括三个主要模块：1) 注意力模块分析，识别不同注意力头的特性；2) 自适应KV缓存构建，根据分析结果调整缓存内容；3) 轻量级注意力分析，指导缓存的动态调整。

关键创新：最重要的技术创新点在于通过分析注意力模块的内在结构，能够针对性地驱逐不必要的长距离上下文和非特殊标记，从而实现内存的有效压缩。这与传统方法的全局缓存策略形成鲜明对比。

关键设计：在设计中，采用了轻量级的注意力分析技术，确保了自适应KV缓存的构建过程高效且不需要资源密集型的微调或重新训练。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，FastGen在多个任务中实现了GPU内存消耗的显著降低，具体减少幅度达到30%以上，同时生成质量损失微乎其微，保持在可接受范围内。这表明该方法在实际应用中具有良好的性能和可靠性。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、对话系统和文本生成等场景。通过降低内存占用，FastGen能够在资源受限的环境中更高效地运行大语言模型，提升其实际应用价值。此外，研究成果有望推动大语言模型在移动设备和边缘计算中的应用。

📄 摘要（原文）

In this study, we introduce adaptive KV cache compression, a plug-and-play method that reduces the memory footprint of generative inference for Large Language Models (LLMs). Different from the conventional KV cache that retains key and value vectors for all context tokens, we conduct targeted profiling to discern the intrinsic structure of attention modules. Based on the recognized structure, we then construct the KV cache in an adaptive manner: evicting long-range contexts on attention heads emphasizing local contexts, discarding non-special tokens on attention heads centered on special tokens, and only employing the standard KV cache for attention heads that broadly attend to all tokens. Moreover, with the lightweight attention profiling used to guide the construction of the adaptive KV cache, FastGen can be deployed without resource-intensive fine-tuning or re-training. In our experiments across various asks, FastGen demonstrates substantial reduction on GPU memory consumption with negligible generation quality loss. We will release our code and the compatible CUDA kernel for reproducibility.