Route to Remember: Adaptive Memory Selection for LLM Agents

AI summary

Overview

Research area: Long-term memory for LLM agents — specifically, whether factual memory can be stored in model weights (a per-user LoRA adapter) rather than in a text store that is retrieved into the prompt.

Technical level: Advanced. The paper assumes familiarity with on-policy distillation, KL-divergence objectives, LoRA adapters, and retrieval-augmented memory systems.

Scope: The paper trains a context-free "weight memory" via on-policy self-distillation (OPSD), compares it against three text-memory baselines on HaluMem and LoCoMo, characterizes how it forgets, and adds a verify-then-regenerate (VtR) router that combines weight and text memory.

What This Paper Is About

Production agents store user memory almost entirely as text: key points are extracted from dialogue, stored externally, and retrieved into the context window at query time — so every query pays a cost that grows with accumulated history. The authors ask whether episodic and semantic memory (user-specific facts) can instead be internalized into model parameters, so that a query can be answered with no context at all, and whether such a "weight memory" can match or beat text memory on recall. Because the two memory types fail on different questions, the paper also asks how to route between them.

Key Contributions

1. An evaluation of on-policy self-distillation (OPSD) for user-memory recall. On the question set for eight HaluMem-Long users, OPSD reaches 59.9%, versus 47.8% for MemOS-style retrieval, 51.9% for extracted context, and 51.3% for raw dialogue. The advantage over retrieval is positive for all eight users and holds at 4B, 9B, and 27B model scales.

2. A characterization of retention and training effects. Repeated questions distinguish context truncation, gradual forgetting in weights, and stable retrieval. Ablations identify training-objective and refusal-supervision effects within the tested recipes, and cost accounting separates query token usage from unmeasured deployment costs.

3. A routing study and VtR integration. Query-side and confidence signals provide limited correctness information in the authors' probes. Answer verification with text-based regeneration adds 6.2 points over OPSD on native HaluMem questions, improving five of six categories while reducing memory-conflict accuracy.

4. A recipe-level diagnosis of why earlier weight-memory attempts failed. Supervised fine-tuning on identical augmented data reaches only 36.9% — below every text baseline — indicating the training objective, not the parametric medium, was the binding constraint.

Main Findings

• OPSD beats every text baseline tested. Overall accuracy on HaluMem's native question set (n=1438) is 59.9% for OPSD, above raw dialogue (51.3%), extracted memory in context (51.9%), and MemOS retrieval (47.8%). The benchmark-provided ground-truth memory points in context, an oracle, reach 64.0%. OPSD also exceeds its own teacher configuration (the base model reading extracted memory in context, 51.9%).

• The margin over retrieval is large and user-consistent. 12.1 points, with a user-clustered confidence interval of [9.4, 14.8], positive for all eight users individually, ranging from 7.1 to 18.1 points. On matched questions, OPSD is uniquely correct on 309 questions against retrieval's 136.

• The objective, not the storage medium, drives the result. SFT on identical augmented data reaches 36.9% overall, below every text baseline. On LoCoMo, both text and weight memory crush SFT (16.9%), and OPSD edges extracted text memory overall (46.9% vs 46.4%).

• Weight memory is expensive to write and cheap to read. OPSD reads 326K input tokens per query against 1.7M for MemOS and roughly a billion for full-context methods, at 1.0 read GPU-hour versus 0.9 for MemOS and 22.5–23.4 for full-context systems. Against full context, roughly 150M extra write tokens are repaid after about 600 queries per stored memory; against retrieval the break-even is on the order of 10^5 queries.

• The three paradigms fail differently as questions age. Under the retention protocol that re-asks questions 20, 40, ..., 80 sessions later, raw dialogue collapses by age 40 and extracted memory holds until 60, while weight memory decays smoothly from 59.9% to 32.4% by age 80, and retrieval stays nearly flat (47.8% to 40.1%). Retrieval overtakes weight memory beyond roughly twenty sessions.

• Forgetting is non-uniform within weight memory. Between age 0 and age 100 the memory-boundary category drifts from 81.7% to 88.1% while every content category collapses — the model keeps the habit of declining but loses the content it protected.

• Absence is not inspectable in weights. HaluMem's memory-boundary category is the one where text memory wins: a text system can see that its retrieved list contains no relevant field, whereas parameters supply a plausible value. The paper notes about a fifth of boundary "hallucination" verdicts were in fact valid abstentions scored wrong, so the true fabrication rate is roughly 80% of the stated one.

• Refusal supervision is the sole source of abstention. Removing unknown-answer augmentation collapses the memory boundary from 88% to 1% and costs 14–16 points overall, with hallucination rising to 60%. Doubling the unknown/counterfactual budget or adding composite-QA augmentation is at best neutral.

• Other recipe levers matter less than expected. Decoupling the teacher from the accumulating student (using a frozen base model on capped cumulative memory) is worth roughly +14 points over the horizon; anti-forgetting replay of cached answers helps the tail but craters dynamic update; LoRA rank is flat from 16 to 128 (best at 64; 128 mildly overfits); full-parameter fine-tuning does not improve on the LoRA adapter.

• Query-side routing signals fail. The two systems disagree on 445 of 1438 questions, and an oracle picking the better one per question would reach 69.3%. But similarity carries no signal (0.735 vs 0.741 for weight memory; 0.737 vs 0.737 for retrieval), a classifier gains at most 0.9 points, and confidence correlates 0.65 with answer length. A length-constant probe reaches AUC 0.918 for answerability but only 0.516 against the event that the model is wrong.

• Verify-then-regenerate works where selection does not. VtR reaches 66.1%, +6.2 points over OPSD (p = 2×10^-6, interval [3.6, 8.6]) and 18.2 over the best text baseline, improving at every question age. The gate keeps 37.9% of answers and regenerates the rest; it flags every memory-boundary question, lifting that category to 99.4%. On LoCoMo it lifts adversarial accuracy to 73.1% and takes overall best at 60.4%.

• The student beats its teacher. The authors offer three possible explanations rather than a settled answer: consolidation over many paraphrases averages out extraction noise; the student never has to locate its answer among hundreds of distractor statements (full-context systems show a 33–34% hallucination rate); and the persistent adapter can retain facts that have dropped out of any single teacher window.

Methodology in Plain English

The researchers start from a shared intermediate representation. One base language model reads each session's dialogue and extracts key memory points, each a narrative sentence with a timestamp. Text memory stores and retrieves those sentences; this paper instead turns them into training data. For each memory point the model generates three kinds of question-answer pair: paraphrases (the same fact asked several ways), counterfactuals (the fact stated wrongly, requiring a correction), and unknowns (attributes the session never recorded, requiring an explicit refusal).

A per-user LoRA adapter — the student — is then trained to answer with no context in the prompt, while the teacher is the same base model conditioned on the user's memory. Crucially, the teacher sees only a recency-bounded window (the current session plus as many recent prior sessions as fit a fixed character budget), not the whole history. Training is on-policy: the student samples its own answer, the teacher scores exactly those tokens, and the loss minimizes the KL divergence between them on the sampled trajectory. Sampling from the student, rather than teacher-forcing the teacher's text, is what avoids repetition loops and exposure bias.

The adapter persists across sessions, so each session's training accumulates on top of previous ones. Memory is written once per session and read after every session through question answering.

For combining the two memory systems, the authors audit the weight memory's answer rather than pre-selecting a system per query. Given the student's answer, they retrieve from a merged text store twice — once using the question and once using the generated answer — and a cheap three-label verifier judges the answer as supported, conflicting, or unsupported. If supported, the weight memory's answer stands; otherwise the adapter-free base model re-answers from the union of both evidence sets. Retrieving with the generated answer is treated as a separate semantic-reasoning step that maps the question into the region where memory points live; merging superseded values matters because without it about a third of verdicts come back as conflicts, versus about one percent with merging.

Evaluation uses Claude Sonnet 4.6 as a judge with a correct/hallucinated/omitted rubric and no partial credit, paired McNemar tests over matched questions, and bootstrap confidence intervals clustered on users. All experiments run on NVIDIA H200 GPUs.

Why This Matters

Impact on research. The paper reframes a pessimistic prior result — that weight memory underperforms once context is removed — as a training-protocol problem rather than a storage-medium limit. Earlier discouraging findings were obtained almost entirely with supervised fine-tuning; the direct comparison here isolates the objective as the dominant lever. The retention and ablation studies also give the field a vocabulary for what weight memory is actually good at (recent content, stable habits) versus what it cannot do (inspect absence).

Real-world applications:

• Persistent personal assistants that must recall user preferences and project conventions across many sessions without re-reading an entire history on every turn.
• High-volume shared knowledge — a workflow thousands of users repeat, a tool-calling convention, or a recurring interaction style — where the write cost is amortized over very heavy query load.
• Multi-session customer-support or account agents where facts update and sometimes contradict (a changing address, for example), handled by routing updates through the merged text store while weights carry the bulk.
• Cost-sensitive deployments that currently spend over a billion input tokens across an evaluation on full-context prompting and could instead read only query tokens.

Industry relevance. The cost accounting is the practical argument: text systems write for free but read an entire context on every query, while weight memory front-loads cost into writing and then reads four orders of magnitude fewer tokens than full context and a factor of five below retrieval. The paper is explicit that this pays off for shared, repeatedly queried memory rather than a single user's transient history, and that the roughly 10^5-query break-even against retrieval means weight memory dominates full context almost immediately but only overtakes retrieval under very heavy, highly shared query load.

Future Directions

• Pushing weight memory past its two intrinsic limits. Weight memory still forgets older facts along with new updates and cannot tell whether a fact was never recorded. The paper isolates refusal supervision as the lever for the second problem but leaves the forgetting curve itself open.

• Extending the objective beyond text. The authors propose applying the same distillation objective to procedural and skill memory, to tool use, and — with a multimodal teacher — to knowledge that never had a clean textual form, so an agent accumulates capability rather than merely facts.

• Internalizing shared regularities rather than single histories. Because write cost is paid once and read cost is negligible, the paper argues the natural unit to internalize is a pattern mined across many users, with per-user memory layered on top.

• Better routing signals. Every query-side signal tested here — similarity, learned query features, model consistency, answer confidence — failed to predict correctness, and even the clean length-constant answerability probe collapses under the retention protocol once the threshold drifts. The paper leaves open whether any pre-generation signal can work, and offers only hypotheses for why the student outperforms its teacher.

Target Audience

Researchers and engineers working on agent memory, long-context systems, or parameter-efficient adaptation who need to decide between retrieval-based and parametric memory for persistent, multi-session deployments. It is most useful to readers already comfortable with LoRA, distillation objectives, and retrieval pipelines — the paper is dense with per-category tables and ablation sweeps — but the high-level finding (the training objective, not the storage medium, is the constraint) is accessible to practitioners choosing a memory architecture.

Authors’ abstract

The memory of production agents is increasingly becoming a bottleneck and is built almost exclusively from text: key points are extracted from out-of-limit context and retrieved into context at query time. Therefore, every new query incurs an additional cost that grows with the accumulated history. Weight memory is a complementary alternative that internalizes knowledge into model parameters and answers with no context at all, and has recently proven effective in self-evolving agents for internalizing skills and behaviors. However, its application in episodic and semantic memory, such as user-specific facts, remains comparatively unexplored, and a gap remains between earlier attempts at weight memory and text-based systems regarding recall performance. We study how to make weight memory competitive as a long-term factual store by applying on-policy self-distillation (OPSD). It trains a context-free student from a teacher, which is the same model conditioned on the user's memory, together with a verify-then-regenerate router that lets weight and text memory each cover the queries they handle best. On the HaluMem benchmark, OPSD matches or surpasses strong text-memory baselines, including full-context prompting and retrieval systems, while reading fewer tokens per query. We further conduct an extensive study that characterizes the behavior of weight memory and analyzes how different training recipes affect its performance. As agent harnesses move toward persistent, multi-session deployments, these results position weight memory as a practical, low-cost component of agent memory and provide insights into feasible training and evaluation protocols for combining it with retrieval.