CORM: Cache Optimization with Recent Message for Large Language Model Inference

📄 arXiv: 2404.15949v2 📥 PDF

作者: Jincheng Dai, Zhuowei Huang, Haiyun Jiang, Chen Chen, Deng Cai, Wei Bi, Shuming Shi

分类: cs.CL, cs.AI, cs.LG

发布日期: 2024-04-24 (更新: 2024-06-21)

💡 一句话要点

提出CORM以优化大语言模型推理中的KV缓存问题

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

关键词: 大语言模型 KV缓存 推理优化 内存管理 动态驱逐策略 Transformer模型 自然语言处理

📋 核心要点

  1. 现有大语言模型在推理时面临GPU内存和计算资源消耗过大的挑战，尤其是KV缓存的内存使用随序列长度增加而线性增长。
  2. 本文提出CORM，一种动态保留关键键值对的KV缓存驱逐策略，基于相邻token查询向量的相似性和前序查询的注意力信息进行优化。
  3. 实验结果显示，CORM在LongBench的六个任务中将KV缓存的推理内存使用降低了70%，且性能几乎没有下降，展示了其有效性和兼容性。

📝 摘要（中文）

大语言模型（LLMs）在多种任务中表现出色，但其推理过程中的GPU内存消耗和计算资源需求巨大。尤其是KV缓存的内存使用随着序列长度线性增长，成为推理的主要瓶颈。本文提出了一种创新的KV缓存优化方法CORM，显著减少了其内存占用。研究发现，大多数Transformer模型中相邻token的查询向量具有显著相似性，当前查询的注意力计算可以仅依赖于少量前序查询的注意力信息。CORM动态保留推理所需的关键键值对，无需模型微调，实验表明在LongBench的六个任务中，CORM将KV缓存的推理内存使用减少了多达70%，且性能损失微乎其微。此外，CORM与GQA兼容，进一步提升压缩率。

🔬 方法详解

问题定义：本文旨在解决大语言模型推理过程中KV缓存内存使用过高的问题。现有方法在处理长序列时，KV缓存的内存占用随序列长度线性增长，导致推理效率低下。

核心思路：CORM的核心思想是通过动态保留关键的键值对，减少KV缓存的内存占用。研究表明，相邻token的查询向量相似性和前序查询的注意力信息可以有效支持当前查询的计算，从而优化内存使用。

技术框架：CORM的整体架构包括KV缓存的动态管理模块和查询向量相似性分析模块。前者负责根据当前推理需求调整缓存内容，后者则分析相邻token的查询向量以决定保留哪些键值对。

关键创新：CORM的主要创新在于其动态驱逐策略，能够在不进行模型微调的情况下，显著降低KV缓存的内存占用。这一方法与传统的静态缓存管理策略有本质区别，后者通常无法适应不同推理场景的需求。

关键设计：在CORM中，关键参数包括缓存大小和驱逐策略的阈值设置。通过实验验证，选择合适的阈值能够在保持性能的同时，最大化内存压缩效果。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，CORM在LongBench的六个任务中将KV缓存的推理内存使用减少了多达70%，且性能损失微乎其微。这一显著的内存优化效果展示了CORM在实际应用中的潜力，尤其是在资源受限的场景中。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、对话系统和机器翻译等需要高效推理的大语言模型。通过优化KV缓存，CORM能够在资源受限的环境中提升模型的推理效率，降低成本，具有重要的实际价值和广泛的应用前景。

📄 摘要（原文）

Large Language Models (LLMs), despite their remarkable performance across a wide range of tasks, necessitate substantial GPU memory and consume significant computational resources. Beyond the memory taken up by model weights, the memory used by the KV cache rises linearly with sequence length, becoming a primary bottleneck for inference. In this paper, we introduce an innovative method for optimizing the KV cache, which considerably minimizes its memory footprint. Upon thorough investigation, we discover that in most Transformer models, (i) there is a striking similarity between adjacent tokens' query vectors, and (ii) the attention calculation of the current query can rely exclusively on the attention information of a small fraction of preceding queries. Based on these observations, we present CORM, a KV cache eviction policy that dynamically retains essential key-value pairs for inference without the need for model fine-tuning. Our validation shows that CORM reduces the inference memory usage of KV cache by up to 70\% with negligible performance degradation across six tasks in LongBench. Furthermore, we demonstrate that CORM is compatible with GQA for further compression rate.