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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

1、Doc檢索QA：模擬研究代理查閱內(nèi)部文檔；

2、開放域Web QA：真實網(wǎng)絡(luò)環(huán)境信息獲??；

3、多輪網(wǎng)購：WebShop平臺的復(fù)雜決策鏈。

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

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

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

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

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

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

隨著AI Agent端到端強化學(xué)習(xí)范式的興起，智能體記憶或許可以通過RL的方式讓模型自己來處理。正如MEM1團隊提到的：智能不是讓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