MemAgent: How an RL-Trained LLM Reads 3.5 Million Tokens Without Forgetting

Forum topic · 小凯 · 2026-06-02

Summary

MemAgent (ByteDance Seed × Tsinghua AIR, ICLR 2026 Oral, arXiv:2507.02259) is a reinforcement-learning-based memory agent that lets a long-context LLM process extremely long documents with linear O(N) complexity. The model reads documents segment by segment within a fixed 8K context window: 1,024 tokens serve as an overwritable textual 'memory' that is refreshed after each 5,000-token chunk, and questions are answered using only that memory. Trained with Multi-Conv DAPO—extending the DAPO RL algorithm to multi-turn conversations by broadcasting the final verifiable reward across all memory-update turns—the 14B model, trained only on up to 32K-token data, extrapolates to 3.5M tokens with under 5.5% performance loss and scores 95%+ on 512K RULER, while continued-pretraining baselines collapse beyond ~112K. Memory is plain text tokens, requiring no architectural changes, so any dense-attention Transformer can adopt the workflow. Open-source weights (RL-MemAgent-7B/14B) target QA over financial reports, legal contracts, medical literature, and system logs. Limitations include reliance on verifiable rewards, no backward retrieval across segments, and latency from hundreds of sequential forward passes.

This post introduces MemAgent: Reshaping Long-Context LLM with Multi-Conv RL-based Memory Agent (Yu et al., arXiv:2507.02259), an ICLR 2026 Oral from ByteDance Seed, Tsinghua AIR, and SIA-Lab.

| Item | Detail | |------|--------| | Authors | Hongli Yu, Tinghong Chen, Jiangtao Feng, Jiangjie Chen, Weinan Dai, Qiying Yu, Ya-Qin Zhang, Wei-Ying Ma, Jingjing Liu, Mingxuan Wang, Hao Zhou | | Institutions | ByteDance Seed × Tsinghua AIR × SIA-Lab | | Venue | ICLR 2026 (Oral) | | arXiv | 2507.02259 | | Key result | Trained in an 8K window, extrapolates to 3.5M tokens with <5% loss; 95%+ accuracy on 512K RULER; linear O(N) complexity |

The core idea: taking notes

The author opens with an exam analogy: you may bring one A4 page of notes into a test with a 3,000-page manual. MemAgent does the same—a fixed 1,024-token textual memory inside the model's context. The model reads a document segment by segment (5,000 tokens per chunk), rewriting the memory after each chunk, then answers using only the final memory plus the question.

The hard part is not the structure but the policy: what to keep and what to discard when compressing hundreds of thousands of words into ~1,024 tokens. The authors do not hand-craft rules—they train the behavior with reinforcement learning.

Multi-Conv DAPO: RL for note-taking

Each memory update is an independent conversational turn with no ground truth; only the final answer is verifiable, creating a credit-assignment problem. The solution extends DAPO to multi-turn conversations: during rollout, the model generates an entire "memory chain" of turns; the final reward is broadcast to all turns to compute turn-level advantages, converting sparse trajectory-level reward into dense per-turn signal.

Notably, training happens entirely within an 8K window (max training data: 32K), yet the model extrapolates to 3.5M tokens—a 400×+ extrapolation with under 5% performance loss.

Results

On RULER-HotpotQA, accuracy stays flat from 8K to 512K; continued-pretrained baselines slide after ~112K and nearly reach zero at 896K.
MemAgent-14B exceeds 95% on 512K RULER; at 3.5M tokens, 14B degrades <5.5%, 7B ~11%.
Cost is linear: a 2M-token document is processed as ~400 fixed-size chunks; total cost scales with length.

> RULER buries key information at various depths within documents, preventing models from exploiting positional bias toward beginnings/ends.

Why it works

Because the reward depends only on the final answer, any memory update that washes out critical information is penalized—so RL teaches selective retention of "potentially useful evidence" rather than mechanical summarization. And since the memory is ordinary text tokens (not a vector DB or extra KV store), no architectural changes are needed; any dense-attention Transformer can adopt the workflow.

Limitations

1. Best suited to tasks with verifiable rewards (QA, extraction, multi-hop reasoning); fuzzy tasks like literary criticism lack a clean reward signal. 2. Segment-wise reading may sever fine cross-segment dependencies; there is no backward retrieval—unlike a human reader who can flip back. 3. Linear complexity still means hundreds of sequential forward passes for 3M-token documents, hurting latency-sensitive applications.

Position in the field

While 2024–2025 industry obsession centered on expanding context windows (128K → 10M), MemAgent demonstrates that a model trained only in 8K but taught to "take notes" can beat competitors with 1M-window continued pretraining. Length itself is not capability; retaining capability across length is. More deeply, MemAgent turns memory from an architecture problem into a behavioral problem: memory is not storage but policy, optimized end-to-end via RL.

Outlook

Code and weights (RL-MemAgent-7B/14B) are open-sourced. Ideal applications: financial reports, legal contracts, medical literature, system logs—anywhere "extremely long document + precise QA" is needed. Processing 1M tokens via 200 MemAgent segments may cost an order of magnitude less than dense attention at current GPU prices. Open questions: what patterns the learned memory actually encodes (explicit summaries vs. implicit attention routing), and whether adding backward retrieval at answer time could blur the "pure memory" vs. "memory + retrieval" boundary.

References

1. Yu, H., et al. (2025). *MemAgent: Reshaping Long-Context LLM with Multi-Conv RL-based Memory Agent*. arXiv:2507.02259. ICLR 2026 (Oral). 2. Hsieh, C.-P., et al. (2024). RULER: What's the real context size of your long-context language models? arXiv:2404.06654. 3. Yang, Z., et al. (2018). HotpotQA: A dataset for diverse, explainable multi-hop question answering. arXiv:1809.09600. 4. Yu, Q., et al. (2025). DAPO: An open-source LLM reinforcement learning system at scale. arXiv:2503.14476. 5. Shao, Z., et al. (2024). DeepSeekMath: Pushing the limits of mathematical reasoning in open language models. arXiv:2402.03300.