通过分层摘要压缩推理过程,显著降低长文本生成的内存占用。
ZoomR: Memory Efficient Reasoning through Multi-Granularity Key Value Retrieval

- 用摘要作为粗粒度索引,动态选择关键细节进行检索。
- 在数学与推理任务中,内存消耗减少4倍以上,性能不降。
- 适合需要长输出生成的复杂推理场景,如代码生成、数学证明。
大型语言模型在复杂推理任务中表现优异,但常需生成大量中间思考过程才能得出最终答案。推理过程中,模型依赖键值缓存(KV cache)实现自回归解码,而其内存开销随输出长度线性增长。现有优化方法多聚焦于压缩输入上下文,却保留完整的KV缓存用于解码。对于长输出任务,这导致计算与内存成本急剧上升。本文提出ZoomR,一种新方法:将冗长的推理过程自适应压缩为摘要,并采用动态的KV缓存选择策略,在解码时利用这些摘要,同时智能“聚焦”到关键细节。通过摘要键作为粗粒度索引,仅对最重要的思想进行细粒度检索。该分层策略大幅减少每步解码所需的完整缓存注意力,显著降低内存使用。在数学和推理任务上的实验表明,该方法性能接近基线,但推理内存降低超过4倍。结果表明,多粒度的KV选择可有效提升长输出生成的内存效率。
原文摘要 · Abstract (English)
Large language models (LLMs) have shown great performance on complex reasoning tasks but often require generating long intermediate thoughts before reaching a final answer. During generation, LLMs rely on a key-value (KV) cache for autoregressive decoding. However, the memory footprint of the KV cache grows with output length. Prior work on KV cache optimization mostly focus on compressing the long input context, while retaining the full KV cache for decoding. For tasks requiring long output generation, this leads to increased computational and memory costs. In this paper, we introduce ZoomR, a novel approach that enables LLMs to adaptively compress verbose reasoning thoughts into summaries and uses a dynamic KV cache selection policy that leverages these summaries while also strategically "zooming in" on fine-grained details. By using summary keys as a coarse-grained index during decoding, ZoomR uses the query to retrieve details for only the most important thoughts. This hierarchical strategy significantly reduces memory usage by avoiding full-cache attention at each step. Experiments across math and reasoning tasks show that our approach achieves competitive performance compared to baselines, while reducing inference memory requirements by more than $4\times$. These results demonstrate that a multi-granularity KV selection enables more memory efficient decoding, especially for long output generation.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。