ZoomR: Memory Efficient Reasoning through Multi-Granularity Key Value Retrieval

📄 arXiv: 2604.10898v1 📥 PDF

作者: David H. Yang, Yuxuan Zhu, Mohammad Mohammadi Amiri, Keerthiram Murugesan, Tejaswini Pedapati, Subhajit Chaudhury, Pin-Yu Chen

分类: cs.AI, cs.CL

发布日期: 2026-04-13

💡 一句话要点

ZoomR：通过多粒度键值检索实现内存高效的LLM推理

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 键值缓存 内存优化 多粒度检索 推理加速

📋 核心要点

  1. 现有LLM推理依赖KV缓存，长输出导致内存占用显著增加，成为性能瓶颈。
  2. ZoomR通过将推理过程压缩成摘要，并动态选择KV缓存，实现多粒度检索。
  3. 实验表明，ZoomR在保持竞争力的同时，显著降低了推理内存需求，最高降低4倍以上。

📝 摘要（中文）

大型语言模型（LLM）在复杂的推理任务中表现出色，但通常需要在得出最终答案之前生成冗长的中间过程。在生成过程中，LLM依赖于键值（KV）缓存进行自回归解码。然而，KV缓存的内存占用随着输出长度的增加而增长。先前关于KV缓存优化的工作主要集中在压缩长输入上下文，同时保留完整的KV缓存用于解码。对于需要长输出生成的任务，这导致计算和内存成本增加。本文介绍了一种名为ZoomR的新方法，该方法使LLM能够自适应地将冗长的推理过程压缩成摘要，并使用动态KV缓存选择策略，该策略利用这些摘要，同时策略性地“放大”细粒度细节。通过在解码期间使用摘要键作为粗粒度索引，ZoomR仅检索最重要过程的细节。这种分层策略通过避免每一步的完整缓存注意力，显著降低了内存使用。在数学和推理任务上的实验表明，与基线相比，我们的方法实现了具有竞争力的性能，同时将推理内存需求降低了4倍以上。这些结果表明，多粒度KV选择能够实现更内存高效的解码，特别是对于长输出生成。

🔬 方法详解

问题定义：论文旨在解决大型语言模型（LLM）在生成长输出时，由于键值（KV）缓存的内存占用过大而导致的推理效率问题。现有的KV缓存优化方法主要集中在压缩输入上下文，而忽略了长输出生成过程中KV缓存的优化，导致计算和内存成本居高不下。

核心思路：ZoomR的核心思路是利用多粒度键值检索，将冗长的推理过程压缩成摘要，并根据摘要动态选择需要关注的KV缓存部分。通过这种方式，避免了在每一步解码时都对整个KV缓存进行注意力计算，从而显著降低了内存占用。

技术框架：ZoomR的技术框架主要包含两个阶段：摘要生成阶段和动态KV缓存选择阶段。在摘要生成阶段，LLM将推理过程中的关键信息压缩成摘要键。在动态KV缓存选择阶段，模型使用查询（query）检索与当前推理步骤最相关的摘要键，并根据检索结果选择相应的细粒度KV缓存进行注意力计算。整体流程是分层级的，先粗粒度检索摘要，再细粒度关注细节。

关键创新：ZoomR的关键创新在于提出了多粒度KV选择策略，将推理过程分解为不同粒度的表示，并根据当前的需求动态选择合适的粒度进行计算。与现有方法相比，ZoomR不是简单地压缩整个KV缓存，而是有选择性地关注重要的信息，从而在保证性能的同时显著降低了内存占用。

关键设计：ZoomR的关键设计包括摘要键的生成方式、摘要键的检索算法以及细粒度KV缓存的选择策略。摘要键的生成可以采用多种方法，例如抽取式摘要或生成式摘要。摘要键的检索算法可以使用相似度匹配等方法。细粒度KV缓存的选择策略可以基于检索到的摘要键的置信度或相关性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，ZoomR在数学和推理任务上取得了与基线方法相当的性能，同时将推理内存需求降低了4倍以上。这一结果验证了多粒度KV选择策略的有效性，并表明ZoomR能够显著提高LLM在长输出生成任务中的内存效率。

🎯 应用场景

ZoomR技术可应用于需要长文本生成和复杂推理的场景，例如代码生成、长篇小说创作、科学研究报告生成等。通过降低LLM的内存需求，ZoomR使得在资源受限的设备上部署大型语言模型成为可能，加速了LLM在边缘计算和移动设备上的应用。

📄 摘要（原文）

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.