面向运行时智能体记忆的查询感知预算层级路由学习 Learning Query-Aware Budget-Tier Routing for Runtime Agent Memory
通过模块化预算层级和强化学习路由器,实现运行时记忆提取的显式性能-成本控制
前置知识
智能体记忆(Agent Memory)
LLM智能体用来保留和复用超出单次上下文窗口信息的外部存储系统。与传统检索增强生成(RAG)不同,智能体记忆通常涉及长期对话历史、用户偏好、知识图谱等复杂结构。记忆系统需要在查询时从大量历史记录中提取相关信息,这是一个计算密集型过程。
本文的核心问题是运行时记忆提取的计算成本过高,理解智能体记忆的工作机制是理解本文动机的前提
运行时记忆提取(Runtime Memory Extraction)
与离线预处理不同,运行时记忆提取在查询到达时才进行记忆处理。这种方式避免了查询无关的预处理导致的信息丢失,但将计算成本推到了推理阶段,使得延迟和成本成为首要考虑因素。
本文提出的BudgetMem正是针对运行时记忆提取的场景,理解这一概念是理解本文方法的基础
强化学习(Reinforcement Learning)
一种通过与环境交互来学习最优策略的机器学习方法。在本文中,强化学习被用来训练路由器,使其能够在每个模块调用时选择最优的预算层级。路由器通过观察查询和中间状态,输出LOW/MID/HIGH的层级选择,并根据任务性能和成本的综合奖励进行优化。
本文使用PPO算法训练预算层级路由器,理解强化学习的基本概念有助于理解路由器的训练过程
模块化管道(Modular Pipeline)
将复杂的记忆处理过程分解为多个独立的模块,每个模块负责特定的功能(如过滤、实体提取、时间信息提取、主题提取、摘要)。每个模块可以独立配置预算层级,使得计算成本可以细粒度地控制。
BudgetMem的核心架构是模块化管道,理解这一设计有助于理解预算层级如何在不同模块间分配
研究动机
当前智能体记忆系统面临一个根本性矛盾:离线、查询无关的记忆构建方式虽然降低了运行时成本,但会不可逆地丢失对特定查询至关重要的信息。例如,MemoryBank和MemoryOS等系统采用固定方式预处理、压缩或索引历史上下文,这种一次构建永远使用的范式存在两个核心问题:第一,无论当前查询需要什么信息,系统都会花费相同的计算量,造成资源浪费;第二,固定的预处理可能丢弃对特定查询至关重要的细节,导致回答质量下降。另一方面,虽然运行时按需记忆提取是自然的替代方案,但现有工作往往带来显著的开销,并且缺乏对性能-成本权衡的显式控制。在实际工业系统中,越来越多的平台开始提供显式的、通常是分层的计算控制(如OpenAI的思考模式、Anthropic的推理级别),这反映出在质量和运行时成本之间平衡的需求。
本文的目标是本文的具体目标是提出一个能够实现显式、可控性能-成本权衡的运行时智能体记忆框架。具体而言,作者希望:(1)将记忆处理模块化,每个模块提供LOW/MID/HIGH三个预算层级;(2)通过学习一个轻量级路由器,在查询处理过程中为每个模块选择最优的预算层级;(3)在三个基准数据集(LoCoMo、LongMemEval、HotpotQA)上验证框架的有效性,特别是在性能优先设置下超越现有基线,并在预算受限时提供更好的准确率-成本前沿。
与已有工作不同的是,本文的独特切入角度在于:现有工作主要关注离线记忆构建的效率优化,而本文将焦点转向运行时记忆提取的显式性能-成本控制。作者观察到,现有系统的权衡机制通常是离线的,而按需记忆将这些决策推到运行时,每个查询都会引发质量-成本选择。这暴露了两个核心问题:第一,预算应该应用在哪里?现有系统通常将记忆视为单一管道,采用固定的计算设置,使得权衡粗糙且难以控制。第二,预算应该如何实现?现有工作缺乏对运行时记忆权衡的系统性指导。本文通过引入模块化预算层级和学习路由器的方式,填补了这一空白。
核心方法
BudgetMem的核心思想可以用一个类比来理解:想象一个智能助手在回答用户问题时,需要从大量的历史对话中查找相关信息。传统方法就像一个勤奋但不够灵活的助手,每次都把所有可能相关的对话都仔细整理一遍,不管当前问题是否需要这么详细的整理。而BudgetMem则像一个更聪明的助手:它会根据问题的复杂程度,动态决定在每个处理步骤上投入多少精力。对于简单问题,它可能只做基本的关键词匹配;对于复杂问题,它才会启动更深入的语义分析和推理。技术路线上,BudgetMem将记忆处理分解为一个模块化管道:首先过滤检索到的原始文本块,然后并行提取实体关系、时间信息和主题信息,最后将这些信息聚合成最终的记忆表示。每个模块都提供LOW/MID/HIGH三个预算层级,通过一个共享的轻量级路由器在查询处理过程中动态选择。
BudgetMem最本质的创新在于将预算控制从离线预处理转移到运行时提取,并通过学习而非启发式规则来实现预算分配。与已有方法的核心区别体现在三个方面:第一,预算粒度。现有系统通常将记忆视为单一管道,采用固定的计算设置,而BudgetMem在模块级别实现预算控制,使得计算可以针对性地分配。第二,决策方式。现有方法要么使用固定的启发式规则,要么简单地增加计算量,而BudgetMem通过强化学习训练路由器,使其能够根据查询和中间状态做出最优的预算分配决策。第三,权衡维度。本文系统地比较了三种互补的预算实现策略(实现、推理、容量),揭示了不同策略在不同预算机制下的权衡特性,这是前所未有的系统性分析。
方法步骤详情
BudgetMem的完整处理流程如下:首先,给定用户查询q,从分块的历史记录C中检索一组潜在相关的文本块Cq。然后,模块化管道开始执行:(1)过滤模块Mfil接收查询和检索到的文本块,输出更聚焦的子集或重新加权的集合;(2)三个提取模块并行操作:实体模块Ment提取实体关系,时间模块Mtmp提取时间信息,主题模块Mtop提取主题信息;(3)摘要模块Msum将这些输出聚合成最终的提取记忆m。在整个过程中,路由器在每个模块调用时选择预算层级(LOW/MID/HIGH)。路由器的状态由三部分组成:查询q、当前模块输入(前一个模块的输出)和模块描述符。路由器输出动作ak属于{LOW, MID, HIGH},决定当前模块的执行层级。训练时,使用PPO算法优化路由器,奖励函数为r = rtask + lambda * alpha * rcost,其中rtask是任务性能奖励,rcost是成本奖励,lambda控制权衡偏好,alpha是用于对齐奖励尺度的因子。
技术新颖性
BudgetMem的技术新颖性体现在以下几个方面:首先,这是第一个系统性地研究运行时智能体记忆中显式性能-成本控制的工作。虽然推理时的性能-成本权衡在LLM系统中已有广泛研究(如快速解码、早期退出、量化等),但这些工作主要关注答案生成阶段,而本文将焦点转向记忆提取阶段。其次,本文提出了三种互补的预算实现策略,并在一个统一框架下进行比较,这为理解不同权衡机制的特性提供了系统性的视角。实现层级通过改变模块实现(从轻量级启发式到BERT模型再到LLM)来控制成本;推理层级通过改变推理行为(从直接生成到思维链再到多步反思)来控制成本;容量层级通过改变模型大小来控制成本。这种正交的维度设计使得研究者可以针对性地选择最适合其场景的权衡策略。
实验结果
实验结果表明BudgetMem在多个基准数据集上都取得了显著的性能提升。在性能优先设置(lambda=0)下,以LLaMA-3.3-70B-Instruct为骨干模型,BudgetMem在三个数据集上的平均F1分数达到41.84(实现层级)、44.19(推理层级)和45.72(容量层级),而最强基线LightMem的平均F1仅为35.45。在LLM-as-a-Judge评分上,BudgetMem的三个变体分别达到57.36、57.39和59.99,而LightMem为49.21。在HotpotQA数据集上,BudgetMem-CAP达到53.87 F1和64.85 Judge,相比LightMem(45.73 F1,58.37 Judge)提升了8.14个F1点和6.48个Judge点。值得注意的是,即使在性能优先设置下,BudgetMem仍然保持了成本效率,在HotpotQA上以0.22美元的成本达到最佳Judge分数(72.08),而推理层级变体以更低的成本(0.17美元)达到了可比的Judge分数(70.83)。当使用Qwen3-Next-80B-A3B-Instruct作为骨干模型时,BudgetMem在HotpotQA上达到58.70 F1和72.08 Judge,成本仅为0.22美元,而LightMem需要8.56美元才能达到41.29 F1和55.42 Judge。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| LoCoMo | F1 | 43.05 | 33.88 (LightMem) | +9.17 |
| LoCoMo | Judge | 54.62 | 40.76 (LightMem) | +13.86 |
| LongMemEval | F1 | 40.24 | 26.74 (LightMem) | +13.50 |
| LongMemEval | Judge | 60.50 | 48.51 (LightMem) | +11.99 |
| HotpotQA | F1 | 53.87 | 45.73 (LightMem) | +8.14 |
| HotpotQA | Judge | 64.85 | 58.37 (LightMem) | +6.48 |
局限与改进
尽管BudgetMem取得了显著的性能提升,但仍存在一些局限性。首先,本文的实验主要集中在对话场景(LoCoMo、LongMemEval)和长上下文问答(HotpotQA),对于其他类型的记忆任务(如代码智能体、多模态智能体)的适用性尚未验证。其次,路由器的训练需要大量的查询-答案对,在数据稀缺的场景下可能面临挑战。第三,虽然本文比较了三种预算层级策略,但每种策略的具体实现(如使用哪些BERT模型、哪些启发式规则)可能对性能有较大影响,论文中对这些实现细节的敏感性分析有限。第四,当前的模块化管道设计(过滤、并行提取、摘要)是针对对话场景定制的,对于其他类型的记忆任务可能需要重新设计管道结构。第五,论文中提到的延迟分析(Table 2和Table 3)显示,BudgetMem的检索/过滤延迟在某些设置下仍然较高(如实现层级在lambda=0时为1176毫秒),这可能限制其在对延迟敏感的实时应用中的部署。
独立分析的弱点
从独立分析的角度,BudgetMem存在以下几个可以改进的弱点:第一,路由器的状态表示目前仅使用简单的文本嵌入拼接,可能无法充分捕捉查询和模块输出之间的复杂语义关系。改进方向是引入更强大的交叉注意力机制或图神经网络来建模模块间的依赖关系。第二,当前的奖励函数设计(r = rtask + lambda * alpha * rcost)使用固定的线性权衡,但在实际应用中,不同查询可能对性能和成本有不同的权衡偏好。改进方向是学习一个自适应的权衡参数,或者允许用户在查询级别指定权衡偏好。第三,实验中使用的模块化管道是针对对话场景设计的,对于其他类型的记忆任务(如文档理解、多模态推理)可能需要重新设计。改进方向是提供一个可配置的模块库,允许用户根据任务需求组合不同的模块。
未来方向
基于本文的成果,未来研究可以从以下几个方向展开:第一,探索更复杂的管道拓扑。当前BudgetMem使用固定的线性管道(过滤、并行提取、摘要),未来可以研究动态拓扑,允许路由器根据查询特性选择不同的模块组合。第二,将预算控制扩展到检索阶段。当前BudgetMem假设检索阶段是固定的(使用Contriever和固定的top-K),但检索的质量和成本也可以作为预算控制的一部分。第三,研究多智能体场景下的预算分配。在多智能体系统中,不同智能体可能需要不同的记忆预算,如何协调和优化全局预算是一个有前景的方向。第四,探索在线学习场景。当前路由器需要离线训练,但在实际部署中,系统可能需要根据用户反馈在线调整预算策略。
复现评估
从复现评估的角度,BudgetMem具有较好的可复现性。代码已在GitHub开源(https://github.com/ViktorAxelsen/BudgetMem),这为研究者提供了直接的参考实现。数据集方面,LoCoMo、LongMemEval和HotpotQA都是公开可用的基准数据集,且论文中提供了详细的数据统计(Table 4)和评估配置(Tables 5和6)。算力需求方面,BudgetMem使用LLaMA-3.3-70B-Instruct和Qwen3-Next-80B-A3B-Instruct作为骨干模型,这些模型可以通过API服务访问,但训练路由器需要强化学习,可能需要一定的计算资源。论文中提到训练使用Adam优化器,批大小为32,最多600个训练步骤,这在中等规模的GPU集群上应该是可行的。复现难度中等,主要挑战在于强化学习训练的稳定性和奖励尺度对齐的调优。
论文图表