MEMENTO：大模型学会自主管理上下文，KV缓存直降2.5倍！

MEMENTO: Teaching LLMs to Manage Their Own Context

当大语言模型在解决复杂数学或编程问题时，往往会生成数以千计的推理Token。这种长长的思维链虽然带来了惊艳的准确率，但也引发了严重的计算灾难。在传统的自回归生成中，每一个历史Token都会被保存在KV缓存中，消耗极其庞大的显存。

ArXiv URL：http://arxiv.org/abs/2604.09852v1

更致命的是，模型面对如此巨大的上下文，没有任何机制能够筛选或遗忘无效信息。它们只能拖着沉重的历史包袱，一步步艰难前行。为了解决这一痛点，微软研究院提出了一项名为 MEMENTO 的创新技术。该研究教会模型将推理过程切分为独立的块，并为每个块生成高度压缩的摘要。

通过这种机制，MEMENTO 在保持强大推理能力的同时，将峰值KV缓存降低了约$2.5$倍。同时，该研究还通过定制的推理引擎，实现了高达$1.75$倍的吞吐量提升。

分块与压缩：大模型的“会议纪要”机制

要理解 MEMENTO 的运作机制，我们可以将大模型的长推理过程比作一场漫长复杂的闭门研讨会。传统的模型就像是一个不开窍的记录员，必须把全天的每一句会议发言都录音并随时回听。随着会议的推进，这堆录音带（即 KV Cache）堆积如山，检索起来极其耗时。

MEMENTO 则教会了这位记录员一种高效的工作方式：边听边做“会议纪要”。具体而言，模型被训练去识别推理链条中的语义连贯片段，称之为思考块（Thinking Block）。当一个思考块结束时，模型会生成一个高度浓缩的备忘录（Memento）。这个备忘录用最少的Token，记录了该阶段的核心结论、中间变量和关键的推导方向。

一旦备忘录生成完毕，系统就会在后台直接把前面那个冗长的思考块“物理抹除”。在后续的推理步骤中，模型只需要去关注过去的备忘录和当前正在生成的思考块即可。如图所示，这使得模型的注意力范围大幅缩减，显存占用呈现出锯齿状的健康模式。

双重信息流：隐藏在KV缓存中的“记忆残留”

表面上看，MEMENTO 只是在做文本摘要。但这引出了一个极其关键且深奥的技术难题。如果仅仅依赖简短的备忘录文本，模型是否会丢失那些对后续推理至关重要的微妙细节？

这正是本文最核心的发现：双重信息流（Dual Information Stream）。让我们继续使用“会议纪要”的比喻。当记录员写下纪要时，由于他刚刚亲历了完整的讨论，他的脑海中其实还残留着许多未写出的隐性语感。如果换一个完全没参加过会议的新人，仅仅只看这几行字，肯定是无法完美接续工作的。

在 MEMENTO 中，由于思考块的抹除是在原地发生的，奇妙的事情发生了。生成备忘录时，完整的思考块仍在上下文中，因此备忘录Token的KV状态吸收了完整块的深层信息。当思考块被抹除后，这些隐式的信息依然被安全地封存在了备忘录的KV表示中。

为了验证这一点，研究人员进行了一项精妙的“重启消融”实验。他们强行丢弃了原始生成的KV缓存，仅仅根据留下的备忘录文本，重新从头计算一次KV状态。这就好比换了一个只看字面意思的新人来接手工作。结果令人震惊：在 $AIME’24$ 数据集上，切断这个隐式KV通道后，模型的准确率直接暴跌了15个百分点。这无可辩驳地证明了，备忘录的KV状态中携带着大量超越文本本身的隐式推理信息。

两阶段训练与数据构建

为了让模型掌握这项复杂的技能，高质量的训练数据必不可少。该研究构建并开源了包含$22.8$万个样本的 OPENMEMENTOS 数据集。构建这一数据集充满了挑战，因为原始的推理轨迹根本没有明确的天然边界。

研究团队设计了一条精巧的流水线。首先，他们使用一个强大的LLM对推理句子进行打分，寻找潜在的切分点。接着，利用算法优化边界选择，最后再由LLM将每个思考块总结为高质量的备忘录。

在 OPENMEMENTOS 数据集的构建中，其领域涵盖了数学、代码与科学三大硬核场景。其中数学与代码类轨迹往往更加绵长，平均每个样本包含多达$9$个思考块。令人惊讶的是，尽管各个思考块的长短差异巨大，但模型生成的备忘录长度却极其稳定。这充分表明，模型真正学会了“提取核心要义”这项高阶技能，而不是死板地按比例压缩文本。

在训练阶段，该研究采用了一种符合课程学习直觉的两阶段监督微调（Supervised Fine-Tuning, SFT）。第一阶段，模型在拥有完整上下文的普通条件下，先学习“思考块+备忘录”的输出格式。第二阶段，系统引入难度，强迫模型在看不到已被抹除内容的条件下，仅凭备忘录进行准确推理。

系统工程落地与性能验证

优秀的算法离不开底层的系统支持。现有的生产级推理框架，由于缺乏动态遮罩支持，无法处理 MEMENTO 这种依赖生成内容的自定义稀疏注意力。为此，研究团队深度改造了 vLLM 引擎的底层逻辑。他们通过Python层面的补丁，使引擎能够在单次生成调用中，物理移除指定的KV缓存条目。

实验结果证明了 MEMENTO 的巨大潜力。在 $Qwen3$、$Phi-4$ 等主流架构上，该技术展现了极强的泛化能力。以 $Qwen3-32B$ 为例，在 $AIME’26$ 上仅损失了$2.6$个百分点的精度，却换来了数倍的资源节省。更令人振奋的是，随着模型规模的增加，这种精度损失在不断缩小，表明大模型更擅长利用压缩后的上下文。

为了进一步提升模型在极端压缩下的表现，该研究还引入了基于反馈的强化学习。他们设计了专门的奖励机制，公式如下所示，对正确且符合格式的推理轨迹给予正向激励。

\[L = -\operatorname{sg}(\operatorname{clip}(r_t(\theta), 1 - \epsilon_{\text{low}}, 1 + \epsilon_{\text{high}})) \cdot A_t \cdot \log \pi_{\theta}(a_t \mid s_t)\]

在 $Qwen3-8B$ 上的实验表明，强化学习微调能够有效弥补因压缩导致的精度下降，使其表现看齐甚至超越基线模型。

技术局限与未来启示

尽管 MEMENTO 开辟了一条极具前景的上下文管理新路径，但它目前仍有一定的局限性。首先，生成备忘录本身也会消耗额外的计算资源。对于非常简单的短推理任务，这种开销可能得不偿失。其次，强依赖底层框架定制，提高了普通开发者直接复用的技术门槛。

但毫无疑问，这项研究为未来的智能体演进提供了重要启示。未来的AI不应仅仅是一个能够被动吞噬百万Token的庞然大物，更应当具备主动整理、压缩和遗忘自身思维的能力。对于拥有超长运行周期的Agent任务而言，MEMENTO 所展现的“边想边忘”机制，或许正是通向真正通用人工智能的一把关键钥匙。