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 training framework from Peking University and Samsung Research Beijing that teaches large language model agents general-purpose memory skills using only unlabeled Wikipedia text. Unlike prior memory-agent approaches that require expensive task-specific annotated data, MemTrain trains on 30k Wikipedia documents (24k-40k tokens each) via two coupled proxy tasks optimized with GRPO: end-to-end masked reconstruction, which rewards memories that let the model recover masked entities from the final memory state, and intermediate memory recall, which requires every intermediate memory to stay complete and retrievable throughout multi-turn document reading. Built on MemAgent's fixed-length (1024-token) memory architecture with a constant 8192-token context, the method lets Qwen3-4B and Qwen2.5-7B models process documents up to 896k tokens. On HotpotQA long-context benchmarks, MemTrain alone lifts Qwen3-4B from 22.0 to 56.2 average score, and combining MemTrain with task-specific MemAgent training yields up to 17.67 points of additional gain on Qwen2.5-7B, with the largest improvements at extreme context lengths. The paper positions general memory ability as a trainable, transferable capability analogous to language pretraining.

This forum post discusses MemTrain: Self-Supervised Context Memory Training (arXiv: https://arxiv.org/abs/2606.03197), a paper by Ziheng Li, Xingrun Xing, Haoqing Wang, Zhi-Hong Deng (Peking University), and Yehui Tang (Samsung Research, Beijing).

Core insight: memory learning is not the same as reasoning learning. Reasoning is single-turn and explicit; memory is multi-turn and latent — the model must continuously decide what to keep, compress, forget, and retrieve, so both outcomes *and* intermediate memory states need supervision.

Key points

The problem

Long-horizon LLM agents suffer from an "amnesia" problem: context windows overflow and early information gets pushed out.
Mainstream approaches collect task-specific annotated data and train with end-to-end RL. This is expensive, covers few memory behaviors (saving, compressing, forgetting, retrieving, integrating), and generalizes poorly.

The MemTrain framework

Built on MemAgent's context-memory architecture: a fixed-length memory state $$m_t$$ (1024 tokens) updated each turn from (m_{t-1}, a_{t-1}, i_t), with a constant 8192-token context (instructions 1024 + input chunk 5120 + memory 1024 + response 1024). This gives constant compute cost and lets models handle documents up to 896k tokens.
Two coupled proxy tasks trained from unlabeled Wikipedia:

End-to-end masked reconstruction

Intermediate memory recall

Used alone, task 1 encourages last-minute cramming; task 2 alone encourages bloated memories. Together they force concise, complete memories throughout.
Optimization: GRPO with total reward = reconstruction EM + 0.5 × recall reward, KL coefficient 1e-3, degenerate samples filtered. Training: 30k Wikipedia docs, Qwen3-4B-Instruct-2507 and Qwen2.5-7B-Instruct, 300 steps, lr 1e-6, batch 32, G1=8 end-to-end trajectories × G2=8 recall trajectories per group.

Results (HotpotQA long-context)

Qwen3-4B-Instruct, average over 7k–896k tokens:

| Method | 7k | 56k | 448k | 896k | Avg | |:---|:---:|:---:|:---:|:---:|:---:| | Base | 57.8 | 10.9 | 3.9 | 3.9 | 22.0 | | +MemTrain | 63.3 | 57.0 | 48.4 | 40.6 | 56.2 | | +MemAgent | 70.3 | 62.5 | 64.1 | 57.0 | 65.1 | | +MemTrain+MemAgent | 79.7 | 73.4 | 61.7 | 62.5 | 70.3 |

Qwen2.5-7B: base 20.8 → 45.4 with MemTrain alone; MemTrain+MemAgent reaches 73.5 avg (+17.67 over MemAgent alone).
Gains grow with context length: at 896k tokens base models score near zero while MemTrain+MemAgent keeps 62.5 / 68.8.
Search-based QA also improves: +10.58 (Qwen3-4B) and +8.50 (Qwen2.5-7B) avg over MemAgent alone.
Qualitatively, MemTrain-trained models retain key entities across multi-turn memory updates where baseline models forget them.

Limitations and future work

Only numeric/named entities are masked; only Wikipedia is used; memory length is fixed at 1024 tokens; GRPO is sample-hungry (72 trajectories per step).
Future directions: multi-domain corpora, hierarchical short/long-term memory, joint training with reasoning/planning/tool use, and more interpretable memory representations.

Takeaway

MemTrain frames general memory ability as a task-agnostic, self-supervised pretrained capability — like language pretraining, but for "how to remember." It provides a strong initialization that composes with task-specific memory training (MemAgent), at constant compute cost, without expensive downstream annotations.