Task-KV：通过注意力头的语义差异进行任务感知的KV Cache优化
Task-KV: Task-aware KV Cache Optimization via Semantic Differentiation of Attention Heads
AI总结：
提出Task-KV方法，利用注意力头的语义差异为不同任务动态分配KV cache预算，为异质头保留完整预算并为非异质头保留关键信息，显著优于现有基线。
摘要：
KV cache是大语言模型（LLMs）推理中广泛使用的加速技术。然而，其内存需求随输入长度快速增长。以往的研究通过要么为所有注意力头移除相同数量的不重要token，要么为预先识别的注意力头分配差异化的KV cache预算，来减小KV cache的大小。然而，由于注意力头的重要性在不同任务间存在差异，预先识别的注意力头无法有效适应各种下游任务。为了解决这个问题，我们提出了Task-KV，一种利用注意力头的语义差异来为不同任务分配差异化KV cache预算的方法。我们证明，远离语义中心的注意力头（称为异质头）对任务输出和语义理解做出了显著贡献。相反，其他注意力头扮演着聚合重要信息和聚焦推理的角色。Task-KV为异质头分配完整的KV cache预算以保留全面的语义信息，同时为非异质头保留少量最近的token和注意力汇。此外，我们创新性地引入中间激活，以保留从非异质头聚合的关键上下文信息。为了动态感知注意力头之间的语义差异，我们设计了一个语义分离器，根据它们与语义中心的距离来区分异质头和非异质头。在多个基准和不同模型架构上的实验结果表明，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.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
