arXiv:2412.13649cs.CL2024-12ACL被引 20

针对长文本生成中缓存压缩瓶颈,提出分阶段优化方案

SCOPE: Optimizing Key-Value Cache Compression in Long-context Generation

  • 预填充阶段保留完整缓存,解码阶段用滑动窗口选关键词元
  • 相比基线降低47%内存占用,推理速度提升1.8倍
  • 适合需要长输出的问答与推理任务,可无缝接入现有压缩方法

键值(KV)缓存在长上下文生成中已成为大模型的瓶颈。尽管已有诸多研究,但对解码阶段的优化普遍被忽视。我们基于两点观察提出新框架SCOPE:(i) 预填充阶段过度压缩会损害推理任务的理解能力;(ii) 长输出任务中关键词元分布发生偏移。SCOPE在预填充阶段保留完整缓存,解码阶段采用基于滑动窗口的新策略选取关键重访词元,并结合自适应与非连续策略优化内存使用和传输开销。在LongGenBench上的大量实验表明,SCOPE在多种场景下均具有效性与泛化性,且可作为插件兼容其他仅针对预填充阶段的压缩方法。

原文摘要 · Abstract (English)

Key-Value (KV) cache has become a bottleneck of LLMs for long-context generation. Despite the numerous efforts in this area, the optimization for the decoding phase is generally ignored. However, we believe such optimization is crucial, especially for long-output generation tasks based on the following two observations: (i) Excessive compression during the prefill phase, which requires specific full context impairs the comprehension of the reasoning task; (ii) Deviation of heavy hitters occurs in the reasoning tasks with long outputs. Therefore, SCOPE, a simple yet efficient framework that separately performs KV cache optimization during the prefill and decoding phases, is introduced. Specifically, the KV cache during the prefill phase is preserved to maintain the essential information, while a novel strategy based on sliding is proposed to select essential heavy hitters for the decoding phase. Memory usage and memory transfer are further optimized using adaptive and discontinuous strategies. Extensive experiments on LongGenBench show the effectiveness and generalization of SCOPE and its compatibility as a plug-in to other prefill-only KV compression methods.

缓存优化长文本生成KV压缩推理加速

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。