MemAgent: An RL Memory Agent That Reads 3.5M Tokens Without Forgetting

Forum topic · 小凯 · 2026-06-02

Summary

MemAgent, a collaboration between ByteDance Seed, Tsinghua AIR, and SIA-Lab accepted at ICLR 2026 (Oral), tackles long-context LLM inference by turning memory management into a learnable behavior rather than an architecture problem. The model reads documents segment by segment (about 5,000 tokens per segment) and overwrites a fixed 1,024-token textual memory after each segment, then answers questions using only that memory plus the query. Training uses Multi-Conv DAPO, extending DAPO reinforcement learning to multi-turn conversations: trajectory-level verifiable rewards are broadcast back to every memory-update turn, teaching the model what to keep and what to discard. Trained with only an 8K context window, MemAgent extrapolates to 3.5M tokens with under 5% performance loss (14B model), scores 95%+ on 512K RULER, and achieves linear O(N) computational complexity. This post explains the method, the credit assignment challenge, why memory overwrite does not destroy crucial information, comparisons against long-context continual pretraining baselines that collapse beyond 112K, limitations such as no backward retrieval, and cost advantages. Code and RL-MemAgent-7B/14B weights are open-sourced (arXiv:2507.02259).

Paper at a glance

| Item | Content | |------|---------| | Title | MemAgent: Reshaping Long-Context LLM with Multi-Conv RL-based Memory Agent | | 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 | | Date | July 3, 2025 | | Key results | Trained with an 8K context window, extrapolates to 3.5M tokens with <5% performance loss; 95%+ accuracy on 512K RULER; linear O(N) complexity |

The core idea: take notes like a student

Imagine sitting in an exam with a 3,000-page technical manual, but you're only allowed one sheet of A4 paper of notes. You'd read chapter by chapter, condensing the essentials onto that sheet — striking out what's irrelevant, adding what matters — then answer using only the sheet. That is essentially MemAgent: the "sheet" is a fixed 1,024-token segment of ordinary text in the model's context, and the "student" is an LLM trained via reinforcement learning.

Why long context is hard

Transformer attention is quadratic: doubling tokens quadruples compute. Existing approaches each have drawbacks:

Position extrapolation (RoPE scaling, interpolation) helps, but models never saw truly long-range dependencies during training, so quality degrades far out of distribution.
Sparse attention reduces complexity toward O(N), but deciding which tokens attend to which is itself an information problem.
New architectures (RNNs, Mamba, state-space models) abandon the entire pretrained ecosystem.

MemAgent's team asked a different question: without changing the model, can we *teach it to take notes*?

How it works

The model reads documents segment by segment. Each segment (~5,000 tokens) is read together with a fixed 1,024-token memory; after each segment, the model overwrites the memory with an updated one. When reading finishes, it discards the original text and answers using only the memory and the question. Memory is not a vector database or external store — it is plain text tokens, so no architecture changes are needed; any dense-attention Transformer can adopt the workflow.

Training with Multi-Conv DAPO

Each memory update is a separate conversational turn with no ground-truth label — only the final answer is verifiable. The team extends DAPO (Direct Advantage Policy Optimization) to the multi-turn setting: rollouts produce a full "memory chain," the final answer's reward is broadcast to all turns, and advantages are computed uniformly across the trajectory — turning trajectory-level sparse rewards into turn-level dense signals. Correct answers reinforce every note-taking step; wrong answers penalize them all.

A key trick: training uses only 8K context (1,024 for memory, 5,000 for the current segment, the rest for question/output), with training data capped at 32K — yet at test time the model handles 3.5M tokens, over 400× extrapolation.

Results

RULER-HotpotQA: accuracy stays essentially flat from 8K to 512K, while continual-pretraining baselines degrade sharply past 112K and approach zero near 896K.
512K RULER: MemAgent-14B exceeds 95% accuracy.
3.5M tokens: performance drop <5.5% for the 14B model, ~11% for 7B; many baselines collapse at 112K.
Complexity: processing is linear — a 2M-token document becomes ~400 fixed-cost segments.

> Note: RULER deliberately buries key information at varying depths in the middle of documents, preventing models from exploiting positional bias (answer-in-beginning/end shortcuts).

Why overwriting memory works

Overwriting risks losing details needed later. But RL rewards depend solely on final answer correctness, so any behavior that washes away crucial information gets penalized. The model learns selective compression — keeping "potentially useful evidence" rather than mechanically summarizing — similar to how humans take notes at a two-hour lecture. The paper does not, however, deeply analyze *what patterns* the learned memory encodes (explicit summarization vs. implicit attention routing), leaving that question open.

Limitations

Best suited to verifiable-reward tasks (QA, extraction, multi-hop reasoning); fuzzy tasks like literary criticism are hard to reward.
Segmented reading can sever cross-segment dependencies; there is no backward retrieval — the model cannot re-read earlier pages.
Linear complexity still carries a constant factor: hundreds of forward passes introduce latency unsuitable for real-time applications.

Where it sits in the field

MemAgent exemplifies a shift from "make the model *longer*" to "teach the model to *handle long*": an 8K-trained model that learned note-taking beats 1M-window continual-pretrained baselines that never learned to prioritize. It reframes memory from an architecture problem into a behavior problem — memory is not storage, but policy.

Practical outlook

Code and weights (RL-MemAgent-7B and 14B) are open-sourced. The approach fits scenarios with very long documents plus precise Q&A: financial reports, legal contracts, medical literature, system logs. Cost-wise, processing 1M tokens in 200 MemAgent segments may be roughly an order of magnitude cheaper than single-pass dense attention at current GPU rental prices.

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.