基本信息

• 📝 原文链接: https://arxiv.org/abs/2412.13649
• 👥 作者: Jialong Wu, Zhenglin Wang, Linhai Zhang, Yilong Lai, Yulan He, Deyu Zhou
• 🏷️ 关键词: large language models, Key-Value Cache
• 📚 分类: 机器学习

摘要

中文摘要

键值（KV）缓存已成为LLMs在长文本生成中的瓶颈。尽管在这个领域已经做了很多努力，但解码阶段的优化通常被忽视。然而，我们认为这种优化至关重要，尤其是在以下两个观察的基础上进行长输出生成任务：(i) 预填充阶段过度的压缩，需要特定的完整上下文，这损害了对推理任务的 comprehension；(ii) 在具有长输出的推理任务中，频繁访问的数据项（heavy hitters）发生偏差。因此，我们引入了SCOPE，这是一个简单而高效的框架，在预填充和解码阶段分别执行KV缓存优化。具体来说，预填充阶段的KV缓存被保留以维护关键信息，同时提出了一种基于滑动窗口的新策略来选择解码阶段的关键频繁访问数据项。通过自适应和断续策略进一步优化内存使用和内存传输。在LongGenBench上的大量实验表明，SCOPE的有效性和泛化能力，以及它作为插件与其他仅预填充的KV压缩方法的兼容性。

原文摘要

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.

论文解读

一句话总结

提出了一种名为SCOPE的简单而高效的框架，通过在预填充和解码阶段分别进行KV缓存优化，以解决长文本生成中KV缓存压缩的问题。

问题1：这篇论文想要解决什么具体问题？

• 问题背景：在长文本生成中，LLMs的KV缓存成为瓶颈，现有方法在解码阶段的优化通常被忽视。
• 现有方案不足：预填充阶段过度压缩会影响推理任务的理解；解码阶段的重热点偏差会导致KV缓存分配不均。
• 研究目标：提出一种框架，在预填充和解码阶段分别进行KV缓存优化，以提高性能和内存效率。

问题2：论文的核心创新点是什么？

• 技术创新：SCOPE框架通过分别优化预填充和解码阶段的KV缓存来提高性能。
• 方法改进：提出滑动策略、自适应策略和断续策略来优化内存使用和内存传输。
• 优势：与现有方法相比，SCOPE在保持性能的同时，减少了内存消耗，并提高了内存传输效率。

问题3：实验结果如何验证了方法的有效性？

• 关键实验：在LONGGENBENCH基准上进行的实验。
• 性能提升：SCOPE在整体压缩率为35%时，性能与完整KV缓存相当。
• 对比结果：与SnapKV、PyramidKV、H2O和PyramidInfer等基线方法相比，SCOPE在大多数任务上都取得了更好的性能。

问题4：这个研究的实际应用价值是什么？

• 应用场景：适用于长文本生成任务，如长文本摘要和多问题回答。
• 实施建议：将SCOPE作为插件与其他预填充仅KV压缩方法结合使用。
• 局限与展望：未来研究可以探索更有效的重热点选择算法，并优化I/O操作以降低延迟。此外，SCOPE可以扩展到视觉模态中的长输出任务。