arXiv 2502.040772025-10-26cs.CLcs.LG

AttentionPredictor：时间模式对KV缓存压缩至关重要

AttentionPredictor: Temporal Patterns Matter for KV Cache Compression

MoE Key Laboratory of Brain-inspired Intelligent Perception and Cognition, University of Science and Technology of China（脑启发智能感知与认知联合实验室，中国科学技术大学）；Huawei Noah’s Ark Lab（华为诺亚实验室）；College of Intelligence and Computing, Tianjin University（智能与计算学院，天津大学）

Qingyue Yang, Jie Wang, Xing Li, Zhihai Wang, Chen Chen, Lei Chen, Xianzhi Yu, Wulong Liu, Jianye Hao, Mingxuan Yuan, Bin Li

备注：: NeurIPS 2025

阅读 PDF

arXiv 页面 ↗HTML 原文 ↗

AI总结：

提出AttentionPredictor，首个基于学习的KV缓存压缩方法，通过统一卷积模型预测注意力模式，实现13倍压缩和5.6倍加速，显著提升LLM推理效率。

摘要：

随着大型语言模型（LLMs）的发展，通过键值（KV）缓存压缩实现高效推理已引起广泛关注，尤其是在长上下文生成场景中。为了压缩KV缓存，近期方法通过静态建模注意力分数来识别关键KV令牌。然而，这些方法往往难以准确确定关键令牌，因为它们忽略了注意力分数中的时间模式，导致LLM性能明显下降。为解决这一挑战，我们提出了AttentionPredictor，这是首个基于学习的方法，直接预测注意力模式以进行KV缓存压缩和关键令牌识别。具体而言，AttentionPredictor学习一个轻量级、统一的卷积模型，以动态捕获时空模式并预测下一令牌的注意力分数。AttentionPredictor的一个吸引人的特点是，它能准确预测注意力分数，并在所有Transformer层之间共享统一预测模型，该模型消耗可忽略不计的内存。此外，我们提出了一种跨令牌关键缓存预取框架，该框架隐藏了令牌估计时间开销以加速解码阶段。通过保留大部分注意力信息，AttentionPredictor在缓存卸载场景中实现了13倍的KV缓存压缩和5.6倍的加速，同时保持可比的LLM性能，显著优于现有最先进方法。代码可在https://github.com/MIRALab-USTC/LLM-AttentionPredictor获取。

英文摘要：

With the development of large language models (LLMs), efficient inference through Key-Value (KV) cache compression has attracted considerable attention, especially for long-context generation. To compress the KV cache, recent methods identify critical KV tokens through static modeling of attention scores. However, these methods often struggle to accurately determine critical tokens as they neglect the temporal patterns in attention scores, resulting in a noticeable degradation in LLM performance. To address this challenge, we propose AttentionPredictor, which is the first learning-based method to directly predict attention patterns for KV cache compression and critical token identification. Specifically, AttentionPredictor learns a lightweight, unified convolution model to dynamically capture spatiotemporal patterns and predict the next-token attention scores. An appealing feature of AttentionPredictor is that it accurately predicts the attention score and shares the unified prediction model, which consumes negligible memory, among all transformer layers. Moreover, we propose a cross-token critical cache prefetching framework that hides the token estimation time overhead to accelerate the decoding stage. By retaining most of the attention information, AttentionPredictor achieves 13$\times$ KV cache compression and 5.6$\times$ speedup in a cache offloading scenario with comparable LLM performance, significantly outperforming the state-of-the-arts. The code is available at https://github.com/MIRALab-USTC/LLM-AttentionPredictor.

机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。