Cross-Self KV Cache Pruning for Efficient Vision-Language Inference

📄 arXiv: 2412.04652v1 📥 PDF

作者: Xiaohuan Pei, Tao Huang, Chang Xu

分类: cs.CV, cs.AI

发布日期: 2024-12-05

🔗 代码/项目: GITHUB

💡 一句话要点

提出Cross-Self Pruning (CSP)方法，用于高效视觉-语言模型推理中的KV缓存剪枝。

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

关键词: 视觉-语言模型 KV缓存剪枝 注意力机制 模态差异 模型压缩

📋 核心要点

  1. 现有VLM的KV缓存剪枝方法忽略了视觉和语言模态间的分布差异，导致token重要性评估不准确和视觉token的过度剪枝。
  2. CSP方法将注意力分数分解为模态内和模态间注意力，并独立管理，从而实现更精确的KV缓存剪枝。
  3. CSP在MileBench基准测试中，在降低KV缓存预算的同时，显著提升了性能，尤其在会话式具身对话等任务上。

📝 摘要（中文）

KV缓存剪枝是降低长上下文自回归生成中内存和计算成本的一种有前景的技术。现有的视觉-语言模型(VLM)方法通常依赖于大型语言模型(LLM)的自注意力分数来识别和剪枝不相关的token。然而，这些方法忽略了模态之间固有的分布差异，常常导致不准确的token重要性估计和关键视觉token的过度剪枝。为了解决这个问题，我们提出将注意力分数分解为模态内注意力（同一模态内）和模态间注意力（跨模态），通过独立管理这些不同的注意力类型来实现更精确的KV缓存剪枝。此外，我们引入了一个n-softmax函数来抵消剪枝引起的分布偏移，保持注意力分数的原始平滑性，并确保稳定的性能。我们最终的免训练方法Cross-Self Pruning (CSP)在与具有完整KV缓存的模型相比，实现了具有竞争力的性能，同时显著优于之前的剪枝方法。在包含29个多模态数据集的MileBench基准上的大量评估表明了CSP的有效性，在诸如会话式具身对话等具有挑战性的任务上实现了高达41%的性能提升，同时将KV缓存预算减少了13.6%。代码可在https://github.com/TerryPei/CSP 获取。

🔬 方法详解

问题定义：论文旨在解决视觉-语言模型（VLM）中，由于模态差异导致的KV缓存剪枝不准确问题。现有方法直接使用LLM的自注意力分数进行剪枝，忽略了视觉和语言模态在数据分布上的差异，导致关键视觉token被过度剪枝，影响模型性能。

核心思路：论文的核心思路是将注意力分数分解为模态内注意力（intra-modality attention）和模态间注意力（inter-modality attention），分别评估和处理不同模态的token重要性。通过独立管理这两种注意力类型，可以更准确地识别和保留重要的视觉token。

技术框架：CSP方法主要包含两个关键模块：注意力分解和n-softmax函数。首先，将原始的注意力分数分解为模态内和模态间两部分。然后，使用独立的剪枝策略处理这两部分注意力。为了缓解剪枝带来的分布偏移，引入n-softmax函数来平滑注意力分数，保持模型性能的稳定性。整个过程无需额外的训练。

关键创新：CSP的关键创新在于其对注意力分数的分解处理，以及针对剪枝引入的n-softmax函数。与现有方法直接使用LLM的注意力分数不同，CSP考虑了模态差异，更准确地评估了token的重要性。n-softmax函数则有效缓解了剪枝带来的负面影响，保证了模型性能。

关键设计：注意力分解的具体实现方式是根据query和key的模态类型进行划分。如果query和key来自同一模态，则属于模态内注意力；否则，属于模态间注意力。n-softmax函数的具体形式未知，但其目的是在剪枝后保持注意力分数的平滑性，可能涉及到对softmax输出的归一化或调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

CSP方法在MileBench基准测试中表现出色，在KV缓存预算减少13.6%的情况下，在会话式具身对话等具有挑战性的任务上实现了高达41%的性能提升。实验结果表明，CSP方法能够有效解决模态差异带来的剪枝问题，显著优于现有的剪枝方法。

🎯 应用场景

该研究成果可应用于各种需要高效视觉-语言推理的场景，例如：机器人导航、智能客服、图像描述、视频理解等。通过降低模型的内存和计算成本，CSP方法能够使VLM在资源受限的设备上运行，并加速推理过程，具有重要的实际应用价值。

📄 摘要（原文）

KV cache pruning has emerged as a promising technique for reducing memory and computation costs in long-context auto-regressive generation. Existing methods for vision-language models (VLMs) typically rely on self-attention scores from large language models (LLMs) to identify and prune irrelevant tokens. However, these approaches overlook the inherent distributional discrepancies between modalities, often leading to inaccurate token importance estimation and the over-pruning of critical visual tokens. To address this, we propose decomposing attention scores into intra-modality attention (within the same modality) and inter-modality attention (across modalities), enabling more precise KV cache pruning by independently managing these distinct attention types. Additionally, we introduce an n-softmax function to counteract distribution shifts caused by pruning, preserving the original smoothness of attention scores and ensuring stable performance. Our final training-free method, \textbf{C}ross-\textbf{S}elf \textbf{P}runing (CSP), achieves competitive performance compared to models with full KV caches while significantly outperforming previous pruning methods. Extensive evaluations on MileBench, a benchmark encompassing 29 multimodal datasets, demonstrate CSP's effectiveness, achieving up to a 41\% performance improvement on challenging tasks like conversational embodied dialogue while reducing the KV cache budget by 13.6\%. The code is available at https://github.com/TerryPei/CSP