Eigen Attention: Attention in Low-Rank Space for KV Cache Compression

📄 arXiv: 2408.05646v2 📥 PDF

作者: Utkarsh Saxena, Gobinda Saha, Sakshi Choudhary, Kaushik Roy

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

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

备注: 12 page, 6 figures, 6 tables

🔗 代码/项目: GITHUB

💡 一句话要点

提出Eigen Attention，通过低秩空间注意力机制压缩KV缓存，降低推理延迟。

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

关键词: 大语言模型 KV缓存压缩 低秩注意力 推理加速 长文本处理

📋 核心要点

  1. 长文本和高并发场景下，KV缓存成为LLM推理的内存瓶颈，限制了模型应用。
  2. Eigen Attention在低秩空间执行注意力计算，降低KV缓存的存储需求，减少计算量。
  3. 实验表明，Eigen Attention能有效压缩KV缓存并降低延迟，同时保持模型性能。

📝 摘要（中文）

大型语言模型（LLMs）凭借其卓越的推理能力，代表了自然语言处理领域的突破性进展。 近来，人们对增加这些模型的上下文长度以增强其在复杂任务中的适用性产生了浓厚的兴趣。 然而，在长上下文长度和大批量大小的情况下，存储注意力键和值的键值（KV）缓存成为推理过程中内存使用的新瓶颈。 为了解决这个问题，我们提出了Eigen Attention，它在低秩空间中执行注意力操作，从而减少了KV缓存的内存开销。 我们提出的方法与现有的KV缓存压缩技术是正交的，可以与它们协同使用。 通过对OPT、MPT和Llama模型系列的广泛实验，我们证明了Eigen Attention可以在KV缓存大小上减少高达40%，在注意力操作延迟上减少高达60%，同时性能下降最小。 代码可在https://github.com/UtkarshSaxena1/EigenAttn 获取。

🔬 方法详解

问题定义：大型语言模型在处理长文本或高并发请求时，需要存储大量的Key和Value向量用于Attention计算，这些向量构成了KV缓存。KV缓存的内存占用随着上下文长度和批量大小线性增长，成为推理过程中的主要瓶颈，限制了模型的部署和应用。现有的KV缓存压缩方法虽然可以缓解这个问题，但仍然存在性能损失或计算复杂度增加的问题。

核心思路：Eigen Attention的核心思想是利用低秩矩阵分解来近似Attention计算。具体来说，它将Key和Value向量投影到一个低维空间，在这个低维空间中执行Attention操作，从而减少了KV缓存的存储需求和计算量。这种方法基于一个假设：Attention矩阵通常是低秩的，因此可以用低秩矩阵来近似。

技术框架：Eigen Attention的整体框架可以分为以下几个步骤：1. 对Key和Value向量进行线性投影，将其映射到低维空间。2. 在低维空间中执行Attention计算。3. 将Attention结果投影回原始高维空间。这个过程可以看作是在原始高维空间中进行低秩近似的Attention计算。Eigen Attention可以无缝集成到现有的Transformer架构中，只需要修改Attention模块即可。

关键创新：Eigen Attention的关键创新在于它将Attention计算限制在一个低秩空间中，从而实现了KV缓存的压缩和计算加速。与现有的KV缓存压缩方法相比，Eigen Attention不需要对KV缓存进行量化或剪枝，因此可以避免性能损失。此外，Eigen Attention的计算复杂度较低，可以有效地降低推理延迟。

关键设计：Eigen Attention的关键设计包括：1. 低维空间的维度选择：需要根据具体的任务和模型来选择合适的维度，以在性能和效率之间取得平衡。2. 投影矩阵的初始化：可以使用随机初始化或预训练的嵌入向量来初始化投影矩阵。3. 正则化：为了避免过拟合，可以对投影矩阵进行正则化。4. 与现有KV缓存压缩技术的结合：Eigen Attention可以与现有的KV缓存压缩技术（如量化、剪枝）结合使用，以进一步提高压缩率和降低延迟。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，Eigen Attention在OPT、MPT和Llama等模型上实现了显著的KV缓存压缩和延迟降低。具体来说，KV缓存大小最多减少40%，注意力操作延迟最多减少60%，同时性能下降最小。这些结果表明Eigen Attention是一种有效的KV缓存压缩方法，可以显著提高LLM的推理效率。

🎯 应用场景

Eigen Attention可广泛应用于需要处理长文本或高并发请求的场景，例如：在线对话系统、机器翻译、文本摘要、代码生成等。通过降低KV缓存的内存占用和推理延迟，Eigen Attention可以提高LLM的部署效率和用户体验，使其能够更好地服务于各种实际应用。

📄 摘要（原文）

Large language models (LLMs) represent a groundbreaking advancement in the domain of natural language processing due to their impressive reasoning abilities. Recently, there has been considerable interest in increasing the context lengths for these models to enhance their applicability to complex tasks. However, at long context lengths and large batch sizes, the key-value (KV) cache, which stores the attention keys and values, emerges as the new bottleneck in memory usage during inference. To address this, we propose Eigen Attention, which performs the attention operation in a low-rank space, thereby reducing the KV cache memory overhead. Our proposed approach is orthogonal to existing KV cache compression techniques and can be used synergistically with them. Through extensive experiments over OPT, MPT, and Llama model families, we demonstrate that Eigen Attention results in up to 40% reduction in KV cache sizes and up to 60% reduction in attention operation latency with minimal drop in performance. Code is available at https://github.com/UtkarshSaxena1/EigenAttn.