Task-KV: Task-aware KV Cache Optimization via Semantic Differentiation of Attention Heads
作者: Xingyang He, Jie Liu, Shaowei Chen
分类: cs.CL
发布日期: 2025-01-25
💡 一句话要点
Task-KV：通过注意力头语义区分实现任务感知的KV缓存优化
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: KV缓存优化 注意力机制 大语言模型 推理加速 语义区分 任务感知 异构计算
📋 核心要点
- 现有KV缓存优化方法未能有效适应不同任务，因为注意力头的重要性随任务变化。
- Task-KV通过语义区分注意力头，为异构头分配完整KV缓存，为非异构头保留少量token和中间激活。
- 实验表明，Task-KV在多个基准测试和模型架构上显著优于现有方法，提升推理效率。
📝 摘要（中文）
KV缓存是大语言模型（LLMs）推理中广泛使用的加速技术。然而，其内存需求随着输入长度迅速增长。以往的研究通过移除所有注意力头相同数量的不重要token，或者为预先确定的注意力头分配不同的KV缓存预算来减少KV缓存的大小。然而，由于注意力头的重要性因不同任务而异，预先确定的注意力头无法有效地适应各种下游任务。为了解决这个问题，我们提出了Task-KV，一种利用注意力头的语义区分来为各种任务分配不同的KV缓存预算的方法。我们证明了远离语义中心的注意力头（称为异构头）对任务输出和语义理解做出了重大贡献。相比之下，其他注意力头则起到聚合重要信息和集中推理的作用。Task-KV为异构头分配完整的KV缓存预算，以保留全面的语义信息，同时为非异构头保留少量最近的token和注意力sink。此外，我们创新性地引入了中间激活来保留从非异构头聚合的关键上下文信息。为了动态感知注意力头之间的语义差异，我们设计了一个语义分离器，根据注意力头与语义中心的距离来区分异构头和非异构头。在多个基准和不同模型架构上的实验结果表明，Task-KV显著优于现有的基线方法。
🔬 方法详解
问题定义：现有KV缓存优化方法，如均匀缩减或预定义注意力头重要性，无法有效适应不同下游任务。这是因为不同任务对不同注意力头的依赖程度不同，预先设定的重要性排序可能失效，导致性能下降。
核心思路：Task-KV的核心思想是根据注意力头的语义角色动态分配KV缓存。它观察到远离语义中心的“异构头”对语义理解贡献更大，而靠近中心的头则负责信息聚合和推理。因此，为异构头分配更多资源，保证语义信息的完整性。
技术框架：Task-KV主要包含三个模块：1) 语义分离器：用于区分异构头和非异构头；2) KV缓存分配器：根据语义分离器的结果，为不同类型的注意力头分配不同的KV缓存预算；3) 中间激活保存器：为非异构头保存中间激活，以保留关键上下文信息。整体流程是，输入文本经过模型后，语义分离器判断每个注意力头的类型，KV缓存分配器据此分配缓存，最后模型输出结果。
关键创新：Task-KV的关键创新在于任务感知的注意力头语义区分。它不再像以往方法那样静态地处理所有注意力头，而是根据每个任务的特点，动态地识别出重要的异构头，并优先保证它们的KV缓存。此外，引入中间激活保存机制，弥补了非异构头缓存缩减带来的信息损失。
关键设计：语义分离器通过计算每个注意力头输出与语义中心（所有注意力头输出的平均值）的距离来判断其类型。距离阈值是一个关键参数，需要根据具体任务和模型进行调整。中间激活保存器保存非异构头在特定层的输出，这些层通常是模型中负责信息聚合的关键层。损失函数的设计目标是最小化Task-KV引入的性能损失，同时鼓励语义分离器准确区分异构头和非异构头。
🖼️ 关键图片
📊 实验亮点
实验结果表明，Task-KV在多个基准测试和不同模型架构上均优于现有方法。例如，在某些任务上，Task-KV能够在保持性能基本不变的情况下，将KV缓存大小减少高达50%。与基线方法相比，Task-KV通常能够取得1%-3%的性能提升，同时显著降低内存占用。
🎯 应用场景
Task-KV可应用于各种需要加速大语言模型推理的场景，例如移动设备上的本地部署、低延迟的在线服务、以及资源受限的边缘计算环境。通过降低KV缓存的内存需求，Task-KV能够显著提升推理速度和吞吐量，并降低部署成本，使得大语言模型能够更广泛地应用。
📄 摘要（原文）
KV cache is a widely used acceleration technique for large language models (LLMs) inference. However, its memory requirement grows rapidly with input length. Previous studies have reduced the size of KV cache by either removing the same number of unimportant tokens for all attention heads or by allocating differentiated KV cache budgets for pre-identified attention heads. However, due to the importance of attention heads varies across different tasks, the pre-identified attention heads fail to adapt effectively to various downstream tasks. To address this issue, we propose Task-KV, a method that leverages the semantic differentiation of attention heads to allocate differentiated KV cache budgets across various tasks. We demonstrate that attention heads far from the semantic center (called heterogeneous heads) make an significant contribution to task outputs and semantic understanding. In contrast, other attention heads play the role of aggregating important information and focusing reasoning. Task-KV allocates full KV cache budget to heterogeneous heads to preserve comprehensive semantic information, while reserving a small number of recent tokens and attention sinks for non-heterogeneous heads. Furthermore, we innovatively introduce middle activations to preserve key contextual information aggregated from non-heterogeneous heads. To dynamically perceive semantic differences among attention heads, we design a semantic separator to distinguish heterogeneous heads from non-heterogeneous ones based on their distances from the semantic center. Experimental results on multiple benchmarks and different model architectures demonstrate that Task-KV significantly outperforms existing baseline methods.