Technical Summary: BudgetMem

Problem Statement

Large Language Model (LLM) agents increasingly rely on memory to operate beyond single context windows for long-horizon interactions and knowledge-intensive reasoning. However, existing systems predominantly utilize offline, query-agnostic memory construction. In these paradigms, past context is preprocessed, compressed, or indexed in a fixed manner without conditioning on the specific downstream query. This "build once, use always" approach is often inefficient, as it expends computation regardless of a query's specific needs and risks discarding information critical to particular queries.

While runtime memory utilization (on-demand extraction) offers a flexible alternative, it introduces significant overhead and latency. Current industrial LLM systems increasingly offer explicit, tiered compute controls (e.g., "thinking" modes or reasoning levels) to balance quality against cost, yet agent memory systems lack a systematic framework for explicit, controllable performance–cost trade-offs during runtime extraction. The core challenges are:

1. Budget Allocation: Defining the appropriate unit for budgeting in a runtime setting, as treating memory as a monolithic pipeline with fixed compute makes trade-offs coarse and difficult to control.
2. Budget Realization: Determining how to operationalize budget control systematically, rather than relying on ad hoc cost reduction or simple compute escalation.

Methodology: BudgetMem

The authors propose BudgetMem, a runtime agent memory framework designed to enable explicit, query-aware performance–cost control. The framework structures memory processing as a multi-stage modular pipeline where computation is controllable at the module level.

1. Modular Pipeline Architecture

BudgetMem implements memory extraction ($f_{mem}$) as a sequence of modules. In their dialogue-centric evaluation, the pipeline follows:
$$M_{fil} \rightarrow \{M_{ent}, M_{tmp}, M_{top}\} \rightarrow M_{sum}$$

• Filter Module ($M_{fil}$): Refines retrieved raw chunks.
• Parallel Extraction Modules: Extract complementary contexts:
  • $M_{ent}$: Entity-centric facts.
  • $M_{tmp}$: Temporal cues.
  • $M_{top}$: Topical summaries.
• Summary Module ($M_{sum}$): Aggregates intermediate contexts into the final extracted memory ($m$).

Crucially, the framework does not assume a specific module set but requires each module to expose a common budget-tier interface.

2. Budget Tiers and Realization Strategies

Each module exposes three budget tiers: LOW, MID, and HIGH. These tiers correspond to increasing computational budgets and typically improved extraction quality. The paper investigates three complementary strategies to realize these tiers:

• Implementation Tiering: Varies the module implementation.
  • LOW: Lightweight heuristics (rules, regex, spaCy).
  • MID: Compact learned models (e.g., BERT-based).
  • HIGH: LLM-based processing.
• Reasoning Tiering: Varies inference behavior while keeping the model backbone fixed.
  • LOW: Direct generation.
  • MID: Chain-of-Thought (CoT) style.
  • HIGH: Multi-step or reflection-style inference.
• Capacity Tiering: Varies the model capacity (size) used within the module.
  • LOW: Small models (e.g., 3B parameters).
  • MID: Medium models (e.g., 8B–32B parameters).
  • HIGH: Large models (e.g., 70B+ parameters).

3. Query-Aware Routing via Reinforcement Learning

To navigate these tiered choices dynamically, BudgetMem employs a shared lightweight router trained with Reinforcement Learning (RL).

• Policy ($\pi_\theta$): At each module invocation step $k$, the router observes a state $s_k$ (comprising the query, current module input, and a module descriptor) and selects an action $a_k \in \{LOW, MID, HIGH\}$.
• Objective: The router is optimized to trade off task performance against memory extraction cost. The reward function is defined as:
  $$r = r_{task} + \lambda \cdot \alpha \cdot r_{cost}$$
  Where $r_{task}$ is the task performance reward (F1 or LLM-Judge), $r_{cost}$ is the normalized cost reward, $\lambda$ controls the trade-off preference, and $\alpha$ is a variance-based alignment factor to stabilize training.
• Training: The router is trained using Proximal Policy Optimization (PPO). The cost is modeled based on token usage and API pricing, normalized via a sliding window to ensure commensurability with task rewards.

Key Contributions

1. Framework Introduction: The authors introduce BudgetMem, a modular runtime agent memory framework that enables explicit performance–cost control for on-demand memory extraction via budget-tiered modules.
2. Routing Mechanism: They propose budget-tier routing, where a shared lightweight router, trained via RL, dynamically selects budget tiers during extraction. This is instantiated with three complementary strategies (implementation, reasoning, and capacity) within a unified framework.
3. Empirical Analysis: Through experiments on LoCoMo, LongMemEval, and HotpotQA, the paper demonstrates that BudgetMem achieves strong performance in performance-first settings and establishes clear, controllable performance–cost trade-off curves under tighter budgets. The analysis disentangles the strengths of different tiering strategies, offering insights into which mechanisms yield the best returns under varying budget regimes.

Experimental Results

The authors evaluated BudgetMem against strong baselines (e.g., ReadAgent, MemoryBank, LightMem) using LLaMA-3.3-70B and Qwen3-Next-80B backbones.

• Performance-First (High Budget): In settings prioritizing performance ($\lambda=0$), BudgetMem variants consistently outperformed baselines. For instance, on LongMemEval with LLaMA, BudgetMem-CAP achieved a Judge score of 60.50, significantly surpassing the best baseline, LightMem (48.51).
• Performance–Cost Trade-offs: By varying the cost weight $\lambda$, BudgetMem traces smooth frontiers that envelope baselines.
  • Implementation and Capacity Tiering: Span a broader cost range. Implementation tiering offers rapid gains under moderate budgets, while capacity tiering pushes the frontier outward in high-budget regimes.
  • Reasoning Tiering: Exhibits pronounced cost concentration. It acts as a fine-grained quality knob within a limited cost bandwidth, offering improvements at similar costs but less flexibility for extreme low/high budget extremes compared to the other strategies.
• Cost Efficiency: Even in performance-first regimes, BudgetMem remains cost-efficient by processing only query-relevant raw chunks on-demand, avoiding the overhead of offline processing entire histories.

Significance and Claims

The paper claims that BudgetMem addresses a critical gap in agent memory systems by moving beyond "one-off heuristics" or "compute-heavy escalation" toward a systematic view of performance–cost control.

• Controllability: It provides a practical mechanism for explicit, controllable trade-offs, aligning with the industry trend of tiered compute options (e.g., "thinking" modes).
• Unified Comparison: By using a unified testbed, the work clarifies when different tiering strategies (implementation, reasoning, capacity) deliver the most favorable trade-offs, offering guidance for deployment under varying budget constraints.
• Runtime Focus: Unlike prior work focused on offline construction, BudgetMem focuses on runtime extraction, ensuring that computation is spent only when necessary and in proportion to the query's requirements.

The authors conclude that BudgetMem improves the usability and controllability of agent memory under real-world budget constraints, making memory computation explicitly manageable without sacrificing the quality of long-horizon reasoning.