PaperlayerView exact revision on arXiv

Source-linked AI summary

Memory-R1: Enhancing Large Language Model Agents to Manage and Utilize Memories via Reinforcement Learning

Sikuan Yan, Xiufeng Yang, Zuchao Huang, Ercong Nie, Zifeng Ding, Zonggen Li, Xiaowen Ma, Jinhe Bi, Kristian Kersting, Jeff Z. Pan, Hinrich Schütze, Volker Tresp, Yunpu Ma

arXiv:2508.19828v5cs.CLcs.MA

TL;DR

LLMs are stateless and external-memory pipelines generally lack learned decisions about what to store, update, or retrieve. Memory-R1 trains a Memory Manager and Answer Agent with outcome-driven PPO or GRPO, achieving strong benchmark performance with 152 training QA pairs and broader generalization. The approach remains scoped to dialogue-centric data and separately trained agents.

  • Problem

    LLMs have finite context windows, while most external-memory pipelines rely on static heuristics for storage, updating, retrieval, and selection.

  • Method

    Memory-R1 uses a Memory Manager for structured memory operations and an Answer Agent for memory distillation and reasoning, fine-tuned with PPO or GRPO.

  • Results

    Memory-R1 achieves state-of-the-art LoCoMo results with 152 training QA pairs, generalizes to MSC and LongMemEval without retraining, and scales across model sizes.

  • Takeaways & Limitations

    Outcome-driven reinforcement learning is a promising direction for adaptive and agentic memory in LLMs.

  • Takeaways & Limitations

    Evaluation focuses on dialogue-centric datasets, and separately training the two agents makes the process less straightforward.

Abstract

from arXiv · show

Large Language Models (LLMs) have demonstrated impressive capabilities across a wide range of NLP tasks, but they remain fundamentally stateless, constrained by limited context windows that hinder long-horizon reasoning. Recent efforts to address this limitation often augment LLMs with an external memory bank, yet most existing pipelines are static and heuristic-driven, lacking a learned mechanism for deciding what to store, update, or retrieve. We present Memory-R1, a reinforcement learning (RL) framework that equips LLMs with the ability to actively manage and utilize external memory through two specialized agents: a Memory Manager that learns structured operations, including ADD, UPDATE, DELETE, and NOOP; and an Answer Agent that pre-selects and reasons over relevant entries. Both agents are fine-tuned with outcome-driven RL (PPO and GRPO), enabling adaptive memory management with minimal supervision. With only 152 training QA pairs, Memory-R1 outperforms strong baselines and generalizes across diverse question types, three benchmarks (LoCoMo, MSC, LongMemEval), and multiple model scales (3B-14B).

1 Introduction

Memory-R1 addresses the limits of stateless LLMs by learning to manage external memory and select relevant retrieved entries for reasoning. Its RL-trained agents improve memory consolidation and answer quality with minimal supervision.

  • LLMs forget information outside finite context windows, limiting knowledge retention across long conversations and evolving tasks.
  • Heuristic memory systems struggle both to retrieve the right amount of context and to decide what information to remember, update, or discard.
  • In the two-dog example, the RL-trained manager consolidates information with UPDATE while the Answer Agent reduces 60 retrieved memories to the relevant fact.
  • Memory-R1 uses a Memory Manager for structured memory operations and an Answer Agent that filters RAG-retrieved memories before reasoning over them.
  • With only 152 question–answer pairs, LLaMA-3.1-8B-Instruct Memory-R1-GRPO improves relative F1 by 28%, BLEU-1 by 34%, and LLM-as-a-Judge by 30% over Mem0 on LoCoMo.
  • The paper contributes an RL framework using PPO and GRPO for adaptive memory operations, retrieval distillation, and answer reasoning.

2 Related Work

Prior work combines external memory with LLMs and increasingly applies RL to structured agent behavior, but memory operation selection and utilization remain comparatively underexplored.

  • External-memory LLM systems support persistent knowledge through selective storage, retrieval, and updating for long-horizon reasoning.
  • RL research has extended beyond supervised fine-tuning to tool use, web search, and reasoning-path selection through dynamic feedback.
  • Existing memory-augmented LLM systems typically use heuristics for memory operations, leaving adaptability and long-term optimization underexplored in RL settings.

3 Method

Memory-R1 uses two RL-fine-tuned agents to manage an external memory bank and answer questions from retrieved memories. The Memory Manager selects memory operations, while the Answer Agent distills retrieved entries before generating answers.

  • Framework overview: The framework separates memory construction and updating from question answering into two stages handled by specialized agents.Stage 1 uses the Memory Manager; Stage 2 uses the Answer Agent.
  • Memory Manager: The Memory Manager takes new information and the current memory bank as input, then outputs an operation and updated memory content.Its operation set is ADD, UPDATE, DELETE, or NOOP.
  • Memory Manager: PPO trains the Memory Manager with answer-based rewards, advantages, importance ratios, and clipped updates for stable optimization.The importance ratio compares current and old policy probabilities, while clipping limits update size.
  • Memory Manager: GRPO alternatively samples grouped candidate actions and computes relative advantages without an explicit value function.It retains PPO-style stability while comparing candidate actions within each group.
  • Memory Manager: Memory Manager operations are trained from their downstream effect on answer correctness rather than from manual operation labels.After applying an operation, the updated memory is passed to a frozen Answer Agent and rewarded using exact match.
  • Answer Agent: The Answer Agent retrieves 60 candidate memories, distills the most relevant entries, and generates an answer using PPO or GRPO fine-tuning.Its reward is the exact-match score against the ground-truth answer, directly tying optimization to final answer correctness.

4 Experiments

Memory-R1 is evaluated across models, benchmarks, metrics, components, and reward designs. It consistently improves performance, generalizes zero-shot beyond LoCoMo, and benefits from reinforcement-learning-based memory management and answer selection.

  • Evaluation Setup: The evaluation covers LoCoMo, MSC, and LongMemEval, using LLaMA-3.1-8B-Instruct and Qwen-2.5 Instruct backbones from 3B to 14B.LoCoMo includes single-hop, multi-hop, open-domain, and temporal questions; models are trained only on LoCoMo and evaluated zero-shot on MSC and LongMemEval.
  • Main Results: Memory-R1 consistently achieves new state-of-the-art performance across both model families on LoCoMo.The comparison includes RAG, A-Mem, Mem0, MemoryOS, and Memory-SFT baselines across F1, BLEU-1, and LLM-as-a-Judge metrics.
  • Generalization and Scalability: Memory-R1 remains effective across Qwen-2.5 model sizes, with PPO and GRPO outperforming base models on F1, BLEU-1, and J.The gains persist at 3B, 7B, and 14B, and the LoCoMo-trained pipeline improves all three metrics on MSC and LongMemEval without retraining.
  • Ablation Studies: Removing the Memory Manager, Answer Agent, or Memory Distillation reduces performance, indicating that each component contributes to the full pipeline.For example, removing the RL-fine-tuned Memory Manager reduces PPO scores from 41.0, 32.9, and 57.5 to 34.5, 28.1, and 49.0 on F1, BLEU-1, and J.
  • Ablation Studies: GRPO converges faster initially than PPO, but both methods reach comparable final reward levels.An EM reward is adopted because J-based rewards encourage verbose answers that perform poorly on F1 and BLEU-1; learned distillation also improves accuracy with lower latency than reranking.

5 Conclusion

Memory-R1 uses reinforcement learning to manage external memory and distill memories for answering. With 152 training examples, it achieves state-of-the-art LoCoMo results, scales across model sizes, and generalizes to MSC and LongMemEval without retraining.

  • Conclusion: Memory-R1 learns memory operations, memory distillation, and memory usage for answering rather than relying on heuristic pipelines.Its framework contains specialized mechanisms for managing and utilizing external memory.
  • Conclusion: With only 152 training examples, Memory-R1 achieves state-of-the-art results on LoCoMo and generalizes to MSC and LongMemEval without retraining.The conclusion also reports scaling across model sizes.
  • Conclusion: The results identify reinforcement learning as a promising direction for adaptive and agentic memory in LLMs.Ablation studies confirm improvements across every component of the system.

Limitations

The evaluation centers on dialogue-centric datasets, and extending Memory-R1 to multimodal data may introduce challenges beyond this work’s scope. The two agents are trained separately for stability under sparse rewards, making the process less straightforward.

  • The evaluation focuses on dialogue-centric datasets, limiting the demonstrated scope of the evaluation.
  • Extending Memory-R1 to multimodal data may introduce challenges beyond the scope of this work.
  • The Memory Manager and Answer Agent are trained separately to ensure stability under sparse rewards, making the process less straightforward.

A.1 From In-context Memory Manager to RL fine-tuned Memory Manager

Case studies show that RL fine-tuning improves Memory Manager updates and Answer Agent retrieval by preserving complementary information and filtering distractors. The resulting memory operations and selected entries support more coherent memories and correct answers.

  • Memory Manager: The vanilla manager fragmented Andrew’s dog memories by deleting Buddy and adding Scout after misinterpreting the second adoption as a contradiction.
  • Memory Manager: RL fine-tuning consolidated separate Buddy and Scout adoptions with one UPDATE instead of deleting the earlier dog memory.The vanilla manager treated differing names as contradictory, while the RL manager recognized the adoptions as complementary.
  • Memory Manager: RL fine-tuning preserved Joanna’s fondness for turtles while incorporating more specific allergy information through UPDATE operations.The vanilla manager treated the specific allergy statement as contradictory and discarded emotional context.
  • Answer Agent: Memory Distillation selected beach-related memories and helped the Answer Agent answer that John lives near a beach rather than mountains.The original model consumed retrieved memories indiscriminately and defaulted to mountains, likely influenced by irrelevant mountaineering mentions.
  • Answer Agent: The Answer Agent’s Memory Distillation filters retrieved distractors, focuses on relevant signals, and improves factual accuracy in the representative case.

B.1 Test Data

The test data span three long-term memory benchmarks covering multi-session dialogue, factual and temporal reasoning, entity tracking, and broader generalization beyond dialogue-centric settings.

  • LoCoMo is a long-term multi-session dialogue benchmark averaging 300 turns and 9k tokens across up to 35 sessions.It is the primary experimental dataset for detailed results.
  • MSC contains open-domain, multi-session dialogues whose questions depend on information distributed across earlier sessions.The modified evaluation setting tests continuity across temporally separated interactions.
  • LongMemEval tests factual recall, temporal reasoning, and entity tracking using questions requiring integration across long and sparse contexts.It complements LoCoMo and MSC by emphasizing broader generalization beyond dialogue-centric settings.

B.2 Training Data

Training data are constructed separately for the Memory Manager and Answer Agent from LoCoMo dialogues. The pipeline uses temporal memory snapshots, retrieved candidates, prompts, and gold answers to support RL fine-tuning without explicit operation labels.

  • Memory Manager Training Data: The Memory Manager receives no explicit ADD, UPDATE, DELETE, or NOOP labels and instead uses downstream answer correctness as its reinforcement-learning signal.
  • Memory Manager Training Data: The Memory Manager training tuples contain a dialogue turn, temporal memory bank, and QA information.The data-construction algorithm builds a temporal memory bank from preceding dialogue turns and stores combined packages as single tuples.
  • Answer Agent Training Data: The Answer Agent is instructed to distill relevant entries, reason over them, and generate concise, correct responses.Its prompt specifies reasoning, memory-selection criteria, and response formatting.
  • Prompting and Evaluation: The Memory Manager is prompted to perform four operations: ADD, UPDATE, DELETE, and NOOP.The full instructions are distributed across multiple figures, including separate ADD/UPDATE and DELETE/NO_OPERATION sections.
  • Prompting and Evaluation: An LLM-as-a-Judge prompt labels generated answers CORRECT or WRONG by comparing them with the gold answer.
  • Answer Agent Training Data: The Answer Agent training tuples pair questions with retrieved memories and gold answers after the Memory Manager maintains the memory bank.Candidate memories are retrieved using the question as a query and paired with the gold answer.

D Implementation Details

Memory-R1 is evaluated across LLaMA and Qwen backbones using PPO and GRPO, with separate procedures for constructing the memory bank and generating memory-augmented answers. The implementation uses staged reinforcement learning, retrieved memory candidates, and an Answer Agent’s distillation step.

  • Model and Training Setup: PPO jointly trains actor and critic networks, whereas GRPO updates only the actor through grouped return normalization.The PPO learning rates are 1 × 10^-6 for the actor and 1 × 10^-5 for the critic.
  • Model and Training Setup: During training, temperature τ = 1.0 encourages exploration, while validation and testing use greedy decoding with τ = 0.The different decoding settings support diverse reward signals during training and deterministic metric evaluation afterward.
  • Memory Bank Construction: Memory Bank Construction extracts information, retrieves related entries, and uses the Memory Manager to select ADD, UPDATE, DELETE, or NOOP operations.The selected operation inserts, merges, prunes, or leaves the memory store unchanged.
  • Memory-Augmented Answer Generation: Answer generation retrieves top-k memory candidates, concatenates them with the question, distills relevant facts, and produces the final response.The resulting answer is added to the answer set.
  • Reinforcement Learning Procedure: The Memory Manager and Answer Agent are optimized separately, with downstream answer correctness providing outcome-based rewards for memory operations.The Answer Agent is frozen during Memory Manager training, and the Memory Manager is fixed during Answer Agent training.

F Extended Results and Type-Level Analysis

Extended evaluations show consistent Memory-R1 gains across reasoning types, especially for multi-hop, temporal, and multi-session questions. GRPO generally outperforms PPO, and GRPO transfers zero-shot from LoCoMo to LongMemEval across model families.

  • LoCoMo Type-Level Results: Memory-R1 improves across all LoCoMo reasoning types, with the largest gains on multi-hop and temporal questions.These results indicate stronger handling of long-range information integration across question types.
  • LongMemEval Type-Level Results: LongMemEval gains are most pronounced in multi-session scenarios requiring continuity across temporally distant interactions.The strongest improvements occur where information must be maintained across separate sessions.
  • LongMemEval Type-Level Results: Memory-R1 substantially improves factual recall and temporal reasoning while also steadily improving knowledge update and open-domain QA.The task abbreviations are SSU, TR, KU, and open-domain QA in the reported analysis.
  • RL Method Comparison: GRPO generally outperforms PPO, particularly when reasoning over multiple or noisy memory entries.The pattern appears across the evaluated reasoning types.
  • Zero-Shot Transfer: After fine-tuning only on LoCoMo, Memory-R1-GRPO outperforms all baseline systems on LongMemEval with both LLaMA-3.1-8B and Qwen-2.5-7B backbones.This evaluation uses zero-shot transfer without additional training on LongMemEval.

G Latency Analysis

Latency analysis compares p50 and p95 inference time across Memory Manager, Memory Search, and Answer Agent components. Memory-R1 generally avoids prohibitive overhead, while GRPO can improve Answer Agent efficiency alongside accuracy.

  • Evaluation Setup: Latency is measured with median (p50) and tail (p95) inference times across Memory Manager, Memory Search, and Answer Agent components.The comparison covers base, PPO-trained, and GRPO-trained variants on LLaMA-3.1-8B and Qwen-2.5-7B.
  • Overall Trends: Memory-R1 does not introduce prohibitive latency overhead across both model families, and GRPO often achieves lower tail latency than base and PPO variants.The reported trend links reinforcement learning with improvements in inference efficiency.
  • Memory Manager: Memory Manager latency remains stable across variants, ranging from 1.98 s to 2.17 s in p50 latency on LLaMA-3.1-8B.Its p95 latency is approximately 3.4–3.6 s, while Qwen-2.5-7B p50 latency remains below 1.4 s.
  • Memory Search: Memory Search remains below 0.35 s in p50 latency and under 0.65 s in p95 latency on both backbones.Differences among base, PPO, and GRPO variants are minimal.
  • Answer Agent: On LLaMA-3.1-8B, the GRPO Answer Agent reaches 0.34 s p50 and 0.67 s p95 latency, below the base model and PPO variant.The corresponding base-model values are 0.65 s and 3.07 s, while PPO reaches 0.91 s and 4.67 s.
  • Accuracy–Latency Relationship: Memory-R1 achieves higher accuracy while reducing both median and tail latency, which the analysis characterizes as a Pareto improvement.The stated explanation emphasizes learned memory distillation and policy optimization.
  • Accuracy–Latency Relationship: Figures 13 and 14 plot median and tail latency against F1, BLEU-1, and LLM-as-a-Judge accuracy for base, PPO, and GRPO variants.The figures compare these accuracy–latency points across LLaMA-3.1-8B-Instruct and Qwen-2.5-7B-Instruct.
Loading 2508.19828v5…