预期注意力:通过估计未来查询分布的注意力实现KV缓存压缩
Expected Attention: KV Cache Compression by Estimating Attention from Future Queries Distribution
AI总结:
针对大语言模型KV缓存内存瓶颈及现有剪枝方法的实际局限,提出无需训练的预期注意力方法,通过预测未来查询注意力估计KV对重要性,在两阶段均优于基线,还发布了含20余种技术的KVPress库。
摘要:
键值(KV)缓存的内存消耗是大语言模型高效推理的主要瓶颈。尽管基于注意力分数的KV缓存剪枝展现出潜力,但它存在关键的实际局限:压缩时无法获取未来token的注意力分数,且Flash Attention等现代实现不会实例化完整注意力矩阵,导致无法获取历史分数。为解决这些挑战,我们提出**预期注意力(Expected Attention)**,这是一种无需训练的压缩方法,通过预测未来查询对KV对的注意力来估计其重要性。该方法利用大语言模型激活的分布特性,以解析形式计算每个KV对的预期注意力分数。这些分数可用于对KV对进行有原则的排序和剪枝,对残差流的影响极小,在不降低性能的前提下实现有效压缩。重要的是,该方法可在预填充和解码阶段无缝运行,在两种场景下均始终优于现有最优基线。最后,我们发布了**KVPress**——一个综合性库,供研究人员实现和基准测试KV缓存压缩方法,目前已包含20余种技术。
英文摘要:
Memory consumption of the Key-Value (KV) cache represents a major bottleneck for efficient large language model inference. While attention-score-based KV cache pruning shows promise, it faces critical practical limitations: attention scores from future tokens are unavailable during compression, and modern implementations like Flash Attention do not materialize the full attention matrix, making past scores inaccessible. To overcome these challenges, we introduce $\textbf{Expected Attention, a training-free compression method}$ that estimates KV pairs importance by predicting how future queries will attend to them. Our approach leverages the distributional properties of LLM activations to compute expected attention scores in closed form for each KV pair. These scores enable principled ranking and pruning of KV pairs with minimal impact on the residual stream, achieving effective compression without performance degradation. Importantly, our method operates seamlessly across both prefilling and decoding phases, consistently outperforming state-of-the-art baselines in both scenarios. Finally, $\textbf{we release KVPress, a comprehensive library to enable researchers to implement and benchmark KV cache compression methods, already including more than 20 techniques}$.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。