CAOTE: KV Cache Selection for LLMs via Attention Output Error-Based Token Eviction

📄 arXiv: 2504.14051v6 📥 PDF

作者: Raghavv Goel, Junyoung Park, Mukul Gagrani, Dalton Jones, Matthew Morse, Harper Langston, Mingu Lee, Chris Lott

分类: cs.LG, cs.CL

发布日期: 2025-04-18 (更新: 2025-10-05)

备注: 15 pages, 3 figures, 13 tables

💡 一句话要点

CAOTE：基于Attention输出误差的KV缓存选择，提升LLM在资源受限设备上的性能。

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

关键词: 大型语言模型 KV缓存 Token驱逐 注意力机制 资源受限设备

📋 核心要点

  1. 现有token驱逐方法依赖注意力分数，但忽略了token对最终注意力输出的实际贡献。
  2. CAOTE方法通过集成注意力分数和值向量，直接优化token驱逐造成的误差，提升驱逐策略的准确性。
  3. 实验表明，CAOTE与现有方法结合使用时，能稳定提升下游任务的准确性，验证了值信息的重要性。

📝 摘要（中文）

大型语言模型（LLM）对长上下文的支持扩展了其能力，但也带来了内存和计算方面的挑战，这在资源受限的设备中成为关键瓶颈。Token驱逐是一种广泛采用的后训练方法，旨在通过从缓存中驱逐不太重要的token来缓解瓶颈。通常，token驱逐使用注意力分数作为token重要性的代理指标。然而，注意力分数作为token级别重要性度量的一个主要限制是，它缺乏关于token对注意力输出贡献的信息。本文提出了一种基于缓存token对注意力输出贡献的简单驱逐准则。我们的方法CAOTE通过无缝集成注意力分数和值向量，优化了由于token驱逐而产生的驱逐误差。这是第一个在闭式中使用基于注意力驱逐分数之上的值token的方法。此外，CAOTE可以作为一种元启发式方法，灵活地与任何token驱逐方法一起使用。我们表明，CAOTE与最先进的基于注意力分数的方法相结合时，总是能提高下游任务的准确性，这表明在token驱逐过程中利用来自值的信息的重要性。

🔬 方法详解

问题定义：现有的大型语言模型在处理长文本时，KV缓存会占用大量内存，尤其是在资源受限的设备上。传统的token驱逐策略通常使用注意力分数作为token重要性的代理，但这种方法忽略了每个token对最终注意力输出的实际贡献，导致驱逐策略不够优化。

核心思路：CAOTE的核心思路是直接优化token驱逐所造成的误差。它认为，应该优先保留那些对最终注意力输出影响最大的token。通过结合注意力分数和值向量，CAOTE能够更准确地评估每个token的重要性，从而做出更明智的驱逐决策。

技术框架：CAOTE可以作为一个独立的token驱逐方法，也可以作为一个元启发式方法，与现有的token驱逐方法结合使用。其主要流程包括：1）计算每个token的注意力分数；2）利用值向量计算每个token对注意力输出的贡献；3）结合注意力分数和贡献度，计算每个token的驱逐优先级；4）根据优先级驱逐token。

关键创新：CAOTE的关键创新在于它首次将值向量信息纳入到token驱逐的决策过程中。与仅仅依赖注意力分数的方法相比，CAOTE能够更准确地评估token的重要性，从而减少驱逐误差。此外，CAOTE采用闭式解的形式，计算效率高，易于实现。

关键设计：CAOTE的关键设计在于如何将注意力分数和值向量信息有效地结合起来。具体来说，CAOTE通过计算每个token的值向量对最终注意力输出的影响，来衡量该token的重要性。这种方法可以有效地捕捉到那些注意力分数不高，但对最终结果至关重要的token。此外，CAOTE可以灵活地与不同的注意力分数计算方法和驱逐策略结合使用。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，CAOTE与最先进的基于注意力分数的token驱逐方法结合使用时，在多个下游任务上均能稳定提升准确性。例如，在长文本问答任务中，CAOTE能够将准确率提高1-3个百分点。这验证了CAOTE方法在token驱逐过程中利用值信息的重要性。

🎯 应用场景

CAOTE方法可以广泛应用于资源受限设备上的大型语言模型部署，例如移动设备、嵌入式系统等。通过更有效地管理KV缓存，CAOTE可以降低内存占用，提高推理速度，从而使这些设备能够运行更大规模的语言模型，并支持更长的上下文处理。这对于智能助手、机器翻译、文本摘要等应用具有重要意义。

📄 摘要（原文）

While long context support of large language models has extended their abilities, it also incurs challenges in memory and compute which becomes crucial bottlenecks in resource-restricted devices. Token eviction, a widely adopted post-training methodology designed to alleviate the bottlenecks by evicting less important tokens from the cache, typically uses attention scores as proxy metrics for token importance. However, one major limitation of attention score as a token-wise importance metrics is that it lacks the information about contribution of tokens to the attention output. In this paper, we propose a simple eviction criterion based on the contribution of cached tokens to attention outputs. Our method, CAOTE, optimizes for eviction error due to token eviction, by seamlessly integrating attention scores and value vectors. This is the first method which uses value tokens on top of attention-based eviction scores in closed-form. Additionally, CAOTE can act as a meta-heuristic method with flexible usage with any token eviction method. We show that CAOTE, when combined with the state-of-the-art attention score-based methods, always improves accuracies on the downstream task, indicating the importance of leveraging information from values during token eviction process.