用动态注意力模式预测压缩KV缓存,显著提速且保持模型效果
AttentionPredictor: Temporal Patterns Matter for KV Cache Compression
- 通过统一卷积模型捕捉注意力的时空变化,动态预测关键缓存
- 实现13倍缓存压缩与5.6倍解码加速,性能接近现有最佳方法
- 轻量级设计适合部署,尤其适合长文本生成场景
随着大语言模型(LLMs)的发展,通过键值(KV)缓存压缩实现高效推理日益重要,尤其在长上下文生成中。现有方法依赖静态注意力分数建模来识别关键KV标记,但因忽略注意力分数的时间模式,常导致性能下降。为此,我们提出AttentionPredictor,首个基于学习的直接预测注意力模式以实现KV缓存压缩与关键标记识别的方法。具体地,AttentionPredictor采用轻量级统一卷积模型,动态捕捉时空模式并预测下一标记的注意力分数。其优势在于准确预测注意力分数,并在所有Transformer层间共享统一预测模型,内存开销极小。此外,我们提出跨标记关键缓存预取框架,隐藏标记估计时间开销以加速解码。通过保留大部分注意力信息,AttentionPredictor在缓存卸载场景下实现13×缓存压缩和5.6×速度提升,性能与当前最先进方法相当,显著超越已有方案。代码已开源。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。