MemAgent: Reshaping Long-Context LLM with Multi-Conv RL-based Memory Agent
Paper: arXiv:2507.02259 Code: BytedTsinghua-SIA/MemAgent Code reference:
main@ef4219b2(2026-05-12)
1. Motivation(研究动机)
现有长上下文方案主要有三类:长度外推/继续预训练、稀疏/线性注意力、上下文压缩或外部记忆插件。它们分别卡在三类瓶颈上:外推方法常在超出训练长度后性能明显掉点;稀疏/线性注意力往往需要改架构或重训,且不一定适配通用 LLM;压缩/插件式方案会打断标准生成流程,兼容性和并行性都差。 论文的问题意识并不只是“窗口不够大”,而是“真正有用的信息远少于原文长度,但模型并不会自动把这些高价值信息稳定保存下来”。在 HotpotQA 这类多跳 QA 中,答案往往依赖跨段实体链接、别名对齐、时间关系或因果链条;如果模型在某个早期 chunk 里漏掉关键实体,后面即使继续读更多内容,也无法弥补这次遗失。作者因此把长上下文失败的根因理解为记忆管理失败,而不只是注意力覆盖范围不足。 从论文的比较框架看,长度外推方法默认要尽量维持 Transformer 在更长 token 序列上的工作方式不变,因此重点放在 RoPE 缩放、继续预训练或长序列 curriculum 上。这些方向确实能延后性能崩塌点,但它们追求的仍是“让模型继续看更多 token”,而不是“让模型学会只保留任务真正需要的 token”。MemAgent 则反过来假设:对于问答任务,模型没有必要忠实记住整篇文章,真正重要的是在阅读过程中持续维护一份 task-conditioned memory。
稀疏/线性注意力方案也有类似局限。它们解决的是计算复杂度,让模型可以更廉价地访问更长序列;但这不等于模型会形成稳定的选择性记忆习惯。对超长输入来说,能“看见”所有 token 与能“保住”关键证据是两回事。MemAgent 刻意给出固定记忆预算,逼迫模型在每一轮决定什么应该被压进 memory,什么可以丢弃,这比单纯扩展可见范围更贴近真实 agent 在有限状态下处理流式信息的场景。压缩/插件式方法表面上更接近 MemAgent,因为两者都承认原文必须被压缩。但作者指出,很多现有方案依赖额外 summarizer、retriever 或外接 memory module,这会带来两类问题:一是训练与推理路径偏离 base LLM 的标准文本生成接口,兼容性差;
二是 memory write 动作往往不在同一个可微或可强化学习的闭环里,难以端到端优化。MemAgent 的动机正是把 memory 作为普通文本输出,让“写记忆”本身成为可被 RL 直接优化的策略。本文要解决的是:能否不改基础 Transformer 结构,只通过训练一个“会写记忆”的 RL agent,让 LLM 在固定窗口下处理任意长文档,并且在百万级 token 上仍然保持接近不降的准确率。更细一点地说,论文其实在回答三个连在一起的问题。第一,memory 能不能被表示成普通 token,而不是专用 latent state。第二,memory update 能不能只靠最终任务回报训练出来,而不需要人工标注摘要。
第三,多轮、上下文彼此独立的 memory conversations 能不能被纳入统一的 RL 目标,而不是退化成只优化最后回答的一次监督学习。MemAgent 的方法设计基本都是围绕这三个子问题展开的。
这个问题值得研究,因为一旦成立,长文档 QA、超长记忆 agent、长链推理等任务就可以用标准 LLM 推理范式直接完成,而不必依赖复杂架构改造或显式检索系统。 这件事的工业含义也很直接。若 MemAgent 的假设成立,那么很多已经存在的 base LLM 不需要重新做昂贵的长序列预训练,就可以通过后训练获得处理百万级上下文的能力;而部署侧几乎不需要新系统,输入输出仍是标准 chat completion,memory 只是一段普通文本。论文强调“兼容标准 LLM workflow”,本质上就是在追求这种低迁移成本的能力增强路径。 从 agent 研究角度看,这篇论文还有一个额外价值:它把“记忆”从经验性 heuristics 变成了可验证、可优化的 workflow。很多 agent 系统里的 memory 仍停留在工程技巧层面,例如会话摘要、scratchpad 压缩或 vector store 回写;MemAgent 则试图证明,只要把 memory 写作动作纳入明确的 RL objective,并把最终任务成败回传给所有中间轮次,memory 本身就能像策略一样被系统性训练。
2. Idea(核心思想)
核心 insight 是把长上下文理解成“流式证据 + 固定长度记忆”的决策问题,而不是单次吃完整段文本的问题。模型每次只看一个 chunk 和当前 memory,然后用 RL 学会覆盖/保留哪些信息,最终只靠 memory 回答。 这里最核心的观念变化是:memory 不是被动缓存,而是主动生成的中间状态。模型在每一轮并不是简单做局部摘要,而是根据问题、已有 memory 和当前 chunk 三者的组合,决定哪些旧信息必须保留,哪些细节已经可以舍弃,哪些新证据值得进入有限记忆预算。换句话说,MemAgent 把“摘要”升级成了“带任务目标的状态更新”。 论文还刻意把 final answer 阶段设计成只能访问最终 memory、不能回看原文 chunk。这是一个非常强的约束,因为它迫使所有真正决定答案的证据都必须在 earlier turns 被搬运进 memory。若模型最终仍能保持高准确率,就说明它学到的不是长上下文模型常见的隐式缓存能力,而是显式、跨轮可持续的记忆策略。 关键创新有两点:一是把 memory 直接做成上下文中的普通 token,保持 base LLM 的 attention/tokenization 不变;二是把多轮独立上下文的整条轨迹都纳入优化目标,用 multi-conv RL 训练“读 chunk、写 memory、最后答题”的完整 agent workflow。
“memory 作为普通 token” 的意义比表面更大。很多 memory-augmented 方法会引入额外 slots、外部数据库或专门模块,这些设计虽然表达力更强,但会把方法绑定到特定 backbone 或推理系统。MemAgent 刻意选择最保守、也最兼容的接口:memory 就是一段自然语言文本,被填回标准 chat template,和问题、chunk 一起送入原始 LLM。这样做的好处是训练目标完全落在常规 next-token 概率上,也便于直接复用已有模型与推理框架。第二个创新点是 multi-conv RL 的目标构造。普通 tool-using agent 的 RL 处理的是一条串行对话,所有 token 共享同一条上下文轨迹;
MemAgent 的每轮 memory update 却是上下文独立的 conversation,输入里只保留当前 memory 状态和当前 chunk,而不是前面所有原始 token。作者因此不能直接套用标准 multi-turn RL mask,而是必须重新定义样本内多轮 conversation 如何共享 reward、如何平均 token loss。可以把 MemAgent 的 idea 压缩成一句话:它不试图让模型“持续看到整个世界”,而是训练模型“在有限状态里持续留下真正重要的世界”。
这也是它和 length extrapolation、RAG、普通多轮 agent 三类方法的根本区别:前者追求更大的可见范围,RAG 追求可反复访问外部证据,普通 agent 追求串行工具交互;MemAgent 追求的是单次流入、不可回看条件下的在线记忆维护。与现有 length extrapolation 的根本区别在于:它不试图让模型“记住所有 token”,而是训练模型“只记住该记住的 token”;与 tool-using multi-turn RL 的区别在于:这里每轮 conversation 是独立上下文,不是简单拼接的对话轨迹。
与 RAG 的差别也值得单列出来。RAG 假设需要时可以重新检索原始证据,因此核心问题是“找得到”;MemAgent 假设证据流只经过一次,过去 chunk 不再可见,因此核心问题是“第一次看见时有没有留下足够的信息”。它更像 streaming setting 下的在线压缩器,而不是开放语料上的搜索器。这也解释了为什么论文在长文档 QA 上强调 nearly lossless extrapolation:它证明的是一次性阅读场景中的记忆保持能力,而非检索能力。
3. Method(方法)
3.1 总体框架
Figure 2:把长上下文改成固定预算状态机
Figure 2 解读:左边是传统长上下文 LLM,必须把整段文本一次性放进上下文;右边是 MemAgent,把文档切成多个 chunk,反复执行“读一段文本 + 更新固定长度 memory”,最后只用 memory 回答。这个设计的关键不是扩大窗口,而是把超长输入转成多轮固定预算决策。
两段式流程与状态转移
MemAgent 的流程分成两个模块。Context-Processing 模块逐 chunk 处理文本并写入 memory;Answer-Generation 模块在文档读完后,只看问题和 memory 输出最终答案。由于 memory 长度固定,整体复杂度对输入长度是线性的;由于没有改 positional embedding 或 attention 结构,模型仍然是标准 LLM。
如果把论文流程写成状态转移,可以记为:给定问题 、chunk 序列 与 memory 状态 ,模型在第 轮执行 m_t = f_\\theta(q, c_t, m_{t-1}),最终再基于 a = g_\\theta(q, m_K) 输出答案。这个写法强调的是,MemAgent 处理超长文档的关键不是单次前向能容纳多少 token,而是是否能在有限状态 上稳定滚动更新。对论文来说,memory 是真正被优化的对象,final answer 只是检验 memory 是否足够好的读出步骤。这个框架天然包含两个硬约束。第一,memory 必须短,否则单轮输入很快重新膨胀成长上下文。
为什么 memory 必须稳定
第二,memory 必须稳,因为任何一轮错误覆盖都可能让后续所有 chunk 建立在错误状态之上。作者没有手工设计复杂 merge rule,而是让 RL 去学习在“保留旧证据”和“接纳新证据”之间做权衡。也因此,MemAgent 不是一般意义上的摘要系统,它追求的不是忠实复述原文,而是最大化最终答题正确率。released code 让这一点更清晰。
released code 中的固定窗口约束
MemoryConfig 直接定义 chunk_size、max_memorization_length、max_chunks 和 max_final_response_length,并把 max_input_length 写成 max_prompt_length + chunk_size + max_memorization_length + template_length。这表明系统不会因为样本更长就动态扩窗,而是把“单轮预算固定、总轮次可变”作为一等设计原则。文档越长,增加的是 recurrent turns,而不是单轮可见上下文。
代码还揭示了一个重要训练细节:数据集阶段只提供 prompt_ids、context_ids 和 context_length,不会提供任何 gold memory。也就是说,memory 轨迹不是预先构造好的监督信号,而是在 rollout 期间由当前策略在线生成。这样训练路径和推理路径保持一致,避免了“训练时看老师写的摘要、推理时自己写摘要”的暴露偏差。把这一点和论文的 nearly lossless 结果放在一起看,含义其实很强:作者并不是在证明“人工设计一个很好的摘要协议可以支持长文档 QA”,而是在证明“标准 LLM 经过合适的 RL 训练后,可以自己学会一套足够好的摘要/记忆协议”。
方法贡献的边界
这让 MemAgent 的贡献从单个 benchmark 提升,变成了一种关于 LLM 是否能自发形成可持续记忆策略的实证结论。另外,Context-Processing 与 Answer-Generation 的拆分也让误差来源更容易分析。若最终答错,至少可以区分两种失败:一种是 memory 在前面轮次已经丢失或扭曲了关键信息;另一种是 memory 本身足够,但 final answer 阶段没有正确读出。论文主结果说明前一种失败被显著压低,而 released code 的显式 memory 文本又让这两类失败在实践中可被人工区分。这种可解释性是很多纯长上下文模型不具备的。
这类模块化拆分还有一个隐含好处:它把“阅读”与“作答”在功能上分离,但没有把它们拆成两个模型。于是,同一个策略既要学会在阅读阶段写出未来可用的 memory,又要在作答阶段消费自己先前写下的 memory。这个闭环让 memory 的价值定义得非常清楚:不是像不像摘要,而是能否被同一模型在最后稳定读出并转化为正确答案。
因此,MemAgent 的总体框架其实可以看成一种“显式可读状态上的自监督决策闭环”:状态由模型自己写,未来也由模型自己读,最终正确性再反过来约束整条状态演化轨迹。 这正是它区别于“更长上下文但同样一次性前向”的根本方法论差异。
3.2 训练目标:GRPO + 多轮独立上下文
标准 GRPO 目标
论文把每个样本的多轮 memory-update conversation 当成一个优化对象。对第 个样本,GRPO 先为同一个 query 采样一组响应 ,再用组内 reward 做归一化 advantage: 其中 是样本级回报, 是 token 级重要性权重。标准 GRPO 目标为:
真正难点:reward 属于哪一轮
其中 。如果只看这套公式,容易把 MemAgent 理解成“在新任务上套标准 GRPO”。真正的难点其实在 reward 属于谁。标准 GRPO 假设每个样本对应一段连续输出,而 MemAgent 一个样本内部包含多次 memory update 加一次 final answer:中间轮次并不直接接受 verifier 检查,但它们又决定最终能否答对。因此训练目标必须把最终任务成败重新分配到所有中间 memory writing 行为上。
论文采用的是一种很典型但很适合 agent workflow 的 delayed credit assignment:final turn 负责产生可验证答案,样本级 reward 则传播回同样本下的所有 memory conversations。这样一来,中间 memory 文本虽然没有逐轮人工标签,仍能因其是否帮助最终答对而获得正负反馈。对长文档任务来说,这比对每轮 memory 做 imitation learning 更自然,因为作者真正关心的是 memory 是否“有用”,而不是是否“像人工摘要”。released code 也能帮忙理解这一点。
released training config
训练脚本 run_memory_7B.sh 虽然入口是 verl.trainer.main_ppo,但显式设置 algorithm.adv_estimator=grpo,并让 actor_rollout_ref.rollout.n=16,说明同一 query 会采样一组 response group 进行组内相对比较。同时保留 kl_coef=0.001,表示策略更新仍受到 reference model 约束,避免 memory 生成因 reward hacking 而过快漂移到异常格式或投机模式。README 还明确提示训练 verifier 比测试 verifier 更严格:训练时答案必须写进 \\boxed{},且大小写、额外字符都更受限;
测试时则会忽略大小写、标点和冠词。这意味着训练目标不只是“最终答对”,还包含一种强格式约束,以减少 reward hacking。方法层面的含义是,MemAgent 学到的并非宽泛的 free-form QA,而是围绕可验证答案构建的一整套 memory-to-answer policy。
不需要 gold memory 的原因
这套设计还有一个隐含优点:因为 reward 来自最终可验证答案,作者不需要构造“什么样的 memory 才算好”的额外判别器。任何 memory 只要最终支持正确答案,就会获得正向训练信号;哪怕 memory 看起来不像人工摘要,只要它保留了足够证据,训练也会接受它。这让方法避免了把 memory 学习额外绑死在某种人类偏好的摘要风格上,而更纯粹地围绕任务成功定义价值。与此同时,这种 reward 设计也有明显挑战:监督非常稀疏。模型可能在前面很多轮都写出“看起来合理”的 memory,但只要最终少一个关键关系,整条轨迹就会失败。
因此 MemAgent 采用 group-based 相对比较和 sample-level reward 传播,本质上是在尽量放大这类稀疏成功信号,让中间 memory turns 也能从最终成败中受益。若没有这套 credit assignment,中间轮次几乎只能依赖 KL 维持语言流畅,而难以学出真正有任务导向的写记忆习惯。token-level averaging 在这里也不只是技术细节。memory update 往往比 final answer 长,若简单把所有 token loss 混在一起,模型可能更倾向于优化“更容易生成长而安全的 memory 文本”,而不是更关键的 final answer correctness。
论文借鉴 DAPO/DrGRPO 系路线索,把 conversation 与 token 两层都纳入平均,目的就是减少长度差异带来的梯度偏置,让训练更关注轨迹整体价值,而不是某一类长输出在 token 数上的天然优势。
从泛化角度看,这套训练目标并不依赖 HotpotQA 的语义结构本身。它真正依赖的是两个条件:一是最终任务可以通过某种 verifier 给出相对可靠的成败信号;二是任务允许把完整输入拆成逐段读取、最后再回答的流程。因此 MemAgent 的方法不只适用于长文档 QA,理论上也可迁移到长日志分析、超长代码库问答、长时序事件归因等场景,只要这些任务的答案可验证且需要跨段保留关键证据。 论文没有把这种迁移展开成大规模实验,但方法形式本身已经给出方向:一旦把 verifier 换成领域对应的可验证规则,memory workflow 本身无需大改。也就是说,MemAgent 更像一个“如何训练可写记忆 agent”的通用 recipe,而不仅是一套针对 HotpotQA 的专用技巧。这也是它相比很多只在单一 benchmark 上报告提升的方法更有启发性的地方。
3.3 multi-conv 关键改造
Figure 3:从 multi-turn 到 multi-conv
Figure 3 解读:这里展示的是为什么不能把 MemAgent 当作普通多轮对话来做 mask。因为同一个样本会生成多个彼此独立的 conversation,最终答案只由最后一轮决定,但中间的 memory 更新轮次也必须获得训练信号。论文的核心做法是把“样本内的多轮 conversation”提升成一个三维优化结构 ,而不是只做 。
reward broadcast 公式
论文给出的做法是:最终 reward 来自包含最终答案的那一轮 conversation,然后把这个样本级 reward 传播到同一样本下的所有 conversation;再把 token-level loss 做均匀平均。形式上,论文写成: 这里 表示第 个样本下的第 个 conversation。论文特别强调,它借鉴 DAPO 的 token-level averaging loss,但把标准的 维度扩展成 ;同时按照 DrGRPO 的处理,不再用 reward 的标准差做归一化。
信用分配直觉
直觉上,这一步解决的是“中间轮次看不到最终答案”的信用分配问题。若只把 final turn 当监督,模型会学成只在最后一轮用 memory;而把 reward 传播到所有 conversation 后,模型才会把“怎么写 memory”当成策略的一部分来优化。
conversation independence 是实现约束
Figure 3 还有一层容易忽略的含义:conversation independence 是实现约束,不是抽象说法。普通多轮对话会把前面所有 user/assistant/tool token 显式留在上下文里;MemAgent 则只把“上一轮生成出的 memory 文本”作为跨轮共享状态,而不是把完整历史 transcript 一起带过去。因此,如果训练时沿用普通 multi-turn RL 的 mask,模型其实会在优化一个推理时根本不存在的条件分布。
论文提出的 multi-conv 视角,本质上是在样本级别建立两套索引:一套是“这一轮 conversation 属于哪个原始样本”,另一套是“这一轮是否是 final turn、应该被 verifier 检查”。released code 中 final_mask_list 和 sample_index_list 正是为此存在。它们把 rollout 期间生成的多轮输出重新映射回样本级 reward,使得系统能在保持 conversation 张量化并行的同时,仍然做 sample-level credit assignment。
released code 中的状态机
这也解释了为什么作者要专门实现 generation_manager.py / async_generation_manager.py。单轮 RL 里,tokenization、padding、attention mask 通常一次构建完毕;multi-conv memory agent 却需要在每轮结束后剔除已完成样本、保留 active 样本、回写 memory、再生成下一轮 prompt。它更像一个小型状态机,而不是一次性前向。论文主打算法贡献,但 released code 表明,真正让 multi-conv RL 可运行的关键之一,是这套 rollout orchestration。从学习信号的角度看,这种设计尤其适合处理长文档里的稀疏延迟监督。
memory writing policy 的训练信号
某轮若错误地把实体 A/B 混淆,当下并不会立刻报错,但最终答错后,负 reward 会回传到所有相关 memory turns;相反,若某轮成功把关键别名、时间线或关系链压进 memory,即使那轮没有显式标签,也能通过最终正确答案获得正反馈。这正是 MemAgent 试图把“记忆质量”内化进策略优化的方式。还有一点很重要:multi-conv 让 memory-writing policy 可以和 answer policy 共享同一个基础模型参数,而不需要拆成“一个模型写摘要、另一个模型读摘要”。这会让策略学习更自洽,因为写 memory 的模型知道未来也是自己来读这份 memory 并作答;从博弈角度看,writer 和 reader 没有分布错配。
很多 pipeline 式系统的问题恰恰在于,摘要器保留下来的格式或细节不一定最利于下游 reader 使用,而 MemAgent 通过 end-to-end RL 把这两者绑成了同一主体。
这也解释了论文为何强调 agent workflow 而不是单纯 summarization module。若把 memory 写作外包给独立摘要器,再把摘要交给 QA 模型,两个模块之间的接口质量很难只靠最终 reward 调整;而 MemAgent 把“写给未来的自己看”作为内在训练目标,memory 自然更倾向保留模型自己在最终推理阶段最需要的线索、命名方式和关系结构。
3.4 实现层:固定窗口 memory token
文本化 memory protocol
论文的实现把 memory 直接塞进普通 chat template 里:每轮 prompt 由 problem + previous memory + current chunk 组成,输出是 Updated memory;最终轮则只保留 problem + previous memory 并要求输出 \boxed{} 答案。这样做的好处是,模型无需额外 memory module,训练/推理都还是标准文本生成。
MemoryAgent 如何固定输入预算
从代码看,recurrent/impls/memory.py 里的 MemoryAgent 会把 chunk_size、max_memorization_length、max_final_response_length 组合成固定的 max_input_length,并按 context_length 逐步推进;recurrent/generation_manager.py 负责循环调用 agent、补齐 tensor、生成、再回写 memory。
这个设计和论文描述一致:每一轮只处理当前 chunk,active sample 会随着 context_length > step * chunk_size 的条件不断减少,直到进入 final turn。released code 中的模板比论文摘要式描述更具体。更新模板 TEMPLATE 明确要求模型读取 problem、当前 section of an article 和 previous memory,并“retain all relevant details from the previous memory while adding any new, useful information”。
这说明 memory update 的目标不是从零重写,也不是 append-only 累积,而是在有限长度预算内做选择性保留与增量整合。最终模板 TEMPLATE_FINAL_BOXED 则拿掉当前 chunk,只保留 problem + previous memory,并强制答案放进 \\boxed{}。首轮没有历史记忆时,代码不会把 memory 槽位留空,而是注入预先 tokenized 的 "No previous memory"。这个细节很有价值,因为它让每一轮 prompt 结构保持完全一致,模型总是在同一个位置读取 memory,从而降低 prompt 形态变化对 RL 训练稳定性的影响。
训练脚本对应的规模
对于 recurrent agent 来说,这种结构稳定性往往比单轮 prompt 写得多漂亮更重要。
数据切块逻辑也非常贴近论文的“流式阅读”设定。MemoryDataset 会把 data_config.max_prompt_length 直接设成 max_chunks * chunk_size,再通过 center truncation 处理原始 context,返回 context_ids 和 context_length。随后 rollout 阶段用 context_ids[:, step * chunk_size : (step + 1) * chunk_size] 切出当前 chunk。
也就是说,chunk 不是预先离线拆成多条训练样本,而是在同一条原始文档上按轮次在线切片,这使每轮 memory update 都是同一文档阅读过程的一部分。7B 训练脚本还能帮助把这个实现量级具体化。run_memory_7B.sh 里给出的 chunk_size=5000、rollout.n=16、data.train_batch_size=128、trainer.total_epochs=30 说明作者不是在句子级或段落级玩具任务上验证想法,而是在较大单轮阅读预算下训练模型学会把 5K token 级别的局部证据压缩进更短 memory。方法真正困难的地方因此不是“生成一段摘要”,而是“在大块证据流中筛出对最终答案最关键的信息”。
固定 memory 长度的硬约束
系统层还有一个与论文主张完全对齐的硬约束:gen_max_tokens_memorization 直接等于 max_memorization_length。换言之,模型无法通过无限拉长 memory 来规避压缩任务;一旦记忆预算固定,唯一出路就是学会更高效地表达和维护证据。这正是 MemAgent 与“只是把更多文本搬进下一轮上下文”之间的根本边界。这也意味着 memory 的好坏不能只用“信息量大不大”来判断,还要看它是否具备跨轮可维护性。若一轮 memory 写得过长、过散或过于贴近原文措辞,下一轮模型继续在其上编辑时就更容易发生覆盖、重复或冲突。
MemAgent 通过固定预算和 recurrent editing,实际上在训练模型形成一种更稳定的内部记号系统:既要保留事实,又要便于后续轮次继续增量修改。
把 hidden state 外显成文本
从这个角度看,MemAgent 的固定 memory protocol 有点像把传统 sequence model 里的 hidden state 外显化成文本,并让这个文本同时接受语言模型先验、RL 目标和长度约束。它不是最紧凑的状态表示,但它是最兼容、最可审计、也最容易与现有 LLM 基础设施对接的表示方式。论文的工程价值,恰恰来自这种“不是最花哨,但非常可落地”的设计取舍。还有一点值得强调:由于每轮都把 memory 放回 prompt,MemAgent 实际训练的是“如何写出未来自己最容易继续编辑和最终读取的记忆”。这与面向人类读者的摘要不同,后者追求可读性和覆盖面;MemAgent 的 memory 更像一种任务内部语言,目标是最大化后续推理可用性。
也正因此,最终学到的 memory 形式未必总是最自然的摘要,但它可能更接近对模型自身最有效的压缩表示。这也给论文中的“兼容标准 LLM”表述补了一层含义。兼容不仅是指不用改 attention 和 tokenizer,也指整个方法把最关键的状态操作都降解成了标准文本输入输出:读 chunk、写 memory、读 memory、答题。于是训练、调试、可视化、错误分析都能沿用成熟的 LLM 工具链完成,而不需要为一个专用 memory tensor 额外搭建观测与诊断体系。如果把 MemAgent 和许多“长上下文 benchmark 上涨点不错、但很难接入真实系统”的方法对比,它的优势就在这里:它把能力提升建立在一个工程上非常朴素的接口之上。
代价是 memory 容量受文本长度严格限制,表达效率未必最优;收益是方法几乎可以立即迁移到任何支持 chat completion 的基础设施中。这种 trade-off 很符合论文所针对的实际落地场景。
rollout 深度由文档长度决定
从 rollout 角度再看,active_mask = context_length > step * chunk_size 这个条件也非常关键。它意味着系统不是机械地让所有样本跑满 max_chunks 轮,而是根据每个样本实际剩余上下文长度动态决定是否继续生成 memory。换句话说,MemAgent 的 recurrent 深度由文档长度决定,而不是由一个固定回合数决定;同一 batch 内部,不同样本会在不同轮次退出,并最终汇入 final turn。这种“异步结束、同步并行”的执行模式正是长文档 agent 落地时最现实的调度方式。这也解释了为什么代码里要保存 sample_index_list。
由于一轮轮筛掉已完成样本后,当前张量中的第 个位置已经不再稳定对应原始 batch 的第 个样本;若不额外保存映射,就无法在 reward 回传时把多轮 conversations 重新归并到正确的样本级轨迹。论文在算法层面强调的是 sample-level optimization,代码在工程层面实现的正是这层“轨迹身份保持”。另一个值得注意的实现选择是,memory 既是下一轮输入的一部分,也是策略自己生成的文本输出。这让 MemAgent 与很多“内部状态不可见”的 recurrent 模型不同:它的状态空间是人类可读、可检查、可审计的。对论文读者而言,这点很重要,因为它意味着方法不仅能在指标上工作,还天然支持 failure analysis。
若最终答错,可以直接检查哪一轮 memory 漏掉了关键实体、哪一轮把旧信息覆盖坏了,而不是只能从隐藏态中猜测出错位置。
这种显式文本状态还有一个副作用:memory token 本身会消耗上下文预算,因此模型必须在“可读性”“完整性”和“压缩率”之间找到平衡。论文虽然没有把这点展开成单独理论,但方法实际上在逼迫 LLM 学习一种 task-specific notation:既要保留足够线索,又要足够短,以便后续轮次还能继续携带它。对 agent memory 研究来说,这是一个比单纯摘要更接近真实系统约束的训练环境。 若把 3.4 和 3.3 连起来看,会发现 MemAgent 的实现哲学其实非常统一:算法层通过 multi-conv RL 解决 delayed credit assignment,协议层通过固定模板与固定预算保证输入输出形式稳定,系统层通过 generation manager 维护跨轮状态映射。三者缺一不可。只有算法没有稳定协议,memory 写作会漂;只有协议没有 sample-level reward 传播,中间 memory turns 学不到东西;只有两者没有状态调度实现,则整条 recurrent workflow 无法在真实 batch 上高效运行。
3.5 论文公式与 released code 实现差异
论文把训练动机描述为 multi-conv RL + DAPO 风格的 token 平均损失;released code 里实际训练入口是 run_memory_7B.sh / run_memory_14B.sh,调用 verl.trainer.main_ppo,并显式设置 algorithm.adv_estimator=grpo、algorithm.grpo_use_adv=False、actor_rollout_ref.actor.use_kl_loss=True、actor_rollout_ref.actor.kl_loss_coef=0.001。
因此,笔记里把方法理解为“以 GRPO 为底座、以 recurrent memory workflow 为外壳”的实现更准确。另一个值得强调的差异是,论文里 memory 被抽象成状态 ,而 released code 中它其实就是上一轮 assistant 输出出的 token 序列。也就是说,MemAgent 的状态传递不是隐藏态复用,也不是额外 memory module 的内部缓存,而是显式文本状态复用。这个实现选择让方法更朴素,但也更透明:任何 memory update 的成功或失败都直接体现在可读文本上。
论文公式与 released code 实现差异:论文在介绍标准 GRPO 时先给出基于组内均值和标准差的 advantage 归一化,但在 Figure 3 的 multi-conv 目标说明以及当前 released code 体现的训练设定中,作者更接近 DrGRPO 风格,强调不再使用 reward 标准差归一化,并把 final conversation 的奖励传播给同样本的全部 conversations。理解这一点很重要,因为它决定了 MemAgent 优化的核心不是单轮 response 的相对排序,而是整条 memory-writing trajectory 是否最终支持正确答案。
这类“论文先给标准公式、代码再做针对性改造”的现象在 RL for LLM 论文里很常见,但在 MemAgent 这里尤其需要记下来,因为方法真正的新意不在某个单独 loss 项,而在 sample-level trajectory 的重组方式。若只按标准 GRPO 理解,容易低估 multi-conv bookkeeping 的必要性;若把代码实现一起看,就会发现作者的主要贡献其实是把 RLVR 稳定地接到了 recurrent memory workflow 上。
3.6 Pseudocode(基于 released code)
Memory 数据与切 chunk
def load_memory_example(row, tokenizer, cfg):
chat = row.pop("prompt")
context = row.pop(cfg.context_key)
context_ids = tokenizer(context, return_tensors="pt", add_special_tokens=False)["input_ids"]
context_ids = postprocess_data(context_ids, max_length=cfg.max_prompt_length, truncation="center")
return {
"context_ids": context_ids[0],
"context_length": int((context_ids != tokenizer.pad_token_id).sum()),
"prompt_ids": tokenizer.encode(chat[0]["content"], add_special_tokens=False),
}逐轮 memory 更新
def memory_step(agent, batch, step):
active = batch["context_length"] > step * agent.config.chunk_size
if active.sum() == 0:
return agent.final_turn(batch)
prompt = batch["prompt_ids"][active]
chunk = batch["context_ids"][active, step * agent.config.chunk_size:(step + 1) * agent.config.chunk_size]
memory = agent.memory[active]
messages = [agent.format_update(p, m, c) for p, m, c in zip(prompt, memory, chunk)]
return messages, {"max_tokens": agent.config.gen_max_tokens_memorization}最终答案生成
def final_turn(agent, batch):
messages = [agent.format_final(p, m) for p, m in zip(batch["prompt_ids"], agent.memory)]
return messages, {"max_tokens": agent.config.gen_max_tokens_final_response}Multi-conv GRPO 训练
def grpo_multiconv_loss(policy, ref_policy, conversations, rewards, epsilon, beta):
adv = rewards - rewards.mean(dim=0)
adv = adv / (rewards.std(dim=0) + 1e-8)
loss = 0.0
for conv in conversations:
ratio = policy.logprob(conv) - policy.old_logprob(conv)
ratio = ratio.exp()
clipped = torch.clamp(ratio, 1 - epsilon, 1 + epsilon)
loss += -(torch.min(ratio * adv, clipped * adv).mean())
kl = kl_div(policy, ref_policy)
return loss / len(conversations) + beta * klCode to paper mapping
Code reference:
main@ef4219b2(2026-05-12) — pseudocode and mapping based on this commit
| Paper concept | Source file | Key class / function |
|---|---|---|
| Fixed-length memory config | recurrent/impls/memory.py | MemoryConfig |
| Dataset → context chunking | recurrent/impls/memory.py | MemoryDataset.__getitem__ |
| Update / final templates | recurrent/impls/memory.py | TEMPLATE, TEMPLATE_FINAL_BOXED |
| Memory agent loop | recurrent/impls/memory.py | MemoryAgent.start, action |
| Multi-turn generation orchestration | recurrent/generation_manager.py | AsyncLLMGenerationManager.run_llm_loop |
| Async rollout integration | recurrent/async_generation_manager.py | AsyncLLMGenerationManager.run_llm_loop |
| Training launcher 7B | run_memory_7B.sh | python3 -m verl.trainer.main_ppo ... |
| Training launcher 14B | run_memory_14B.sh | python3 -m verl.trainer.main_ppo ... |
4. Experimental Setup(实验设置)
数据集
- 主任务:HotpotQA 派生的长上下文 QA,训练集/验证集由
taskutils/memory_data/processing.py、filter.py、filter2.py处理得到;README里明确提到hotpotqa_train.parquet和hotpotqa_dev.parquet。 - 主评测:RULER-HotpotQA,输入长度覆盖 7K 到 3.5M token。
- OOD 评测:RULER 子任务,包括 needle-in-a-haystack、variable tracking、frequent words extraction、SQuAD-based QA,论文中说明长度覆盖 8K 到 512K;SQuAD 由于原文长度限制只到 256K。
Baseline
QwenLong-L1-32BQwen2.5-Instruct-14B-1MQwen2.5-Instruct-7B-1MDS-Distill-Qwen-32BDS-Distill-Qwen-14BDS-Distill-Qwen-7B
指标
accuracy (%):主表所有数值均为准确率。RULER average:多子任务平均分,用于图 6 的总体比较。
训练配置
- 7B 脚本:
run_memory_7B.shMODEL_PATH=Qwen/Qwen2.5-7B-InstructMAXLEN=8192MAX_NEW_TOKEN=1024chunk_size=5000data.train_batch_size=128actor_rollout_ref.actor.optim.lr=1e-6actor_rollout_ref.actor.ppo_mini_batch_size=8actor_rollout_ref.rollout.n=16actor_rollout_ref.actor.kl_loss_coef=0.001actor_rollout_ref.rollout.gpu_memory_utilization=0.6trainer.total_epochs=30trainer.n_gpus_per_node=8trainer.nnodes=4
- 14B 脚本:
run_memory_14B.shMODEL_PATH=Qwen/Qwen2.5-14B-InstructMAXLEN=8192MAX_NEW_TOKEN=1024chunk_size=5000data.train_batch_size=256actor_rollout_ref.actor.optim.lr=1e-6actor_rollout_ref.actor.ppo_mini_batch_size=16actor_rollout_ref.rollout.n=16actor_rollout_ref.actor.kl_loss_coef=0.001actor_rollout_ref.rollout.gpu_memory_utilization=0.60trainer.total_epochs=30trainer.n_gpus_per_node=8trainer.nnodes=16
5. Experimental Results(实验结果)
主结果
Figure 1 解读:这张图是论文最核心的结果。横轴从 7K 拉到 3.5M token,纵轴是 RULER-HotpotQA 准确率。MemAgent-7B/14B 在长度继续增加时只出现很小波动,而基线模型要么在 112K 左右明显下降,要么在 896K 附近直接掉到 0,说明“固定窗口 memory + RL”确实带来了接近无损的长度外推。
主表中 RL-MemAgent 的准确率如下:
| Model | 7K | 14K | 28K | 56K | 112K | 224K | 448K | 896K | 1.75M | 3.5M |
|---|---|---|---|---|---|---|---|---|---|---|
| RL-MemAgent-14B | 83.59 | 82.03 | 84.38 | 80.47 | 76.56 | 81.25 | 75.00 | 77.34 | 76.56 | 78.12 |
| RL-MemAgent-7B | 82.03 | 79.69 | 78.91 | 77.34 | 79.69 | 72.66 | 74.22 | 76.56 | 75.78 | 71.09 |
消融
Figure 5 解读:消融说明 memory 结构只能解决“能装下更长上下文”的问题,真正决定性能是否稳定的是 RL。没有 RL 的 memory 模型虽然比纯 truncation 强,但随着长度增长仍然下降;加上 RL 后,曲线在 28K 到 896K 都保持稳定,说明策略学会了更有效地筛选/保留信息。
OOD 与类别平均
Figure 6 解读:平均结果显示 MemAgent 在多种 RULER 任务上都保持稳定优势,尤其在 8K 到 512K 区间,14B 的平均准确率超过 95%。这表明它不是只在 HotpotQA 模板上有效,而是学到了可泛化的“记忆策略”。
Figure FWE 解读:频繁词提取任务用来测试模型是否真的保留了关键信息,而不是只记住表面问答格式。MemAgent 在这类任务上依旧稳,说明它的 memory update 不是专门为某一种 QA 模式过拟合。
Figure VT 解读:variable tracking 检查的是长上下文中的状态追踪能力。该图说明 MemAgent 在需要跨多个位置维护变量状态时仍然表现稳定,支持它的 memory 机制确实能做“状态压缩”。
局限
论文明确指出的限制不多;从结果本身可以看出它依赖长文档训练数据和固定的 chunk / memory 预算,且当前评测仍然集中在 HotpotQA/RULER 风格任务上。换言之,论文证明了“超长上下文的可扩展记忆策略”,但并没有证明这种策略在所有开放域生成任务上都天然最优。
结论
结果支持一个清晰结论:MemAgent 不是单纯扩大上下文窗口,而是通过 RL 学会在固定记忆预算下进行信息选择与压缩,因此能在百万级 token 上保持接近无损的答案质量。