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

AI summary

Overview

• Research area: Natural Language Processing — specifically memory systems for Large Language Model (LLM) agents, with connections to inference-time compute control and reinforcement learning.
• Technical level: Intermediate — assumes familiarity with LLM agents, memory pipelines, and basic reinforcement learning (policy optimization), but the core ideas are explained accessibly.
• Scope: The paper introduces BudgetMem, a framework that makes the performance-versus-cost trade-off in runtime agent memory explicit, query-aware, and controllable through learned budget-tier routing.

What This Paper Is About

Most LLM agent memory systems are built "offline" and query-agnostic: they preprocess, compress, and index past context in a fixed way regardless of what the user will later ask, which wastes computation and can irreversibly discard query-critical details. BudgetMem instead defers memory construction to query time and gives each stage of that process a set of Low/Mid/High budget tiers, so a learned router can spend more compute only when a query actually needs it. The goal is to give practitioners an explicit knob for trading answer quality against memory-extraction cost, and to systematically study how those budget tiers should be realized.

Key Contributions

1. A modular runtime memory framework with a common budget-tier interface. BudgetMem structures on-demand memory extraction as a multi-stage pipeline (filter → parallel entity/temporal/topic extraction → summarization), where each module can be invoked under Low, Mid, or High compute while preserving the same input–output contract.
2. Learned budget-tier routing via reinforcement learning. A single lightweight router policy observes the query, the current module input, and a module descriptor, and selects a tier per module. It is trained end-to-end with a cost-aware reward that trades task performance against token-based extraction cost.
3. A unified testbed for three complementary tiering strategies. The paper defines and compares implementation tiering (varying method complexity, from heuristics to BERT-scale models to LLM processing), reasoning tiering (direct vs. chain-of-thought vs. reflective inference), and capacity tiering (smaller vs. larger backbone models) within one framework.
4. Empirical characterization of when each strategy pays off. Experiments across LoCoMo, LongMemEval, and HotpotQA show strong performance-first results and smoother, controllable performance–cost frontiers, plus analysis disentangling the strengths of each tiering axis under different budget regimes.

Main Findings

• Strong gains in performance-first settings: With no cost pressure (λ = 0), BudgetMem variants outperform strong baselines such as ReadAgent, MemoryBank, A-MEM, LangMem, Mem0, MemoryOS, and LightMem on F1 and LLM-as-a-Judge across all three datasets. For example, on LongMemEval with LLaMA-3.3-70B, BudgetMem-Cap reaches a Judge score of 60.50 versus LightMem's 48.51.
• Controllable performance–cost frontiers: By sweeping the cost weight λ, BudgetMem traces smooth trade-off curves that envelope baselines in both low- and high-cost regimes — achieving higher quality at comparable cost, or lower cost at comparable quality.
• Different tiering axes have different trade-off shapes: Implementation and capacity tiering span a broader cost range (implementation gives rapid quality gains at moderate budgets; capacity keeps pushing the frontier outward at high budgets), while reasoning tiering has pronounced cost concentration and acts more like a fine-grained quality knob within a limited cost bandwidth.
• Reward-scale alignment is essential: Removing the variance-based alignment between task and cost rewards causes the router to collapse onto the Low tier, producing a degenerate low-cost policy with the worst Judge scores — confirming that naive cost weighting biases learning.
• Interpretable routing behavior: Module-level selection ratios shift systematically from Mid/High toward Low as cost pressure increases, giving direct evidence that the router allocates compute in a cost-aware, predictable way.
• Retrieval size has a sweet spot: On LoCoMo, retrieving 5 chunks gives the best balance; retrieving too many chunks introduces noise and lowers Judge scores, while too few provides insufficient evidence.
• Practical latency reduction: Under controlled local deployment, implementation tiering cuts total inference latency from 3881 ms (λ = 0) to 1167 ms (λ = 0.9), while the router itself adds only about 44–52 ms of GPU time per query.
• Transfer works without retraining: A router trained with LLaMA as the extraction backbone transfers directly to Qwen3-Next-80B-A3B, still ranking at the top of its comparison block.

Methodology in Plain English

The researchers keep the agent's raw history intact and only split it into indexed text chunks — no offline summarization or rewriting. When a query arrives, a retriever pulls the top-K relevant chunks, and those chunks flow through a fixed sequence of memory modules: a filter, three parallel extractors (entity, temporal, topical), and a summarizer that produces the final compact memory. Each of these modules is available in three versions — Low, Mid, and High — that differ in cost and expected quality.

The key move is that the choice of which version to use at each step is not hard-coded. A small neural policy examines the query, whatever the previous module produced, and which module is being routed, then picks a tier. Because the modules involve non-differentiable components (rules, LLM calls), the team frames routing as a sequential decision problem and trains the policy with reinforcement learning (PPO). The reward combines two signals: how well the final answer matched the ground truth, and how much the extraction cost (measured in tokens priced at API rates). A tunable weight λ controls how much the system cares about cost versus quality, and a variance-based scaling factor keeps the two reward terms from drowning each other out during training.

To isolate which design axis matters, the authors implement the Low/Mid/High tiers three different ways: by swapping module implementations, by changing inference behavior (direct vs. chain-of-thought vs. reflection), and by changing the underlying model's size. This lets them compare the trade-off shapes of all three axes on equal footing.

Why This Matters

Impact on research. This work reframes agent memory as an explicit resource-allocation problem rather than a fixed preprocessing step. It supplies a unified testbed for comparing how different "compute knobs" — algorithmic complexity, reasoning depth, and model capacity — behave under budget constraints, which prior work studied mostly in isolation or in offline settings. It also shows that RL-trained routing can produce stable, interpretable budget policies when reward scales are properly balanced.

Real-world applications:

• Personalized assistants and chatbots that must recall long conversation histories but cannot afford to reprocess everything on every turn.
• Customer-support agents operating under strict per-query latency or API-cost budgets, where routing can automatically downgrade memory processing for simple questions.
• Enterprise document QA and research tools that need to pull precise evidence from large corpora without paying for full-context processing on every query.
• Resource-constrained deployments (edge devices, small teams, startups) where a predictable accuracy–cost frontier makes agent memory economically viable at all.

Industry relevance. Commercial LLM providers increasingly expose tiered compute controls ("thinking" modes, reasoning levels, heavier-model options). BudgetMem extends this pattern into the memory layer, which is where long-horizon agents spend much of their runtime budget. It offers a concrete, transferable mechanism for building agents whose cost and latency can be planned for rather than discovered after the fact.

Future Directions

• Extending beyond the fixed filter → extraction → summary backbone. The authors note the framework does not assume a specific module set, leaving open how routing behaves in other modular memory pipelines (e.g., graph-based or hierarchical memory).
• System-level acceleration of the full pipeline. The paper explicitly leaves batching, caching, and serving-level optimizations as orthogonal future work, noting that router overhead is small but the overall pipeline can still be sped up.
• Richer or adaptive tiering. All experiments use exactly three tiers; whether finer granularity, learned tier definitions, or mixing tiering axes (rather than comparing them separately) yields better frontiers is unexplored.
• Reducing reliance on reward shaping. The reward-scale alignment ablation shows training is sensitive to how task and cost signals are combined; more robust or principled objective designs could remove that tuning burden and generalize across backbones and datasets.

Target Audience

This paper is most useful for researchers and engineers working on LLM agents, long-context memory systems, and inference-time compute optimization. It will also interest practitioners building production agents under latency or cost constraints, and ML researchers studying reinforcement learning for systems-level resource allocation. Readers should have some background in LLM inference, retrieval pipelines, and basic RL concepts to fully engage with the methodology, though the high-level framing and results are accessible to anyone familiar with agent memory challenges.

Authors’ abstract

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.