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

Abstract: Memory is increasingly central to 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.

• The paper introduces BudgetMem, a modular framework that enables runtime cost-performance trade-offs in memory-augmented LLM agents.
• It employs a lightweight reinforcement learning router to dynamically assign budget tiers across implementation, reasoning, and capacity strategies.
• Experiments on benchmarks like LongMemEval and HotpotQA show that BudgetMem outperforms baselines in both quality and cost-efficiency.

Learning Query-Aware Budget-Tier Routing for Runtime Agent Memory: A Technical Analysis

Problem Formulation and Motivation

This work addresses the operational inefficiencies and rigidity inherent in traditional memory-augmented LLM agent systems, specifically those using offline, query-agnostic memory construction pipelines. These conventional approaches often fail to allocate computational resources dynamically according to query-specific requirements, leading to both surplus cost and loss of crucial information for downstream tasks. The paper introduces the BudgetMem framework, which enables explicit, rigorous performance-cost trade-offs during runtime memory extraction. The objective is to support LLM agents in long-horizon, personalized, and knowledge-intensive tasks by making agent memory both query-aware and budget-controllable.

BudgetMem Framework: Architecture and Core Methodology

BudgetMem emphasizes a modular, multi-stage runtime memory pipeline. Each module exposes a uniform budget-tiered interface (Low/Mid/High), allowing fine-grained cost-quality decisions per module invocation without altering the architectural backbone. At runtime, a shared lightweight neural router, trained using reinforcement learning (specifically PPO), dynamically selects the budget tier per module for each query. The router's policy is conditioned on the query, module descriptor, and intermediate pipeline states, yielding an end-to-end differentiable, query-aware computation allocation.

Three principal budget tiering strategies are standardized across all modules:

• Implementation Tiering: Progressive increase in algorithmic sophistication, from heuristics to compact learned models to full LLM calls.
• Reasoning Tiering: Incremental inference complexity, e.g., direct prediction, CoT prompting, and multi-step reflective reasoning.
• Capacity Tiering: Scaling the underlying model size per module, keeping inference patterns fixed but varying inductive capacity.

The system is designed to optimize a cost-aware objective, balancing task reward and normalized computational cost, with additional alignment via variance-based reward-scaling to prevent degenerate solutions.

Experimental Evaluation and Numerical Results

Rigorous experiments are conducted on three long-context/high-memory QA benchmarks: LoCoMo, LongMemEval, and HotpotQA. Two state-of-the-art LLM backbones (LLaMA-3.3-70B-Instruct, Qwen3-Next-80B-A3B-Instruct) are used to demonstrate robustness and transferability. Baselines include a diversity of competitive memory-augmented agent architectures (ReadAgent, MemoryBank, A-MEM, Mem0, MemoryOS, LangMem, LightMem).

Key quantitative results:

• On LongMemEval with LLaMA-3.3-70B, BudgetMem-CAP achieves a Judge score of 60.50, surpassing the next best baseline (LightMem, 48.51) by a wide margin.
• On HotpotQA (Qwen3-Next-80B-A3B), BudgetMem-CAP obtains the highest Judge score (72.08) at minimal extraction cost ($0.22), while BudgetMem-REA achieves a comparable Judge (70.83) at an even lower cost ($0.17).
• BudgetMem variants consistently outperform all baselines in both F1 and Judge across datasets, demonstrating superiority in both cost-constrained and high-quality regimes.
• The framework exhibits highly controllable performance-cost curves, with smooth transitions as the budget parameter is varied. Notably, no existing baseline achieves comparable efficiency across the entire Pareto frontier.

Ablation experiments confirm the necessity of reward-scale alignment for stable RL optimization, and sensitivity analyses reveal optimal retrieval chunk sizes for cost-performance trade-off.

Comparative and Mechanistic Insights

The comparative analysis reveals distinct characteristics of the tiering strategies:

• Implementation and Capacity Tiering: Both span a broader range of the cost-quality pareto frontier, supporting effective operation under both stringent and lax budgets. Capacity tiering achieves the highest absolute quality in high-budget settings.
• Reasoning Tiering: Provides a fine-grained quality knob within a relatively narrow cost bandwidth; meaningful gains in answer quality often accrue at marginal extra cost, but less effective for extreme budget scenarios.

BudgetMem's unified framework distinguishes itself by allowing these axis-specific trade-offs to be directly compared, yielding detailed insight into when each axis dominates under varying regime constraints. Module-wise analysis of routing behaviors demonstrates interpretable and predictable tier allocation in response to budget signals, validating the expressiveness and reliability of the router.

Theoretical and Practical Implications

The explicit modularization and tiered abstraction in BudgetMem introduce a new canonical structure for runtime memory management in agentic LLM systems. The approach:

• Allows deployment architects to guarantee cost and latency bounds through direct budget specification.
• Enables adaptive resource allocation, improving both user experience in real-world applications (by controlling for unpredictable query complexity) and overall system cost-efficiency.
• Provides a unified testbed to systematically evaluate orthogonal directions for memory optimization in LLM agents.

From a theoretical standpoint, this work brings together reinforcement learning, modular deep learning design, and cost-aware decision-making, framing runtime agent memory as a general multi-agent, sequential RL problem over computation graphs.

Prospects and Future Directions

The BudgetMem paradigm opens avenues for a range of research topics:

• Extension to heterogeneous agent architectures, enabling even richer module-level customization (arbitrary function contracts, multimodal memory).
• Incorporation of learning-based state summarization and long-term planning for the budget router to address adversarial or uncertainty-driven user queries.
• Multi-tiered budget negotiation in multi-agent settings, balancing population-level utility and individual performance.
• Exploration of more nuanced cost models reflecting real-world serving constraints, such as latency, token quotas, or competing workload priorities.

Conclusion

The BudgetMem framework establishes an explicit, learnable approach to runtime, budget-aware agent memory for LLM-driven systems. By introducing a modular, tiered abstraction with RL-based query routing and by empirically demonstrating superior cost-performance frontiers across multiple benchmarks and model families, BudgetMem advances the design, deployment, and evaluation of controllable, efficient memory augmentation for LLM agents. Its systematic treatment of budget-realization axes clarifies design choices for both researchers and practitioners, with direct applicability to scalable, cost-sensitive AI deployments.

Reference: "Learning Query-Aware Budget-Tier Routing for Runtime Agent Memory" (2602.06025)