arXiv:2501.15113cs.CL2025-01被引 4

根据任务特性动态分配注意力头的缓存,提升长文本推理效率。

Task-KV: Task-aware KV Cache Optimization via Semantic Differentiation of Attention Heads

  • 按语义差异区分注意力头,为不同任务分配不同缓存预算。
  • 对关键语义头保留完整缓存,其他头仅存少量最近令牌。
  • 适合需要高效长文本生成的场景,如对话系统、摘要任务。

KV 缓存是大语言模型推理中广泛使用的加速技术,但其内存需求随输入长度迅速增长。以往方法或对所有注意力头统一移除不重要令牌,或为预识别的注意力头分配差异化缓存预算。然而,由于注意力头的重要性在不同任务间变化,预识别方案难以适配多种下游任务。为此,我们提出 Task-KV,利用注意力头的语义差异,为不同任务动态分配差异化 KV 缓存预算。实验表明,远离语义中心(称为异质头)的注意力头对任务输出和语义理解贡献显著,而其他头则用于聚合关键信息并聚焦推理。Task-KV 为异质头分配完整缓存以保留完整语义信息,而非异质头则仅保留少量近期令牌与注意力汇点。此外,我们创新引入中间激活来保存非异质头聚合的关键上下文信息。为动态感知注意力头间的语义差异,设计语义分离器,基于其与语义中心的距离判断是否为异质头。在多个基准测试及不同模型架构上的实验结果表明,Task-KV 显著优于现有基线方法。

原文摘要 · Abstract (English)

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.

KV缓存注意力机制推理优化LLM

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。