PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning Query-Aware Budget-Tier Routing for Runtime Agent Memory

Haozhen Zhang, Haodong Yue, Tao Feng, Quanyu Long, Jianzhu Bao, Bowen Jin, Weizhi Zhang, Xiao Li, Jiaxuan You, Chengwei Qin, Wenya Wang

arXiv:2602.06025v3cs.CLcs.AIcs.LG

TL;DR

Existing agent-memory systems often construct memories offline and query-agnostically, while runtime extraction offers flexibility but limited explicit control over quality–cost trade-offs. BudgetMem introduces modular runtime extraction with three budget tiers and reinforcement-learning-based routing, achieving strong high-budget results and improved performance–cost frontiers under tighter budgets across three benchmarks.

  • Problem

    Existing memory systems often use inefficient query-agnostic construction, while runtime memory extraction offers limited explicit control over performance–cost trade-offs.

  • Method

    BudgetMem uses budget-tiered memory modules, a lightweight router trained with reinforcement learning, and implementation, reasoning, and capacity tiering strategies.

  • Results

    Across LoCoMo, LongMemEval, and HotpotQA, BudgetMem shows strong performance in high-budget settings and improved performance–cost frontiers under tighter budgets.

  • Takeaways & Limitations

    BudgetMem provides a unified way to compare tiering strategies and identify their returns across different runtime memory budget regimes.

  • Takeaways & Limitations

    BudgetMem introduces no new data sources, and its potential risks remain similar to existing memory systems, including irrelevant retrieval or unintended historical details.

Abstract

from arXiv · show

Memory is increasingly central to Large Language Model (LLM) agents operating beyond a single context window, yet most existing systems rely on offline, query-agnostic memory construction that can be inefficient and may discard query-critical information. Although runtime memory utilization is a natural alternative, prior work often incurs substantial overhead and offers limited explicit control over the performance-cost trade-off. In this work, we present \textbf{BudgetMem}, a runtime agent memory framework for explicit, query-aware performance-cost control. BudgetMem structures memory processing as a set of memory modules, each offered in three budget tiers (i.e., \textsc{Low}/\textsc{Mid}/\textsc{High}). A lightweight router performs budget-tier routing across modules to balance task performance and memory construction cost, which is implemented as a compact neural policy trained with reinforcement learning. Using BudgetMem as a unified testbed, we study three complementary strategies for realizing budget tiers: implementation (method complexity), reasoning (inference behavior), and capacity (module model size). Across LoCoMo, LongMemEval, and HotpotQA, BudgetMem surpasses strong baselines when performance is prioritized (i.e., high-budget setting), and delivers better accuracy-cost frontiers under tighter budgets. Moreover, our analysis disentangles the strengths and weaknesses of different tiering strategies, clarifying when each axis delivers the most favorable trade-offs under varying budget regimes.

1. Introduction

BudgetMem addresses the inefficiency and limited controllability of query-agnostic or runtime memory processing with modular, query-aware budget-tier routing. Across three benchmarks, it reports strong high-budget performance and improved performance–cost trade-offs as budgets tighten.

  • Offline, query-agnostic memory construction can waste computation and omit information critical to particular queries.
  • Runtime extraction makes memory processing query-aware but turns cost and latency into explicit concerns requiring controllable performance–cost trade-offs.
  • BudgetMem makes runtime extraction modular and exposes a common LOW/MID/HIGH tier interface so a learned router can control computation per module.
  • Experiments on LoCoMo, LongMemEval, and HotpotQA show strong performance in high-budget settings and clear trade-off curves under tighter budgets.
  • The framework learns budget-tier routing with a shared lightweight router trained using reinforcement learning.

2. Related Work

Prior work studies memory construction and runtime compute control largely as separate problems. BudgetMem brings explicit budget-tier control into runtime agent memory and compares complementary realization strategies in one framework.

  • Memory-augmented agents commonly use external stores, while many systems construct memories offline or ahead of time and retrieve them later.
  • Runtime memory remains costly and is rarely framed as explicit performance–cost control because existing systems often rely on fixed pipelines.
  • Controllable compute research includes algorithmic, systems-level, and reasoning-level mechanisms for trading quality against runtime cost or latency.
  • BudgetMem applies explicit budget-tier control to runtime memory extraction and systematically studies complementary budget realization strategies in a unified setting.

3. Problem Setup and Method Overview

The task keeps historical records intact, retrieves query-relevant chunks, and performs runtime memory extraction before generating an answer conditioned on the query and extracted memory. BudgetMem makes this extraction budget-controllable through module-wise routing.

  • The system defers memory computation until query time, retaining raw historical records instead of applying irreversible query-agnostic preprocessing.
  • A lightweight segmentation step creates a chunk store, after which a retriever returns potentially relevant chunks for the current query.
  • Runtime memory extraction maps the query and retrieved chunks to a compact memory, which an answer generator uses with the query to produce the final answer.
  • BudgetMem uses a modular pipeline with LOW/MID/HIGH tiers and a shared router that selects each module’s tier from the query and intermediate signals.

4. BudgetMem

BudgetMem is a modular runtime memory framework that applies query-aware, module-level budget control through tiered processing and learned routing. It combines implementation, reasoning, and capacity tiering with reinforcement-learning optimization of performance–cost trade-offs.

  • Modular Runtime Memory Pipeline: BudgetMem processes retrieved chunks through a modular pipeline that filters context, extracts entity, temporal, and topical signals, and summarizes them into query-focused memory.The framework uses a fixed pipeline structure while allowing the module set to be adapted to other memory pipelines.
  • Budget-Tier Interface: Each module exposes LOW, MID, and HIGH tiers through a shared interface, preserving its input–output contract while enabling fine-grained runtime budget decisions.The tiers provide increasing computational budgets and typically improved extraction quality.
  • Tiering Strategies: BudgetMem studies implementation, reasoning, and capacity tiering as orthogonal ways to vary computation through procedure, inference behavior, and model size.Implementation tiers range from heuristics to learned models and LLM processing; reasoning tiers vary inference intensity, while capacity tiers vary model capacity.
  • Budget-Tier Routing: A shared router selects a tier for each module from the query, current module input, and module identity, producing sequential routing decisions during pipeline execution.The router observes intermediate states as each selected tier produces the next signal or state.
  • Reinforcement-Learning Objective: The router optimizes task performance and extraction cost jointly, using reinforcement learning because runtime memory extraction can contain non-differentiable components.The objective uses λ to control performance–cost preference, while α aligns reward magnitudes; PPO is used as a standard policy-optimization algorithm.
  • Cost Modeling: Extraction cost sums per-module costs, estimates LLM costs from token usage and prices, treats non-LLM costs as negligible by comparison, and normalizes recent costs for reward use.A sliding-window normalization uses robust quantiles, including Q5 and Q95, while variance-based alignment addresses differing reward variances.

5. Experimental Setup

BudgetMem is evaluated on three memory and long-context question-answering benchmarks using task-performance and monetary extraction-cost measures. The setup compares it with diverse memory-augmented baselines under matched retrieval and evaluation protocols.

  • Datasets: Experiments use LoCoMo, LongMemEval, and HotpotQA to evaluate long-horizon memory and long-context question answering.LoCoMo and LongMemEval target agent-memory evaluation, while HotpotQA represents long-context QA.
  • Metrics: Task performance is measured with F1-score and LLM-as-a-judge, while extraction cost aggregates model-call input and output tokens and converts them into monetary cost.The cost calculation uses the corresponding service pricing.
  • Evaluation Protocol: The evaluation reports results across all three datasets using F1 and LLM-as-a-Judge metrics in the performance-first setting.Table 1 presents the experimental comparisons for LoCoMo, LongMemEval, and HotpotQA.
  • Baselines: BudgetMem is compared against ReadAgent, MemoryBank, A-MEM, LangMem, Mem0, MemoryOS, and LightMem.These baselines cover a broad range of memory architectures.
  • Models and Retrieval: The router is trained with LLaMA-3.3-70B-Instruct and directly tested with Qwen3-Next-80B-A3B-Instruct without retraining.Retrieval uses Contriever, and BudgetMem and baselines share matched top-K retrieval with K=5.
  • Training and Fair Comparison: All methods use the same data splits and evaluation protocols, with a 6/2/2 train-validation-test split across datasets.PPO is the default reinforcement-learning algorithm, with Adam optimization and batch size 32 for up to 600 steps.

6. Experimental Analysis

BudgetMem shows strong performance across three datasets while exposing controllable performance–cost trade-offs through tiered routing. Analyses clarify how reward alignment, routing behavior, retrieval size, and latency shape these trade-offs.

  • 6.2. Exploring Trade-off Across Tiering Strategies: Varying λ produces smooth, controllable performance–cost frontiers that envelop baselines in both low- and high-cost regimes.Implementation tiering provides rapid Judge gains under moderate budgets, capacity tiering reaches the best quality at high budgets, and reasoning tiering offers finer improvements within a narrower cost range.
  • 6.1. Main Results: BudgetMem consistently outperforms prior methods across LoCoMo, LongMemEval, and HotpotQA in both F1 and LLM-Judge.On LongMemEval with LLaMA-3.3-70B, BudgetMem-CAP achieves a Judge score of 60.50 versus 48.51 for LightMem.
  • 6.1. Main Results: BudgetMem maintains strong quality under controlled cost, with BudgetMem-CAP reaching a Judge score of 72.08 at cost 0.22 on HotpotQA and BudgetMem-REA reaching 70.83 at cost 0.17.The paper attributes this efficiency to query-relevant retrieval and on-demand extraction, while noting smaller cost differences on LoCoMo.
  • 6.1. Main Results: Across datasets, BudgetMem variants consistently rank at the top within each LLaMA and Qwen backbone block, indicating broad performance rather than a single-benchmark gain.This aggregation supports strong effectiveness across diverse evaluation settings in the performance-first regime.
  • 6.3. Ablating Reward-scale Alignment: Removing reward-scale alignment causes unstable optimization and a degenerate LOW-tier policy, substantially reducing answer quality and driving Judge to its lowest level.The ablation shows that aligned task and cost rewards are needed for meaningful cost–performance control rather than trivial cost minimization.
  • 6.4. Discussion: The router shifts probability mass from HIGH and MID toward LOW as λ increases, while retrieval-size changes trade evidence coverage against noise, cost, and Judge score.Retrieving 5 chunks provides the reported best balance; implementation tiering also reduces total latency from 3881 ms at λ = 0 to 1167 ms at λ = 0.9.

7. Conclusion

BudgetMem provides a unified framework for query-aware performance–cost control in runtime agent memory. Across three benchmarks, it improves high-budget performance and tight-budget performance–cost frontiers while comparing three tiering strategies.

  • BudgetMem enables explicit performance–cost control for on-demand memory extraction through a modular pipeline with budget-tier routing.A lightweight router selects tiers as queries are processed.
  • Experiments on LoCoMo, LongMemEval, and HotpotQA show strong performance when performance is prioritized.
  • BudgetMem improves performance–cost frontiers under tighter budgets across the evaluated benchmarks.
  • The framework compares implementation, reasoning, and capacity tiering strategies and characterizes their trade-off behavior across budgets.

Impact Statement

BudgetMem targets practical deployment by making memory computation controllable under budget constraints. Its expected impact is improved usability and controllability of agent memory, with standard safeguards addressing risks shared with existing memory systems.

  • Explicitly controllable memory computation can benefit applications requiring predictable latency or cost and resource-limited settings.
  • BudgetMem introduces no new data sources, so its potential risks resemble those of existing memory systems.These include irrelevant retrieval and exposure of unintended historical details.
  • Retention policies, access control, and routine failure-case evaluation are identified as safeguards for these risks.
  • The framework’s primary expected impact is improving agent-memory usability and controllability under real-world budget constraints.

A.1. Dataset Statistics

The experiments use LoCoMo, LongMemEval, and HotpotQA, with dataset statistics covering split sizes and average context lengths. HotpotQA follows a prior evaluation setup that concatenates evidence and distractor documents into long contexts.

  • Table 4 reports train, validation, and test split sizes together with average context length for LoCoMo, LongMemEval, and HotpotQA.
  • HotpotQA contexts concatenate evidence and distractor documents, using preprocessed test queries from the prior evaluation setup.

A.2. Evaluation Details

Evaluation details include shared and dataset-specific hyperparameters, repeated runs with consistent protocols, and a structured router state representation. The router combines query, module-input, and module-descriptor information to predict budget tiers.

  • Tables 5 and 6 specify hyperparameters shared across datasets and additional dataset-dependent evaluation settings.
  • Each BudgetMem experiment uses three random seeds and reports averaged results under consistent data splits, retrieval settings, and evaluation protocols.This reduces effects from initialization and reinforcement-learning router stochasticity.
  • The router state represents the user query, current module input, and module descriptor before predicting the budget tier.
  • The final router state is a 512-dimensional vector combining projected context and module representations.

A.4. Details on PPO Training Objectives

BudgetMem trains a shared PPO router to select LOW, MID, or HIGH budget tiers for each module in a fixed runtime memory pipeline. Its objective combines answer quality and normalized extraction cost, while tier implementations vary computation, reasoning behavior, and model capacity.

  • Router and tier design: The router selects tiers across the fixed pipeline Mfil → {Ment, Mtmp, Mtop} → Msum, preserving the overall extraction structure.The pipeline filters retrieved context, extracts complementary entity, temporal, and topic contexts, and summarizes them into final memory.
  • Router and tier design: Each query forms one episode in which a shared actor–critic router sequentially selects a three-way budget tier for every module.The state combines the query, current module input, and module descriptor, enabling shared parameters with module-specific routing behavior.
  • PPO objective: The reward combines task quality, normalized extraction cost, a performance–cost weight λ, and a reward-scale alignment factor α.Task reward is based on F1 or LLM-as-a-judge, while cost reward is derived from routed module-call cost.
  • PPO objective: PPO optimizes the joint likelihood of the full sequence of module actions, alongside value-function and entropy terms.The likelihood ratio compares the updated policy with the rollout behavior policy, and the entropy coefficient regularizes routing decisions.
  • PPO objective: Entropy is averaged across module decisions, preventing its contribution from scaling with the number of routed modules.Advantages use a one-step Monte-Carlo return because each episode ends after the routed pipeline completes, with the critic as baseline.
  • Router and tier design: Budget tiers preserve module input–output contracts while changing internal computation or inferential strength across implementation, reasoning, and capacity strategies.Implementation tiers move from symbolic pipelines to compact learned specialists and LLM-based processing; reasoning tiers progress from direct inference to CoT and reflection-style inference.

B. More Experimental Results

The supplementary material details category-level evaluations, answer prompts, and relevance-scoring prompts across LoCoMo, LongMemEval, and HotpotQA. These materials report F1, Judge score, and cost while specifying dataset- and tier-specific evaluation procedures.

  • More Experimental Results: LoCoMo results report per-category F1-score, Judge score, and Cost, separately for LLaMA-based and Qwen-based models.Category 5, the adversarial category, is excluded from evaluation.
  • More Experimental Results: LongMemEval results likewise provide category-level F1-score, Judge score, and Cost for LLaMA-based and Qwen-based model families.Each category is evaluated independently to examine retention, retrieval, and reasoning over long-term contextual information.
  • More Experimental Results: HotpotQA results report per-category F1-score, Judge score, and Cost, with separate tables for LLaMA-based and Qwen-based models.The categories correspond to predefined multi-hop reasoning types.
  • Answer prompts: The LoCoMo answer prompt requests a short phrase using exact context words whenever possible, while HotpotQA requires concise answers enclosed in <answer> tags.These prompts define dataset-specific answer formatting rather than changing the reported metrics.
  • Filter prompts: The filter module scores each memory for query relevance on an integer scale, with higher values indicating more directly useful information.The prompts require one score per memory in input order and return an empty array when no memories are present.
  • Filter prompts: The high-tier filter prompt adds explicit planning, action, and reflection steps before producing relevance scores.Its scoring rubric distinguishes essential facts, strong relevance, partial usefulness, weak overlap, and irrelevance.
Loading 2602.06025v3…