突破Agent長程推理效率瓶頸！MIT&NUS聯(lián)合推出強(qiáng)化學(xué)習(xí)新訓(xùn)練方法

時(shí)令 2025-08-22 16:35:00 來源：量子位

MEM1團(tuán)隊(duì) 投稿

量子位 | 公眾號 QbitAI

AI Agent正在被要求處理越來越多復(fù)雜的任務(wù)。

但當(dāng)它要不停地查資料、跳頁面、篩選信息時(shí)，顯存狂飆、算力吃緊的問題就來了。

針對這一難題，MIT和新加坡國立大學(xué)聯(lián)合提出了MEM1框架。實(shí)驗(yàn)結(jié)果顯示，7B參數(shù)MEM1模型的推理速度能達(dá)到傳統(tǒng)14B模型的3.5倍，同時(shí)峰值上下文token（peak token）數(shù)量約為后者1/4。

該框架提出了一種基于強(qiáng)化學(xué)習(xí)的推理方案，通過讓智能體學(xué)會將記憶更新作為上下文的一部分，從而實(shí)現(xiàn)近似常量級的顯存開銷（near constant memory）。

目前相關(guān)論文已被大語言模型會議COLM 2025 Reasoning，Attention & Memory：RAM 2 workshop收錄為口頭報(bào)告（Oral），會議將于今年10月份在加拿大蒙特利爾舉辦。

MEM1：讓AI學(xué)會自主整理思緒+管理工作記憶

想象一位科研工作者連續(xù)工作一周后的大腦——充斥著各種公式、實(shí)驗(yàn)數(shù)據(jù)和臨時(shí)靈感，但真正用于解決問題的只是其中核心片段。

對于人來說，每過一段時(shí)間起身喝杯咖啡醒醒腦，可以整理思緒，但是對于AI agent來說，這就難了。

經(jīng)典的大語言模型采用全上下文提示技術(shù)，每輪交互都完整附上所有歷史記錄。隨著對話輪次增加，計(jì)算成本和內(nèi)存需求都呈線性增長（O(N)）。

無限線性增長的上下文導(dǎo)致三個(gè)嚴(yán)重問題:

• 無限增長的計(jì)算代價(jià)；
• 超出訓(xùn)練長度的泛化失效；
• 因信息過載導(dǎo)致的推理能力退化，即使相關(guān)信息仍在上下文中，模型也無法有效利用它們。

MEM1的核心創(chuàng)新在于通過訓(xùn)練的方式讓模型自主將記憶與推理統(tǒng)一。

不同于簡單地添加外部記憶模塊（如RAG類的方法），MEM1通過端到端強(qiáng)化學(xué)習(xí)訓(xùn)練代理，使其在每一步自動(dòng)完成三個(gè)關(guān)鍵操作：

1、提取——從新觀察中識別關(guān)鍵信息；
2、整合——將新信息與內(nèi)部記憶狀態(tài)融合；
3、修剪——丟棄冗余或無關(guān)內(nèi)容。

通過這種訓(xùn)練方法，AI Agent僅需維護(hù)自己上下文中的一個(gè)<IS>內(nèi)部狀態(tài)（Internal State），這個(gè)<IS>包含之前上下文中所有模型自己認(rèn)為需要保留的重要信息。

通過引入<IS>，模型的上下文不會隨交互輪次增加而膨脹。

不僅如此，<IS>的引入使得模型的推理過程與記憶整合可以有機(jī)地結(jié)合在一起，它不僅提供對當(dāng)前查詢的深度洞察，還充當(dāng)“工作記憶”，從收集信息中提取關(guān)鍵組件構(gòu)建下一步推理。

這種過程也十分符合人們自己整理思緒的過程：雜亂陳舊的記憶被遺忘，同時(shí)隨著記憶被整合，下一步的工作重點(diǎn)也隨之顯現(xiàn)出來。

MEM1采用端到端的基于結(jié)果獎(jiǎng)勵(lì)（Outcome Reward）的強(qiáng)化學(xué)習(xí)訓(xùn)練方式，并引入一種特殊的注意力掩碼機(jī)制（Attention Mask）。

該機(jī)制限定Agent在生成下一步輸出時(shí)只能關(guān)注到上一輪交互的內(nèi)容，從而迫使其學(xué)會對歷史信息進(jìn)行高效壓縮，同時(shí)提升推理與問題求解能力。

相對于傳統(tǒng)的RL訓(xùn)練，MEM1在rollout期間引入了多個(gè)trajectories來訓(xùn)練compressed reasoning。為了提升訓(xùn)練效率，MEM1使用attention masking技巧將面向同一個(gè)任務(wù)的多個(gè)trajectories壓縮成一個(gè)進(jìn)行高效訓(xùn)練。

在推理階段，MEM1 agent會在推理時(shí)由agent自主地不斷整合自己的context。整合完之后，之前的memory會自動(dòng)從context中移出，從而達(dá)到在長程環(huán)境交互任務(wù)下控制context長度的目的。

上圖為MEM1（下）和現(xiàn)有推理模型（上）在長程任務(wù)上上下文占用的對比，可以看出現(xiàn)有推理模型上下文會隨著輪次增加線性增長，但MEM1每一輪后將被壓縮過的信息移出context使得上下文保持近似不變。

推理與信息整合的雙重能力

經(jīng)典的RL環(huán)境通常是為短交互設(shè)計(jì)，并不能很好的鼓勵(lì)模型進(jìn)行長期多輪的推理。大家比較熟知的多輪對話數(shù)據(jù)集，例如hotpot QA等通常也只需要模型進(jìn)行兩輪推理。

為了驗(yàn)證MEM1訓(xùn)練方法的效果，MEM1團(tuán)隊(duì)基于現(xiàn)有的數(shù)據(jù)集構(gòu)造了一個(gè)高難度多目標(biāo)問答任務(wù)的訓(xùn)練環(huán)境。團(tuán)隊(duì)基于以下不同領(lǐng)域的數(shù)據(jù)集構(gòu)建了訓(xùn)練環(huán)境并進(jìn)行訓(xùn)練：

1、Doc檢索QA：模擬研究代理查閱內(nèi)部文檔；
2、開放域Web QA：真實(shí)網(wǎng)絡(luò)環(huán)境信息獲??；
3、多輪網(wǎng)購：WebShop平臺的復(fù)雜決策鏈。

團(tuán)隊(duì)在2目標(biāo)任務(wù)上訓(xùn)練MEM1 agent，然后在復(fù)雜的16目標(biāo)任務(wù)上測試。

MEM1展現(xiàn)出了超越訓(xùn)練范圍的強(qiáng)大泛化能力，在16目標(biāo)任務(wù)上，MEM1準(zhǔn)確率，上下文長度，推理速度三個(gè)維度上全方面超過比他更大的模型以及帶有外部記憶模塊的模型。

隨后，研究團(tuán)隊(duì)對MEM1 agent的推理行為進(jìn)行了定性分析，發(fā)現(xiàn)MEM1在處理多輪搜索推理任務(wù)時(shí)成功展現(xiàn)出了推理與信息整合的雙重能力。

記憶方面，MEM1展現(xiàn)出了分問題獨(dú)立存儲以及信息更新能力。在推理方面，MEM1能夠掌握自適應(yīng)的搜索策略，例如拆解子問題，調(diào)整關(guān)鍵詞查詢，以及任務(wù)規(guī)劃等等。

MEM1為處理AI Agent長推理上下文這一重大挑戰(zhàn)提供了一個(gè)非常有意思的思路。

當(dāng)下工業(yè)界處理上下文仍是以接入外部記憶模塊作為主流方法。但是這種做法通常需要很繁瑣的工程，而且效果難以掌控。

隨著AI Agent端到端強(qiáng)化學(xué)習(xí)范式的興起，智能體記憶或許可以通過RL的方式讓模型自己來處理。正如MEM1團(tuán)隊(duì)提到的：智能不是讓AI記住一切，而是教會它自己決定應(yīng)該如何記憶。

論文地址：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