Haozhen Zhang 現為 Nanyang Technological University（NTU）博士生，導師 Wenya Wang 為 NTU 計算機與數據科學學院助理教授。團隊研究聚焦 LLM Agent Memory、LLM Interpretability 與高效智能體系統，關注大語言模型在 reasoning、memory 與 multimodal grounding 等方面的能力構建，致力于發展更可解釋、高效且具備泛化能力的大模型系統。本工作同時與香港科技大學（廣州）Chengwei Qin 老師合作完成；Chengwei Qin 現為香港科技大學（廣州）人工智能學域助理教授，主要關注大語言模型智能體、多模態學習以及模型的高效學習。

BudgetMem：給 Runtime Agent Memory 裝上 “預算路由器”

讓記憶系統學會按需分配運行成本

當 LLM Agent 處理長期對話、多輪交互和復雜文檔時，Memory 已經成為不可或缺的核心模塊。它幫助智能體保存歷史、檢索信息、維持個性化上下文，并支撐跨時間的推理能力。

但一個現實問題常常被忽視：Agent Memory 到底應該花多少成本來處理歷史？

如果一個 query 只需要簡單事實，是否有必要調用昂貴的 LLM 進行復雜總結？如果一個 query 依賴跨時間、跨實體、跨主題的多條證據，低成本檢索和粗粒度摘要又是否足夠？

對此，研究團隊提出BudgetMem: Learning Query-Aware Budget-Tier Routing for Runtime Agent Memory，旨在讓 Agent Memory 從固定的記憶處理流水線，轉向 query-aware 的動態預算分配機制：對于簡單 query 使用低成本處理路徑，對于復雜 query 則自動調用更高質量的記憶模塊。

• 論文：《Learning Query-Aware Budget-Tier Routing for Runtime Agent Memory》
• 鏈接：https://arxiv.org/abs/2602.06025

背景：固定記憶流水線難以適應不同 Query

現有很多 Agent Memory System 仍然遵循一種固定范式：先離線構建記憶，再在未來統一檢索使用。

這種build once, use always的方式雖然直觀，但存在兩個問題。首先，它是 query-agnostic 的。系統在不知道未來問題的情況下提前壓縮歷史，可能會丟掉后續 query 真正需要的細節。其次，它缺乏顯式的 performance-cost control。簡單問題和復雜問題往往被同一套流程處理，前者可能浪費成本，后者又可能預算不足。

因此，本文關注一個更貼近真實部署的問題：當 query 到來時，記憶系統應該花多少計算去處理歷史？

Runtime Query-Aware Memory Extraction

BudgetMem 將記憶構建從固定離線流程轉向runtime query-aware extraction。系統首先將歷史保留為原始 chunks，當用戶 query 到來時，再檢索相關片段，并通過模塊化流水線構建 query-focused memory。

在本文實驗中，團隊采用了一個簡單且可解釋的模塊化實例：

Filtering → Entity / Temporal / Topic Extraction → Summarization

其中，Filtering 負責篩選相關 chunks；Entity、Temporal、Topic 模塊分別提取實體、時間和主題相關信息；Summary 模塊最終整合為面向當前 query 的記憶。

需要強調的是，這一模塊組合只是 BudgetMem 在實驗中的具體實例，并不是框架本身的限制。BudgetMem 的核心在于為模塊化 memory pipeline 提供統一的 budget-tier interface，因此也可以適配到其他 memory modules 或不同形式的記憶處理流水線中。

關鍵在于，每個 memory module 都提供LOW / MID / HIGH三種 budget tiers。也就是說，同一個模塊既可以用低成本方式執行，也可以用更高質量但更昂貴的方式執行。這樣，記憶系統不再固定調用同一套流程，而是根據當前 query 的復雜度和信息需求，動態選擇合適的記憶處理路徑。

三種 Budget-Tier 策略

為了系統研究 runtime memory 中的性能成本權衡，BudgetMem 比較了三種 budget-tier realization strategies。

Implementation Tiering：改變模塊實現方式，從規則 / 啟發式方法，到輕量模型，再到 LLM-based 模塊。

Reasoning Tiering：改變推理行為，從 direct extraction，到 CoT-style reasoning，再到 multi-step 或 reflection-style processing。

Capacity Tiering：改變模型大小，用不同規模的模型實現同一個 memory module。

這三種策略分別對應算法復雜度、推理深度和模型規模三個預算軸，使 BudgetMem 不只是一個省 token trick，而是一個系統研究 runtime memory performance-cost trade-off 的統一框架。

用強化學習訓練 Budget Router

BudgetMem 設計了一個輕量級Budget Router，在 runtime memory extraction 的每一步為當前模塊選擇合適的預算檔位。Router 會根據 query、當前模塊輸入和模塊描述，動態決定該模塊使用LOW / MID / HIGH哪一檔，從而在低成本處理和高質量提取之間做出取舍。

由于整個 memory extraction 過程包含離散的檢索、規則、小模型和 LLM 調用，難以直接端到端求導，本文將 budget-tier selection 建模為一個sequential decision problem，并使用強化學習訓練 Router。

每個 query 的處理過程對應一個 episode。最終回答質量作為task reward，記憶提取成本作為cost reward。通過調節 cost weight，BudgetMem 可以在performance-first和cost-sensitive設置之間靈活切換：預算寬松時優先提升記憶質量和回答性能，預算緊張時則主動選擇低成本路徑，降低整體調用開銷。

實驗結果

在效果與成本之間取得更優權衡

本文在 LoCoMo、LongMemEval 和 HotpotQA 上評估 BudgetMem，并與 ReadAgent、MemoryBank、A-MEM、Mem0、MemoryOS、LightMem 等強基線比較。實驗顯示，在 performance-first 設置下，BudgetMem 在 F1 和 LLM-Judge 上整體優于現有 memory baselines。

同時，當我們調節 cost weight 時，BudgetMem 能形成連續、可控的 performance-cost frontier：在相近成本下取得更好效果，在相近效果下降低記憶提取成本。

進一步分析表明，不同 tiering strategy 適用于不同場景。Implementation 和 Capacity Tiering 覆蓋更寬的預算范圍，更適合從低成本到高性能的部署需求；Reasoning Tiering 更像細粒度質量調節旋鈕，適合在相近成本區間內進一步提升記憶質量。

總結：讓 Agent Memory 學會 “按需計算”

BudgetMem 的核心觀點是：未來的 Agent Memory 不應只是固定地存儲、檢索和壓縮歷史，而應根據當前 query 的需求，動態決定投入多少計算來提取和組織記憶。

對于簡單 query，系統可以走低成本路徑快速處理；對于復雜 query，則可以主動調用更強模塊、更深推理或更大模型，獲得更可靠的記憶支持。

因此，BudgetMem 希望推動 Agent Memory 從 “固定流程” 走向 “按需計算”：讓記憶系統不只是會記住歷史，還能判斷什么時候該 “快速掃一眼”，什么時候該 “認真回看”、整理證據并投入更多計算。這種能力也將成為長期對話、個性化智能體和真實部署場景中的關鍵基礎。