Source-linked AI summary
MemAgent: Reshaping Long-Context LLM with Multi-Conv RL-based Memory Agent
Hongli Yu, Tinghong Chen, Jiangtao Feng, Jiangjie Chen, Weinan Dai, Qiying Yu, Ya-Qin Zhang, Wei-Ying Ma, Jingjing Liu, Mingxuan Wang, Hao Zhou
TL;DR
Long-context LLMs still need to process extremely long documents without extrapolation degradation or superlinear cost. MEMAGENT addresses this with an RL-trained, segment-based memory workflow and Multi-Conv DAPO, achieving strong extrapolation from limited context to million-token QA tasks. Its memory capacity introduces a trade-off between retaining useful information and managing redundancy or insufficient storage.
Problem
Long-context processing still lacks a demonstrated way to handle infinitely long documents with linear complexity and without performance degradation during extrapolation.
Method
MEMAGENT uses reinforcement learning to train a fixed-length memory that is selectively updated while the LLM processes text segments, with Multi-Conv DAPO optimizing independent conversations from verifiable outcome rewards.
Results
MEMAGENT extrapolates from 8K context to QA documents up to 3.5 million tokens with no performance drop and linear computation cost, while showing strong performance across long-context tasks.
Takeaways & Limitations
The results support memory-based reinforcement learning as a scalable approach for long-context tasks and show that memory capacity influences performance across task types.
Takeaways & Limitations
Memory size creates a trade-off: larger memories can store more useful information but increase management challenges and redundancy, while smaller memories may lack necessary references.
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
MEMAGENT addresses long-context processing by iteratively reading document chunks, selectively updating fixed-length memory, and generating answers from the accumulated information. Its multi-conversation DAPO training supports linear-time processing and strong extrapolation to extremely long documents.
- 1 INTRODUCTION: Current LLMs struggle to process entire books, long reasoning chains, and agent histories because generated text can exceed typical context windows.
- 1 INTRODUCTION: Long-context systems must process arbitrarily long text, avoid significant performance loss during scaling, and decode efficiently with linear complexity.
- 1 INTRODUCTION: MEMAGENT processes text segment-by-segment while selectively updating a dynamically maintained fixed-length memory for final answer generation.
- 1 INTRODUCTION: Multi-Conv DAPO generates multiple conversations per sample and uses the final conversation’s reward and advantage to optimize preceding conversations.
- 1 INTRODUCTION: 8K context and 1024-token memory models trained on 60K-length documents answered questions on documents up to 3.5 million tokens with no performance drop and linear computation cost.
- 1 INTRODUCTION: The approach is designed to process arbitrarily long inputs within limited context windows, using end-to-end multi-conversation training and minimal performance degradation during extrapolation.
2 METHODOLOGY
MemAgent processes long documents as a stream of chunks, repeatedly reading and overwriting a fixed-length memory, while RL trains memory updates for answer quality. This design separates context processing from answer generation and yields linear scaling with input length.
- Workflow: The workflow separates chunk-level context processing, which updates memory, from final answer generation, which consults the problem and accumulated memory.The base LLM's tokenization and attention layout remain unchanged across both modules.
- Training: Multi-conversation DAPO extends loss computation from group–token structure to group–conversation–token structure for independent-context generation.The final-answer conversation supplies the advantage, which is uniformly applied to conversations from the same sample.
- Workflow: MemAgent processes arbitrarily long documents by iterating over chunks and maintaining a fixed-length memory that summarizes prior evidence.The memory is overwritten after each chunk, so its length remains constant.
- Scalability: O(N) end-to-end complexity follows because each chunk uses constant-size memory and bounded local context, keeping per-step compute and memory at O(C + M).Here, C bounds chunk length and M is the constant memory length.
- Training: Reinforcement learning treats memory updates as a policy and optimizes the read–write trajectory using the final task reward.The formulation models reading and writing as a Markov Decision Process.
3 EXPERIMENTS
Experiments evaluate MEMAGENT across long-context QA, retrieval, summarization, ablations, and context-distribution probes. RL-MEMAGENT maintains strong performance at extreme lengths, while ablations show that reinforcement learning is central and memory length involves a robustness–capacity trade-off.
- Main Results: MEMAGENT shows marginal performance decay through 3.5M-token RULER-HQA contexts, whereas several baselines degrade substantially or reach zero earlier.Qwen2.5-Instruct-1M reaches zero at 896K tokens, while MEMAGENT is evaluated at 1.75M and 3.5M.
- Main Results: MEMAGENT demonstrates superior overall performance on LongBench-QA, outperforming larger long-context or reasoning models.The result supports generalization of reinforcement-learning-acquired memory capabilities to diverse, information-dense tasks.
- Main Results: RL-MEMAGENT achieves only minimal performance loss of less than 5% at 512K on NIAH, despite evaluation involving more than 100 dialogue turns.
- Main Results: RL-MEMAGENT achieves state-of-the-art performance on almost all LongBench-SUM ROUGE-{1,2,L} recall metrics.This indicates memory and context management capabilities extend beyond question answering.
- RL Training: MEMAGENT without reinforcement learning declines as input length increases, while RL-MEMAGENT improves substantially across RULER-HQA and LongBench-QA.Directly applying MEMAGENT yields only marginal or negative improvements, indicating reinforcement learning is essential for robust memory abilities.
- Memory-Length Ablation: Memory lengths from 256 to 4096 show that the chosen configuration is a reasonable sweet spot and that MEMAGENT remains robust across memory sizes.Larger memory increases storage capacity but also raises management and redundancy challenges; smaller memory can provide insufficient references.
- Context Distribution: MEMAGENT remains consistently robust across tested key-information position patterns without catastrophic performance degradation.The probe uses five context distributions ranging from beginning- and end-position groups to intermediate placements.
4 RELATED WORK
Related work addresses long-context processing through positional extrapolation, alternative attention or sequence architectures, memory mechanisms, and reinforcement-learning training. These lines of work provide the main technical context for MEMAGENT.
- Long Context LLMs: RoPE extrapolation methods modify positional embeddings to capture long-range semantic dependencies.Examples include NTK, PI, YaRN, and DCA.
- Long Context LLMs: Linear attention, recurrent and state-space models, and sparse attention pursue alternative architectures for long-context processing.
- Memory Mechanism: Transformer memory mechanisms commonly add external memory modules or extend the model’s memory-related components.
- Reinforcement Learning for LLMs: Recent LLM reinforcement learning increasingly uses rule-based feedback alongside PPO- or GRPO-based optimization.
5 CONCLUSION
MEMAGENT is an RL-trained memory-module method that selectively records relevant information and extrapolates long-context capability from 8K contexts to 3.5M tokens. Its ablations identify RL training and memory capacity as important factors.
- MEMAGENT uses an RL-trained memory module to selectively record relevant information while disregarding extraneous details.
- Training on 60K-length sequences enables MEMAGENT to extend its effective context to 3.5M tokens with an 8K context window.
- MEMAGENT achieves state-of-the-art performance across a diverse range of long-context tasks.
- Ablations identify a critical role for RL-based training and show how memory capacity influences performance across task types.
REPRODUCIBILITY STATEMENT
The paper supplies implementation, algorithm, training, evaluation, and template details for reproducing MEMAGENT. Its training procedure uses Multi-conv DAPO with independent-context conversations and rule-based rewards.
- Reproducibility Statement: Reproducibility materials include the prompt template, pseudocode, training recipe, algorithm hyperparameters, and evaluation settings.The authors state that code, datasets, and model weights will be available through open-source platforms.
- Templates: MEMAGENT’s template details cover both context processing and final answer generation, with placeholders replaced by actual content.
- Algorithm: Multi-conv DAPO samples multiple context-independent conversations for one prompt and scores each final response with a rule-based verifier.
- Algorithm: The algorithm updates the policy using clipped-objective gradients after generating grouped conversations and computing verifier rewards.
A.3 TRAINING
Training uses staged long-document QA data, progressing from synthetic 32K-token HotpotQA-based instances to difficult 60K-token examples while applying off-policy DAPO optimization.
- Optimization: DAPO training uses off-policy optimization with a fixed sample-batch to backpropagation-batch ratio of 16.The rollout batch size is 256 with group size 16, while multi-conversation generation makes the actual mini-batch size differ from rollout batchsize/16.
- Optimization: Training shifts to Stage II after Stage I fully converges, which takes about 400 steps.
- Stage I: 32,768 synthetic QA instances of approximately 32K tokens form Stage I training data, based on HotpotQA and RULER-style distractor construction.Golden answer paragraphs are embedded within distractor content sampled from the same dataset.
- Stage II: 2,560 training instances with maximum length 60K tokens form Stage II, combining DocQA-RL-1.6K data with Stage I examples.
- Data preparation: Stage I samples contain 200 HotpotQA articles totaling approximately 28K tokens, after filtering questions already answered without context.The filtering targets questions likely internalized as common knowledge by the base models.
- Data preparation: Extrapolation test sets span approximately 7K to 3.5M tokens by varying the number of wiki items from 50 to 6,400.
A.4 EVALUATION
Evaluation combines normalized sub-EM scoring, RULER needle-in-a-haystack tests, and FLOP analysis across context lengths to compare long-context behavior and computational scaling.
- Scoring: Answers receive sub-EM scores after normalization, with partial credit proportional to the correctly provided ground-truth parts.Normalization removes definite articles, ignores case, and follows prior evaluation practice.
- RULER evaluation: NIAH single tests at RULER levels 1–3 use increasingly complex haystacks and needles, including seven-digit numbers and 36-character UUID strings.
- RULER evaluation: MEMAGENT uses a workflow-incompatible RULER format adjustment by omitting the original benchmark’s answer prefix.
- Computational analysis: O(n2) complexity characterizes the baseline FLOP cost, whereas MEMAGENT achieves O(n) complexity.
- Computational analysis: Figure 10 compares floating-point operations for baseline and MEMAGENT across context lengths from 8K to 4M.
- Computational analysis: MEMAGENT FLOPs sum initialization, repeated memory updates, and final answering stages.The evaluation uses q = 1024, o = 1024, N = 5000, with c ranging from 8K to 4M.
C CASE STUDY
The case study traces MEMAGENT retaining, revising, and using memory across context chunks to solve a two-hop question despite an irrelevant intervening round.
- Setup: The case study asks which New York City location is associated with the director of Big Stone Gap, a two-hop question over relevant Wikipedia entries.
- Memory trajectory: The model first retains Ghost’s New York City association, even though that production team is not the queried film’s director.
- Memory trajectory: During an irrelevant second round, the model maintains its agent state and remains robust against distraction.
- Answer formation: After receiving both relevant entries, the model updates memory and correctly answers Greenwich Village, New York City.The entries identify Big Stone Gap’s director as Adriana Trigiani and her base as Greenwich Village.
- Interpretation: The trajectory illustrates proactive estimation, immediate updating on relevant context, and insensitivity to irrelevant information.The paper attributes these behaviors to text-generation abilities reinforced through RL rather than architectural attention mechanisms.
D.1 ABLATION STUDY ON CONTEXT SIZE AND MEMORY SIZE
The ablation and comparison studies vary memory and context sizes, evaluate retrieval-based baselines, and examine both successful and failed memory updates.
- Ablation design: The ablation keeps total memory-plus-chunk length constant while linearly varying memory size to control context length per conversation turn.Results are reported for NIAH and Longbench-QA.
- Baseline comparisons: MEMAGENT is compared with Mem0 and RAG agents using multiple retrieval configurations and top-K settings.Mem0 uses top-30 memory retrieval, while RAG varies chunk size and top-K.
- Baseline comparisons: The results report that MEMAGENT outperforms the evaluated Mem0 and RAG methods.The authors characterize end-to-end RL-trained memory as more flexible and coherent than retrieval-based strategies.
- Information overwriting: A failure case shows irrelevant early memories consuming capacity, causing a later crucial item to be truncated and the final answer to be incorrect.Later summarization creates space for a second key item, but the first key item has already been lost.
- Information retention: A successful QA example combines information about Sachin Warrier’s employment at TCS with TCS headquarters in Mumbai, Maharashtra.
F.2 MISSING CRITICAL INFORMATION
The examples show that MemAgent can lose critical information when prerequisite evidence arrives later, causing incorrect answers despite later relevant context. Its memory may also preserve irrelevant information when the required connection is absent.
- F.2 MISSING CRITICAL INFORMATION: The model fails a multi-hop question because it does not recognize the first key information before observing the prerequisite evidence.Although it recognizes the problem when the second key piece appears at Turn 215, it still fails to produce the correct final answer.
- F.2 MISSING CRITICAL INFORMATION: At Turn 143, the memory lacks information about The Hard Easy and its guest stars, so Brian Doyle-Murray’s family information is not directly relevant.The memory therefore remains unchanged with respect to the question until later context supplies information about the episode.
- F.2 MISSING CRITICAL INFORMATION: The later context identifies Ryan Truex as Martin Lee Truex Jr.’s younger brother, but this fact does not answer the question about The Hard Easy’s guest stars.The example contrasts retained relevant information about Martin Lee Truex Jr. with the still-missing relation involving the episode guest stars.
F.3 PRIMACY BIAS
The primacy-bias example shows the model committing early to an interpretation of “country” and reinforcing it with later related information. Even after encountering the critical evidence, the model still answers incorrectly.
- F.3 PRIMACY BIAS: The model initially interprets “country” as a nation and reinforces that interpretation when it later encounters additional information about countries.This early commitment is identified as a primacy bias in the example.
- F.3 PRIMACY BIAS: The critical evidence eventually identifies Brown State Fishing Lake as being in the United States, but the model still fails to answer correctly.The question asks for the population of the country containing the lake, while the relevant country is established as the United States.
- F.3 PRIMACY BIAS: The updated memory retains information about Kunming that is unrelated to the question while preserving the lake’s location in the United States.The memory includes Kunming’s population and other facts, but these do not directly address the requested country population.
- F.3 PRIMACY BIAS: The answer reported for the United States is around 327,167,434 inhabitants as of 2018, whereas Brown County’s 2010 population is 9,984.The example distinguishes the requested country-level population from the county-level population.