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 need external memory because limited context windows make long information streams difficult to retain, while fixed memory instructions do not teach agents how to construct complex memories. Mem-α trains memory management with reinforcement learning, specialized multi-turn data, and downstream question-answering rewards, achieving gains over memory-augmented baselines and generalizing from 30k-token training instances to sequences exceeding 400k tokens.

  • Problem

    Limited context windows require external memory, but existing memory-augmented agents typically rely on fixed instructions and tools that do not teach effective construction of complex memories.

  • Method

    Mem-α trains agents with reinforcement learning to process sequential information, choose memory operations, and construct core, episodic, and semantic memories using downstream question-answering rewards.

  • Results

    Mem-α significantly improves over existing memory-augmented baselines and generalizes from 30k-token training instances to sequences exceeding 400k tokens, over 13× the training length.

  • Takeaways & Limitations

    The results indicate that reinforcement learning can produce robust memory-management strategies for long-context retention across diverse benchmarks and unseen distributions.

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 construct and manage complex memories, addressing the limits of fixed memory instructions and tools. It improves memory-augmented performance and generalizes from 30k-token training instances to sequences exceeding 400k tokens.

  • 1 INTRODUCTION: Reinforcement learning teaches agents to select appropriate memory tools and types, replacing pre-defined memory-management behavior with learned policies.Before training, agents struggle with tool selection; after training, they learn effective memory-management policies.
  • 1 INTRODUCTION: Mem-α trains agents to learn memory construction through sequential decisions, downstream question-answering rewards, and a dataset of diverse multi-turn interactions.Agents process information chunks, choose memory operations, and receive rewards based on answering questions over the full interaction history.
  • 1 INTRODUCTION: The framework uses a memory architecture with core, episodic, and semantic components, each equipped with specialized memory-operation tools.This architecture is designed to represent diverse information types while remaining learnable through reinforcement.
  • 1 INTRODUCTION: 30k-token training instances generalize to sequences exceeding 400k tokens, over 13× the training length.The authors describe this as robust length generalization and evidence that agents learn memory-management principles rather than merely memorizing patterns.
  • 1 INTRODUCTION: Mem-α achieves significant improvements over existing memory-augmented agent baselines across diverse benchmarks.The evaluation includes validation datasets and out-of-distribution MemoryAgentBench tests, with particularly substantial gains on Accurate Retrieval and Long-Range Understanding.

2 RELATED WORK

Related work spans latent-space memory, external memory systems, and reinforcement-learning approaches for memory construction. Existing reinforcement-learning methods generally use simpler memory representations or tool-calling settings than the complex memory management targeted by Mem-α.

  • 2 RELATED WORK: Latent-space memory stores information in model internals such as hidden states, key-value caches, prompts, parameters, or external matrices.Its main advantage is efficient compression.
  • 2 RELATED WORK: External-memory agents use databases or vector stores to provide persistent memory systems for language models.The related systems include MemGPT, Mem0, Zep, MIRIX, and other external-memory frameworks.
  • 2 RELATED WORK: Reinforcement-learning approaches to memory construction remain preliminary, with early systems often training models to update simple text-only memories.Later methods introduce somewhat richer representations or simplified tool-calling settings.

3 METHOD

Mem-α formulates memory construction as reinforcement learning over sequential conversation chunks, optimizing memory-writing actions with downstream question-answering and auxiliary rewards. Its architecture combines core, semantic, and episodic memory components with specialized operations and evaluates the resulting memory through a fixed retrieval-and-generation pipeline.

  • 3.1.1 TASK SETUP: Mem-α treats memory construction as an RL problem in which an agent processes conversation chunks, observes the current memory, and issues write operations before continuing.The action space includes structured memory insert, update, and delete calls, applied sequentially to the evolving memory.
  • 3.1.2 REWARD FUNCTIONS: The training reward combines question-answering correctness, tool-call success, compression, and semantic validity of memory operations.Correctness is evaluated from the final memory through QA, while other components assess execution, compactness, and operation quality.
  • 3.1.3 MEMORY COMPREHENSIVENESS EVALUATION VIA RAG: Memory comprehensiveness is evaluated after processing all chunks using a decoupled RAG pipeline with fixed retrieval and generation components.The pipeline retrieves entries from semantic and episodic memory, generates an answer, and compares it with the reference answer to obtain correctness indicators.
  • 3.2 POLICY OPTIMIZATION: Mem-α uses GRPO to optimize the expected reward over memory actions, with group-normalized advantages derived from the four-part reward.The method discards GRPO’s KL term to encourage policy exploration.
  • 3.3 MEMORY INSTANTIATION: The memory architecture contains a continuously accessible core summary plus expandable semantic and episodic stores, whose operations reflect their distinct update patterns.Core memory is limited to one 512-token paragraph and supports rewriting, whereas semantic and episodic memories support insert, update, and delete operations.

4 EXPERIMENTS

Mem-α is evaluated against long-context, retrieval, and memory-agent baselines across validation and out-of-distribution benchmarks. Results emphasize superior task performance, memory efficiency, reinforcement-learning gains, and structured memory management.

  • 4.2 OVERALL PERFORMANCE COMPARISON: Mem-α significantly outperforms existing baselines across metrics, with especially substantial gains on Accurate Retrieval and Long-Range Understanding tasks.The evaluation includes validation datasets and the out-of-distribution MemoryAgentBench benchmark.
  • 4.2 OVERALL PERFORMANCE COMPARISON: Training on sequences below 20k tokens generalizes to documents exceeding 400k tokens, demonstrating strong length generalization.The reported test length exceeds the training-document scale by more than 20× in this experiment description.
  • 4.2 OVERALL PERFORMANCE COMPARISON: Mem-α reduces memory footprint by approximately 50% compared with Long-Context and RAG-Top2 while maintaining stronger task performance.The comparison is reported among the principal baselines in the overall evaluation.
  • 4.3 PERFORMANCE BOOST FROM REINFORCEMENT LEARNING: Reinforcement learning raises average performance from 0.389 for base Qwen3-4B to 0.642 for Mem-α, surpassing gpt-4.1-mini with the same memory framework.The ablation compares models using the proposed memory framework and attributes the gain to RL optimization rather than memory structure alone.
  • 4.4 ABLATION STUDIES: The memory-content reward is critical, whereas stronger compression rewards shorten memories at the cost of reduced performance.With γ = 0, memory construction quality collapses; increasing β produces shorter memories but can reduce task performance.
  • 4.5 CASE STUDIES: Mem-α organizes informative core and semantic memories, consolidates duplicate-timestamp episodic events, and retains both user behavior and assistant responses.Compared with baseline traces, this organization preserves more information while using memory space more efficiently.

5 CONCLUSION, LIMITATION AND FUTURE WORK

The paper concludes that reinforcement learning enables LLM agents to learn memory-management strategies that generalize to longer interaction patterns. It identifies richer memory architectures and real-world deployment as future directions.

  • CONCLUSION: Mem-α trains LLM agents to learn memory-management strategies through interaction, feedback, and question-answering-based rewards.The framework moves beyond predefined heuristics toward learned memory operations for diverse scenarios.
  • FUTURE WORK: The authors propose integrating more sophisticated memory systems to provide additional structural advantages for complex reasoning tasks.MIRIX is given as an example of a possible future integration.
  • LIMITATIONS AND FUTURE WORK: Applying Mem-α to real-world systems would require connecting reinforcement learning with databases and production environments while addressing latency, scalability, and safety.The paper identifies this transition from simulated environments as an open direction.

A.1 TRAINING DATASET

The training dataset targets accurate retrieval, test-time learning, and long-range understanding through varied chunked interactions. Its tasks require retrieval, cross-document synthesis, adaptive classification, or summarization from sequential information.

  • DATASET ORGANIZATION: The dataset is organized around Accurate Retrieval, Test-Time Learning, and Long-Range Understanding capabilities.Dataset statistics are reported across eight data sources.
  • ACCURATE RETRIEVAL: Accurate Retrieval tasks require storing and retrieving facts across single documents, multiple documents, or episodic and semantic user information.HotPotQA tests cross-document synthesis, while PerLTQA tests memory integration, profile consistency, and multi-hop reasoning.
  • ACCURATE RETRIEVAL: LongMemEval-Train uses 10K–30K-token dialogue contexts paired with four to five questions per context without overlap with MemoryAgentBench evaluation data.The resulting training subset contains 50 samples.
  • TEST-TIME LEARNING: Test-Time Learning trains agents to infer new classification patterns from examples and apply them to later instances across medical, natural-language, and intent datasets.The category includes PubMed-RCT, NLU, and TREC-C adaptations.
  • LONG-RANGE UNDERSTANDING: BookSum simulates incremental long-context processing by splitting chapters into 10–20 conversational chunks and evaluating summaries by ground-truth keyword recovery.The metric is the ratio of correctly identified generated-summary keywords to the ground-truth keyword set.

A.2 EVALUATION DATASET

The evaluation suite covers accurate retrieval, test-time learning, and long-range understanding across nine datasets, while reward definitions assess answer correctness, tool use, compression, and memory content.

  • Evaluation framework: Nine datasets evaluate accurate retrieval, test-time learning, and long-range understanding using task-specific metrics.The suite contains 112 test instances, and Table 8 reports dataset statistics.
  • Accurate Retrieval: Accurate retrieval includes single-hop and multi-hop question answering plus extended-interaction memory querying.RULER-QA1 tests direct retrieval, RULER-QA2 requires synthesis across memory chunks, and LME(S*) tests querying complex memory representations.
  • Test-Time Learning: Test-time learning measures whether models infer classification patterns from thousands of labeled examples and apply them to new instances.The listed datasets include coarse and fine-grained question classification, natural-language understanding, clinical intent classification, and banking intent detection.
  • Reward evaluation: The overall reward combines correctness, tool-call format, compression, and memory-content components, with correctness computed from dataset-specific answer metrics.Correctness uses the final memory and generated answers; compression rewards compact memory, while tool and content rewards are computed per time step.

C.1 JUSTIFICATION OF BACKBONE MODEL SELECTION

The authors selected Qwen3-4B because Qwen3-8B repeatedly violated the required function-call API and received lower rewards even after the interface was adapted.

  • Backbone selection: Qwen3-8B was excluded after consistently generating malformed memory-tool calls that appended an invalid “ memory” suffix to memory-type arguments.The required values were “semantic,” “core,” and “episodic,” but the model generated forms such as “semantic memory.”
  • Backbone selection: Qwen3-8B continued to obtain lower rewards than Qwen3-4B after the function signatures were changed to accept its preferred argument format.Qwen3-4B handled both formats correctly, whereas the larger model remained lower-performing.
  • Backbone selection: The selection therefore reflects observed instruction-following reliability and experimental reward performance rather than parameter count alone.The passage describes poorer API adherence and lower overall performance for Qwen3-8B than for Qwen3-4B.

C.2 BASELINE INTRODUCTION AND IMPLEMENTATION DETAILS

The experiments compare Mem-α with long-context, retrieval, MemAgent, and MEM1 baselines, while standardizing chunk processing, answer generation, and memory-content evaluation procedures.

  • Baseline comparison: The baseline suite includes a 32k-token long-context model and BM25 retrieval of the two most relevant prior chunks.Long-Context truncates longer inputs to the last 32k tokens, while RAG-Top2 answers using Qwen3-32B after retrieval.
  • Memory evaluation: Core memory is invalid when it contains literal labels such as “core memory” or remains an unspecified placeholder.The validity criteria reject statements that mention a summary or rules without providing their actual content.
  • Baseline comparison: MemAgent and MEM1 process all chunks to construct memory before answering questions from the resulting memory.MemAgent uses its original model for answering, whereas MEM1 uses Qwen3-32B after memory construction.
  • Memory construction: Training uses a universal prompt that asks the agent to memorize each chunk through core-memory updates, summaries, or classification rules.The prompt requires brief updates tailored to the task type.
  • Evaluation procedure: Memory-content rewards are judged by Qwen3-32B with separate prompts for core, episodic, and semantic memory.Question answering also uses a Qwen3-32B prompt shown separately from the memory-content prompts.

C.4 ADDITIONAL ABLATION STUDY

The additional study examines reward-weight settings on validation and test data, alongside prompts for judging memory components and answering questions on MemoryAgentBench.

  • Reward-weight ablation: Validation and MemoryAgentBench test results show consistent observations across different β and γ reward-weight settings.The test comparison is reported in Table 9, while the validation comparison appears in Section 4.4.
  • Memory-content evaluation: The ablation evaluates memory quality separately for core, episodic, and semantic memory using dedicated analyst prompts.Each memory component has its own content-measurement prompt.
  • Question answering: The final question-answering stage uses a reasoning-assistant prompt that accesses structured memory to produce accurate and comprehensive responses.Figure 10 presents the corresponding question-answering prompt.
Loading 2509.25911v1…