🤖 AI

Did You Check the Right Pocket? Cost-Sensitive Store Routing for Memory-Augmented Agents

This paper proposes formulating memory retrieval as a cost-sensitive store-routing problem to enable memory-augmented agents to selectively query specialized stores, thereby improving both answer accuracy and token efficiency compared to uniform retrieval.

Original authors: Madhava Gaikwad

Published 2026-03-18

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Madhava Gaikwad

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a brilliant but very expensive personal assistant (the AI) who is trying to answer your questions. To do their job, this assistant has access to four different filing cabinets in your office:

The Sticky Note (Short-Term): What you just said five minutes ago.
The ID Card (Summary): Your name, phone number, and favorite pizza topping.
The Diary (Long-Term): A summary of what you talked about last week or last year.
The Tape Recorder (Episodic): The raw, unedited audio transcripts of every conversation you've ever had.

The Problem: "Did You Check the Right Pocket?"

Currently, most AI systems are like a clumsy assistant who, whenever you ask a question, dumps all four filing cabinets onto your desk at once.

If you ask: "What's my phone number?"
The Clumsy Assistant: Opens the Diary, the Tape Recorder, the Sticky Note, and the ID Card.
The Result: Your desk is now covered in thousands of pages of irrelevant text. Your assistant has to read through all of it to find the phone number. This is slow, costs a lot of money (because reading takes time and computing power), and the sheer amount of noise makes it harder for the assistant to find the right answer.

This paper asks a simple question: Why don't we just check the right pocket first?

The Solution: The "Smart Router"

The authors propose a new system called Cost-Sensitive Store Routing. Think of this as a smart receptionist who stands at the door before the assistant starts working.

When you ask a question, the receptionist looks at it and decides, "Okay, this is about your phone number. We only need the ID Card. Let's ignore the Diary and the Tape Recorder."

They tested three approaches:

The "Dump Everything" Method: Always open all cabinets. (High cost, often messy).
The "Oracle" Method: A magical receptionist who knows exactly which cabinet holds the answer every single time. (Perfect efficiency, but impossible in real life).
The "Smart Heuristic" Method: A receptionist using simple rules (e.g., "If the question says 'yesterday,' check the Diary").

What They Found (The "Aha!" Moments)

1. Less is More (The "Needle in a Haystack" Effect)
When the assistant is forced to read 787 pages of text to find one fact, they get confused. It's like trying to find a specific grain of sand on a beach.

The Result: When the system only opened the necessary cabinets, the assistant got more accurate answers (86% vs 81%) while using 62% less "reading time" (tokens).
The Metaphor: It's easier to find your keys if you only look in the kitchen, rather than searching the kitchen, the garage, the car, and the attic simultaneously.

2. Long Contexts are Dangerous
The problem gets worse when the "Diary" and "Tape Recorder" are huge. If you ask about something from last year, and the system dumps 1,000 pages of old conversations onto the desk, the assistant might get distracted by conflicting stories.

Example: The "Summary" cabinet says your manager is "Jennifer." The "Diary" says she used to be "Michael" before a promotion. If you show the assistant both, they might get confused and guess "Michael." If you only show the "Summary," they get it right.

3. Simple Rules Work, But Learning is Better
The authors built a simple rule-based receptionist (e.g., "If the word 'before' is used, check the Diary"). This worked pretty well (94% success rate in picking the right cabinet).

However: Even when the receptionist picked the right cabinet, the assistant didn't always get the answer right. This suggests that while picking the right files is crucial, the assistant also needs to be better at reading them.
The Future: The paper argues that instead of hard-coding rules, we should train an AI to learn how to pick the right cabinets automatically, just like a human learns which drawer to open based on experience.

The Bottom Line

This paper proves that being selective is smarter than being thorough.

In the world of AI, throwing more data at the problem doesn't always make it better. Sometimes, it just makes it slower, more expensive, and more confusing. By teaching AI to "check the right pocket" before digging in, we can build systems that are cheaper to run, faster to respond, and actually smarter at answering your questions.

In short: Don't read the whole library to find a phone number. Just open the directory.

1. Problem Statement

Memory-augmented agents typically maintain multiple specialized memory stores (e.g., Short-Term Memory, Long-Term Memory, Summaries, Episodic Transcripts). However, most existing frameworks employ uniform retrieval, where the system queries all stores for every user query regardless of relevance.

This approach introduces two critical issues:

Computational Waste: Retrieving from irrelevant stores consumes tokens and infrastructure resources unnecessarily.
Accuracy Degradation: Irrelevant or noisy context reduces the signal-to-noise ratio, hindering the Large Language Model's (LLM) ability to locate the correct answer, particularly in long-context scenarios.

The paper posits that store selection should be treated as a distinct, first-class decision problem (routing) that occurs before retrieval, rather than relying on the LLM to filter noise during generation.

2. Methodology

2.1 Memory Architecture & Problem Formulation

The authors define a memory system with four distinct stores:

STM (Short-Term Memory): Current conversation turns.
Sum (Summary Store): Compressed user facts (preferences, bio).
LTM (Long-Term Memory): Summaries of past conversations.
Epi (Episodic Memory): Raw transcripts.

The Routing Problem is formalized as selecting a subset of stores $\hat{G} \subseteq S$ for a given query $q$ . The goal is to maximize downstream answer accuracy while minimizing retrieval cost (tokens).

2.2 Decision-Theoretic Framework

The authors formalize store routing as a cost-sensitive subset-selection problem:
$\pi^*(q) = \arg \max_{G \subseteq S} \left[ \mathbb{E}[\text{Acc}(q, G)] - \lambda \sum_{s \in G} c_s \right]$
Where:

$\text{Acc}(q, G)$ is the expected accuracy given retrieved stores $G$ .
$c_s$ is the cost of accessing store $s$ .
$\lambda$ controls the trade-off between accuracy and cost.

2.3 Evaluation Metrics

Three primary metrics are used to evaluate routing policies:

Coverage: Did the router include all necessary stores? (Binary: 1 if $G_{true} \subseteq \hat{G}$ ).
Exact Match (EM): Did the router select exactly the required stores? (Binary: 1 if $G_{true} = \hat{G}$ ).
Waste: The number of unnecessary stores retrieved (proxy for token cost).

2.4 Experimental Setup

The evaluation is conducted in two stages:

Synthetic Routing Evaluation: Using 1,000 synthetic queries with ground-truth labels derived from query taxonomies (e.g., "temporal" queries map to LTM+Epi). This isolates the routing decision from retrieval quality.
LLM-Based QA Evaluation: Testing 12 routing policies on 150 real questions using GPT-3.5-turbo and GPT-4o-mini. The dataset includes "Short" (200 tokens/store) and "Long" (1000 tokens/store) context regimes.

Policies Tested:

Uniform: Retrieve from all stores.
Oracle: Retrieve only ground-truth necessary stores (Upper Bound).
Fixed Subset: e.g., STM + Sum + LTM.
Hybrid Heuristic: A rule-based system using semantic signals (quantity, temporal, multi-hop) with a conservative fallback (Sum + LTM).

3. Key Contributions

Formalization of Store Routing: Defined memory access as a cost-sensitive routing problem, decoupling store selection from within-store ranking.
New Metrics: Introduced Coverage, Exact Match, and Waste metrics specifically for evaluating multi-store selection quality.
Empirical Evidence: Demonstrated that selective retrieval outperforms uniform retrieval in both efficiency and accuracy.
Decision Framework: Provided a theoretical framework explaining the accuracy-cost trade-off, showing that uniform retrieval is a special case where cost sensitivity is zero.

4. Key Results

4.1 Synthetic Routing Performance

Uniform Retrieval: Achieved 100% coverage but only 8% Exact Match (high waste).
Hybrid Heuristic: Achieved 94% coverage and 58% Exact Match, significantly outperforming simple linguistic rules (57% coverage) while maintaining low waste.
Key Insight: Semantic signals (e.g., "list all," "before," "compare") are far more effective than shallow linguistic features (pronouns, tense) for routing.

4.2 Downstream QA Performance (LLM Evaluation)

Accuracy vs. Efficiency: The Oracle policy achieved 86.7% accuracy using only 299 tokens on average.
Uniform Retrieval: Achieved only 81.3% accuracy while using 787 tokens (a 62% increase in cost).
Long-Context Penalty: In the "Long" condition (1000 tokens/store), the gap widened. Oracle reached 72% accuracy vs. Uniform's 60%. Irrelevant context in long windows significantly distracts the model ("Needle in a Haystack" effect).
Fixed Subsets: A simple fixed policy (STM+Sum+LTM) achieved 84.7% accuracy with moderate cost, suggesting that even non-adaptive strategies can capture most benefits if the architecture is well-structured.
Heuristic Gap: While the Hybrid Heuristic had high routing coverage (94%), its downstream QA accuracy (70.7%) lagged behind the Oracle (86.7%). This indicates that routing errors (missing a store) and extraction errors (model failing to use correct context due to noise) both contribute to performance gaps.

4.3 Why Uniform Underperforms

Signal Dilution: Adding irrelevant stores increases token count, diluting the signal.
Conflicting Information: Different stores may contain outdated or contradictory facts (e.g., old manager vs. new manager), causing the model to hallucinate or select the wrong source.

5. Significance and Future Directions

Paradigm Shift: The paper argues that "checking the right pocket" (selecting the right memory store) is as critical as the retrieval or generation steps themselves.
Cost Reduction: Selective routing can reduce context tokens by ~62% without sacrificing (and often improving) accuracy, directly lowering inference costs.
Scalability: As agents accumulate more memory, uniform retrieval becomes computationally prohibitive. Routing is essential for scalable multi-store systems.
Future Work: The gap between heuristic routing (70%) and oracle routing (86%) suggests that learned routing policies (end-to-end optimization) are necessary to fully realize the potential of memory-augmented agents, moving beyond rule-based heuristics.

In conclusion, the paper establishes that cost-sensitive store routing is a fundamental component of agent design, offering a principled way to balance retrieval cost against answer quality, and demonstrating that "less context" (when selected correctly) often leads to "better answers."