Mem-{\alpha}: Learning Memory Construction via Reinforcement Learning | Summary

Contents

Summary
1 INTRODUCTION
2 Related Work
3 Method
4 Experiments
5 Conclusion, Limitation and Future Work
Appendix
Brief Thoughts

This article explains the key points of Mem-{\alpha}: Learning Memory Construction via Reinforcement Learning.

2025-09-30 (arXiv)
Wang, Yu, Takanobu, Ryuichi, Liang, Zhiqi, Mao, Yuzhen, Hu, Yuanzhe, McAuley, Julian, Wu, Xiaojian.
Anuttacon, University of California San Diego, Stanford University
Paper
Github

Summary

Large language model (LLM) agents need to preserve information outside their limited context windows, but providing storage tools does not ensure that they retain and organize useful information. With the memory architecture, retrieval system, and frozen answer generator held fixed, Mem-α uses reinforcement learning, optimization through reward feedback, to raise Qwen3-4B’s average validation score from 0.389 to 0.642, compared with 0.517 for gpt-4.1-mini as the memory writer.
The agent learns to update core memory, an always-accessible summary; semantic memory, a collection of factual statements; and episodic memory, a collection of timestamped events. Training combines downstream answer correctness, successful tool execution, compression, and model-judged memory validity, using a reported 562 sampled training instances with a maximum input length of 30k tokens.
On MemoryAgentBench, a benchmark for incremental memory processing, Mem-α averages 0.592 versus 0.502 for retrieval over the original chunks and 0.461 for a model restricted to a 32k context window, with evaluation inputs reaching 474K tokens. Results vary by task, conflict resolution and production performance remain unevaluated, and the paper contains inconsistent reward-weight and dataset-count reports.

1 INTRODUCTION

Existing external-memory agents provide persistent storage and update tools but often rely on prompts to determine what to retain and how to organize it. Mem-α learns a sequential memory-writing policy from rewards without requiring reference memory-construction traces.

Its three contributions are a sequential memory-writing formulation, training data covering several interaction patterns, and a memory architecture with distinct summary, fact, and event stores.
Figure 1 illustrates the motivation: the untrained writer leaves core memory empty and stores an incomplete semantic entry, while the trained example uses all three stores.

Memory-tool selection and stored information before and after reinforcement learning.

The paper distinguishes Latent-Space Memory, information encoded in model internals, from LLM Agents with External Memory, information maintained in databases or other external stores. It positions Learning Memory Construction with Reinforcement Learning as a way to improve external-memory tool use with richer storage structures than paragraph rewriting or simple fact lists.

The authors describe latent-space methods as efficient compressors whose bounded capacity and need for access to model internals can limit deployment with proprietary models.
Earlier reinforcement-learning approaches include MEM1, MemAgent, Memory-R1, Learn-to-Memorize, and REMEMBER. The paper contrasts its evaluation with the relatively short conversational settings used by some of these methods.

3 Method

The method separates learning to write memory from retrieving information and generating answers. It combines a sequential writing environment, a four-part reward, policy optimization, a three-store memory implementation, and task-specific data converted into conversational chunks.

3.1 Reinforcement Learning Framework

In 3.1.1 TASK SETUP, the agent starts with empty memory and processes conversations one chunk at a time. For each chunk, an action comprises a sequence of structured function calls using memory_insert, memory_update, or memory_delete, with arguments such as record identifiers, memory types, and content.

The calls modify the current memory before the agent advances to the next chunk; the final memory and recorded actions determine the rewards.
Figure 2 separates the trainable memory writer from the fixed downstream process that evaluates whether its memory supports answering questions.

Sequential memory construction and the four reward signals in the Mem-α training framework.

In 3.1.2 REWARD FUNCTIONS, the combined reward is \(r_t = r_1 + r_{2,t} + \beta r_3 + \gamma r_{4,t}\). Correctness measures downstream answers, tool-call format measures successful execution, compression measures the reduction from input length to stored-memory length, and memory content measures whether Qwen3-32B judges operations semantically valid.

Correctness and compression depend on the final memory and are shared across all actions in a sequence; execution and content validity are computed separately for each action.
The compression reward is \(r_3 = 1 - \frac{l_m}{l_c}\), where l_m is total stored-memory length and l_c is total input-chunk length.
The weights on correctness and tool execution are fixed at 1; β and γ control compression and content validity.

In 3.1.3 MEMORY COMPREHENSIVENESS EVALUATION VIA RAG, retrieval-augmented generation (RAG), answering with retrieved memory as supporting context, evaluates the terminal memory. BM25, the retrieval method used to rank memory entries for a question, selects top-k entries, the k highest-ranked entries, from semantic and episodic memory; a frozen generator then answers questions and its outputs are compared with references using dataset-specific metrics.

Only the writing policy is learned. Retrieval and generation remain fixed, so training evaluates the writer through the information this downstream pipeline can retrieve and use.

3.2 Policy Optimization

Mem-α uses Group Relative Policy Optimization (GRPO), which normalizes rewards against the mean and standard deviation of a sampled action group. Equation (2) applies a clipped probability-ratio objective, limiting changes in action-token probabilities relative to the previous policy, across action tokens and sequence steps.

The authors remove the Kullback–Leibler (KL) divergence term, a penalty for policy deviation, to encourage exploration.
Both sequence-level and action-level rewards enter the normalized advantage, the action’s reward relative to its sampled group.

3.3 MEMORY Instantiation

Core Memory is a persistent summary of at most 512 tokens that remains available in the agent’s context. Semantic Memory stores individually retrievable factual statements, while Episodic Memory stores chronologically organized events with timestamps.

Semantic and episodic stores support memory_insert, memory_update, and memory_delete; core memory supports only memory_update and is rewritten as a complete summary.
Figure 3 distinguishes the bounded core paragraph from the expandable fact and event lists.
The architecture is described as separable from the training framework, but the experiments evaluate this particular implementation.

Core, semantic, and episodic memory storage formats.

3.4 TRAINING DATASET Preparation

Training targets Accurate Retrieval, locating information in prior context; Test-Time Learning, learning classification patterns from examples during deployment; and Long-Range Understanding, integrating information across extended sequences. Conflict Resolution, revising stored information when evidence contradicts it, is excluded because the authors consider available benchmarks insufficiently realistic.

The compiled dataset contains 4,139 training instances across eight sources; stratified sampling, sampling separately across data sources to reduce imbalance, produces a reported 562-instance training subset.
Each instance contains multiple chunks and associated evaluation questions, creating a sequence of memory-writing decisions.
The data include document questions, personal conversations, labeled classification examples, and book chapters.

4 Experiments

The experiments compare task performance and stored-memory token counts on validation data and the longer MemoryAgentBench test suite. Further comparisons examine writer training, reward-weight choices, and an illustrative memory-construction trace.

4.1 EXPERIMENTAL SETUP

The test suite covers Single-Doc and Multi-Doc document retrieval and LME(S*), the adapted LongMemEval conversational-memory task; TREC-C and TREC-F, coarse and fine question classification; NLU, intent classification; CLINIC and BANKING77, domain-specific intent classification; and InfBench-Sum, long-novel summarization. Baselines are Long-Context, Qwen3-32B restricted to the last 32k tokens; RAG-Top2, BM25 retrieval of two original chunks followed by Qwen3-32B; and the learned memory systems MemAgent and MEM1.

Training uses verl, the training framework, with Qwen3-4B, 32 H100 graphics processing units (GPUs), learning_rate 1e-6, batch_size 32, and grpo_rollout_n 8 for three days and 205 steps. The best checkpoint is selected by validation performance.
Section 4.1 reports β = 0.05 and γ = 1, whereas Section 4.4 and Tables 4 and 9 identify β = 0.05 and γ = 0.1 for the configuration matching the main results.
The reported average combines dataset-specific scores, including classification accuracy and summarization keyword coverage, the fraction of reference keywords present in generated summaries. It is not a single uniform metric.

4.2 Overall Performance Comparison

Table 1 reports an average validation score of 0.642 for Mem-α, compared with 0.588 for Long-Context, 0.567 for RAG-Top2, 0.236 for MemAgent, and 0.111 for MEM1. Table 2 reports test averages of 0.592, 0.461, 0.502, 0.198, and 0.071, respectively; Mem-α leads on five of the nine test tasks.

On Multi-Doc, Mem-α scores 0.680 versus 0.450 for RAG-Top2 and 0.270 for Long-Context, while storing 323K tokens versus RAG-Top2’s 474K.
On LME(S), RAG-Top2 leads with 0.581 versus Mem-α’s 0.520; Long-Context leads on NLU, Clinic, and Banking77 with 0.740, 0.860, and 0.770.
Average test storage is 129K tokens for Mem-α versus 207K for RAG-Top2, but exceeds Long-Context’s reported 33K. The storage advantage therefore depends on the baseline.
The maximum reported training length is 30k tokens, while Multi-Doc evaluation reaches 474K, supporting length extrapolation under this chunk-processing setup.

Validation task scores and stored-memory token counts across five methods.

MemoryAgentBench task scores and stored-memory token counts.

4.3 Performance Boost from Reinforcement Learning

Table 3 holds the memory architecture, BM25 retrieval, and Qwen3-32B answer generator fixed while changing the writer. Mem-α raises the average validation score from Qwen3-4B’s 0.389 to 0.642, exceeding gpt-4.1-mini’s 0.517 and supporting the benefit of training within this architecture.

Average stored memory increases from 3.9K tokens for the untrained Qwen3-4B writer to 7.9K for Mem-α, so the gain accompanies greater information retention.
gpt-4.1-mini performs better on BookSum, the book-chapter summarization task, with 0.246 versus 0.187.
This comparison does not independently establish the advantage of the chosen architecture over alternative memory designs.

Comparison of untrained Qwen3-4B, gpt-4.1-mini, and the trained Mem-α memory writer.

4.4 ABLATION Studies

Table 4 shows a validation average of 0.642 at β = 0.05, γ = 0.1, compared with 0.543 when γ = 0 and 0.630 when β = 0. Increasing β to 0.2 or 0.4 reduces average memory to 4.7K or 3.6K tokens but lowers average performance to 0.525 or 0.509.

Removing the content-validity reward lowers average performance, although retrieval scores remain substantial; the table supports degradation without establishing complete failure to learn useful memory construction.
On BookSum, adding β = 0.05 at γ = 0.1 reduces memory from 4.5K to 2.2K tokens while changing performance from 0.183 to 0.187.
The selected setting has the best reported average among these configurations, although other settings lead on individual tasks.

Validation performance and memory consumption for compression and content-reward weight settings.

4.5 Case Studies

Table 5 compares memory written from a conversation about condo living and noise reduction. The untrained Qwen3-4B writer leaves core memory empty and stores one generic semantic entry; gpt-4.1-mini separates semantic content but splits same-timestamp events and omits assistant behavior from episodic memory.

Mem-α writes an informative core summary, two distinct semantic entries, and a consolidated event covering both the user’s request and the assistant’s response.
This selected trace illustrates differences in memory organization but does not establish how frequently these behaviors occur.

Core, semantic, and episodic memory entries produced by three writers for a condo-advice conversation.

5 Conclusion, Limitation and Future Work

The paper concludes that training memory-writing policies improves downstream performance and supports processing sequences much longer than the training inputs. It proposes richer memory architectures and integration with real databases as future work, identifying latency, scalability, and safety as unresolved deployment concerns.

The experiments use simulated interaction formats and do not measure production database behavior, end-to-end deployment latency, or conflict resolution.
Training requires three days on 32 H100 GPUs despite the small sampled dataset.

Appendix

A DATASETS DETAILS and A.1 TRAINING DATASET describe combining SQuAD, a document question-answering dataset; distributing HotpotQA evidence across chunks for questions requiring information from multiple documents; preserving personal facts and events in PerLTQA; and constructing LongMemEval training contexts of 10K–30K tokens with reported exclusion of MemoryAgentBench evaluation data. PubMed-RCT, sentence-role classification in medical abstracts, replaces semantic labels with numeric labels 0–4, while NLU and TREC-C preserve label coverage within each instance. BookSum chapters are converted into conversational chunks and scored through coverage of keywords extracted from reference summaries.
A.2 EVALUATION DATASET reports 112 test instances across nine sources: 100 InfBench-Sum novels, five LongMemEval contexts, and one context for each of the other seven sources, each of those seven paired with 100 questions. InfBench-Sum is described as averaging 172k tokens per novel; keyword coverage measures selected content retention rather than all aspects of summary quality. Table 7’s listed training counts sum to 567 rather than its reported 562, and its listed validation counts sum to 458 rather than 463; it omits the five LME-Train validation instances in Table 6 and lists 50 LME-Train training instances rather than Table 6’s 45.
B FORMAL DEFINITIONS OF REWARD COMPONENTS restates correctness over reference answers, the fractions of executable and semantically valid operations, and compression from total input length to final-memory length. C EXPERIMENTAL DETAILS includes C.1 JUSTIFICATION OF BACKBONE MODEL SELECTION, which reports that Qwen3-8B repeatedly produced invalid memory_type values such as ‘semantic_memory’ and remained below Qwen3-4B in reward after the accepted argument values were adjusted. These observations explain backbone selection under this interface but do not establish a general ranking of the models.
C.2 BASELINE INTRODUCTION AND IMPLEMENTATION DETAILS identifies BytedTsinghua-SIA/RL-MemoryAgent-14B for MemAgent and https://huggingface.co/Mem-Lab/Qwen2.5-7B-RL-RAG-Q2-EM-Release for MEM1; MemAgent answers with its original model, whereas MEM1 uses Qwen3-32B after memory construction. C.3 PROMPTS USED IN TRAINING documents the shared writing prompt, Qwen3-32B validity checks for each memory type, and the answer-generation prompt; the checks reject placeholder core summaries and assess whether event and fact entries match their intended memory types. C.4 ADDITIONAL ABLATION STUDY reports test averages of 0.445 without content reward, 0.580 without compression reward, and 0.592 at β = 0.05, γ = 0.1; stronger compression weights lower the averages to 0.351 and 0.375.

Brief Thoughts

The comparison using the same architecture and downstream answer generator provides the most direct evidence for improved memory writing after training, accompanied by increased storage. The longer test inputs support length extrapolation, but several tasks contain only one test context, no uncertainty estimates are reported, and the long-context baseline is restricted to 32k tokens. The memory-content weight and dataset-count discrepancies need resolution for precise reproduction. The results show a task-dependent tradeoff between retained information and storage: Mem-α stores more than the truncated baseline and, on some tasks, more than the original chunks. Contradictory updates and production database behavior remain evaluation gaps.