Keyformer: KV Cache Reduction through Key Tokens Selection for Efficient Generative Inference

📄 arXiv: 2403.09054v2 📥 PDF

作者: Muhammad Adnan, Akhil Arunkumar, Gaurav Jain, Prashant J. Nair, Ilya Soloveychik, Purushotham Kamath

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

发布日期: 2024-03-14 (更新: 2024-04-06)

期刊: Proceedings of the 7th Annual Conference on Machine Learning and Systems (MLSys), 2024

💡 一句话要点

提出Keyformer以解决KV缓存大小和内存带宽问题

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

关键词: 生成语言模型 KV缓存 内存带宽 注意力机制 长上下文处理 推理优化 文本生成

📋 核心要点

  1. 现有生成语言模型在推理过程中面临KV缓存大小和内存带宽的瓶颈，影响性能。
  2. Keyformer通过识别和保留关键token，减少KV缓存的大小，从而优化内存带宽利用率。
  3. 实验结果显示，Keyformer在多个模型上实现了推理延迟降低2.1倍和生成吞吐量提高2.4倍。

📝 摘要（中文）

随着Transformer架构在大型语言模型（LLMs）中的广泛应用，生成语言模型的推理过程面临着内存带宽瓶颈，尤其是在长上下文和大规模文本生成的场景中。本文提出了Keyformer，一种创新的推理时方法，通过识别和保留关键token，显著减少KV缓存的大小和内存带宽使用，同时保持模型的准确性。实验表明，Keyformer在多个基础模型上表现出色，推理延迟降低了2.1倍，token生成吞吐量提高了2.4倍。

🔬 方法详解

问题定义：本文旨在解决生成语言模型推理过程中KV缓存大小和内存带宽的限制，现有方法在处理长上下文时效率低下。

核心思路：Keyformer的核心思想是识别并保留在推理中占据大部分注意力权重的关键token，从而减少KV缓存的使用。

技术框架：Keyformer的整体架构包括一个新的评分函数，用于识别关键token，并在推理阶段仅保留这些token，从而优化内存带宽的使用。

关键创新：最重要的技术创新在于提出了一种新的评分函数，能够有效识别关键token，与传统方法相比，显著减少了KV缓存的大小。

关键设计：在设计中，Keyformer通过分析注意力权重分布，确定保留的token，并在不同的基础模型上进行了验证，确保了模型的准确性不受影响。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，Keyformer在推理延迟方面降低了2.1倍，生成吞吐量提高了2.4倍，且在保持模型准确性的前提下，显著优化了KV缓存的使用。这些结果表明Keyformer在实际应用中的有效性和潜力。

🎯 应用场景

Keyformer的研究成果在需要长上下文处理和高效文本生成的应用场景中具有广泛的潜在价值，如对话系统、文本摘要和内容生成等。通过优化内存带宽和缓存使用，Keyformer能够提升生成模型的响应速度和处理能力，推动相关技术的进一步发展。

📄 摘要（原文）

Transformers have emerged as the underpinning architecture for Large Language Models (LLMs). In generative language models, the inference process involves two primary phases: prompt processing and token generation. Token generation, which constitutes the majority of the computational workload, primarily entails vector-matrix multiplications and interactions with the Key-Value (KV) Cache. This phase is constrained by memory bandwidth due to the overhead of transferring weights and KV cache values from the memory system to the computing units. This memory bottleneck becomes particularly pronounced in applications that require long-context and extensive text generation, both of which are increasingly crucial for LLMs. This paper introduces "Keyformer", an innovative inference-time approach, to mitigate the challenges associated with KV cache size and memory bandwidth utilization. Keyformer leverages the observation that approximately 90% of the attention weight in generative inference focuses on a specific subset of tokens, referred to as "key" tokens. Keyformer retains only the key tokens in the KV cache by identifying these crucial tokens using a novel score function. This approach effectively reduces both the KV cache size and memory bandwidth usage without compromising model accuracy. We evaluate Keyformer's performance across three foundational models: GPT-J, Cerebras-GPT, and MPT, which employ various positional embedding algorithms. Our assessment encompasses a variety of tasks, with a particular emphasis on summarization and conversation tasks involving extended contexts. Keyformer's reduction of KV cache reduces inference latency by 2.1x and improves token generation throughput by 2.4x, while preserving the model's accuracy.