MemTrain：不靠下游标注，光啃维基百科就把"记忆代理"训出来了
核心摘要
如果你最近在做 long-horizon Agent，多半会撞上一个非常具体的麻烦：把所有历史塞进 context 太贵，做 context memory 又训不动——因为高质量、记忆密集型的标注数据太少，下游 RL 一上去就过拟合到固定任务上，换个 benchmark 就废。这篇 MemTrain 的思路非常对我胃口：别再卷下游标注了，用维基百科的纯文本，自监督地把"维护记忆"这件事先练熟。具体做法是在分块阅读上设计两个耦合的代理任务——一个是端到端的掩码实体重建（强迫模型在多轮记忆更新后还能把 [MASK] 填回来），一个是中间记忆回忆（IMR，随机抽一个中间记忆状态去恢复之前看过的某个块），两者用 GRPO 联合优化。结果：Qwen2.5-7B 在 long-text QA 上比直接做 MemAgent 后训练多涨 17.67 个点，而且训练时只用 32k 的上下文，能稳定泛化到 896k。arXiv ID：2606.03197。
这是一篇我愿意推荐的论文——它不在"又一个新 RL 算法"上做花活，而是把"记忆能力本身可以被自监督预训"这件事讲清楚了。
论文信息
- 标题：MemTrain: Self-Supervised Context Memory Training
- 作者：Ziheng Li, Xingrun Xing, Haoqing Wang, Zhi-Hong Deng, Yehui Tang
- 机构：北京大学智能学院通用人工智能国家重点实验室 / 三星研究院（北京）
- arXiv：2606.03197
- 日期：2026/06/02
一、为什么要给"记忆"专门训一道？
先说一个我自己被困过的场景。
之前在做一个长文档问答的 Agent，要它边读边写笔记、最后回答跨段落的问题。最朴素的做法是直接把全文塞进 context 然后 RL，这条路在文档不长、模型上下文够用时还行。但等文档长度从 32k 涨到 200k，再涨到 800k，事情就开始失控——KV cache 涨爆显存只是表面问题，更要命的是注意力被稀释，模型在 100k 之后基本就在装样子。
Context memory agent 是这两年逐渐成型的一类方案：不要无限往后堆 context，而是在每一轮维护一个固定长度的 memory state \(m_t\)，每读一块新输入就把 memory 更新一次。MemAgent、MEM1 这些工作都属于这条路。它的好处是计算开销恒定、能突破原生上下文长度，缺点也很明显：
模型必须学会"把刚看到的关键信息塞进固定大小的 memory，并且在下一轮还能调出来用"。这个能力不能凭空获得，得训。
图1：传统长上下文代理与上下文记忆代理的对比。这张图把 MemTrain 想解决的问题摆得很直接——左侧灰色的传统范式上下文越堆越长，右侧绿色的记忆范式每轮把"记忆"当成一个可写可改的 slot 持续更新。
之前主流做法是怎么训这个能力的？基本都是端到端 RL on 下游任务：拿 HotpotQA、2Wiki 这种多跳 QA 数据，给一个 outcome reward，让模型在做题过程中自己学会怎么用 memory。
听起来挺合理。但实际跑过你就会发现两个问题：
- 数据稀缺：记忆密集型的高质量标注题（比如真正需要跨大段落整合信息的多跳 QA）非常贵，且容易被覆盖、被记住答案；
- 行为偏窄：模型学到的是"针对这个任务怎么调 memory"，换个任务就漏。说白了它学的是 task-specific 的记忆策略，不是通用的记忆能力。
我之前调一个类似的 setup 时，最直观的体感就是——一旦换个 benchmark，性能立刻塌一半。这种 brittleness 用再多 RL 也补不上来，因为问题不在算法，在数据本身的覆盖度上。
MemTrain 给的方案是：别管下游了，先用无标注的维基百科把"记忆"这个肌肉练扎实。下游再做任何 RL 都能站在更高的起点上。
二、方法：两个代理任务，联合优化
2.1 把"长文档阅读"改造成自监督任务
核心 trick 在数据构造上。MemTrain 不需要任何 QA 标注，它是这样把维基百科改造成训练信号的：
- 从维基里随机选一个 pivot passage；
- 用语义检索拉出 \(n_1\) 篇相关段落，再随机塞 \(N - n_1 - 1\) 篇无关段落，凑成一篇"长文档"；
- 在 pivot 段落里随机选一个实体 \(y\)（数字、地名、人名这类用 spaCy NER 抽出来的），把它在整篇文档中所有出现位置都换成
[MASK]； - 把这篇带 mask 的长文档切成固定长度的块 \(\{c_1, c_2, \dots, c_T\}\)。
为什么这个构造是关键？
因为 \(y\) 的所有副本都被 mask 了，模型没法靠"在某一块里直接抄答案"来作弊——它必须把分散在不同块里的上下文线索（"……出生于 [MASK]……"、"……[MASK] 的人口约 50 万……"）整合起来，才能反推出被 mask 的实体到底是什么。
这就把"长距离信息聚合"变成了一个有客观答案、可自动校验的任务。这件事的妙处不在于多新颖（BERT 那一套 MLM 思想都能对上），而在于它把 mask 跨到了多个 chunk、需要 memory 才能解的尺度上。
2.2 任务一：端到端掩码重建（End-to-End Reconstruction）
LLM 在 episode 里按顺序读这些 chunk：
- 第 \(t\) 步看到 \((m_{t-1}, c_t)\)，输出新的 memory \(m_t\) 和（可选的）一段思考内容；
- 到了最后一步 \(T\)，模型只看 \(m_{T-1}\)（注意，原始 chunk 已经看不到了），生成对实体 \(y\) 的预测 \(\hat{y}^E\)；
- 如果 \(\hat{y}^E = y\) 就给 reward 1，否则 0。
形式化的 reward 长这样：
后面那项是 IMR 的奖励叠加，下一节会展开。
这个任务给的是结果级监督——它告诉模型"你最终能不能从 memory 里把答案推出来"，但不告诉中间每一步该怎么写 memory。这就引出了第二个任务的必要性。
2.3 任务二：中间记忆回忆（IMR）
光给 outcome reward 有个老问题：信用分配。模型可能在第 3 步把关键信息丢了，但因为最后一步靠运气猜对了，整条轨迹都被打了高分。这种伪正样本用多了，memory 的写入策略就会越学越随机。
IMR 的设计很直接——给中间步骤也加监督：
- 跑完一条完整 episode、得到所有中间 memory \(\{m_1, \dots, m_{T-1}\}\) 后，随机抽一个中间步骤 \(k\)；
- 取出当时的 memory 状态 \(m_k^E\)；
- 再随机挑一个 \(k\) 之前的 chunk \(c_l\)（\(l \lt k\)），把里面的某个实体 \(\tilde{y}_i\) 也 mask 掉，得到 \(\tilde{c}_l\)；
- 让模型只看 \(m_k^E\) 和 \(\tilde{c}_l\)（没有原始历史），单步内把 \(\tilde{y}_i\) 恢复出来；
- 恢复对了就给 reward \(R_{i,j}^I = \mathbb{I}[\hat{y}_{i,j}^I = \tilde{y}_i]\)。
我盯着这个设计看了一会儿，觉得它真的挺漂亮——
它没有去构造一个新的 token 级 reward 模型，而是直接复用了 mask reconstruction 的检验机制，把"中间 memory 是否信息充分"翻译成了"能不能从 memory 里再 decode 出之前的内容"。
这等于在说：memory 的好坏 = memory 的可逆性。如果 \(m_k\) 里没保留之前 chunk 的关键信息，那再 decode 一次必然失败。这种 self-consistency 的味道很对。
2.4 联合优化：GRPO 一勺烩
两个任务的轨迹用 GRPO 一起算 advantage。具体来说，每条端到端 episode 会衍生出 \(T\) 步交互；每条 episode 又会再 sample \(G_2\) 条 IMR 轨迹。所以一个 batch 里总的交互数是 \(G_1 T + G_1 G_2\)。
图2：MemTrain 的 rollout pipeline。这张图值得仔细看——左半部分是端到端的多轮记忆更新，右半部分是从某个中间记忆状态出发的 IMR 分支。两类轨迹的 reward 在 GRPO 里被合并归一化，所以 IMR 的信号会反向传给端到端轨迹中"应该写好这一步 memory"的那个动作。
advantage 用的是非归一化形式 \(\hat{A}_{i,j,k} = R_i - \text{mean}\{R_i\}_{i=1}^G\)，没做标准差归一。这个细节其实是 DAPO/Dr.GRPO 这一脉的取舍——长度不均的轨迹做 std 归一会引入 bias，干脆不归一更稳。论文这里没展开论证，但跟近期的趋势是吻合的。
整体目标函数就是标准 GRPO：策略比率 + 裁剪 + KL 正则。\(\lambda = 0.5\) 平衡两类奖励，KL 系数 \(1\mathrm{e}{-3}\)，学习率 \(1\mathrm{e}{-6}\)，跑 300 步。
三、实验：先看主表，再看长度泛化
3.1 长文本 QA：训练 32k，测到 896k
主实验放在 long-context HotpotQA 上，把上下文从 7k 一路拉到 896k——注意训练时上下文上限只设到 8192 tokens（chunk 5120 + memory 1024 + 指令 1024 + 响应 1024），相当于用 8k 的训练长度去测 100 倍长的推理。
| 模型 | 7k | 14k | 28k | 56k | 112k | 224k | 448k | 896k | 平均 |
|---|---|---|---|---|---|---|---|---|---|
| Qwen3-4B-Instruct | |||||||||
| 基础 | 57.81 | 51.56 | 34.38 | 10.94 | 8.59 | 4.69 | 3.91 | 3.91 | 21.97 |
| +MemTrain | 63.28 | 60.16 | 60.16 | 57.03 | 60.94 | 58.59 | 48.44 | 40.62 | 56.15 |
| +MemAgent | 70.31 | 64.06 | 71.88 | 62.50 | 64.84 | 66.41 | 64.06 | 57.03 | 65.14 |
| +MemTrain+MemAgent | 79.69 | 73.44 | 75.78 | 73.44 | 68.75 | 67.19 | 61.72 | 62.50 | 70.31 |
| Qwen2.5-7B-Instruct | |||||||||
| 基础 | 53.12 | 51.56 | 35.16 | 13.28 | 10.16 | 1.56 | 1.56 | 0.00 | 20.80 |
| +MemTrain | 59.38 | 55.47 | 48.44 | 46.09 | 42.19 | 38.28 | 39.84 | 33.59 | 45.41 |
| +MemAgent | 64.06 | 67.19 | 62.50 | 59.38 | 55.47 | 50.00 | 46.88 | 41.41 | 55.86 |
| +MemTrain+MemAgent | 76.56 | 79.69 | 77.34 | 75.00 | 70.31 | 75.78 | 64.84 | 68.75 | 73.53 |
几个点值得拎出来讲：
第一，MemTrain 单独用就已经能打。 Qwen2.5-7B 上，纯 MemTrain（没做下游 RL）就把平均分从 20.80 拉到 45.41，涨了 24.6 个点。这意味着维基百科上的自监督训练已经给了模型一套通用的记忆动作。
第二，叠加下游 RL 是真"叠加"，不是替代。 MemTrain + MemAgent 在 Qwen2.5-7B 上达到 73.53，比单独 MemAgent 的 55.86 高 17.67 个点。这说明 MemTrain 学到的能力跟 MemAgent 学到的能力在表征层面没有冲突，前者是更通用的 memory schema，后者是任务特定的 reasoning policy，两者是正交的。
第三，越长越凸显价值。 看 896k 那一列：Qwen2.5-7B 基础模型直接跌到 0，MemAgent 单独也就 41.41，叠加 MemTrain 之后能稳到 68.75。这才是这套方法真正想证明的事——记忆能力可以被自监督地、长度无关地学到。
3.2 搜索 QA：换个赛道也成立
为了证明不是只在 HotpotQA 这一个数据集上 overfit，作者还在 search-based QA（基于工具调用的多跳检索 QA）上测了一遍，跟 MEM1 这套方案叠加：
| 模型 | TriviaQA | Bamboogle | HotpotQA | NQ | PopQA | 2Wiki | MusiQUE | 平均 |
|---|---|---|---|---|---|---|---|---|
| Qwen3-4B-Instruct-2507 | ||||||||
| 基础 | 42.71 | 21.78 | 18.94 | 19.92 | 21.81 | 14.36 | 4.76 | 20.61 |
| +MEM1 | 44.29 | 23.39 | 18.80 | 21.97 | 23.62 | 12.80 | 5.63 | 21.50 |
| +MemTrain+MEM1 | 55.63 | 34.68 | 27.85 | 32.24 | 37.91 | 25.84 | 10.43 | 32.08 |
| Qwen2.5-7B-Instruct | ||||||||
| 基础 | 18.84 | 8.87 | 11.15 | 12.22 | 12.59 | 10.45 | 4.43 | 11.22 |
| +MEM1 | 49.08 | 22.58 | 19.79 | 24.21 | 27.13 | 17.81 | 6.96 | 23.94 |
| +MemTrain+MEM1 | 57.21 | 30.65 | 27.73 | 35.18 | 38.36 | 27.32 | 10.64 | 32.44 |
我特别留意了 MEM1 单独和 MEM1+MemTrain 的差距——4B 上从 21.50 到 32.08（+10.58），7B 上从 23.94 到 32.44（+8.50）。两个不同体量的模型、七个不同数据集，几乎全线一致提升。这种一致性在论文里其实比单点 SOTA 更有说服力。
注意 MusiQUE 这一列特别低（10 左右），那是因为 MusiQUE 是一个特意设计来对抗"组合泛化"的多跳数据集，主流方法在上面普遍跌得很惨。MemTrain 也没把这个数字救起来——这是它的 honest weakness，不是所有"难题"都能靠记忆训练化解，有些是 reasoning 本身的问题。
3.3 消融：IMR 到底重不重要？
图3：消融实验。三条曲线分别是完整方法、去掉 IMR 的端到端版、把端到端 reward 与 IMR 解耦的版本。完整方法在所有上下文长度上都更稳，尤其是右半段长上下文区域优势拉开。
两个变体：
- End-to-End：把 IMR 分支拿掉，只保留 outcome reward。平均从 70.31 跌到 63.28。
- Decoupled：保留 IMR，但端到端 reward 不再加上 IMR 那一项（即两类轨迹的 reward 互相独立）。在短上下文（≤56k）反而稍好，但长上下文掉得很厉害。
第二个变体的现象挺有意思——它说明把 IMR 的 reward 反向加权到端到端轨迹上这件事，对长上下文鲁棒性是有用的。直觉上你可以这样理解：当端到端轨迹自己结果对错并不能完全代表 memory 写得好不好（长上下文下噪声大），让 IMR 来给"中间步骤"打分，可以有效降低信用分配的方差。
这个细节在原论文里只用了几行字带过，但我觉得是整篇里最值得 hands-on 实验验证的设计点。
3.4 案例研究：MemAgent 把人名忘了
Table 3 给了一个具体的失败-成功对比：问题是 "Who directed Big Stone Gap?"，需要跨多个 chunk 抓取导演的名字。
- MemAgent（无 MemTrain）：第一个 chunk 里写出了导演的关键信息，但 memory update 时把关键的人名给丢了，到后面几轮 memory 里只剩下 "the film was set in Virginia" 这种无关细节，最终答错。
- MemTrain+MemAgent：从第一轮就把导演姓名稳稳塞进 memory，并且在后续 4 轮更新里始终保留这个 entity，最后正确推出答案。
这个 case 把 IMR 在做什么讲得很形象——它就是逼模型"该记的人名你必须记住"。
3.5 持续后训练真的不如预训记忆？
图4：把 MemAgent checkpoint 继续后训练 300 步，对比从 MemTrain 初始化再做 MemAgent 后训练。前者只能拿到边际收益甚至下降，后者保持 2.64 个百分点的稳定优势。
这张图回答了一个很自然的质疑——"你光是多训了 300 步而已，会不会 baseline 多训 300 步也能涨上去？"作者直接补了实验：MemAgent 自己继续训 300 步，平均分基本横盘甚至略降，没有任何 free lunch。
也就是说，MemTrain 用维基百科训出来的能力，在下游 RL 阶段是补不回来的——不是简单的"训练步数差异"。
四、我的判断：这篇值不值得跟？
直说我的看法。这是一篇方向选得很对、执行也扎实的工作。它没有去发明新算法（GRPO 是现成的、mask reconstruction 是现成的、IMR 的检验机制也是 mask reconstruction 的复用），但它选对了一个被忽视的位置：
大家都在卷"下游 RL 算法"和"reward shaping"，没人认真问过"记忆能力本身能不能预训"。MemTrain 把这件事摊开了，而且证明它不仅 work，还跟下游 RL 正交。
几个我觉得值得继续挖的点：
第一，IMR 这个 self-consistency 思想可以泛化。 你完全可以把它推广到 tool use、code execution 这些场景——只要中间状态有"可被验证恢复"的属性，就能造一个细粒度 reward。我隐约觉得这是一个 underexplored 的训练信号源。
第二，"用维基百科训通用能力"这条路还能再深挖。 不止记忆，长程 reasoning、检索决策、甚至 tool routing 都可能可以在 wiki 这种结构化纯文本上自监督地训出来。MemTrain 是个早期样板，但思路可以延伸。
第三，几个 honest weakness。 MusiQUE 上提升有限（多跳推理本身的瓶颈）；训练块长度固定 5120，对真正变长的 chunking 场景泛化没测；IMR 的 \(G_2 = 8\) 让单步训练成本翻倍，工程上需要权衡。
最后给做工程的同学一个实操建议：如果你正在做长文档 Agent 或者多轮 memory-based agent，MemTrain 这套自监督预训完全可以加在你现有的 RL pipeline 之前——成本相对可控（300 步、维基语料免费、复用 GRPO），收益在 5-17 个点这个量级。这种 "免费的菜"在 2026 年的 LLM 训练里已经不多见了。
参考文献
- MemTrain: Self-Supervised Context Memory Training, arXiv:2606.03197
- MemAgent: Reshaping Long-Context LLM with Multi-Conv RL-based Memory Agent
- MEM1: Learning to Synergize Memory and Reasoning for Efficient Long-Horizon Agents
- DeepSeek GRPO 原论文（Group Relative Policy Optimization 算法基础）
- veRL：本文实现使用的 RL 训练框架
觉得有启发的话，欢迎点赞、在看、转发。跟进最新AI前沿，关注我