PaperlayerView exact revision on arXiv

Source-linked AI summary

Mem-α: Learning Memory Construction via Reinforcement Learning

Yu Wang, Ryuichi Takanobu, Zhiqi Liang, Yuzhen Mao, Yuanzhe Hu, Julian McAuley, Xiaojian Wu

arXiv:2509.25911v1cs.CL

TL;DR

LLM agents have limited context windows, while existing memory systems typically rely on predefined instructions and tools rather than training agents to construct memory. Mem-α trains agents to manage memory with reinforcement learning and improves over existing memory-augmented baselines, with generalization to sequences exceeding 400k tokens despite training only up to 30k tokens.

  • Problem

    Existing memory systems typically rely on predefined instructions and fixed tool sets without training to optimize memory construction.

    Source: §1 · p. 1

  • Method

    Mem-α uses reinforcement learning to train agents to manage complex memory systems through interaction, rewarding downstream question-answering accuracy over the interaction history.

    Source: §1 · p. 2

  • Results

    Across diverse benchmarks, Mem-α significantly improves over existing memory-augmented baselines, and agents trained on instances up to 30k tokens generalize to sequences exceeding 400k tokens, over 13× the training length.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The results support learning-based memory management beyond predefined heuristics for diverse scenarios.

    Source: §5 · p. 9

Abstract

from arXiv · show

Large language model (LLM) agents are constrained by limited context windows, necessitating external memory systems for long-term information understanding. Current memory-augmented agents typically depend on pre-defined instructions and tools for memory updates. However, language models may lack the ability to determine which information to store, how to structure it, and when to update it, especially as memory systems become more complex. This results in suboptimal memory construction and information loss. To this end, we propose Mem-alpha, a reinforcement learning framework that trains agents to effectively manage complex memory systems through interaction and feedback. We also construct a specialized training dataset spanning diverse multi-turn interaction patterns paired with comprehensive evaluation questions designed to teach effective memory management. During training, agents process sequential information chunks, learn to extract and store relevant content, then update the memory system. The reward signal derives from downstream question-answering accuracy over the full interaction history, directly optimizing for memory construction. To illustrate the effectiveness of our training framework, we design a memory architecture comprising core, episodic, and semantic components, equipped with multiple tools for memory operations. Empirical evaluation demonstrates that Mem-alpha achieves significant improvements over existing memory-augmented agent baselines. Despite being trained exclusively on instances with a maximum length of 30k tokens, our agents exhibit remarkable generalization to sequences exceeding 400k tokens, over 13x the training length, highlighting the robustness of Mem-alpha.

1 INTRODUCTION

Mem-α uses reinforcement learning to train agents to manage complex memory systems through interaction and feedback, optimizing memory construction for downstream question answering. The introduction reports improvements over memory-augmented baselines and generalization to sequences far longer than the training examples.

Motivation

  • Existing systems such as Mem0, MemGPT, and MIRIX offer predefined instructions and memory tools without training agents to choose what to store, how to structure it, or when to update it.The introduction notes that complicated prompts are difficult to adjust for all scenarios and may confuse smaller models; earlier RL efforts MEM1, MemAgent, and Memory-R1 use simpler memory structures.Source: §1 · pp. 1–2

Training approach

  • Mem-α frames memory construction as sequential decisions over information chunks, rewarding agents with downstream question-answering accuracy across the full interaction history.Unlike supervised fine-tuning that requires ground-truth memory traces, RL lets agents discover strategies through trial and error; rewards also use memory-quality metrics.Source: §1 · pp. 1–2

Training data

  • Its specialized dataset spans conversations, document sharing, pattern recognition, and storytelling, paired with comprehensive questions requiring broad memory to answer correctly.This exposes agents to diverse memory-management scenarios during training.Source: §1 · p. 2

Memory architecture

  • The memory architecture combines core, episodic, and semantic components with specialized tools to handle diverse information types while remaining learnable through reinforcement.Figure 1 illustrates how RL training improves agents’ memory-tool selection and management policies.Source: §1 · p. 2, Figure 1 · p. 2

Results

  • Over 400k-token sequences—more than 13× the 30k-token training maximum—Mem-α agents generalize while achieving significant improvements over memory-augmented agent baselines across diverse benchmarks.The authors suggest this length generalization reflects learned memory-management principles rather than memorized patterns.Source: §1 · p. 2

2 RELATED WORK

The section contrasts latent-space and external memory approaches with reinforcement-learning methods for memory construction, motivating explicit training to manage more capable memory systems.

Latent-Space Memory

  • Latent-space methods compress memory efficiently, but remain capacity-bounded—M+ retains approximately 160k tokens—and require model-internal access, limiting proprietary-system compatibility (Wang et al., 2025a; Wang & Chen, 2025).SELF-PARAM can memorize hundreds of contexts without external storage, but these methods require access to model internals.Source: §2 · p. 3

LLM Agents with External Memory

  • External-memory systems support proprietary models and organize, retrieve, and update large information stores, but their effectiveness depends on tool use and complex systems may still be updated poorly (Zhang et al., 2025a; Wang & Chen, 2025).This limitation motivates training models to manage memory rather than relying purely on prompting.Source: §2 · p. 3

Learning Memory Construction with Reinforcement Learning

  • Prior reinforcement-learning work trains models on simple or richer memories, but focuses on LoCoMo contexts under approximately 26k tokens and same-distribution subsets (Zhou et al., 2025; Yu et al., 2025; Yan et al., 2025; Zhang et al., 2025b; Zhang et al., 2023a).This paper instead trains models to operate a substantially more capable memory system and reports improvements across memory quality and efficiency.Source: §2 · p. 3

3 METHOD

Mem-α trains a memory-writing policy with reinforcement learning, rewarding memories that support accurate question answering while accounting for operation reliability, content quality, and compression. Its method specifies sequential memory updates, fixed RAG-based evaluation, GRPO optimization, a three-part memory architecture, and a balanced training subset.

3.1 REINFORCEMENT LEARNING FRAMEWORK

  • Mem-α formulates memory construction as reinforcement learning, optimizing a write policy whose memory quality is evaluated separately through retrieval-augmented question answering.The framework treats the write policy as learnable and evaluates its constructed memory through a separate QA process.Source: §3.1 · p. 3

3.1.1 TASK SETUP

  • At each conversation chunk, the agent observes the current memory and may issue insert, update, or delete operations before proceeding.The interaction sequence spans varied conversation formats, and the memory starts empty.Source: §3.1.1 · pp. 3–4

3.1.2 REWARD FUNCTIONS

  • Four rewards measure final-memory QA correctness, successful tool-call execution, memory compression, and semantic validity of memory operations.Correctness and compression are global rewards based on the final memory; tool-call success and content quality are action-level rewards. The tool-call reward weight is fixed at 1, while other weights require tuning.Source: §3.1.2 · pp. 4–5

3.1.3 MEMORY COMPREHENSIVENESS EVALUATION VIA RAG

  • A fixed RAG pipeline evaluates memory comprehensiveness by retrieving top-k semantic and episodic entries with BM25, generating answers with a frozen model, and scoring them against references.Only the write policy is learnable; retrieval and generation remain fixed.Source: §3.1.3 · p. 5

3.2 POLICY OPTIMIZATION

  • The policy is optimized with GRPO (Shao et al., 2024) to maximize expected reward across actions, discarding the KL term to encourage exploration.GRPO computes advantages from rewards within sampled action groups.Source: §3.2 · p. 5

3.3 MEMORY INSTANTIATION

  • The architecture combines a persistent core summary of at most 512 tokens with discrete factual semantic memory and timestamped episodic memory.Semantic and episodic memories support insertion, updating, and deletion; core memory supports only full updates. The architecture is modular and can be replaced without changing the RL framework.Source: §3.3 · pp. 5–6

3.4 TRAINING DATASET PREPARATION

  • Training uses 4,139 instances and a stratified, balanced subset of 562, focusing on accurate retrieval, test-time learning, and long-range understanding.Conflict resolution is excluded because the authors report that available benchmarks are predominantly synthetic and do not adequately capture real-world complexity.Source: §3.4 · p. 6

4 EXPERIMENTS

Experiments evaluate Mem-α across retrieval, classification, and summarization tasks, compare it with memory-augmented baselines, and examine RL, reward hyperparameters, and memory-construction traces. The reported results show strong performance and generalization to documents much longer than those used in training.

4.1 EXPERIMENTAL SETUP

  • Evaluation covers accurate retrieval, test-time classification, and long-range summarization, comparing Mem-α with long-context, BM25 RAG, MemAgent, and MEM1 baselines.Mem-α uses Qwen3-4B and was trained for 205 steps on 32 H100 GPUs; the Long-Context baseline uses Qwen3-32B with a 32K context window.Source: §4.1 · pp. 6–7

4.2 OVERALL PERFORMANCE COMPARISON

  • Mem-α significantly outperforms existing baselines across metrics, particularly on Accurate Retrieval and Long-Range Understanding, and generalizes from training documents averaging <20K tokens to documents up to 474K tokens.The 474K-token documents are in MemoryAgentBench’s Multi-Doc dataset.Source: §4.2 · p. 7, Table 2 · p. 8

4.3 PERFORMANCE BOOST FROM REINFORCEMENT LEARNING

  • Mem-α reaches 0.642 average performance versus 0.389 for base Qwen3-4B, surpassing gpt-4.1-mini and showing gains from RL optimization.Base Qwen3-4B also falls below RAG-Top2 (0.567) and Long-Context (0.588); all three configurations use the memory framework.Source: §4.3 · pp. 7–8, Table 3 · p. 8

4.4 ABLATION STUDIES

  • Setting the memory-content reward γ to 0 causes severe degradation, while increasing compression weight β shortens memories at the cost of performance.With γ = 0.1, β = 0.05 reduces BookSum memory from 4.5K to 2.2K tokens versus β = 0, while maintaining comparable memory lengths on other datasets.Source: §4.4 · p. 8

4.5 CASE STUDIES

  • Mem-α builds informative core memory, distinct semantic entries, and consolidated episodic events while retaining both user behavior and assistant responses.Qwen3-4B leaves core memory empty and compresses concepts into one semantic entry; GPT-4.1-mini creates duplicate-timestamp episodic entries and ignores assistant responses.Source: §4.5 · pp. 8–9, Table 5 · p. 9

5 CONCLUSION, LIMITATION AND FUTURE WORK

Mem-α trains LLM agents to manage memory through reinforcement-learning interaction and question-answering feedback, improving over existing memory-augmented baselines and generalizing to longer interactions. The conclusion identifies architectural integration and real-world deployment as future directions.

Conclusion

  • Mem-α learns memory-management strategies through interaction and feedback, using a training dataset and rewards based on question-answering correctness.The approach moves beyond pre-defined heuristics to let agents discover memory operations for diverse scenarios.Source: §5 · p. 9
  • Experiments show significant improvements over existing memory-augmented baselines, with strategies that generalize to much longer interaction patterns.The passage reports no specific numerical result in this conclusion.Source: §5 · p. 9

Future Work

  • Future work could integrate more sophisticated memory systems such as MIRIX to provide structural advantages for complex reasoning tasks.The conclusion presents this integration as a potential direction, not an evaluated result.Source: §5 · p. 9
  • Moving from simulated environments to real-world applications would require database and production-system connections, with latency, scalability, and safety challenges to investigate.The stated goal is practical deployment of memory-augmented LLM agents.Source: §5 · p. 9

A DATASETS DETAILS · A.1 TRAINING DATASET

The training dataset spans Accurate Retrieval, Test-Time Learning, and Long Range Understanding, using tasks designed to exercise different memory capabilities. Dataset size was capped because of computational constraints and imbalance.

A.1 TRAINING DATASET

  • The dataset targets three memory capabilities—accurate retrieval, test-time learning, and long-range understanding—with statistics reported across eight data sources.Table 6 lists category, instance, chunk, and question statistics with metrics suited to each task type.Source: §A.1 · p. 14, Table 6 · p. 14
  • SQuAD (Rajpurkar et al., 2016), HotPotQA (Yang et al., 2018), and PerLTQA train retrieval across documents, chunks, and user-memory types.The tasks require recalling specific facts, synthesizing relationships across documents, and integrating episodic and semantic user memories while maintaining profile consistency.Source: §A.1 · p. 14
  • LongMemEval-Train (Wu et al., 2024) uses 50 samples of 10K–30K-token contexts, each paired with 4–5 questions, with no overlap with MemoryAgentBench evaluation data.The subset collects 200 questions from longmemeval oracle.json3 and concatenates haystack dialogues into contexts.Source: §A.1 · p. 14
  • PubMed-RCT (Dernoncourt & Lee, 2017) trains test-time classification by replacing five semantic-role labels with numeric labels and asking agents to classify new examples.Conversational chunks contain sentence-label examples, and each instance covers all five categories across chunks.Source: §A.1 · p. 14
  • NLU and TREC-C (Hu et al., 2025) provide test-time learning over 68 and 6 classes, respectively, using 200 instances per dataset with 10 chunks each.The original approximately 100K-token instances are partitioned into chunks of roughly 500–2,000 tokens, with labels redistributed to ensure complete label coverage per instance.Source: §A.1 · p. 15
  • BookSum (Kry´sci´nski et al., 2021) supports long-range understanding by splitting each book chapter into 10–20 conversational chunks and evaluating summary keywords against ground truth.The metric is the ratio of correctly identified keywords in generated summaries to the ground-truth keyword set.Source: §A.1 · p. 15
  • Each dataset was limited to at most 100 instances because of computational constraints and dataset imbalance, and every chunk was formatted as a conversation.After three days of training with 32 H100 GPUs, only a small portion of the complete datasets could be processed; final composition and statistics appear in Table 7.Source: §A.1 · p. 15, Table 7 · p. 15

A.2 EVALUATION DATASET · B FORMAL DEFINITIONS OF REWARD COMPONENTS

The evaluation suite tests retrieval, test-time learning, and long-range understanding across nine datasets and 112 instances. The reward formalization combines answer correctness with tool-call validity, memory compression, and content validity.

A.2 EVALUATION DATASET

  • The evaluation suite uses nine datasets and 112 test instances to assess accurate retrieval, test-time learning, and long-range understanding.Dataset statistics are summarized in Table 8.Source: §A.2 · p. 15, Table 8 · p. 16
  • Accurate Retrieval covers RULER-QA1 single-hop retrieval, RULER-QA2 multi-hop synthesis, and LME(S*) queries over fewer contexts from extended interactions.LME(S*) derives from LongMemEval (Wu et al., 2024) and was reformatted by Hu et al. (2025); RULER-QA1 is attributed to Hsieh et al. (2024).Source: §A.2 · pp. 15–16
  • Test-Time Learning uses TREC-Coerse’s six categories, TREC-Fine’s 50 types, NLU’s 68 intents, CLINIC150’s 150 medical intents, and Banking77’s 77 financial intents.TREC-Coerse and TREC-Fine each have 5,452 training and 500 test questions; the other datasets assess semantic, healthcare, and banking-domain classification patterns.Source: §A.2 · pp. 16–17
  • Long-range understanding is evaluated with InfBench-Sum, which requires summaries of 100 novels averaging 172k tokens.The task tests synthesis across multiple chunks and distributed memory representations (Zhang et al., 2024).Source: §A.2 · p. 18

B FORMAL DEFINITIONS OF REWARD COMPONENTS

  • Correctness reward evaluates predicted answers from the final memory state against ground truth using each dataset’s metric, such as exact match or F1.The prediction is generated by the RAG pipeline for each question.Source: B FORMAL DEFINITIONS OF REWARD… · p. 18
  • Tool-call format reward measures whether each function call executes without error at each time step.The passages define a correctness indicator for individual calls and a reward at each step.Source: B FORMAL DEFINITIONS OF REWARD… · p. 18
  • Compression reward compares total input-chunk length with final memory length, encouraging compact representations that preserve essential information.It approaches 1 for highly compressed memory and 0 when memory size equals input size.Source: B FORMAL DEFINITIONS OF REWARD… · p. 18
  • Memory content reward uses a language-model judge to assess semantic validity, and the overall reward adds correctness, tool-call, weighted compression, and weighted content components.Correctness and compression are global rewards shared across time steps; tool-call and content rewards are computed per time step.Source: B FORMAL DEFINITIONS OF REWARD… · pp. 18–19

C EXPERIMENTAL DETAILS · C.1 JUSTIFICATION OF BACKBONE MODEL SELECTION

Qwen3-8B was excluded from the final experiments after repeatedly violating the specified function-call format and receiving lower reward than Qwen3-4B, even after its preferred argument format was accommodated.

C.1 JUSTIFICATION OF BACKBONE MODEL SELECTION

  • Qwen3-8B repeatedly generated malformed function calls by appending “ memory” to allowed memory-type arguments, despite explicit signature specifications.The failures occurred reliably across multiple trials.Source: §C.1 · p. 19
  • The researchers changed the valid arguments to “semantic memory,” “core memory,” and “episodic memory” to accommodate Qwen3-8B’s apparent formatting preference.Source: §C.1 · p. 19
  • Qwen3-4B handled both the original and modified function-argument formats correctly, and the adaptation did not impact its performance.Source: §C.1 · p. 19
  • Even with the accommodated argument format, Qwen3-8B continued to exhibit lower reward values than Qwen3-4B.Source: §C.1 · p. 19
  • The larger Qwen3-8B variant showed poorer instruction-following capabilities and lower overall performance than the 4B variant.Source: §C.1 · p. 19
  • The researchers excluded Qwen3-8B from their final experiments because of its instruction-following failures and lower reward relative to Qwen3-4B.Source: §C.1 · p. 19

C.2 BASELINE INTRODUCTION AND IMPLEMENTATION DETAILS · C.3 PROMPTS USED IN TRAINING

The appendix specifies four baseline implementations and their answering procedures, then documents the training prompts for chunk memorization, memory-content reward judging, and final question answering.

C.2 BASELINE INTRODUCTION AND IMPLEMENTATION DETAILS

  • The Long-Context baseline uses Qwen3-32B with a 32k-token maximum context and truncates longer combined chunks to their last 32k tokens.Source: §C.2 · p. 19
  • RAG-Top2 retrieves the top two previous chunks with BM25 using the question as query, then answers with Qwen3-32B.Source: §C.2 · p. 19
  • MemAgent uses RL-MemoryAgent-14B, while MEM1 uses Qwen2.5-7B-RL-RAG-Q2-EM-Release to construct memory.Source: §C.2 · p. 20
  • Both memory baselines process all chunks with the task instruction before answering from memory; MemAgent answers with its original model, whereas MEM1 uses Qwen3-32B.Source: §C.2 · p. 20

C.3 PROMPTS USED IN TRAINING

  • A universal prompt asks the agent to memorize each chunk during memory updates across the training dataset.Source: §C.3 · p. 20
  • For memory-content reward r4, Qwen3-32B judges Core, Episodic, and Semantic Memory using prompts shown in Figures 7, 8, and 9.Source: §C.3 · p. 20, Figure 7 · p. 20
  • The final Qwen3-32B question-answering model uses the prompt shown in Figure 10.Source: §C.3 · p. 20

C.4 ADDITIONAL ABLATION STUDY

The additional ablation compares β and γ settings on MemoryAgentBench and reports observations consistent with Section 4.4. The section also includes prompts for measuring episodic and semantic memory during training and answering questions in Mem-α.

C.4 ADDITIONAL ABLATION STUDY

  • On MemoryAgentBench, the β and γ settings show observations consistent with Section 4.4; Table 9 reports task-specific performance and memory consumption.Performance uses task-specific F1 or Accuracy metrics, and memory is measured in thousands of tokens.Source: §C.4 · p. 20, Table 9 · p. 21
  • Figure 8 presents the prompt used to measure episodic memory content during training, while Figure 9 presents the corresponding semantic-memory prompt.Source: Figure 8 · p. 21, Figure 9 · p. 21
  • Figure 10 presents the prompt used to answer questions in Mem-α.Source: Figure 10 · p. 22

Related papers

Loading 2509.25911v1…