用未来预测过去,动态清理冗余缓存,提升大模型推理速度
Back from the Future: Key-Value Cache Management by Counter-Causal Surprise
- 基于反因果注意力判断哪些历史token可被预测,优先淘汰
- 在多个模型上实现更快推理,且生成质量不下降
- 无需训练,仅需一次轻量级后处理,适合部署优化
近年来,通过压缩和淘汰策略管理键值(KV)缓存成为重要研究方向。大语言模型(LLMs)及其多模态变体在生成输出时的计算开销可通过缓存先前计算的键值来缓解,但随之带来缓存大小随上下文长度线性增长的问题,导致长文本生成时迅速耗尽GPU内存。通过周期性淘汰缓存条目可降低内存占用并提升推理速度。本文提出一种简单有效的缓存淘汰机制:若过去某词元能由更近期词元准确预测,则其对应的键值可被移除。通过在原顺序上运行模型,复用缓存中的键值表示,并施加反因果注意力掩码,使每个位置仅关注后续上下文。该方法为分布内、直接依赖实际缓存内容,无需额外训练。为进一步降低开销,还提出仅在最后一层Transformer进行反因果传播的快速近似,显著加快刷新周期,精度损失极小。在多个开源大模型和基准数据集上的评估显示,本方法性能优于或媲美现有最先进方法。参考代码已公开于 https://github.com/metacognitionai/counter_causal。
原文摘要 · Abstract (English)
Key-value (KV) cache management through compression and eviction strategies has emerged as an important research direction in recent years. Computational demands of large language models (LLMs) and their multi-modal variants during output generation can be partially alleviated by caching previous key and value calculations needed by subsequent scaled dot-product attention operations. However, this leads to another problem: the size of the resulting KV cache grows linearly with context length and quickly consumes all available GPU memory when either the prompt or the generated output are long. KV cache management periodically prunes entries from the cache thereby reducing its memory footprint while attempting to retain sufficient information for accurate generation. A by-product is faster inference speed. We propose a simple yet effective KV eviction scheme motivated by the insight that past tokens which can be well-predicted from more recent tokens are redundant and their associated keys and values can be removed from the cache. To score entries for eviction we run the model on the tokens in their original order, reusing the key and value representations already stored in the KV cache, and applying a counter-causal attention mask so that each position attends only to its future context. This is in-distribution, tied directly to the actual cache contents, and requires no additional training. To further reduce cost, we additionally propose a fast single-layer approximation that restricts the counter-causal pass to the last transformer layer, achieving a significant speedup per refresh cycle at marginal accuracy cost. We evaluate our strategy on various open-source LLMs and benchmark datasets showing competitive or improved performance over other state-of-the-art methods. Reference code is available at https://github.com/metacognitionai/counter_causal.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。