MemAgent: Reshaping Long-Context LLM with Multi-Conv RL-based Memory Agent

TL;DR

MemAgent introduces a fixed-length, overwrite-based memory for LLMs to process arbitrarily long inputs with linear processing costs. By training with multi-conversation DAPO, MemAgent learns to compress and preserve task-relevant information across segmented reads, enabling end-to-end RL optimization of context memory. Empirical results show strong extrapolation capabilities, maintaining high accuracy up to 3.5 million tokens and robust OOD performance, outperforming long-context baselines and arguing for memory-augmented RL as a scalable path for long-context reasoning. The work provides a practical framework for turning moderate-context LLMs into efficient, long-context reasoners with minimal architectural changes.

Abstract

Despite improvements by length extrapolation, efficient attention and memory modules, handling infinitely long documents with linear complexity without performance degradation during extrapolation remains the ultimate challenge in long-text processing. We directly optimize for long-text tasks in an end-to-end fashion and introduce a novel agent workflow, MemAgent, which reads text in segments and updates the memory using an overwrite strategy. We extend the DAPO algorithm to facilitate training via independent-context multi-conversation generation. MemAgent has demonstrated superb long-context capabilities, being able to extrapolate from an 8K context trained on 32K text to a 3.5M QA task with performance loss < 5% and achieves 95%+ in 512K RULER test.

Figures (6)

• Figure 1: Accuracy scores of RULER-HotpotQA hsieh2024ruleryang2018hotpotqa . Even models that employ long-context continual pretraining and extrapolation techniques fail to maintain consistent performance. In contrast, MemAgent with RL demonstrates nearly lossless performance extrapolation.
• Figure 2: MemAgent is inspired by the way humans process long documents. It divides the document into multiple chunks and allows LLMs to process them iteratively, recording relevant information in memory. Finally, LLMs generate answers based on the information stored in the memory.
• Figure 3: Comparison between vanilla GRPO and Multi-Conv DAPO. During the rollout phase of Multi-conv DAPO, each sample generates multiple conversations. The answer contained in the final conversation is used to compute the reward and advantage, which are then employed to optimize all preceding conversations.
• Figure 4: The architecture and graphic model of MemAgent. The memory is modeled as a latent memory variable, thereby enabling the decomposition of the autoregressive language model into multiple steps of reading from and writing to the memory.
• Figure 5: Ablation study on RULER-HotpotQA comparing models with and without RL training across context lengths from 28K to 896K tokens.