Learning to Remember: End-to-End Training of Memory Agents for Long-Context Reasoning

AI summary

Overview

Research area: Long-context reasoning with memory-augmented LLM agents, reinforcement learning for agent training, and long-horizon state tracking benchmarks.

Technical level: Advanced. The paper formulates memory management as a Markov Decision Process, introduces a custom tool API (CRUD-style memory operations), and trains a single policy with a modified GRPO algorithm (Task-Stratified GRPO) on 32 NVIDIA H200 GPUs. Familiarity with RL fine-tuning (GRPO/PPO-style objectives, advantage normalization) is assumed.

Scope in one sentence: The paper introduces Unified Memory Agent (UMA), one reinforcement-learned policy that builds a query-agnostic structured Memory Bank from a chunked stream and then answers many future questions from that shared memory, together with Ledger-QA, a synthetic benchmark for long-horizon state tracking over accumulated updates.

What This Paper Is About

Long-context LLMs and Retrieval-Augmented Generation (RAG) both defer reasoning to query time: they reprocess context for every question, which the authors argue is brittle when facts evolve and answers depend on latent states. The goal is to instead construct a reusable external memory once from a stream of chunks and share it across multiple future question-answering sessions — the "one-to-many" setting where memory is built before the questions are revealed. To make this trainable, the authors need a learning signal that connects query-agnostic memory construction to how useful that memory turns out to be downstream.

Key Contributions

1. Unified Memory Agent (UMA): A single policy that operates in two phases — Phase I incrementally maintains a structured Memory Bank plus a core summary through CRUD operations (Add, Update, Delete, Retrieve, List, UpdateCore) over chunks, and Phase II answers queries using both structured bank retrieval and raw-context retrieval (BM25, Embedding).
2. Task-Stratified GRPO: A training algorithm that samples memory trajectories, branches each resulting memory state into multiple QA trajectories, uses the mean QA reward from each memory state to supervise memory maintenance, and normalizes memory rollouts and per-question QA rollouts in separate groups.
3. Ledger-QA: A controllable synthetic diagnostic benchmark for continuous state tracking over accumulated updates, built from date-stamped consumption scenes between 2023-01-01 and 2024-12-31, with questions across 8 categories that require multi-step aggregation rather than local span matching.
4. Empirical demonstration across 13 benchmarks plus Ledger-QA, including ablation variants (w/o Phase I, w/o RL, w/o Raw-Context, Global Group, Two-Stage) and session-budget scaling (8k/16k/32k).

Main Findings

• Highest aggregate score at the 16k budget: UMA-Generalist (16k) reaches an average of 75.86 across the 13 test-time-learning (TTL) and accurate-retrieval (AR) benchmarks, the best among compared methods. For reference, Mem-α scores 63.02, A-MEM 58.75, RAG 53.52, and Concat 50.30.
• TTL performance: UMA-Generalist averages 82.70 across the six TTL tasks, the highest six-task average at the 16k budget, leading four tasks. It also performs well on Banking77 (87.00) and Clinic (91.00), which the authors state are unseen during training, indicating transfer to new classification domains.
• AR performance: UMA-Generalist obtains the highest seven-task average of 70.00 and ranks among the top three methods on five tasks, leading ConvoMem (87.05) and ranking top-three on HotpotQA (76.56), MSC (66.20), PerLTQA (74.50), and SQuAD (87.53).
• Session-budget scaling: Moving from 8k (71.45 average) to 16k (75.86) clearly improves performance; 32k (76.76) gives only modest additional gains at higher compute and memory cost, so 16k is retained as the default.
• Memory maintenance matters: After RL, full UMA exceeds the w/o Phase I variant by 13.3 points on average; before RL, enabling write/update operations already outperforms the retrieval-only variant.
• RL training matters: RL raises full UMA by 16.2 points on average, and also strengthens the retrieval-only configuration.
• Stratified grouping matters: The Global Group variant, which normalizes all memory and QA sessions in one global group, is 5.9 points lower on average.
• End-to-end beats two-stage: Two-stage training, which optimizes memory and QA separately, is 11.2 points lower on average.
• Ledger-QA transfer vs. adaptation: Without Ledger-QA training, UMA-Generalist reaches 38.18 LLM-as-a-Judge and 34.39 EM (vs. 30.45 and 19.09 for UMA w/o RL). MemAgent records a higher Judge score of 40.52 under its query-conditioned protocol, but UMA-Generalist reuses one query-hidden memory state across all questions and achieves higher EM, F1, and ROUGE-L. Task adaptation raises UMA-Specialist to 53.02 Judge and 49.39 EM.
• Long-horizon decay of baselines: RAG and MemAgent-woq are competitive at H=2 but fall to 5.77 and 2.88 at H=500; even query-conditioned MemAgent drops to 11.54. UMA-Specialist scores 25.00 at 500 sessions, matching UMA (w/o Raw-Context) and exceeding the strongest non-UMA baseline.
• Arithmetic control: In a 100,000-example arithmetic-only control where the model sums five random two-decimal numbers below 1000, accuracy changed little after UMA-Specialist training (base 25.41% vs. UMA-Specialist 25.14%), suggesting Ledger-QA gains come from memory maintenance and state access rather than improved standalone arithmetic.
• Structured memory is the default access path: Across 952 Ledger-QA QA traces, UMA-Specialist accesses structured memory in every trace: 878 (92.2%) use only Memory Bank tools, 74 (7.8%) additionally invoke raw-context retrieval, and no trace uses raw-context retrieval alone.

Methodology in Plain English

The authors treat long-document reasoning as a sequential decision problem. A long input is cut into a stream of chunks. The agent's state at any step consists of two memory components — a high-level summary ("core memory") and a structured key-value Memory Bank — plus whatever it is currently focused on (either the current chunk or a user query) and a short buffer of recent tool calls and results. Because only these pieces are fed to the model, the system does not need the entire document in the context window.

In the first phase, the agent reads one chunk at a time and calls tools to add, update, or delete entries in the Memory Bank, optionally reading them back. When it stops calling tools and just writes plain text, the controller treats that text as a commit: the core summary is updated and the agent advances to the next chunk. In the second phase, focus switches to a query, and the agent can use lexical search (BM25), dense embedding search, or direct lookups into the Memory Bank to gather evidence before producing a final answer.

Training is where the distinctive idea lies. For each context, the system samples 16 memory trajectories, each producing a final memory state. It then samples a set of questions once per context and reuses them across all memory states, generating one multi-turn QA session per (memory state, question) pair. Each QA session gets a reward combining a tool-use quality term and answer correctness. Crucially, each memory session is credited with the average outcome of all QA sessions branching from its state — so good memory is rewarded for supporting many future questions, not one known target. The GRPO advantages are then normalized in two separate groups: all memory sessions together, and QA sessions grouped by the specific question they answer (so cross-question difficulty differences are removed). A single objective combines both session types back into one update. The tool reward term is weighted by a fixed λ = 0.5, and a small ε_std = 10⁻⁶ is used to avoid division by zero.

Implementation: dense retrieval uses sentence-transformers/all-MiniLM-L6-v2 served via infinity; training builds on customized veRL, initializes from Qwen/Qwen3-4B-Instruct-2507, and runs on 32 NVIDIA H200 GPUs. Evaluation uses an LLM-as-a-Judge protocol with DeepSeek-V3.2, plus deterministic lexical metrics via a shared final-answer extraction and normalization pipeline.

Why This Matters

Impact on research. The paper reframes memory construction as an optimization problem with a one-to-many credit assignment structure: a memory state is good if it helps many unknown future questions. Task-Stratified GRPO offers a concrete template for supervising upstream state-building from downstream outcomes, which is directly relevant to other agentic pipelines where intermediate artifacts serve multiple later tasks. Ledger-QA also adds a diagnostic that isolates accumulated-state tracking, a capability that retrieval benchmarks and LongMemEval-style evaluations do not directly stress.

Real-world applications.

• Personal financial or expense tracking, where a system must maintain a running ledger and answer aggregation questions ("total spending") whose answers appear in no single entry — the example the paper uses in Figure 1.
• Assistive agents for clinical or customer-support histories, where records accumulate and answers depend on the current consolidated state rather than one passage.
• Enterprise knowledge bases over evolving documents (policies, contracts, tickets) where facts are updated and superseded over time.
• Long-running conversational assistants that must remember preferences and prior decisions across many sessions without re-reading everything at each turn.

Industry relevance. The efficiency argument is practical: RAG reprocesses retrieved passages per query, while a maintained memory is built once and reused. The paper's budget scaling results (8k → 16k → 32k) give a concrete cost/accuracy trade-off for deployment, and the observation that a 4B-parameter backbone can be trained for this role is relevant to teams that cannot serve very large long-context models.

Future Directions

• Multi-agent and evolving tool interfaces: The limitations section states the work studies single-agent settings with a fixed tool interface; collaborative agents and evolving APIs are explicitly left for future work.
• Broader real-world validation: Ledger-QA is a controlled synthetic diagnostic with transparent generation rules; the authors state that validation on real-world proactive state-maintenance tasks is still needed and that strong synthetic performance does not guarantee robustness in noisier scenarios.
• Finer-grained credit assignment within sessions: Because downstream outcomes supervise memory maintenance at the session level, all actions within one tool-calling session share an advantage. Attributing contribution to individual CRUD operations would require finer-grained or counterfactual attribution.
• Latency and scaling under other conditions: The authors note that latency and scaling measurements reflect the reported hardware and workloads, leaving characterization on other hardware and workloads open.
• Privacy and retention mechanisms: The ethical discussion raises the need for informed consent, content minimization, retention limits, access controls, and user-facing inspection, correction, and deletion — including deletion of derived retrieval artifacts such as indexes and caches.

Target Audience

This paper is aimed at machine learning researchers and engineers working on long-context modeling, retrieval-augmented generation, and RL-trained LLM agents. It will be most useful to readers already comfortable with group-relative policy optimization, agentic tool-calling loops, and benchmark design, and to practitioners evaluating whether to replace query-time retrieval with a maintained external memory. Readers looking for a beginner-friendly introduction to memory agents will find the architecture description accessible but the training formulation demanding.

Funding note as reported: The work was supported by the National Natural Science Foundation of China (General Program, Grant No. 62376260) and the Beijing Natural Science Foundation (Grant No. L257006). The authors state ChatGPT (OpenAI) was used for language polishing and manuscript consistency checks, and as an auxiliary tool during reference review, with all references manually verified. Code is stated to be available at https://github.com/ictnlp/unified-memory-agent.git.

Authors’ abstract

Long-context LLMs and Retrieval-Augmented Generation defer state tracking and evidence consolidation to query time, which is brittle when facts evolve and answers depend on latent states. We introduce Unified Memory Agent (UMA) for a one-to-many setting: query-agnostic external memory is constructed once from a stream and reused across multiple future QA sessions. A single policy maintains a structured Memory Bank through CRUD operations and answers using both the Memory Bank and raw context. Task-Stratified GRPO uses the mean reward of QA trajectories branching from each sampled memory state to supervise memory maintenance, while normalizing memory and per-question QA groups separately. We also introduce Ledger-QA, a diagnostic benchmark for long-horizon state tracking over accumulated updates. At the 16k budget, UMA-Generalist achieves the highest average score among compared methods across the test-time-learning and accurate-retrieval benchmarks and transfers to Ledger-QA without task-specific training; UMA-Specialist further improves long-horizon tracking after task adaptation. These results support learned proactive memory management for long-context reasoning.