MemOPD:动作在线不等于状态在线,记忆对齐让F1最高提升416.2%

MemOPD: On-Policy Distillation through Memory State Alignment for Long-Horizon Agents

论文原文 ↗ 论文发布 解读发布 解读:AI前沿分享

MemOPD:动作在线不等于状态在线,记忆对齐让F1最高提升416.2% 论文图示

在多轮交互、信息检索和复杂工具调用的长程任务(Long-horizon tasks)中,大模型智能体(Agent)面临一个难以回避的物理瓶颈:随着交互轮数增加,上下文长度迅速膨胀,推理成本呈二次方上升,模型的有效注意力也会被海量无关历史稀释。为了让智能体能够长期运行,学术界和工业界逐渐达成共识——必须引入“紧凑记忆”(Compact Memory)机制,让模型在每一轮交互后自主压缩、重写并更新上下文,只把精炼后的关键记忆传递给下一步。

ArXiv URL:https://arxiv.org/abs/2608.07068v1

然而,如何教会模型在每一步里“保留什么、丢弃什么”,一直是个训练难题。主流做法通常依靠近端策略优化(PPO)等强化学习算法,但现实任务的最终结果反馈极其稀疏,走完十几轮检索后如果答案错误,模型根本无法判断究竟是第几步的记忆压缩出了差错。为了引入密集的监督信号,在线策略蒸馏(On-Policy Distillation, OPD)被寄予厚望:让更强的教师模型对学生模型自主探索生成的动作进行逐步打分指导。

来自哈尔滨工业大学、北京大学和浙江大学的研究团队在一篇新工作中指出了一个被长期忽视的基础漏洞:在引入紧凑记忆后,智能体的“动作在线”不再等于“状态在线”。当训练框架像处理常规对话一样把交互历史展平成单条长序列时,教师模型给学生动作打分所依赖的上下文,早已不是学生当初做决策时看到的真实上下文。这种因记忆重写诱发的状态失真,导致传统蒸馏信号严重失真甚至有害。为此,研究团队提出了 MemOPD(Memory-Aligned On-Policy Distillation)框架,在长程问答基准上使 3B 模型的 F1 值相比纯 PPO 基线最高跃升 416.2%,并在保证数值一致性的前提下带来了最高 1.63 倍的训练加速。

记忆重写背后的隐秘陷阱:动作在线,但状态脱节

要理解这一问题的根源,首先需要审视紧凑记忆的工作机制。在标准的无状态多轮对话中,第 $t$ 步的输入包含了过去所有步骤的拼接序列,上下文随时间单调递增。但在采用紧凑记忆的 Agent(如 MEM1 框架)中,模型在第 $t$ 步接收到的输入是系统前缀 $q$ 与当前可变记忆上下文 $x_t$ 的组合,即 $z_t = [q, x_t]$。模型在此基础上自回归采样生成响应 $a_t$,$a_t$ 中不仅包含了要执行的环境工具动作(如检索关键词),还包含用于下一步更新的记忆内容。随后,环境执行动作并返回观测 $o_t$,记忆更新函数 $U$ 则将旧上下文、模型输出和新观测融合,重写并压缩出下一步的输入:

\[x_{t+1} = U(x_t, a_t, o_t)\]

这种机制的典型执行流表现为:$[q] \rightarrow a_0$;随后历史被重写为 $[q, m_0, o_0] \rightarrow a_1$;再下一步重写为 $[q, m_1, o_1] \rightarrow a_2$。这里,前一步响应中被保留的信息 $m_0$ 换了位置、换了前后缀,被丢弃的信息则彻底从模型的视野中消失。

MemOPD 方法总览

常规的强化学习和蒸馏训练为了凑批次、打满显卡利用率,惯常的做法是将整条交互轨迹拉平(Flatten)为一个连续的持久序列 $y = [q, x_0, a_0, x_1, a_1, \dots]$。在无记忆重写的普通轨迹中,这种做法没有问题;但在紧凑记忆场景下,灾难发生了。

当整个轨迹被平铺后,第 $t$ 步生成的动作 $a_t$,在重构后的序列中可能处于完全不同的 Token 绝对位置,其前面可能会出现其在生成时不曾见过的“未来上下文”,或者那些在记忆重写阶段早就被删掉的过期观测。更隐蔽的是,同一个响应 $a_t$ 的部分片段可能在后面的 $x_{t+1}$ 中以记忆形式再次出现。

这就产生了一个概念分离:动作的采样来源(Provenance)确实是在线的(On-policy),因为这串 Token 确由当前学生策略采样而出;但它被训练框架重构出来供模型评分的状态(State),却从未在采样过程中真实存在过。

作者团队在 3B 参数模型上的实测审计揭示了这种脱节的破坏力:直接平铺历史进行打分,导致模型对同一动作对数概率计算的 p99 误差高达 1.774,改变了 651 个采样动作位置上的顶层预测结果,并且因为概率比值严重漂移,无故触发了 13.29% 动作位置的 PPO 截断(Clipping)。如果此时直接引入教师模型进行反向 KL 散度蒸馏,教师模型本质上是在强行指导一个“学生根本没见过的幻觉状态”,不仅无法提供高质量的密集反馈,反而会破坏策略的收敛。

状态对齐:MemOPD 如何让时空倒流?

MemOPD 的核心设计哲学在于:如果要对采样动作做策略梯度更新和教师知识蒸馏,训练时呈现给前向计算的必须是该动作在生成时刻的精确物理状态。

研究团队将第 $t$ 次模型调用(Invocation)时预测第 $k$ 个动作 Token $a_{t,k}$ 所依赖的采样状态严格定义为四元组:

\[\sigma_{t,k} = (z_t, a_{t,<k}, \rho_t, V_t, \gamma_t(k))\]

其中 $z_t$ 是当时的输入 Token 序列,$a_{t,<k}$ 是当前轮次此前已自回归生成的动作前缀,$\rho_t$ 是当时每个 Token 的绝对位置编码,$V_t$ 是因果可见性矩阵(Casual Visibility / Attention Mask),而 $\gamma_t(k)$ 则定义了自回归预测的目标物理位置。记忆状态对齐(Memory State Alignment)的要求极其严苛:训练重构出的状态 $\widetilde{\sigma}{t,k}$ 必须在每一个动作 Token 上恒等于采样状态 $\sigma{t,k}$。

为了实现这一目标,MemOPD 采取了三步走重构策略:

第一,精确调用录制与物理隔离。在环境交互生成(Rollout)阶段,系统不仅记录文本,而是直接记录每一次独立调用的输入 Token ID 和采样的输出 Token ID。这避免了任何后续分词器重新解码再编码导致的切分边界偏移。

第二,解耦采样动作域(Sampled Action Domain)。在紧凑记忆机制下,前一步生成的响应内容可能在下一步作为记忆上下文重新输入给模型。此时,这些 Token 在当前步只属于“给定的观测环境”,而绝不能被视作“当前正在采样的决策动作”。如果训练框架不加区分,同一个决策就会在 Actor、Critic 和 GAE(广义优势估计)中被重复计费。MemOPD 显式构造了动作掩码 $m^{\mathrm{act}}_j$,只对属于真正采样发生的物理位置打上决策标签。

第三,紧凑调用打包(Packed Call Execution)与 RCE 验证。如果将每一次调用完全独立地拆开进行单次前向计算,显卡计算中会有大量公用的系统前缀 $q$ 被重复计算,显存吞吐也会暴跌。MemOPD 设计了专用的批次打包编译器,将同一交互轨迹内的多次调用在保持状态隔离的前提下拼接打包。为了防止注意力掩码或相对位置编码在打包过程中引发隐式泄漏,研究团队设立了“采样上下文等价性”(Rollout Context Equivalence, RCE)标准:

\[\left\|\boldsymbol{\ell}^{\mathrm{pack}}_{g(t,k)} - \boldsymbol{\ell}^{\mathrm{ind}}_{t,k}\right\|_{\infty} \le \epsilon_{\mathrm{num}}\]

该标准要求打包前向计算输出的全词表 Logits,与彻底独立隔离调用执行输出的 Logits 之间的无穷范数误差,必须严格控制在计算精度(如 BF16/FP16)的数值公差 $\epsilon_{\mathrm{num}}$ 范围内。只有通过 RCE 审计的打包方案,才被认定为既恢复了真实状态、又实现了高效批处理。

结合密集引导与全局目标:双驱动的优化目标

在确保状态完全对齐后,MemOPD 在训练目标上打出了一套组合拳,巧妙平衡了教师模型的局部引导与强化学习的任务级全局反馈。

对于每一个真正属于采样动作的位置 $j$(由掩码 $m^{\mathrm{T}}_j$ 指定),固定权重的更强教师模型与学生模型在完全一致的上下文、位置编码和可见性下同步计算。模型采用全词表反向 KL 散度进行在线策略蒸馏:

\[\mathcal{L}_{\mathrm{OPD}} = \frac{1}{\sum_j m^{\mathrm{T}}_j} \sum_j m^{\mathrm{T}}_j \sum_{v \in \mathcal{V}} p_{\theta,j}(v) \log \frac{p_{\theta,j}(v)}{p_{T,j}(v)}\]

这里的核心价值在于提供稠密的 Token 级监督。在多步检索中,学生即使最后没有得到最高任务奖励,只要其在某一具体步骤里提取实体或生成检索语句的动作符合更强教师模型的知识倾向,依然能得到及时的正向梯度拉动。

然而,仅有教师蒸馏是不够的。教师模型擅长评估单步动作的局部合理性,但面对高度动态的环境反馈和极其复杂的记忆更新链路,教师并不能预见这套记忆压缩策略在第十几轮后是否能成功回答最终问题。因此,MemOPD 依然保留了任务级 PPO 目标,以最终环境给出的任务完成奖励为核心,配合参考模型(Reference Policy)的 KL 惩罚与 GAE 优势计算:

\[\mathcal{L}_{\mathrm{actor}} = \mathcal{L}_{\mathrm{PPO}} - \alpha \mathcal{H} + \lambda \mathcal{L}_{\mathrm{OPD}}\]

其中 $\mathcal{H}$ 代表策略熵正则化,$\lambda$ 和 $\alpha$ 为权衡系数。PPO 拥有决定交互全局成败的“一票否决权”,一旦局部的教师偏好导致了全局任务的失败,PPO 的全局梯度便能及时介入并矫正这一偏差;反之,在全局奖励稀疏的广阔探索空间中,$\mathcal{L}_{\mathrm{OPD}}$ 则如明灯一般为学生指引微观动作的搜索方向。

实验评测:极限长任务上的惊人跃升

为了验证 MemOPD 的有效性,研究人员以 Qwen2.5-3B 为基础学生模型,7B 模型为固定教师,在极具挑战性的多目标长程检索基准上进行了五次独立随机种子的严格评测。任务基于 HotpotQA 和 Natural Questions 构造,模型需要在一个庞大的维基百科语料库中多轮检索证据。基准分为 Q2、Q8 和 Q16,分别代表单次交互需要同时解决 2 个、8 个乃至 16 个子问题的极限长程检索。

所有模型首先在由强模型生成的 20,036 条单轮优质格式轨迹上进行一轮监督微调(SFT),学会记忆、搜索和作答的标准语法,随后统一在 Q2 任务上进行强化训练,并在 Q8 和 Q16 上检验向极长交互时程的泛化迁移能力。

实验结果展现出了惊人的性能分化。在基础的 Q2 任务上,MemOPD-3B 相对纯 PPO 基线,Token F1 提升了 14.3%;随着交互轮数增加、任务难度上升,当扩展到极度考验长期记忆维持的 Q8 和 Q16 时,纯 PPO 训练的模型由于缺乏细粒度记忆更新指导,策略迅速崩溃,而 MemOPD-3B 展现出极强的韧性:

在通用开放域单目标检索任务 Wiki-RAG 的跨任务零样本迁移评测中,MemOPD 同样全面胜出。相较于 PPO,它在 Exact Match(EM)和 F1 上分别提升了 6.1% 和 7.4%,同时推理依赖度(Dependency)降低 31.6%,推理用时缩短 28.5%,证明了这种对齐记忆学习不仅没有过拟合特定协议,反而培养出了更加通用且高效的上下文筛选习惯。

状态对齐究竟有多重要?拆解控制实验

这篇论文最精彩的部分之一,在于作者团队设计了一个极其干净的控制对照实验,用来直接回答一个核心追问:性能的大幅提升,究竟有多少来自“教师蒸馏本身”,又有多少来自“状态对齐”?

在 Q2 任务的严格受控对比中,研究人员设置了三组实验:

  1. 纯 PPO 基线:使用重构的独立状态调用,但不引入教师指导;

  2. 持久历史蒸馏(Persistent-History OPD):加入完全相同的 7B 冻结教师模型,但教师在打分时直接读取展平拼接的整条交互历史(即传统做法);

  3. MemOPD:在相同教师和系数下,教师严格在状态对齐重构后的调用下打分。

对照数据给出了极具穿透力的结论:

这一消融结果无可辩驳地证明:过去直接将长程轨迹展平打分的研究,实际上白白浪费了一半以上的蒸馏收益;状态不对齐不仅仅是一个微小的工程妥协,而是会实打实截断模型性能上限的核心瓶颈。

进一步的误差归因实验揭示了背后的微观机理:如果单独人为破坏 Token 绝对位置编码,会让 260 个动作的顶层预测发生改变,引发 4.97% 的误剪枝;而如果进一步破坏可见性掩码(Visibility),误差则会呈数量级放大。另外,如果不能精确区分采样动作域,而错误地把整条轨迹中重复出现的记忆内容当成动作算入梯度,会导致多引入 66,830 到 74,357 个虚假决策位置,严重稀释训练信号。

最后,在工程效率方面,经过状态对齐编译打包后,MemOPD 在 9 种不同批处理配置下均实现了对无打包基准的显著加速,在维持 Logits 完全处于数值容差范围内的同时,Actor 的训练计算取得了最高 1.63 倍 的真实加速。

通向更鲁棒的长程智能体架构

长期以来,在 Agent 的强化学习与知识蒸馏训练中,工程实现往往遵循着无状态自回归模型的惯性思维:只要数据是模型自己采出来的,打包塞进 Transformer 就能照常反向传播。

MemOPD 从机理层面打破了这种盲目的工程惯性。它清晰地表明,紧凑记忆与记忆重写并不是免费的午餐——它打破了 Transformer 训练中最底层的“自回归前缀不变性”假定。当智能体开始学会遗忘和覆写历史时,训练系统必须具备在时空上“还原现场”的能力。

更具实用价值的是,MemOPD 提出的编译器接口并不与某一种具体的记忆算法(如 MEM1)深度绑定。因为重构器消费的是每次调用真实发生的物理 Token 上下文,而非抽象的语义标签,研究人员在全量响应保留、后缀保留、摘要替换、滑动窗口(Sliding Windows)以及检索动态刷新等多种常见更新策略下,均验证了 RCE 等价性的成功保持。

随着大型语言模型走向处理更长时程、更深层次的复杂任务,Agent 自主管理上下文必然成为标配。MemOPD 的工作为这一演进方向补齐了一块关键的训练基础设施:在动态可变的记忆拓扑中,唯有守住状态对齐的底线,密集的知识蒸馏与任务级的强化学习才能真正形成合力。