Mem-π:通过学习何时生成和生成什么内容来实现自适应记忆 Mem-π: Adaptive Memory through Learning When and What to Generate
提出生成式记忆策略,让LLM智能体学会何时需要记忆辅助以及如何生成任务相关指导
前置知识
检索增强生成(RAG)
RAG是一种通过从外部知识库检索相关文档来增强大语言模型生成能力的技术。它首先将用户查询转换为向量表示,然后在向量数据库中找到最相似的文档片段,最后将这些检索到的内容作为上下文输入给语言模型,帮助模型生成更准确、更有依据的回答。这种方法特别适合需要访问最新信息或特定领域知识的场景,能够在不重新训练模型的情况下扩展模型的知识范围。
本文对比的基线方法RAG代表了现有的主流检索式记忆范式,理解它的工作原理和局限性是理解Mem-π创新价值的基础
情景记忆(Episodic Memory)
情景记忆来源于认知心理学,指对个人经历过的具体事件和情景的记忆,包括事件发生的时间、地点、情境等细节。在人工智能领域,情景记忆通常指智能体通过与环境交互积累的经验记录,每个记录包含任务上下文、采取的行动、观察到的状态和获得的结果。这些经验可以组织成记忆库,在未来的任务中通过相似性检索来复用成功的策略或避免重复错误。
本文对比的现有记忆系统大多基于情景记忆库进行检索,理解这个概念有助于理解Mem-π为什么要抛弃检索范式而转向生成范式
策略梯度强化学习
策略梯度是一类直接优化策略参数的强化学习算法,它不通过学习价值函数来间接改进策略,而是直接对策略参数进行梯度更新,最大化期望累积奖励。算法的核心是计算策略梯度,即奖励对策略参数的导数,指示如何调整参数以增加获得高奖励的概率。在离散动作空间中,通常使用REINFORCE算法或其变体,它们基于采样的轨迹估计梯度,通过增加高奖励轨迹对应的动作概率、降低低奖励轨迹对应的动作概率来优化策略。
本文第二阶段适应蒸馏使用的GRPO算法是策略梯度方法的一种变体,理解策略梯度有助于理解Mem-π如何通过下游任务奖励来优化记忆生成策略
Abstention机制
Abstention是指模型在面对不确定性或认为输出可能有害时,选择不输出或拒绝回答的机制。在分类任务中,可以添加一个额外的不确定类别;在生成任务中,可以设置特殊的退出token或判断条件。这种机制可以提高模型在遇到分布外输入、知识不足或容易产生幻觉的情况下的可靠性,避免给出错误或误导性的答案。
本文的创新点之一就是将abstention机制引入记忆生成,让模型学会何时需要生成记忆指导、何时应该让基础智能体独立工作,这是提高可靠性的关键设计
研究动机
现有的记忆增强智能体系统面临严重的检索对齐问题。这些系统通常从情景记忆库或技能库中基于相似度检索相关条目,但这些检索到的条目是静态的、固定的,经常与当前的查询上下文不对齐。具体表现为检索到的记忆可能包含不相关的内容、部分对齐或过度具体的信息,无法适应智能体当前的上下文需求。例如,论文中提到在WebArena任务中,当查询要求列出前3个搜索术语时,RAG检索到的源任务是前2个搜索术语,导致提示中说读取前两行,这与查询中的3不匹配,造成失败。此外,检索式记忆的另一个问题是记忆库会不断增长,带来记忆管理开销,包括合并新经验和防止旧经验被遗忘的挑战。
本文的目标是本文的目标是提出一个新的框架Mem-π,将智能体的记忆建模为一个生成式策略πmem,而不是对显式存储条目的检索。这个生成式策略应该能够根据当前的智能体上下文动态生成简洁的、任务自适应的指导,学会判断何时应该生成记忆指导、何时应该选择abstain让基础智能体独立工作,通过强化学习利用下游任务结果作为奖励信号来对齐记忆生成与任务成功,将可重用的经验编码到模型参数中,避免记忆库无限增长带来的管理开销。
与已有工作不同的是,本文的独特切入角度是从认知科学中的人类记忆构建性过程中获得启发。认知科学研究表明,人类的记忆不是字面回放机制,而是一种构建性过程,回忆是从先验知识和当前上下文中动态重建的。基于这个洞见,Mem-π抛弃了检索范式,转而将记忆建模为一个独立的、有自己参数的生成式策略。与现有工作相比,Mem-π的创新在于不是无条件生成辅助指导,而是联合学习何时生成和生成什么,引入决策内容解耦的强化学习目标,将路由决策和内容生成分开优化,通过两阶段蒸馏将离线经验转化为参数知识,再用RL对齐任务效用。这与ParaMem和SEAM等近期工作不同,它们要么没有条件化到当前上下文,要么将生成作为始终开启的辅助步骤。
核心方法
Mem-π的核心思想是将智能体的记忆建模为一个独立的生成式策略πmem,该策略由一个专门的语言或视觉语言模型实现,有自己的参数,与下游智能体分离。给定智能体上下文,πmem生成一个输出y等于d连接m,其中d是决策token,m是记忆指导内容。如果选择GENERATE,则将m前缀到下游智能体的输入中;如果选择ABSTAIN,则智能体使用原始上下文。训练分为两个阶段:经验蒸馏阶段通过监督学习将离线经验库压缩到πmem参数中,学习从上下文到指导的映射;适应蒸馏阶段通过强化学习使用下游任务结果作为奖励信号,进一步优化πmem,使其生成对任务成功真正有用的指导。为了可靠性,πmem可以学习在生成没有必要或不确定时选择abstain。
Mem-π的核心创新点是提出决策内容解耦的强化学习目标。标准的策略梯度算法将决策token和内容token放在同一个扁平目标下优化,但由于决策token只占输出序列的很少部分,内容层面的梯度会主导决策层面的学习。此外,第二阶段是从第一阶段的有监督策略初始化的,该策略倾向于生成,因此标准采样可能产生没有abstain分支的组,无法比较生成和abstain的相对价值。为了解决这两个问题,Mem-π设计了结构化counterfactual采样,确保每个采样组包含一个abstain分支和多个generate分支,然后优势分解,将学习信号分解为跨分支的决策优势和分支内的内容优势,最终通过pertoken credit assignment将不同的学习信号路由到相应的token位置。
方法步骤详情
Mem-π的训练过程分为两个阶段。第一阶段是经验蒸馏,给定离线经验库E包含上下文指导对,优化自回归监督目标,通过最大似然学习将显式的离线经验转化为参数知识。第二阶段是适应蒸馏,将模型词汇表扩展两个特殊决策tokenGENERATE和ABSTAIN,对称初始化它们的嵌入以确保两个决策有可比的初始概率。采用GRPO算法作为基础RL算法,对于每个上下文x,从πmem采样一组G个输出并计算组相对优势。奖励定义为R,包括下游任务奖励和长度正则化器。为了解决决策和内容的imbalance,使用结构化采样构造一个abstain分支和G减1个generate分支。然后计算决策优势,对于abstain分支为正Delta,对于generate分支为负Delta,其中Delta等于Vabs减去Vgen,Vabs和Vgen分别是abstain分支和generate分支的平均奖励。内容优势在generate分支内通过组归一化排名,对于abstain分支内容优势为零。最终构造pertoken优势,将决策优势路由到决策token,将Deltagated后的内容优势路由到内容token,实现决策内容解耦优化。
技术新颖性
Mem-π的技术新颖性体现在多个方面。首先,它提出了将记忆建模为生成式策略的新范式,与现有的检索式记忆有本质区别。这个范式有几个优势:记忆足迹由模型大小而非累积经验数量决定;可以融合许多过去经验的信号到单个上下文特定的提示中;分离专业化和执行,可以将较小的本地模型微调为πmem,插入更大的前沿智能体模型中。其次,决策内容解耦的强化学习目标是技术创新,它通过结构化counterfactual rollouts和pertoken credit assignment将路由决策和内容生成分开优化,这是首次在记忆生成场景中处理这种特殊的输出结构。第三,两阶段蒸馏的设计既有创新性又实用,第一阶段通过监督学习提供强初始化,第二阶段通过RL对齐任务效用,实验证明两个阶段都是必要的。最后,引入abstention机制并学习何时生成记忆是智能体记忆系统中的一个新颖设计,提高了可靠性。
实验结果
Mem-π在四个智能体基准测试中取得了最先进的性能。在WebArena上,Mem-π达到43.1%的成功率,相比基线智能体27.1%有16.0个百分点的绝对提升,相对提升接近50%。在各个子域都领先,最大的绝对提升在Reddit加23.8个百分点和CMS加28.2个百分点,这些领域的结构化导航模式从记忆的经验中受益最大。在WorkArena上,Mem-π将基线智能体从42.0%提高到50.3%,平均提升8.3个百分点,在Form类别上有强劲的14.9个百分点提升。在ALFWorld上,Mem-π达到91.6%,比已经很强的GPT5.4mini基线有6.3个百分点提升。体验蒸馏本身就已经匹配或超过基于RL的基线:Mem-π Stage 1在WebArena上达到35.0%,与MemoryR1的33.2%和MemRL的34.0%相当,不需要任何RL训练。这验证了离线参数知识作为强初始化策略的有效性。RL阶段在WebArena上提供了显著的额外增益,从Stage 1到完整模型有8.1个百分点的提升,最大的跳跃在CMS加25.4个百分点、Reddit加4.8个百分点和Maps加3.0个百分点。ALFWorld的增益更 modest,为1.6个百分点,这与前沿智能体的高基线留给改进空间较少一致。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| WebArena网络导航 | 任务成功率(SR %) | 43.1% | MemRL 34.0% | +9.1 pp (相对提升26.8%) |
| WorkArena企业软件导航 | 任务成功率(SR %) | 50.3% | 基线智能体 42.0% | +8.3 pp (相对提升19.8%) |
| ALFWorld具身环境 | 任务成功率(SR %) | 91.6% | GPT-5.4-mini 85.3% | +6.3 pp (相对提升7.4%) |
| LifelongAgentBench终端工具使用 | 任务成功率(SR %) | 36.7% (DB+OS平均) | 基线智能体 26.8% | +9.9 pp (相对提升37.0%) |
| WebArena跨智能体泛化(Qwen2.5-7B) | 任务成功率(SR %) | 46.1% | RAG 32.1% | +14.0 pp (3.4× RAG增益) |
局限与改进
作者在论文中承认了几个局限性。首先,Mem-π依赖于高质量的离线经验库,论文中使用JEFHinter从原始交互轨迹中提炼,但如果经验库质量差或覆盖不足,会影响性能。其次,在某些环境下下游任务奖励可能很稀疏,如多步骤交互中的延迟奖励,这会让第二阶段RL训练变得困难。论文提到这可以通过更细粒度的奖励设计或使用内在奖励来缓解。第三,虽然Mem-π展示了跨智能体泛化能力,但增益会随着基线智能体变强而缩小,这在论文中观察到。最后,当前的Mem-π是离线学习后固定的,不支持持续从新经验中更新参数记忆,这与人类记忆的持续学习能力不同。论文提到未来工作可以扩展到闭环记忆学习。我自己观察到的局限性包括:训练两阶段需要大量计算资源;视觉语言变体需要额外的视觉特征提取器;在某些基线智能体已经很强的任务如ALFWorld,提升空间有限;abstention率的最优阈值可能因任务而异,需要调参。
独立分析的弱点
Mem-π的第一个独立分析弱点是在奖励稀疏环境下的性能可能受限。在多步骤交互任务中,最终的二进制成功或失败信号提供的信息很稀疏,这使得第二阶段RL训练变得困难。改进方向包括:设计更细粒度的中间奖励,如基于子目标完成的奖励;使用内在奖励机制,基于新颖性或探索进展;采用hierarchical RL将长期目标分解为短期目标。第二个弱点是离线经验库的覆盖和偏见问题。Mem-π的性能依赖于经验库的质量和覆盖,如果经验库缺乏某些任务类型的样本或包含有偏见的内容,生成的指导也会有偏见。改进方向包括:主动学习策略,选择性地收集对当前策略最有价值的经验;多样性正则化,确保经验库覆盖任务分布的不同区域;在线持续学习,让πmem能够从新经验中更新参数。第三个弱点是计算开销。两阶段训练和RL采样需要大量计算,特别是在大型模型上。改进方向包括:参数高效微调,如LoRA或适配器;蒸馏到更小的模型;更高效的RL算法,如离线RL或基于模型的RL。
未来方向
作者提出的未来工作方向包括扩展到闭环记忆学习,让智能体持续收集新经验并更新参数记忆;探索可归因的参数记忆,使生成的指导可以追溯到支持它的经验,同时保持生成记忆的灵活性。基于论文成果可以延伸的方向包括:多智能体协作记忆,多个智能体共享或协商一个共同的生成式记忆;层次化记忆生成,不同粒度的指导由不同的策略生成;跨模态记忆生成,支持视觉、语言、代码等多模态指导的统一生成;元学习适应,快速适应新任务领域或新的下游智能体;可解释性研究,分析πmem学到了什么知识,如何在不同上下文下生成不同的指导;安全性和鲁棒性,确保生成的指导不会产生有害或误导性行为,特别是在对抗性输入或分布外输入下;实际应用部署,将Mem-π集成到实际的智能体系统中,如自动化研究助手、计算机使用智能体等。
复现评估
论文提供了相对详细的实验设置和实现细节。论文提到使用Qwen2.57BInstruct作为记忆模型πmem的骨干,在视觉设置下使用Qwen2.5VL7BInstruct。下游智能体包括Qwen2.57BInstruct按照WebAgentR1设置微调和GPT5.4mini。训练使用离线经验库,通过JEFHinter从原始交互轨迹中提炼。论文提供了架构图和训练流程的伪代码,以及主要的超参数设置。但是,论文没有提到代码或模型是否开源,这是一个复现的障碍。数据方面,WebArena、WorkArena、ALFWorld和LifelongAgentBench都是公开的基准测试数据集,但论文使用的具体划分和JEFHinter提炼的经验库可能不公开。算力方面,训练7B参数的模型需要多个GPU和大量计算时间,普通研究者可能难以复现。总的来说,论文提供了足够的技术细节让研究者理解方法,但完全复现实验需要算力和数据支持,开源代码和模型将大大降低复现难度。
论文图表
Figure 1对比了三种记忆系统的架构。左边的工作流式记忆系统使用预定义的工作流来控制记忆构造、检索和更新,从记忆库检索后提供给智能体。中间的学习式记忆系统通过学习优化记忆操作,与下游智能体结果联合优化,有可学习的控制器。右边的Mem-π将记忆建模为可学习的生成器,通过离线经验蒸馏和在线适应蒸馏将可重用经验内化到模型参数中,实现自适应生成。这张图清晰展示了Mem-π与现有方法的本质区别:不是检索固定的记忆条目,而是根据当前上下文生成指导。
这张图对于理解论文的核心创新点至关重要,它直观地对比了Mem-π与现有两种主要记忆范式的架构差异,帮助读者快速把握Mem-π的设计理念和新颖性