Source-linked AI summary
MemTrain: Self-Supervised Context Memory Training
Ziheng Li, Xingrun Xing, Haoqing Wang, Zhi-Hong Deng, Yehui Tang
TL;DR
Long-horizon agents need general memory capabilities, but downstream reinforcement-learning approaches require costly, insufficiently diverse labeled data. MemTrain uses coupled self-supervised reconstruction and recall tasks over Wikipedia, jointly optimized with GRPO, and consistently improves downstream long-text and search-based QA across models, with gains up to 17.67 points. Its evaluation and training evidence does not establish performance beyond the reported tasks and model settings.
Problem
Existing memory-agent training requires costly labeled data that may not cover diverse, general memory behaviors.
Method
MemTrain jointly optimizes end-to-end masked reconstruction and intermediate memory recall over Wikipedia passages using GRPO.
Results
MemTrain consistently improves downstream long-text QA and search-based QA across models, with gains up to 17.67 points over direct task-specific post-training.
Takeaways & Limitations
Outcome-level and process-level supervision improves the general-purpose context-memory capability used in downstream post-training.
Takeaways & Limitations
Poorly constructed memories can force task solving from inadequate information and lead to more severe hallucination under long-horizon settings.
Abstract
from arXiv · showhide
Memory is an indispensable capability for long-horizon LLM agents, enabling them to preserve and utilize information accumulated across extended interactions. Existing memory-agent approaches are typically trained end-to-end with reinforcement learning on downstream tasks. However, collecting high-quality annotated problems for memory-intensive scenarios is costly, and the resulting training data often lack sufficient diversity to cover general memory behaviors. In this work, we propose MemTrain, a self-supervised training framework for generally enhancing the context-memory capability of LLM agents for more effective downstream post-training. MemTrain introduces two coupled proxy tasks over unlabeled Wikipedia corpora: (1) an end-to-end masked reconstruction objective, which requires the model to recover masked entities after multiple rounds of memory updates, thereby encouraging memory maintenance from the final outcome perspective; and (2) an intermediate memory recall objective, which requires the model to reconstruct masked historical information using intermediate memory states, encouraging faithful compression and memory completeness throughout the interaction process. The two objectives are jointly optimized using GRPO. Extensive experiments on long-text QA and search-based QA benchmarks demonstrate that MemTrain consistently improves downstream memory-intensive reasoning performance across different models, achieving gains of up to 17.67 points over direct task-specific post-training.
1 Introduction
Long-horizon agents need fixed-size memory to preserve information across growing interactions, but existing training is costly, domain-specific, and insufficiently diverse. MemTrain addresses this gap with coupled self-supervised proxy tasks that improve downstream memory-intensive QA.
- Fixed-size context memory helps agents preserve historical information while avoiding the cost and attention pressure of retaining full interaction histories.
- Existing memory agents typically rely on reinforcement learning with expensive labeled downstream data, limiting behavioral diversity and generalization.
- Memory proxy-task design is challenging because targets are latent and process-dependent, requiring continuous decisions about preservation, compression, and recall.
- MemTrain jointly trains end-to-end masked reconstruction and intermediate memory recall on Wikipedia passages using GRPO.The tasks provide outcome-level and process-level supervision for memory maintenance, compression, completeness, and utilization.
- MemTrain consistently improves downstream long-text QA and search-based QA performance across 4B and 7B models.Reported average gains include 5.17 and 10.58 points for Qwen3-4B-Instruct-2507, and 17.67 and 8.50 points for Qwen2.5-7B-Instruct.
2 Related Works
Related work contrasts full-history context handling with fixed-length context memory and motivates self-supervised reinforcement training from unlabeled corpora. These approaches address context limits and the scalability and generalization constraints of curated downstream data.
- Conventional agents append environmental observations and responses to the context, which is limited by the finite context window of LLMs.
- External memory systems compress or summarize interaction records for storage outside the model’s context.
- Context memory agents maintain a fixed-length memory state updated at each interaction step instead of retaining the entire history.
- Reinforcement post-training commonly depends on curated question-answer datasets, limiting scalability and generalization.
3 Self-Supervised Memory Training
MemTrain trains context memory through coupled end-to-end reconstruction and intermediate recall tasks over multi-turn Wikipedia interactions, jointly optimized with GRPO. The design targets both final masked-entity recovery and the quality and retrievability of intermediate memory states.
- Context Memory Agent: Context memory maintains a fixed-size state across interaction steps, avoiding dependence on an ever-growing context and enabling long-horizon interaction beyond the native context limit.The model updates memory and action from the previous state and current interaction input at each step.
- Training Sample Construction: MemTrain constructs long documents from Wikipedia passages, masks an entity, segments the document into chunks, and processes them sequentially as interaction steps.Samples combine semantically related and randomly selected passages before masking all occurrences of a selected entity.
- End-to-End Masked Reconstruction: The end-to-end task requires recovering the masked entity after repeated memory updates, with the final prediction generated solely from the resulting memory state.Because the entity is masked throughout the document, success requires long-range information aggregation rather than local copying.
- Intermediate Memory Recall: Intermediate recall supplements coarse end-to-end rewards by constraining intermediate memory quality and reducing the weakness caused by multi-step error accumulation.The objective encourages memory states that preserve sufficient historical information for direct retrieval and downstream reasoning.
- Intermediate Memory Recall: Intermediate Memory Recall samples an intermediate memory state and an earlier chunk, then requires reconstructing the masked entity in a single interaction.This directly tests whether historical information remains available in the current memory representation.
- Joint GRPO Optimization: GRPO jointly optimizes end-to-end and intermediate-recall trajectories, combining final-prediction rewards with associated recall rewards through a balancing coefficient.For each end-to-end trajectory, the method constructs an IMR prompt and samples additional IMR trajectories before computing rewards and advantages.
4 Experiments
MemTrain is evaluated as a self-supervised memory-training initialization for long-context and search-based QA post-training. It consistently improves downstream performance across models, context lengths, and search benchmarks.
- Evaluation: MemTrain is evaluated on long-context multi-hop QA and search-based multi-hop QA, covering in-domain and out-of-domain memory-intensive settings.The long-context evaluation uses HotpotQA with inputs from 7k to 896k tokens, while search-based evaluation spans seven multi-hop QA benchmarks.
- Long-Text QA Results: 17.67% higher average performance is achieved by MemTrain+MemAgent over direct MemAgent on Qwen2.5-7B-Instruct.The corresponding gain on Qwen3-4B-Instruct is 5.17%.
- Long-Text QA Results: The gains remain consistent from 7k to 896k tokens, and MemTrain reduces Qwen2.5-7B-Instruct’s 28k-to-896k drop from 21.09% points to 8.59% points.MemTrain also transfers to shorter contexts such as 7k and 14k.
- Long-Text QA Results: MemTrain alone raises average performance from 21.97% to 56.15% on Qwen3-4B-Instruct and from 20.80% to 45.41% on Qwen2.5-7B-Instruct.These results are obtained without labeled supervision.
- Search-Based QA Results: 10.58 points and 8.50 points are gained over MEM1 on Qwen3-4B-Instruct-2507 and Qwen2.5-7B-Instruct, respectively, across search-based QA benchmarks.Improvements are reported consistently across all benchmarks and are larger on harder multihop tasks.
5 Analysis
The analysis shows that intermediate memory recall and coupled supervision are important for performance across long contexts. MemTrain also outperforms simply extending downstream post-training and preserves critical information in a representative case.
- 5.1 Ablation Study: Removing intermediate memory recall lowers average performance from 70.31% to 63.28% across evaluated context lengths.The Full model outperforms both the End-to-End and Decoupled variants, showing the importance of IMR.
- 5.1 Ablation Study: The Decoupled variant is stronger than End-to-End at contexts ≤56k but deteriorates significantly as context length increases.The reported pattern suggests that IMR supports memory utilization, while decoupling it from end-to-end rewards weakens long-context behavior.
- 5.1 Ablation Study: The analysis attributes severe long-horizon hallucination in the Decoupled variant to insufficient guidance for high-quality memory generation.Poorly constructed memories force the model to solve tasks with inadequate retained information.
- 5.2 Scaling Post-Training: Further post-training yields only marginal gains or degradation, while MemTrain initialization retains a 2.64 percentage-point advantage at the best-performing extended checkpoint.This comparison uses additional post-training beyond step 500.
- 5.3 Case Study: MemTrain preserves Adriana Trigiani’s location from chunk 1, enabling correct answer deduction in chunk 2 where direct MemAgent fails.The case study examines memory retention across a memory update step.
6 Conclusion
MemTrain is a self-supervised framework that trains general-purpose memory capabilities in LLMs through two coupled proxy tasks. It improves downstream performance on long-text and search-based question answering across two models.
- MemTrain trains general-purpose LLM memory using end-to-end masked reconstruction and intermediate memory recall.The coupled objectives encourage memory completeness, faithful compression, and effective utilization.
- The framework performs memory training on Wikipedia corpora rather than relying solely on labeled downstream tasks.
- MemTrain consistently improves downstream performance on long-text and search-based question answering across two models.
End-to-End Memory Generation Prompt
The prompt templates implement end-to-end memory updating and answer generation, alongside intermediate recall from previously formed memory states. They repeatedly process text sections while asking the model to recover masked entities.
- The memory generation prompt asks the model to update previous memory with useful information from each new section.
- The end-to-end answer generation prompt asks the model to identify the masked entity using the memory and section.
- The intermediate memory recall prompt asks the model to predict a target entity from an existing memory and a masked section.
A Prompt Template
MemTrain uses three prompt templates to support iterative memory generation, final answer generation, and intermediate memory recall. The end-to-end task processes all chunks before producing the final answer, while recall uses a masked target.
- MemTrain employs three prompt templates for its end-to-end reconstruction and intermediate recall tasks.
- The end-to-end memory generation prompt is applied iteratively until all text chunks have been processed.
- After processing the chunks, an answer generation prompt produces the final output.
- The intermediate recall task introduces the placeholder [TARGET] for reconstructing masked information.