💬 NLP

PReM: Learning What to Preserve and When to Refresh for Context Compression

PReM 是一种新颖的上下文压缩框架，它通过一个专门的记忆层和特殊标记，动态地学习如何保留并刷新内部层级 KV 记忆，从而在不依赖外部压缩器或做出过早保留决策的情况下，实现高效的长上下文推理。

原作者： Bohan Yu, Lei Shen, Chenxi Zhou, Chen Han, Junlin Liu, Wenbo Su, Yu Cheng, Bo Zheng

发布于 2026-07-17

📖 1 分钟阅读☕ 轻松阅读

原作者： Bohan Yu, Lei Shen, Chenxi Zhou, Chen Han, Junlin Liu, Wenbo Su, Yu Cheng, Bo Zheng

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在试图解决一个庞大且复杂的谜题，但交给你的不是单一的线索，而是一个包含 32,000 页文档的图书馆。你的大脑（或者在这种情况下，一个被称为“大语言模型”的计算机程序）必须阅读所有的内容，才能找到回答问题所需的特定事实。问题在于，试图将整个图书馆同时装进你的“工作记忆”中是非常沉重且缓慢的。这就像是在比赛中背着一个装满图书馆的背包跑步；最终你会感到疲惫，并且可能会丢掉你正需要的那本书。

为了解决这个问题，科学家们到目前为止尝试了两种主要的技巧。第一种是“压缩”，即在你在开始阅读之前，尝试将图书馆缩减为一个微小的摘要。第二种是“缓存”，即保留整个图书馆，但只尝试查看最热门的页面。但两者都有一个缺陷：它们在开始解谜之前就做出了决定。如果谜题在途中发生了变化，而你突然需要书后的一页被扔掉的页面，你就陷入困境了。你无法回头去取它，因为决定做得太早了。这篇论文介绍了一种处理此问题的新方法，它建议我们需要的不是一个静态的摘要，而是一个能够随着故事展开而自我更新的记忆系统。

由此诞生了 PReM（保留并刷新记忆，Preserve and Refresh Memory），这是一个聪明的全新框架，它将模型的记忆视为一个动态的、活生生的笔记本，而不是一个静态的文件柜。其核心思想简单而强大：不要只决定保留什么一次；要决定此时此刻该保留什么，并准备好在故事需要的那一刻进行更换。

以下是 PReM 的工作原理，我们用一个有趣的类比来解释。想象你是一名正在带领团队（AI 的不同层）破案的侦探。在旧的方法中，团队会阅读整个案卷，挑选出他们最喜欢的 10 页，然后将剩下的锁进保险库。如果你稍后需要保险库里的一页，那就太晚了。

PReM 改变了规则。它保留了整个案卷，但将其组织成小的“块”（就像章节一样）。当侦探开始撰写解决方案时，一个特殊的“记忆标记”（我们称之为魔法书签，记作 ⟨m⟩）充当了一个信号。当侦探写下这个书签时，这就像是在说：“等等！我需要重新思考一下我桌子上的东西。”

就在那一刻，一个专门的“记忆管理器”（AI 中的一个特殊层）苏醒了。它观察当前的问题和魔法书签，然后快速扫描所有章节。它决定：“好吧，对于故事的下一步，我们肯定需要第三章和第十二章。我们可以暂时忘记其他的。”然后，它会将桌上的旧章节与新章节进行交换，在保持重要细节清晰锐利的同时，将不需要的部分压缩成一个微小的摘要。这一切都发生在写作过程中，而不是在开始之前。

论文表明，这种“刷新”机制是一个游戏规则的改变者。在针对 32,000 个 token 上下文（那是非常多的文本！）的测试中，PReM 成功地将记忆压缩了 16 倍甚至 32 倍。尽管在任何时刻“桌面上”留存的信息都少得多，但它的回答能力甚至比那些试图保持整个图书馆敞开的模型还要好。例如，在一个 70 亿参数的模型上，与现有的最佳方法相比，PReM 将答案的准确度提高了多达 12.55 个百分点。

研究人员还发现了一些关于这个记忆管理器应该“居住”在哪里有趣的事情。他们测试了将“记忆管理器”放置在 AI 大脑的早期、中期或后期部分的情况。他们发现，中期和后期层在决定保留什么方面表现得更好，而早期层则太慢，无法学会这项技能。这就像是让一名资深侦探（中间层）来做出关于哪些证据相关的决策，而不是让一名新手（早期层）去做。

P-ReM 最令人惊讶的发现之一是，由于它拥有如此出色的“学习型”记忆管理能力，一个较小的模型（30 亿参数）使用 PReM 后，可以超越使用其他压缩技巧的较大模型（70 亿参数）。这表明，知道何时刷新你的记忆与你的大脑有多大同样重要。

论文还排除了一些想法。它表明，仅仅使用 AI 的自然注意力（关注它当前正在思考的内容）不足以决定保留什么；你需要一个专门训练过的系统来进行选择。它还表明，你不能只是在开始时压缩一次文本然后就期望一切顺利；这种“静态”的方法（即在整个回答过程中保持相同的压缩摘要）会导致更差的结果，尤其是在处理复杂的、多步骤的问题时。

简而言之，PReM 表明，高效 AI 的未来不仅仅在于缩小数据，而在于构建一个会呼吸的记忆。它学习如何保留当前对话步骤所需的证据，并准确知道何时按下“刷新”按钮，为下一步换入新的证据。通过这样做，它在保持 AI 快速高效的同时，不会忘记解开谜团所需的线索。

技术摘要：PReM (保留与刷新记忆)

1. 问题陈述

大语言模型（LLM）的高效长上下文推理面临着一个根本性的矛盾：即在降低计算成本与维持对必要上下文证据的访问能力之间进行权衡。随着上下文长度的增加，键值（KV）缓存不断增长，显著增加了解码延迟和内存占用。

现有的缓解方法，如 KV 缓存压缩和上下文压缩，存在两个主要局限性：

静态证据保留： 大多数方法在早期阶段就做出了关于保留哪些上下文的单一决策。这种静态压缩无法适应多步推理中不断变化的需求，因为不同的推理步骤可能需要不同的证据片段。
解耦的证据选择： 依赖外部压缩器的模型将压缩过程与生成过程分离。这引入了额外的计算开销，并阻碍了模型根据特定的推理步骤动态调整证据选择的能力。

本文认为，有效的压缩不应仅仅是简单的规模缩减；它必须能够动态地保留有用的证据，并在后续推理步骤需要不同信息时对其进行刷新。

2. 方法论：PReM

PReM（Preserve and Refresh Memory，保留与刷新记忆）是一个上下文压缩框架，它将长上下文维护为 LLM 内部的层级化 KV 记忆。PReM 不依赖于外部压缩器，而是学习如何根据模型的内部状态来决定保留什么以及何时刷新。

核心机制

层级化 KV 记忆表示： 长上下文被编码为存储在模型内部的层级化 KV 状态（ $M_\ell$ ）。这些状态被划分为连续的块（chunks）。
专用记忆层： 指定一个特定层（ $\ell_m$ $ℓ_{m}$ ）用于做出记忆选择决策。该层使用专门的选择查询投影矩阵（ $W^c_\ell$ $W_{ℓ}^{c}$ ）对上下文块进行评分。
- 块选择： 模型基于当前查询（来自问题或刷新标记）对块进行评分。
- 保留与池化（Preserve-and-Pool）： 被选中的块以 Token 粒度进行保留。未被选中的块则被压缩为一个平均化的 KV 代表值。
记忆标记 $\langle m \rangle$ ： 一个特殊的标记作为显式的内部信号，用于刷新记忆。当模型生成 $\langle m \rangle$ 时，记忆层会重新评估上下文块，选择新的相关证据，并更新用于后续层的压缩 KV 记忆。这使得模型能够根据当前的推理步骤来控制何时进行刷新。

训练：阶段分离刷新训练

为了训练模型使记忆选择与生成保持一致，PReM 引入了阶段分离刷新训练策略。每个推理步骤被分为两个截然不同的前向传播过程：

记忆选择阶段： 模型处理输入（问题或 $\langle m \rangle$ ）以计算块评分并选择前 $k$ 个证据块。该阶段通过记忆选择损失（ $L_{ms}$ ）进行监督，旨在鼓励模型将证据块的排名排在非证据块之上。
记忆条件生成阶段： 模型在由选择阶段产生的新刷新记忆状态的条件下，生成答案片段（ $y_s$ ）。

为了确保训练与推理之间的一致性，训练目标包含了边界交叉熵项（ $L_{bd}$ ）。这些项将选择阶段与第一个答案 Token 以及答案片段末尾与下一个 $\langle m \rangle$ 标记连接起来，确保模型学习到记忆更新与生成之间的自回归联系。

3. 核心贡献

形式化定义： 本文将上下文压缩重新定义为一个学习到的记忆保留与刷新问题，消除了对外部压缩器的需求。
架构设计： PReM 引入了一个专门的记忆层和一个特殊的 $\langle m \rangle$ 标记，从而在 LLM 的内部 KV 状态中实现动态、针对特定步骤的记忆更新。
训练策略： 提出了阶段分离刷新训练，通过在训练期间将记忆选择与生成解耦，确保模型学习在特定步骤所需的特定记忆状态下进行条件生成。
实证验证： 通过全面的实验证明，学习到的动态记忆更新优于静态压缩和全上下文提示。

4. 实验结果

实验在 Qwen2.5-Instruct 的 3B 和 7B 版本上进行，使用 32K Token 上下文，涵盖六个基准测试（TriviaQA, SQuAD, NaturalQuestions, 2WikiMQA, HotpotQA, MuSiQue）。

性能对比基线： 在 16× 和 32× 压缩率下，PReM 均优于八种强力基线模型，包括 KV 缓存方法（SnapKV, StreamingLLM, CAKE）和上下文压缩方法（LongLLMLingua, EXIT, ICAE）。
- 在 3B 主干网络、16× 压缩率下，PReM 的平均精确匹配度（EM）和 F1 值分别比最强基线提高了 5.30 和 5.15 个百分点。
- 在 7B 主干网络、32× 压缩率下，提升幅度达到 8.20 (EM) 和 9.89 (F1) 个百分点。
与全上下文对比： PReM 在 3B 和 7B 模型上均优于直接使用 32K Token 提示（无压缩）的表现，这表明学习到的记忆选择比向模型暴露未经处理的全上下文更为有效。
小模型优势： 3B PReM 模型超越了已发布的 7B 软上下文压缩基线（如 Activation Beacon），表明学习到的保留与刷新可以补偿较小的模型规模。
效率： PReM 显著降低了峰值 GPU 显存（例如，全量记忆为 30.0GB，而 PReM 约为 19.57GB），同时保持了具有竞争力的推理延迟。使用 $\langle m \rangle$ 标记允许在解码计算量（刷新频率）与显存节省之间进行可控的权衡。

5. 重要性与主张

论文声称，PReM 证明了有效的长上下文推理需要与生成耦合的记忆更新，而非静态的上下文缩减。通过将上下文维护为内部 KV 记忆并学习动态刷新，模型可以自适应地访问每个推理步骤所需的特定证据。

作者强调，他们的方法解决了静态保留和解耦选择的局限性，表明模型可以学习“知道要保留什么以及何时刷新”，从而在多跳推理任务中实现更优的性能，并相比现有的压缩技术提供更好的效率-质量权衡。这项工作表明，与生成共同训练的内部记忆管理，是扩展 LLM 至更长上下文且不产生高昂计算成本的一条可行路径。