MemTrain: Self-Supervised Context Memory Training

2026-06-02

AI Summary

This work addresses the limited long-range contextual memory and reliance on costly annotated data in existing large language model (LLM) agents. To overcome these challenges, the authors propose MemTrain, a self-supervised memory training framework that requires no human annotations. MemTrain jointly optimizes two proxy tasks—masked entity reconstruction and intermediate memory recall—on unlabeled Wikipedia corpora to systematically enhance the model’s ability to retain and compress information. This approach constitutes the first end-to-end self-supervised method for training general-purpose memory capabilities in LLMs and is compatible with diverse model architectures. Evaluated on long-form question answering and search-based QA benchmarks, MemTrain significantly improves memory-intensive reasoning performance, achieving gains of up to 17.67 points over standard fine-tuning.

Abstract

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.

Problem

• memory training
• long-horizon LLM agents
• self-supervised learning
• memory-intensive reasoning
• context memory

Innovation

• self-supervised training
• context memory
• masked reconstruction
• memory recall
• GRPO