MemTrain: Self-Supervised Context Memory Training Gives LLM Agents Long-Term Memory Without Labeled Data

Forum topic · 小凯 · 2026-06-07

Summary

MemTrain is a self-supervised context memory training framework from Peking University and Samsung Research Beijing that teaches LLM agents to maintain long-range memory without task-specific labeled data. Built on a fixed-length memory state architecture (MemAgent-style), MemTrain uses two coupled proxy tasks on unlabeled Wikipedia text: end-to-end masked reconstruction, which supervises the final memory's ability to answer questions, and intermediate memory recall, which supervises that memory stays complete and retrievable at every step. The two tasks are jointly optimized with GRPO (lambda=0.5, KL coefficient 1e-3). Trained on 30k Wikipedia documents (24k-40k tokens) with Qwen3-4B-Instruct and Qwen2.5-7B-Instruct, MemTrain alone lifts average HotpotQA performance from 22.0 to 56.2 (Qwen3-4B) and from 20.8 to 45.4 (Qwen2.5-7B) across inputs up to 896k tokens. Combined with task-specific MemAgent training, it adds up to +17.67 points on Qwen2.5-7B and +10.58 on search-based QA, with the largest gains at extreme context lengths. The key insight: memory learning differs from reasoning learning — memory objectives are hidden-state and process-dependent, requiring supervision of both outcomes and intermediate states. MemTrain acts as a general-purpose enhancer, transferable to long-document reading, multi-turn dialogue, and long-horizon agent tasks at constant compute cost.

Overview

Paper: MemTrain: Self-Supervised Context Memory Training Authors: Ziheng Li, Xingrun Xing, Haoqing Wang, Zhi-Hong Deng (Peking University), Yehui Tang (Samsung Research, Beijing) Link: https://arxiv.org/abs/2606.03197 Core insight: Memory learning is not reasoning learning. Memory objectives are inherently hidden-state and process-dependent, requiring supervision of both results and process.

Key points

The problem: Long-horizon LLM agents suffer from an "amnesia" issue — fixed context windows cause earlier key information to be pushed out. Existing approaches rely on expensive downstream task annotations for end-to-end RL, covering few memory behaviors and generalizing poorly.
Memory differs from reasoning: Reasoning is single-turn and explicit (reward from the final answer); memory is multi-turn and implicit — the model must continually decide what to retain, compress, forget, and retrieve, so intermediate memory states must also be evaluated.
Architecture: Fixed-length memory state $$m_t$$ (1024 tokens) updated each turn: input is (previous memory $m_{t-1}$ , previous action $a_{t-1}$ , environment feedback $$i_t$$ ), output is (new memory $$m_t$$ , action $$a_t$$ ). Total context stays fixed at 8192 tokens regardless of interaction length, enabling documents up to 896k tokens with constant compute and no attention dilution.

The two proxy tasks

1. End-to-end masked reconstruction (result supervision)

Pick a center paragraph from Wikipedia, add 29 semantically related + 120 random paragraphs (24k-40k token documents).
Mask one entity (number, place, person name, etc.) everywhere via spaCy NER.
Split into ~5k-token chunks; the model reads over ~8 turns, updating memory, and must reconstruct the masked entity from the final memory only.
Reward is high only if memory is both concise and complete.

2. Intermediate memory recall (process supervision)

At a random mid-trajectory memory state, ask the model to recall a masked entity from earlier chunks directly from that intermediate memory.
Forces memory to remain complete and retrievable at every step, not just at the end.

Complementarity

| Dimension | Masked reconstruction | Intermediate recall | |---|---|---| | Supervision | Result (final memory) | Process (intermediate memory) | | Prevents | Redundant memory, final failure | Mid-trajectory loss, unretrievable memory | | Encourages | Concise, key information retention | Complete, faithful compression throughout |

Alone, task 1 lets the model cram key info only at the end; alone, task 2 encourages hoarding everything. Together the model must keep memory concise and complete throughout.

GRPO optimization

Total reward = end-to-end exact-match reward + λ × intermediate recall reward (λ = 0.5).
GRPO uses group-relative rewards (no critic), KL regularization (coefficient 1e-3), and filters all-0/all-1 reward samples.
Training: 30k Wikipedia docs, 8192-token context (1024 instruction + 5120 input chunk + 1024 memory + 1024 response), batch size 32, G1=8 end-to-end trajectories per group each with G2=8 recall trajectories, 300 steps, LR 1e-6.
Models: Qwen3-4B-Instruct-2507 and Qwen2.5-7B-Instruct.

Results

HotpotQA long-context (7k-896k tokens)

Qwen3-4B-Instruct (average):

| Method | Avg | |---|---| | Base | 22.0 | | +MemTrain | 56.2 | | +MemAgent | 65.1 | | +MemTrain+MemAgent | 70.3 |

Qwen2.5-7B-Instruct (average):

| Method | Avg | |---|---| | Base | 20.8 | | +MemTrain | 45.4 | | +MemAgent | 55.9 | | +MemTrain+MemAgent | 73.5 |

MemTrain alone: +34.2 (Qwen3-4B) and +24.6 (Qwen2.5-7B) average points.
MemTrain+MemAgent vs. MemAgent alone: +5.17 (Qwen3-4B) and +17.67 (Qwen2.5-7B).
At 896k tokens, base models collapse (3.9 / 0.0) while MemTrain+MemAgent retains 62.5 / 68.8 — gains grow with context length.
Search-based QA: MemTrain+MemAgent beats MemAgent alone by +10.58 (Qwen3-4B) and +8.50 (Qwen2.5-7B) average.

Qualitative example

Given the question "What is Adriana Trigiani's profession?" with the key fact split across chunks 1 and 2, the non-MemTrain model drops it during a later memory update, while the MemTrain model carries the entity through multi-turn updates and answers correctly.

Design rationale

Why not next-token prediction RL (RPT/RLPT/PretrainZero)? Those suit explicit single-turn reasoning rewards; memory's hidden intermediate states require the two-angle supervision (end-to-end usability + intermediate usability).
Why Wikipedia? Broad entity types, no annotation needed, retrieved related paragraphs simulate real cross-document entity references, and 30k docs suffice for general memory strategies.
Fixed 8192-token budget (instruction/chunk/memory/response) keeps compute constant regardless of document length.

Limitations and future work

Only numbers and named entities are masked; concepts, relations, and events are not yet covered.
Only Wikipedia; memory behaviors for code, dialogue, or scientific text may differ.
Fixed 1024-token memory may be insufficient for complex tasks; adaptive memory length is future work.
GRPO is sample-heavy (72 trajectories per step).
Future directions: multi-domain training, hierarchical (short+long-term) memory, joint training with reasoning/planning/tool use, and interpretable memory (explicit key-value or summaries).

Takeaways for practitioners

Memory is a general capability that can be pre-trained independently of tasks, like language itself.
One self-supervised training run on Wikipedia improves all memory-intensive downstream tasks — no expensive per-task annotation.
MemTrain provides the initialization ("how to take notes"), MemAgent the task specialization ("how to write meeting minutes"); combined they are strongest, with constant compute cost and context limits far exceeded.

References

Li Z, Xing X, Wang H, et al. MemTrain: Self-Supervised Context Memory Training. arXiv:2606.03197, 2026.
Yu Z, et al. MemAgent: Context Memory for LLM Agents. 2025.
Dong L, et al. RPT: Reasoning Pre-Training. 2025.
Xing X, et al. PretrainZero: Active Pre-Training. 2025.