arXiv:2604.09852cs.AIcs.LG2026-04被引 13

让大模型自动压缩推理过程,提升效率且不丢精度。

MEMENTO: Teaching LLMs to Manage Their Own Context

  • 用记忆块压缩中间推理,只保留关键摘要。
  • 峰值KV缓存减少2.5倍,推理吞吐提升1.75倍。
  • 适合需要高效推理的数学、编程等长序列任务。

推理模型以无结构的长序列方式思考,缺乏对中间状态的压缩或组织机制。我们提出MEMENTO:一种让模型将推理分块,将每块压缩为紧凑的状态摘要(即记忆体),并仅通过关注这些记忆体进行后续推理的方法,从而降低上下文长度、键值缓存和计算量。为训练MEMENTO模型,我们发布了OpenMementos——一个包含22.8万条推理轨迹的公开数据集,源自OpenThoughts-v3,已分段并标注了中间摘要。我们发现,在OpenMementos上采用两阶段监督微调(SFT)策略,对不同模型家族(Qwen3、Phi-4、Olmo 3)及规模(8B–32B参数)均有效。训练后的模型在数学、科学与编码基准上保持高准确率,同时实现约2.5倍的峰值KV缓存减少。我们扩展vLLM以支持该推理方法,带来约1.75倍的吞吐提升,并支持强化学习进一步优化准确率。最后,我们发现信息通过双重通道传递:记忆体文本与对应的键值状态均携带原始推理块的信息;移除后者会使AIME24上的准确率下降15个百分点。

原文摘要 · Abstract (English)

Reasoning models think in long, unstructured streams with no mechanism for compressing or organizing their own intermediate state. We introduce MEMENTO: a method that teaches models to segment reasoning into blocks, compress each block into a memento, i.e., a dense state summary, and reason forward by attending only to mementos, reducing context, KV cache, and compute. To train MEMENTO models, we release OpenMementos, a public dataset of 228K reasoning traces derived from OpenThoughts-v3, segmented and annotated with intermediate summaries. We show that a two-stage SFT recipe on OpenMementos is effective across different model families (Qwen3, Phi-4, Olmo 3) and scales (8B--32B parameters). Trained models maintain strong accuracy on math, science, and coding benchmarks while achieving ${\sim}2.5\times$ peak KV cache reduction. We extend vLLM to support our inference method, achieving ${\sim}1.75\times$ throughput improvement while also enabling us to perform RL and further improve accuracy. Finally, we identify a dual information stream: information from each reasoning block is carried both by the memento text and by the corresponding KV states, which retain implicit information from the original block. Removing this channel drops accuracy by 15\,pp on AIME24.

大模型推理上下文压缩KV缓存优化智能摘要

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