MemTrain：自监督上下文记忆训练

# MemTrain：自监督上下文记忆训练  
来源：https://arxiv.org/html/2606.03197  

李梓恒1,2†，邢星润2†，王浩庆2，邓志鸿1✉\{\}^\{1\{~\\textrm\{\{\\char 0\\relax\}\}\}\}，唐晔辉2✉\{\}^\{2\{~\\textrm\{\{\\char 0\\relax\}\}\}\}  

1北京大学通用人工智能国家重点实验室，智能科学与技术学院  
2三星研究院，北京，中国  
\{liziheng,zhdeng\}@pku\.edu\.cn  
yehui\.tang@samsung\.com  
†同等贡献  
✉\{\}^\{\\textrm\{\{\\char 0\\relax\}\}\}通讯作者  

###### 摘要  

记忆是长周期LLM智能体不可或缺的能力，使其能够保留并利用跨长时间交互累积的信息。现有的记忆型智能体方法通常在下游任务上通过强化学习进行端到端训练。然而，为记忆密集型场景收集高质量标注问题成本高昂，且由此产生的训练数据往往缺乏足够的多样性来覆盖通用的记忆行为。在这项工作中，我们提出了MemTrain，一个自监督训练框架，用于普遍增强LLM智能体的上下文记忆能力，以支持更有效的下游后训练。MemTrain在未标注的维基百科语料上引入了两个耦合的代理任务：(1) 端到端掩码重建目标，要求模型在多次记忆更新后恢复被掩码的实体，从而从最终结果角度鼓励记忆维护；(2) 中间记忆召回目标，要求模型利用中间记忆状态重建被掩码的历史信息，鼓励在整个交互过程中进行忠实的压缩和记忆完整性。这两个目标使用GRPO联合优化。在长文本问答和基于搜索的问答基准上的大量实验表明，MemTrain在不同模型上持续提升下游记忆密集型推理性能，相较直接的任务特定后训练，提升幅度最高达17.67个百分点。  

## 1 引言  

大语言模型（LLM）已迅速发展成为越来越强大的智能体，能够进行推理、规划并与外部环境交互（Singh等人，2025（https://arxiv.org/html/2606.03197#bib.bib26）；Team等人，2025（https://arxiv.org/html/2606.03197#bib.bib27）；DeepSeek-AI等人，2025（https://arxiv.org/html/2606.03197#bib.bib9））。然而，长周期智能体任务的一个关键瓶颈是**记忆**：即保留并利用多个回合前获取的信息的能力。在真实的交互环境中，智能体持续接收新的观测，生成中间思考，并且必须跨回合维护相关的过往信息。一个直接的解决方案是将完整的交互历史附加到提示中（Yao等人，2023（https://arxiv.org/html/2606.03197#bib.bib36）），但随着轨迹增长，这很快变得代价高昂。因此，使智能体能够在**固定大小的持久记忆**下运行，对于可扩展的长周期部署仍然是一个重要挑战。  

最近的工作探索了**上下文记忆**智能体（Zhou等人，2025b（https://arxiv.org/html/2606.03197#bib.bib3）；Yu等人，2025a（https://arxiv.org/html/2606.03197#bib.bib2）；Yan等人，2025（https://arxiv.org/html/2606.03197#bib.bib33）；Yuan等人，2026（https://arxiv.org/html/2606.03197#bib.bib41）），其中每一轮交互都基于一个紧凑的记忆状态，而不是完整历史。在第t轮，模型接收形如[记忆_{t-1};输入_{t}]的输入，生成响应，并将记忆更新为记忆_{t}。这种范式允许在保留历史信息的同时实现近乎恒定的上下文使用，并且可以在语言模型内部进行端到端优化。然而，现有的记忆型智能体通常使用基于可验证奖励的强化学习（RLVR）在下游任务上进行训练。这种方法需要昂贵的标注数据，使得获得足够多样的训练数据变得困难，这些数据需要涵盖广泛的记忆行为。因此，以这种方式学习到的记忆能力往往是领域特定的，并且泛化能力有限。这些局限性凸显了对通用自监督训练范式的需求。  

与此同时，推理方面的最新进展探索了使用预训练数据的强化学习（Dong等人，2025（https://arxiv.org/html/2606.03197#bib.bib10）；Li等人，2025（https://arxiv.org/html/2606.03197#bib.bib17）；Xing等人，2025（https://arxiv.org/html/2606.03197#bib.bib31））。他们通过在未标注语料上构建基于思维链的下一个token预测的自监督代理任务，来普遍提升推理能力。然而，记忆学习与推理面临不同的挑战。记忆目标本质上是隐式的且依赖过程的，因为模型必须持续决定保留、压缩和回忆哪些信息。因此，设计一个能够忠实捕捉底层记忆机制的代理任务仍然是一个重大挑战。  

为了解决这一挑战，我们提出了**MemTrain**，一个自监督训练框架，用于提升LLM智能体的通用上下文记忆能力，以更好地支持下游后训练。MemTrain基于两个耦合的代理任务，这些任务从维基百科段落构建：(1) 端到端掩码重建任务，要求模型在多次记忆更新后恢复被掩码的实体，从而鼓励有效的记忆维护和利用；(2) 中间记忆召回任务，要求模型利用中间记忆状态从较早的交互历史中重建额外的被掩码实体，鼓励在整个记忆更新过程中实现记忆完整性和忠实压缩。这两个目标使用GRPO联合优化。大量实验表明，MemTrain在长文本问答和基于搜索的问答上持续优于直接的任务训练。在Qwen3-4B-Instruct-2507上的平均提升分别达到5.17和10.58个百分点，在Qwen2.5-7B-Instruct上分别达到17.67和8.50个百分点。  

我们的贡献总结如下：  

- • 我们提出了MemTrain，这是第一个旨在普遍提升LLM智能体上下文记忆能力以支持有效下游后训练的自监督训练框架。  
- • 我们引入了一种新颖的面向记忆的代理训练范式，该范式为记忆生成和利用同时提供结果级和过程级的监督信号。  
- • 在长文本问答和基于搜索的问答任务上的大量实验表明，MemTrain在4B和7B模型上持续提升下游后训练的性能上限。  

## 2 相关工作  

#### 长周期LLM智能体的记忆。  
LLM智能体最广泛采用的记忆管理策略是将环境观测和模型响应不断附加到上下文窗口中（Yao等人，2023（https://arxiv.org/html/2606.03197#bib.bib36）），但这从根本上受限于LLM有限的上下文窗口。为了实现无界记忆，人们提出了外部记忆系统，将交互记录压缩或总结后存储在外部。（Yoon等人，2024（https://arxiv.org/html/2606.03197#bib.bib38）；Li等人，2023（https://arxiv.org/html/2606.03197#bib.bib16）；Chhikara等人，2025（https://arxiv.org/html/2606.03197#bib.bib8）；Xu等人，2025（https://arxiv.org/html/2606.03197#bib.bib32））。Qian等人（2026（https://arxiv.org/html/2606.03197#bib.bib23））；Xu等人（2025（https://arxiv.org/html/2606.03197#bib.bib32））；Chen等人（2026（https://arxiv.org/html/2606.03197#bib.bib7））进一步引入了多智能体框架以支持更复杂和高效的记忆管理。然而，外部记忆系统往往忽视了记忆与推理之间的内在协同，同时增加了整体系统的复杂性。  

更近期的研究（Zhou等人，2025b（https://arxiv.org/html/2606.03197#bib.bib3）；Yu等人，2025a（https://arxiv.org/html/2606.03197#bib.bib2）；Wu等人，2026（https://arxiv.org/html/2606.03197#bib.bib30）；Ye等人，2025（https://arxiv.org/html/2606.03197#bib.bib37）；Yuan等人，2026（https://arxiv.org/html/2606.03197#bib.bib41））将记忆构建和利用直接整合到智能体自身的推理过程中，实现了端到端优化。尽管有效，但这些方法通常依赖于昂贵的任务特定标注，严重限制了数据多样性。在这项工作中，我们转而提出一种自监督训练框架，使得能够在通用互联网语料上进行训练，显著增强了数据多样性。  

参见标题下的图1：现有长周期智能体与上下文记忆智能体的比较。传统上，为了处理长上下文文档或多轮环境交互，LLM必须将所有输入保留在上下文中，导致高计算成本和注意力压力。相比之下，上下文记忆智能体维护一个固定长度的上下文记忆，每轮更新，从而允许在可行的资源限制内处理不断增加的输入。  

#### 用于LLM预训练的强化学习。  
在后训练阶段，强化学习已被广泛采用以增强LLM的推理和工具使用能力（DeepSeek-AI等人，2025（https://arxiv.org/html/2606.03197#bib.bib9）；Yu等人，2025c（https://arxiv.org/html/2606.03197#bib.bib42））。然而，后训练方法通常依赖于精心策划的问答数据集，这限制了可扩展性和泛化能力。受自监督语言模型预训练成功的启发，最近的工作探索了利用大规模互联网文本的强化预训练范式。Quiet-STaR（Zelikman等人，2024（https://arxiv.org/html/2606.03197#bib.bib43）；Huang等人，2025（https://arxiv.org/html/2606.03197#bib.bib13））在每个token位置生成潜在理由以更好地预测未来文本。RPT（Dong等人，2025（https://arxiv.org/html/2606.03197#bib.bib10））引入了下一个token推理的RLVR目标，并首次证明了可扩展的强化学习预训练。RLPT（Li等人，2025（https://arxiv.org/html/2606.03197#bib.bib17））采用了类似的公式，同时引入了一个生成式奖励模型。RLP（Hatamizadeh等人，2025（https://arxiv.org/html/2606.03197#bib.bib11））用对比奖励替代下一个token预测以显式诱导推理。PretrainZero（Xing等人，2025（https://arxiv.org/html/2606.03197#bib.bib31））进一步提出了一个主动预训练框架，合成更有信息量和价值的训练样本。然而，现有的基于RL的预训练方法主要集中在单轮推理上，学习有效的多轮记忆维护和利用的问题在很大程度上尚未被探索。  

## 3 自监督记忆训练  

在本节中，我们首先形式化上下文记忆智能体（§3.1（https://arxiv.org/html/2606.03197#S3.SS1））。然后我们介绍两个代理任务——端到端掩码重建（§3.2（https://arxiv.org/html/2606.03197#S3.SS2））和中间记忆召回（§3.3（https://arxiv.org/html/2606.03197#S3.SS3））。最后我们描述如何使用GRPO进行记忆训练（§3.4（https://arxiv.org/html/2606.03197#S3.SS4））。  

### 3.1 问题设定  

我们的研究基于MemAgent（Yu等人，2025a（https://arxiv.org/html/2606.03197#bib.bib2））提出的多轮上下文记忆框架。如图1（https://arxiv.org/html/2606.03197#S2.F1）所示，现有的上下文记忆机制可以抽象为在交互步骤t维护一个固定长度的记忆状态m_t。在每个交互步骤，模型接收输入元组（m_{t-1}, a_{t-1}, i_t），其中a_t表示模型在当前步骤选择的动作。动作空间取决于目标应用。对于长上下文阅读智能体，动作可能对应请求下一个文本块或生成最终答案。对于搜索智能体，动作可能涉及调用外部搜索工具或直接返回答案。对于与环境交互的非终止动作，i_t表示执行所选动作后返回的环境输入或反馈。基于（m_{t-1}, a_{t-1}, i_t），模型生成更新后的记忆状态和动作，即（m_t, a_t），这些将在后续的交互步骤中使用。  

与传统的智能体范式（其中整个交互历史不断附加到上下文窗口）相比，上下文记忆在整个轨迹中保持恒定的上下文大小。这种设计消除了对不断增长的上下文长度的依赖，使得能够超出模型原生上下文限制进行长周期交互，同时减轻注意力稀释，并避免与长上下文处理相关的增加的计算成本。  

参见标题下的图2：MemTrain在GRPO训练期间的rollout流水线说明。首先，我们从维基百科语料中选择N个段落，并构建一个分块的输入集合c_{1:T-1}。然后我们采样G_1个多轮轨迹o^E_{1:T}，用于通过顺序读取c_{1:T-1}并更新上下文记忆来恢复被掩码的词ŷ。对于每个多轮轨迹，我们随机选择一个中间记忆来恢复之前的一个输入块，并生成G_2个中间记忆召回轨迹。最后，我们为所有G_1T + G_1G_2个交互计算奖励和优势。  

### 3.2 端到端掩码重建  

我们从原始的维基百科文本构建训练样本。首先，我们随机选择一个段落作为锚点段落。然后我们从语料中检索n_1个语义相关的段落，以及N - n_1 - 1个随机采样的段落。这N个段落以随机顺序拼接，形成一个长文档。接下来，我们从锚点段落中随机选择一个实体y（例如，一个数字或地点），并将文档中所有出现该实体的地方替换为特殊token [MASK]。遵循上下文记忆研究中的实践（Yu等人，2025b（https://arxiv.org/html/2606.03197#bib.bib39）），我们将长文档分割成固定长度的块{c_1, c_2, ..., c_T}，其中每个块对应一个交互步骤。LLM顺序处理这些块，按照o_{i,t}^E ∼ π_θ(·|q^E, o_{i,t-1}^E, c_t)生成多轮轨迹o_i^E（第i次rollout），其中q^E表示重建提示，详见附录A（https://arxiv.org/html/2606.03197#A1）。对于第i次rollout中的第t步，我们评估是否成功恢复被掩码的实体：r_{i,t}^E = 1如果o_{i,t}^E包含y，否则为0。最终的端到端重建奖励是r_{i,T-1}^E，即最后一次交互步骤的奖励，此时模型已读取所有块并应产生最终答案。对于G_1个rollout中的每一个，我们进一步对最后一步执行G_2次中间记忆召回rollout，如下所述。  

### 3.3 中间记忆召回  

中间记忆召回任务旨在鼓励模型在记忆更新过程中保留完整信息。具体来说，对于每个端到端重建轨迹，我们随机选择一个中间步骤j（1 ≤ j ≤ T-1），并提示模型仅基于截至步骤j的中间记忆状态m_j重建一个之前输入块c_k（其中k ≤ j）中的被掩码实体。该提示表示为q^R，详见附录A（https://arxiv.org/html/2606.03197#A1）。模型生成s个token的rollout o^R，然后我们再次评估o^R是否包含被掩码的实体：r^R = 1如果o^R包含y，否则为0。  

### 3.4 GRPO训练  

我们使用GRPO联合优化这两个目标。对于每个训练样本，我们采样G_1个端到端重建轨迹，并为每个轨迹采样G_2个中间记忆召回轨迹。所有轨迹共享相同的输入文档和掩码实体。GRPO损失为：  

L(θ) = −(1/(G_1T + G_1G_2)) Σ_{i=1}^{G_1} Σ_{t=1}^{T} A_{i,t}^E log π_θ(o_{i,t}^E | ...) − (1/(G_1G_2)) Σ_{i=1}^{G_1} Σ_{k=1}^{G_2} A_{i,k}^R log π_θ(o_{i,k}^R | ...)  

其中A_{i,t}^E和A_{i,k}^R是基于组内归一化奖励计算的优势函数。  

## 4 实验  

我们通过测量后训练后的最终下游性能来评估MemTrain的有效性。我们考虑两个代表性任务：(1) 长上下文多跳问答（§4.2（https://arxiv.org/html/2606.03197#S4.SS2）），这与记忆训练设置高度匹配，模型读取分块的长文档并回答；(2) 带搜索工具的多跳问答（§4.3（https://arxiv.org/html/2606.03197#S4.SS3）），一个领域外任务。