关于ZAKER 合作
量子位 26分钟前

突破 Agent 长程推理效率瓶颈!MIT& 新加坡国立联合推出强化学习新训练方法

AI Agent 正在被要求处理越来越多复杂的任务。

但当它要不停地查资料、跳页面、筛选信息时,显存狂飙、算力吃紧的问题就来了。

针对这一难题,MIT 和新加坡国立大学联合提出了MEM1 框架。实验结果显示,7B 参数 MEM1 模型的推理速度能达到传统 14B 模型的 3.5 倍,同时峰值上下文 token(peak token)数量约为后者 1/4。

该框架提出了一种基于强化学习的推理方案,通过让智能体学会将记忆更新作为上下文的一部分,从而实现近似常量级的显存开销(near constant memory)。

目前相关论文已被大语言模型会议 COLM 2025 Reasoning,Attention & Memory:RAM 2 workshop 收录为口头报告(Oral),会议将于今年 10 月份在加拿大蒙特利尔举办。

MEM1:让 AI 学会自主整理思绪 + 管理工作记忆

想象一位科研工作者连续工作一周后的大脑——充斥着各种公式、实验数据和临时灵感,但真正用于解决问题的只是其中核心片段。

对于人来说,每过一段时间起身喝杯咖啡醒醒脑,可以整理思绪,但是对于 AI agent 来说,这就难了。

经典的大语言模型采用全上下文提示技术,每轮交互都完整附上所有历史记录。随着对话轮次增加,计算成本和内存需求都呈线性增长(O ( N ) )

无限线性增长的上下文导致三个严重问题 :

无限增长的计算代价;

超出训练长度的泛化失效;

因信息过载导致的推理能力退化,即使相关信息仍在上下文中,模型也无法有效利用它们。

MEM1 的核心创新在于通过训练的方式让模型自主将记忆与推理统一

不同于简单地添加外部记忆模块(如 RAG 类的方法),MEM1 通过端到端强化学习训练代理,使其在每一步自动完成三个关键操作:

1、提取——从新观察中识别关键信息;

2、整合——将新信息与内部记忆状态融合;

3、修剪——丢弃冗余或无关内容。

通过这种训练方法,AI Agent 仅需维护自己上下文中的一个 <IS>内部状态(Internal State),这个 <IS> 包含之前上下文中所有模型自己认为需要保留的重要信息。

通过引入 <IS>,模型的上下文不会随交互轮次增加而膨胀

不仅如此,<IS> 的引入使得模型的推理过程与记忆整合可以有机地结合在一起,它不仅提供对当前查询的深度洞察,还充当 " 工作记忆 ",从收集信息中提取关键组件构建下一步推理。

这种过程也十分符合人们自己整理思绪的过程:杂乱陈旧的记忆被遗忘,同时随着记忆被整合,下一步的工作重点也随之显现出来。

MEM1 采用端到端的基于结果奖励(Outcome Reward)的强化学习训练方式,并引入一种特殊的注意力掩码机制(Attention Mask)。

该机制限定 Agent 在生成下一步输出时只能关注到上一轮交互的内容,从而迫使其学会对历史信息进行高效压缩,同时提升推理与问题求解能力。

相对于传统的 RL 训练,MEM1 在 rollout 期间引入了多个 trajectories 来训练 compressed reasoning。为了提升训练效率,MEM1 使用 attention masking 技巧将面向同一个任务的多个 trajectories 压缩成一个进行高效训练。

在推理阶段,MEM1 agent 会在推理时由 agent 自主地不断整合自己的 context。整合完之后,之前的 memory 会自动从 context 中移出,从而达到在长程环境交互任务下控制 context 长度的目的。

上图为 MEM1(下)和现有推理模型(上)在长程任务上上下文占用的对比,可以看出现有推理模型上下文会随着轮次增加线性增长,但 MEM1 每一轮后将被压缩过的信息移出 context 使得上下文保持近似不变。

推理与信息整合的双重能力

经典的 RL 环境通常是为短交互设计,并不能很好的鼓励模型进行长期多轮的推理。大家比较熟知的多轮对话数据集,例如 hotpot QA 等通常也只需要模型进行两轮推理。

为了验证 MEM1 训练方法的效果,MEM1 团队基于现有的数据集构造了一个高难度多目标问答任务的训练环境。团队基于以下不同领域的数据集构建了训练环境并进行训练:

1、Doc 检索 QA:模拟研究代理查阅内部文档;

2、开放域 Web QA:真实网络环境信息获取;

3、多轮网购:WebShop 平台的复杂决策链。

团队在 2 目标任务上训练 MEM1 agent,然后在复杂的16 目标任务上测试。

MEM1 展现出了超越训练范围的强大泛化能力,在 16 目标任务上,MEM1 准确率,上下文长度,推理速度三个维度上全方面超过比他更大的模型以及带有外部记忆模块的模型。

随后,研究团队对 MEM1 agent 的推理行为进行了定性分析,发现 MEM1 在处理多轮搜索推理任务时成功展现出了推理与信息整合的双重能力

记忆方面,MEM1 展现出了分问题独立存储以及信息更新能力。在推理方面,MEM1 能够掌握自适应的搜索策略,例如拆解子问题,调整关键词查询,以及任务规划等等。

MEM1 为处理 AI Agent 长推理上下文这一重大挑战提供了一个非常有意思的思路。

当下工业界处理上下文仍是以接入外部记忆模块作为主流方法。但是这种做法通常需要很繁琐的工程,而且效果难以掌控。

随着 AI Agent 端到端强化学习范式的兴起,智能体记忆或许可以通过 RL 的方式让模型自己来处理。正如 MEM1 团队提到的:智能不是让 AI 记住一切,而是教会它自己决定应该如何记忆。

论文地址:https://arxiv.org/abs/2506.15841

代码地址:https://github.com/MIT-MI/MEM1

开源模型:https://huggingface.co/Mem-Lab/Qwen2.5-7B-RL-RAG-Q2-EM-Release

一键三连「点赞」「转发」「小心心」

欢迎在评论区留下你的想法!

点亮星标

科技前沿进展每日见