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

arXiv: https://arxiv.org/abs/2602.06025 · HTML: https://mlsanity.com/paper/2602.06025/

## Key facts

- Cited by 13 papers in the map (data through 14 Oct 2026).
- Placed in Agent memory systems, one of 16 directions in LLM agents and tool use.
- Builds on 40 of its 47 references that are in the map.
- First posted on arXiv 5 Feb 2026.

## Placed in

- LLM agents and tool use / [Agent memory systems](https://mlsanity.com/agents/agent-memory-systems/)

## Builds on (most cited first, 30 of 40)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [Self-Consistency Improves Chain of Thought Reasoning in Language Models](https://mlsanity.com/paper/2203.11171/) (2022)
- [Reflexion: Language Agents with Verbal Reinforcement Learning](https://mlsanity.com/paper/2303.11366/) (2023)
- [HotpotQA: A Dataset for Diverse, Explainable Multi-hop Question Answering](https://mlsanity.com/paper/1809.09600/) (2018)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Self-Refine: Iterative Refinement with Self-Feedback](https://mlsanity.com/paper/2303.17651/) (2023)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models](https://mlsanity.com/paper/2211.10438/) (2022)
- [Unsupervised Dense Information Retrieval with Contrastive Learning](https://mlsanity.com/paper/2112.09118/) (2021)
- [SparseGPT: Massive Language Models Can Be Accurately Pruned in One-Shot](https://mlsanity.com/paper/2301.00774/) (2023)
- [A-MEM: Agentic Memory for LLM Agents](https://mlsanity.com/paper/2502.12110/) (2025)
- [OpenAI GPT-5 System Card](https://mlsanity.com/paper/2601.03267/) (2025)
- [On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes](https://mlsanity.com/paper/2306.13649/) (2023)
- [SnapKV: LLM Knows What You are Looking for Before Generation](https://mlsanity.com/paper/2404.14469/) (2024)
- [Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads](https://mlsanity.com/paper/2401.10774/) (2024)
- [H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models](https://mlsanity.com/paper/2306.14048/) (2023)
- [Evaluating Very Long-Term Conversational Memory of LLM Agents](https://mlsanity.com/paper/2402.17753/) (2024)
- [A Simple and Effective Pruning Approach for Large Language Models](https://mlsanity.com/paper/2306.11695/) (2023)
- [Mem0: Building Production-Ready AI Agents with Scalable Long-Term Memory](https://mlsanity.com/paper/2504.19413/) (2025)
- [LongMemEval: Benchmarking Chat Assistants on Long-Term Interactive Memory](https://mlsanity.com/paper/2410.10813/) (2024)
- [LLM-Pruner: On the Structural Pruning of Large Language Models](https://mlsanity.com/paper/2305.11627/) (2023)
- [EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty](https://mlsanity.com/paper/2401.15077/) (2024)
- [Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs](https://mlsanity.com/paper/2310.01801/) (2023)
- [LLM-QAT: Data-Free Quantization Aware Training for Large Language Models](https://mlsanity.com/paper/2305.17888/) (2023)
- [Confident Adaptive Language Modeling](https://mlsanity.com/paper/2207.07061/) (2022)
- [Memory in the Age of AI Agents](https://mlsanity.com/paper/2512.13564/) (2025)

## Built on by (most cited first, 13 of 13)

- [Externalization in LLM Agents: A Unified Review of Memory, Skills, Protocols and Harness Engineering](https://mlsanity.com/paper/2604.08224/) (2026)
- [MemoryCD: Benchmarking Long-Context User Memory of LLM Agents for Lifelong Cross-Domain Personalization](https://mlsanity.com/paper/2603.25973/) (2026)
- [EMBER: Efficient Memory via Budgeted Evidence Retention for Long-Horizon Agents](https://mlsanity.com/paper/2606.05894/) (2026)
- [SmartSearch: How Ranking Beats Structure for Conversational Memory Retrieval](https://mlsanity.com/paper/2603.15599/) (2026)
- [Learning What to Remember: Observability-Safe Memory Retention via Constrained Optimization for Long-Horizon Language Agents](https://mlsanity.com/paper/2606.10616/) (2026)
- [Self-Evolving Agents as Dynamic Graph Transformation: A Survey and New Perspective](https://mlsanity.com/paper/2608.18104/) (2026)
- [LLM Agents Are Latent Context Managers: Eliciting Self-Managed Context via State Proprioception](https://mlsanity.com/paper/2606.30005/) (2026)
- [Exploring Cross-Scenario Generality of Agentic Memory Systems: Diagnostics and a Strong Baseline](https://mlsanity.com/paper/2606.04315/) (2026)
- [What Eviction Destroys: A Restore-Counterfactual Audit of Forgetting in Agent Memory](https://mlsanity.com/paper/2609.08279/) (2026)
- [BudgetBench: A Budget-Tiered Protocol and Pilot Harness for Memory Strategy Evaluation in Local Large Language Model Agents](https://mlsanity.com/paper/2609.13149/) (2026)
- [Contract Memory Compiler: Resolve, Then Traverse](https://mlsanity.com/paper/2609.32658/) (2026)
- [When Should Agents Check External State? Budgeting Observations for Stored Intentions](https://mlsanity.com/paper/2609.37125/) (2026)
- [MemPilot: Orchestrating On-Demand Multimodal Memory Curation for LLM Agents](https://mlsanity.com/paper/2610.06830/) (2026)

## Abstract (arXiv)

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 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., textscLow/textscMid/textscHigh). 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.