# MemTrain: Self-Supervised Context Memory Training

Source: [MemTrain: Self-Supervised Context Memory Training](<https://www.sota2.com/research/paper/6a21065ec10484bac23efebf>)

Collection index: [llms.txt](<https://www.sota2.com/research/llms.txt>)

Data updated: 2026-06-07T06:00:01.589Z

Authors: Ziheng Li, Xingrun Xing, Haoqing Wang, Zhi-Hong Deng, Yehui Tang

Published: 2026-06-02T05:56:24.000Z

## 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.

## Recorded benchmark results

A research result applies to its stated dataset, split, metric and evaluation configuration. Compare like-for-like settings, preserve the metric's unit and direction, and distinguish paper-reported results from independent reproduction. The best result recorded here is not a claim of exhaustive, current global coverage.

- [Multi-hop Question Answering on HotpotQA (Size-Cut Evaluation)](<https://www.sota2.com/research/sota/multi-hop-question-answering-on-hotpotqa-size-cut-evaluation.md>): Multi-hop Question Answering; HotpotQA; Score (7k): 79.69.

- [Multi-hop Question Answering on TriviaQA](<https://www.sota2.com/research/sota/multi-hop-question-answering-on-triviaqa.md>): Multi-hop Question Answering; TriviaQA; Accuracy: 57.21.

- [Multi-hop Question Answering on NQ (Accuracy)](<https://www.sota2.com/research/sota/multi-hop-question-answering-on-nq-accuracy.md>): Multi-hop Question Answering; NQ; Accuracy: 35.18.

- [Multi-hop Question Answering on PopQA](<https://www.sota2.com/research/sota/multi-hop-question-answering-on-popqa.md>): Multi-hop Question Answering; PopQA; Accuracy: 38.36.

- [Multi-hop Question Answering on Musique](<https://www.sota2.com/research/sota/multi-hop-question-answering-on-musique.md>): Multi-hop Question Answering; MuSiQue; EM: not recorded.

## Original sources and implementation

[Original paper](<https://arxiv.org/abs/2606.03197>)

[PDF](<https://arxiv.org/pdf/2606.03197.pdf>)
