突破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
