PaperlayerView exact revision on arXiv

Source-linked AI summary

Look Back to Reason Forward: Revisitable Memory for Long-Context LLM Agents

Yaorui Shi, Yuxin Chen, Siyuan Wang, Sihang Li, Hengxing Cai, Qi Gu, Xiang Wang, An Zhang

arXiv:2509.23040v5cs.CLcs.AI

TL;DR

Long-context question answering must connect evidence dispersed across very large documents, while memory agents face premature pruning, overwriting, and sparse training signals. ReMemR1 adds callback retrieval over historical memories and multi-level rewards combining answer outcomes with step-level signals. It consistently surpasses general-purpose and specialized memory agents across benchmark settings, with negligible retrieval latency overhead.

  • Problem

    Long-context QA requires connecting evidence dispersed across millions of tokens, while linear memory agents can prune latent evidence, lose information through overwriting, and receive sparse final-answer supervision.

    Source: §1 · p. 1, §2.1 · p. 4, §2.3 · p. 5

  • Method

    ReMemR1 generates callback queries during memory updates to retrieve historical memories, and RLMLR combines trajectory-level answer rewards with step-level state rewards.

    Source: §1 · pp. 2–3, §4 · p. 10

  • Results

    ReMemR1 consistently surpasses general-purpose LLMs and specialized memory agents across in- and out-of-distribution benchmarks, with retrieval latency below 0.2% time overhead.

    Source: §1 · p. 3

  • Takeaways & Limitations

    The approach trades marginal computational cost for reasoning gains while addressing progressive information loss without prohibitive scalability issues.

    Source: §1 · p. 3

    Storing and recalling long-context information may raise privacy and security risks when private or proprietary data are used without robust safeguards.

    Source: ETHICS STATEMENT · p. 11

Abstract

from arXiv · show

Large language models face challenges in long-context question answering, where key evidence of a query may be dispersed across millions of tokens. Existing works equip large language models with a memory buffer that is dynamically updated via a linear document scan, also known as the "memorize while reading" methods. While this approach scales efficiently, it suffers from pruning of latent evidence, information loss through overwriting, and sparse reinforcement learning signals. To tackle these challenges, we present ReMemR1, which integrates the mechanism of memory retrieval into the memory update process, enabling the agent to selectively callback historical memories for non-linear reasoning. To further strengthen training, we propose a multi-level reward design, which combines final-answer rewards with dense, step-level signals that guide effective memory use. Together, these contributions mitigate information degradation, improve supervision, and support complex multi-hop reasoning. Extensive experiments demonstrate that ReMemR1 significantly outperforms state-of-the-art baselines on long-context question answering while incurring negligible computational overhead, validating its ability to trade marginal cost for robust long-context reasoning.

1 INTRODUCTION

Long-context QA remains difficult when evidence is dispersed across millions of tokens and attention struggles to track long-range dependencies. ReMemR1 adds historical memory callback and step-level rewards to address limitations in linear memory agents.

1 INTRODUCTION

  • Long-context reasoning is difficult when query evidence is scattered across millions of tokens and attention must track long-range dependencies.The introduction notes applications such as synthesizing legal precedents and reviewing scientific literature.Source: §1 · p. 1
  • Early evidence may be pruned before later context reveals its importance, while fixed-length memory compression progressively degrades information needed for distant-evidence queries.Final-answer-only reinforcement learning also provides limited guidance for the many intermediate memory updates.Source: §1 · p. 2
  • ReMemR1 generates callback queries during memory updates to retrieve historical memories, integrating recalled information with the current chunk for non-linear reasoning.The agent can selectively revisit facts from earlier stages as new evidence arrives.Source: §1 · p. 2, Figure 2 · p. 3
  • RLMLR combines trajectory-level answer correctness with fine-grained step-level signals evaluating information gain in memory transformations.Because document chunks remain identical across trajectories at a given step, the training design targets memory updates and callback actions without environmental noise.Source: §1 · pp. 2–3
  • ReMemR1 consistently surpasses general-purpose LLMs and specialized memory agents on in- and out-of-distribution benchmarks, with retrieval latency below 0.2% time overhead.The introduction reports over 20% error rate reduction in reasoning accuracy while describing the computational cost as marginal.Source: §1 · p. 3

2 METHOD

ReMemR1 augments sequential memory updates with retrieval over memory history and trains the agent using trajectory-level and step-level rewards.

2.1 PRELIMINARIES: MDP MEMORY AGENT FOR LONG-CONTEXT QA

  • Conventional memory agents process chunks sequentially, treating the current memory as a sufficient summary, but fixed-length compression and irreversible overwriting can discard evidence needed for later hops.The standard transition updates memory from the question, current chunk, and current memory.Source: §2.1 · p. 4

2.2 MEMORY AGENT WITH HISTORY-AUGMENTED STATE

  • The history-augmented state pairs current memory with a callback query that retrieves relevant earlier memories, allowing iterative recall and non-linear reasoning paths.Retrieval selects prior content by query-word overlap.Source: §2.2 · pp. 4–5

2.3 REINFORCEMENT LEARNING WITH MULTI-LEVEL REWARD SHAPING

  • Multi-level rewards combine terminal answer outcomes with dense state rewards, comparing trajectories that share the same external context at each step but maintain different internal states.The design normalizes outcome rewards across trajectories and state rewards across states at the same step.Source: §2.3 · p. 5, Figure 4 · p. 6

2.3.1 TRAJECTORY-LEVEL OUTCOME REWARDS FOR FINAL CORRECTNESS

  • Terminal outcome reward evaluates whether the predicted answer exactly matches any acceptable ground-truth answer.The reward is calculated from each trajectory’s terminal state.Source: §2.3.1 · p. 5

2.3.2 STEP-LEVEL ACTION REWARDS FOR BEHAVIOR SHAPING

  • Step-level state rewards provide fine-grained supervision by rewarding increases in memory recall of information relevant to the ground-truth answers.The memory update reward compares recall before and after the update.Source: §2.3.2 · p. 5
  • Additional step rewards assess whether callback retrieval adds relevant information beyond current memory and context, and whether outputs use the required tags.The format checks cover callback and memory tags at intermediate steps and the answer tag at the final step.Source: §2.3.2 · pp. 5–6

2.3.3 TRAINING OBJECTIVE

  • The GRPO-variant objective combines trajectory-level outcome advantages with state-level advantages computed relative to states at the same step.The normalization omits the standard-deviation term to avoid difficulty bias.Source: §2.3.3 · pp. 6–7

3 EXPERIMENTS

The experiments test long-context accuracy, distant-evidence recall, computational cost, and the effects of multi-level rewards and learned callbacks. They evaluate on HotpotQA and 2WikiMultiHopQA against general LLMs, long-context models, and memory agents.

3.1 EXPERIMENTAL SETUP

  • Evaluation uses HotpotQA for training and in-distribution testing, 2WikiMultiHopQA for out-of-distribution testing, and contexts padded to as many as 6,400 documents.Baselines include general LLMs, long-context LLMs, and tailored memory agents such as MemAgent.Source: §3.1 · pp. 7–8

3.2 MAIN RESULTS (RQ1)

  • ReMemR1 achieves the best accuracy across model scales, datasets, and context lengths, with up to 7.3% higher accuracy on 3B and 7.6% on 7B models than MemAgent.The paper reports that gains are more pronounced on OOD 2WikiMultiHopQA and grow with document length.Source: §3.2 · p. 8

3.3 DISTANT EVIDENCE CHALLENGE (RQ2)

  • In the distant-evidence setting, ReMemR1 surpasses MemAgent by a large margin on 2Wiki accuracy.Supporting evidence is placed in reverse reasoning order, with successive evidence more than half a context apart.Source: Figure 5 · p. 8, §3.3 · p. 8

3.4.1 INFERENCE-TIME EFFICIENCY

  • ReMemR1's callback operations add less than 2 seconds of latency and under 1MB of memory, even with 6,400 documents.Retrieved states are compact model-generated summaries; ReMemR1 achieves up to 5% absolute accuracy improvement over the baseline, corresponding to a 20% reduction in error rate.Source: §3.4.1 · pp. 8–9, Figure 6 · p. 9

3.4.2 TRAINING-TIME EFFICIENCY

  • Training has higher per-step latency than MemAgent because callback queries must be generated, while peak GPU memory remains similar and is dominated by the LLM backbone.Table 2 reports average per-step time, early/late step latency, and peak GPU memory.Source: §3.4.2 · p. 9, Table 2 · p. 9

3.5.1 EFFECTIVENESS OF MULTI-LEVEL REWARD DESIGN (RQ4)

  • α = 0.8 gives the best accuracy across tested context lengths, balancing final-answer rewards against step-level guidance.Outcome-only rewards underperform, while overly emphasizing step-level shaping distracts from final correctness.Source: §3.5.1 · p. 9, Table 3 · p. 10

3.5.2 RL-DRIVEN V.S. RULE-BASED MEMORY CALLBACK (RQ5)

  • RL-driven callback outperforms MemAgent and a fixed question-based callback on both datasets across context lengths; the rule-based callback can reduce performance by up to 7.9%.The learned strategy determines when and what to recall, and its advantage grows with document length.Source: Table 4 · p. 10, §3.5.2 · p. 10

4 CONCLUSION

The conclusion presents ReMemR1 as a response to forward-only processing, memory overwriting, and sparse supervision in long-context QA. It reports consistent benchmark gains and robustness on distant evidence, while identifying RL training and learned callbacks as necessary components.

4 CONCLUSION

  • ReMemR1 combines historical-memory callbacks with RLMLR, which pairs trajectory-level outcome rewards with step-level state rewards.The conclusion reports consistent gains over general LLMs and prior memory agents on ID and OOD benchmarks, plus robustness under distant evidence.Source: §4 · p. 10

ETHICS STATEMENT

The work is computational research on public QA benchmarks, not an evaluation of high-stakes applications. The ethics statement also flags privacy, security, and fairness considerations for downstream deployment.

ETHICS STATEMENT

  • Experiments are limited to HotpotQA and 2WikiMultiHopQA; the paper does not develop or evaluate medicine, law, or finance applications.The datasets use publicly sourced text without sensitive data, and the work involves no human subjects.Source: ETHICS STATEMENT · p. 11
  • Storing and recalling private or proprietary data could pose privacy and security risks without robust safeguards.The authors call for downstream evaluation of fairness, transparency, and potential discriminatory impacts.Source: ETHICS STATEMENT · p. 11

REPRODUCIBILITY STATEMENT

The reproducibility statement describes code, configurations, and scripts for dataset generation, training, and evaluation. The supplied related-work and extended-results passages review adjacent research areas and compare ReMemR1 with larger long-context baselines.

REPRODUCIBILITY STATEMENT

  • The reproducibility materials include fixed-seed dataset-generation code, environment configurations, ReMemR1 training scripts, and evaluation scripts.The paper also points to setup and hyperparameter descriptions in §3.1 and Appendix C.Source: REPRODUCIBILITY STATEMENT · p. 11

A RELATED WORK

  • Related work covers memory-augmented LLM agents, long-context model approaches, and reinforcement learning for LLM reasoning.The review discusses memory systems, context-length methods, and outcome-based reward approaches.Source: A RELATED WORK · p. 16

B.1 COMPARISON AGAINST MORE BASELINES

  • At 6,400 documents, ReMemR1 reaches 80.8% HotpotQA accuracy versus 38.3% for QwenLong-L1-32B, and 50.3% 2WikiMultiHopQA accuracy versus 29.9% for that baseline.Table 5 also compares ReMemR1 with 14B-level R1-Distill-Qwen and other long-context baselines.Source: §B.1 · p. 16, Table 5 · p. 17

B.2 IMPACT OF RL TRAINING

RL improves long-context reasoning across HotpotQA and 2WikiMultiHopQA, with ReMemR1 strongest across most tested context lengths. Without RL, formatting errors initially hinder ReMemR1, while its advantage becomes clearer beyond 800 documents.

B.2 IMPACT OF RL TRAINING

  • With RL, ReMemR1 consistently achieves the highest scores across most context lengths on HotpotQA and 2WikiMultiHopQA, outperforming MemAgent.Without RL, both methods drop sharply as context grows; above 800 documents, ReMemR1 improves by about 6% on both benchmarks.Source: §B.2 · pp. 17–18
  • Without RL, ReMemR1 initially trails MemAgent at fewer than 800 documents, reflecting imperfect instruction-following and added formatting requirements.The 3B model starts with around 0.6 average format reward, failing to extract updated memory in 40% of steps.Source: §B.2 · pp. 17–18
  • During the first 20 training steps, ReMemR1 rapidly improves as it learns to follow the required format.The additional requirements concern generating both inner memory and callback queries; early parsing errors initially lower success.Source: Figure 7 · p. 18

B.3 DETAILED INFLUENCE OF DIFFERENT ALPHA VALUES

The three α settings show similar early reward growth, but their later training curves diverge. α = 0.8 yields the highest converged outcome reward, while outcome-only training is slower and overly dominant step-level reward is unstable.

B.3 DETAILED INFLUENCE OF DIFFERENT ALPHA VALUES

  • During the first 100 steps, α = 1.0, 0.8, and 0.5 follow similar trajectories as outcome reward rises with formatting ability and coarse-grained reasoning.The curves diverge as training progresses.Source: §B.3 · p. 18
  • After convergence, α = 0.8 consistently achieves the highest outcome reward among the three settings.The paper attributes this to moderate step-level reward helping address sparse and noisy credit assignment.Source: §B.3 · p. 18
  • With α = 1.0, outcome-only training converges more slowly and reaches a lower plateau than the other settings.Without step-level feedback, credit assignment across interacting memory updates is difficult.Source: §B.3 · p. 18
  • With α = 0.5, reward initially tracks the other curves but collapses mid-training as dominant step-level signals make it sensitive to noisy intermediate states.The paper characterizes this as training instability that leads to divergence.Source: §B.3 · p. 18
  • A balanced mix, such as α = 0.8, combines intermediate guidance with final-answer correctness for more stable and effective training dynamics.The final-answer reward keeps optimization grounded in the evaluation metric.Source: §B.3 · p. 18

B.4 ADDITIONAL RESULTS ON COMPUTATIONAL OVERHEAD

The inference-time analysis reports small callback costs alongside accuracy gains, with callback-query generation accounting for most added latency. The section also sets up a qualitative multi-hop case study comparing ReMemR1 with MemAgent.

B.4.1 INFERENCE-TIME PERFORMANCE

  • The inference-time evaluation uses HotpotQA contexts of 50 to 6400 documents and reports accuracy, latency, and GPU memory measures.Latency includes total time, callback time, and their ratio; memory includes total and callback consumption and corresponding ratios.Source: §B.4.1 · p. 19
  • Callback contributes fewer than 0.2% of total inference-time latency and less than 0.001% of total GPU memory across all scales.The reported comparison covers full inference-time performance; ReMemR1 also improves accuracy by up to 5% versus MemAgent, corresponding to a 20% reduction in error rate.Source: §B.4.1 · p. 19, Table 7 · p. 19
  • The additional latency comes primarily from autoregressively generating the <recall> callback query at each step, rather than from retrieval itself.Each stored intermediate memory is a short model-generated summary.Source: §B.4.1 · p. 19

B.5 CASE STUDY

  • The qualitative case study compares ReMemR1 with MemAgent on a 2WikiMultiHopQA question requiring two film-director identifications, death-date retrieval, and temporal comparison.Its visualizations distinguish correct, erroneous, and ignored evidence.Source: §B.5 · pp. 19–20

B.5.1 CASE STUDY OF BASELINE METHOD (MEMAGENT)

In the case study, MemAgent processes chunks sequentially and misses Stuart Paton’s death date before learning he directed Is There Justice? It later uses irrelevant context to form a false answer; ReMemR1 recalls the retained date and answers correctly.

B.5.1 CASE STUDY OF BASELINE METHOD (MEMAGENT)

  • MemAgent misses Paton’s death date in Chunk 10 before learning in Chunk 109 that he directed Is There Justice?, and cannot revisit the earlier chunk.Its final reasoning instead uses irrelevant information about Chris Marker and film release dates, producing an incorrect answer.Source: §B.5.1 · pp. 21–22
  • The baseline’s finite memory bottleneck lets noisy local correlations overwrite valid query states, turning a search for information into a confident false inference.The false memory ultimately contributes to the incorrect answer.Source: §B.5.1 · p. 22

B.5.2 CASE STUDY OF REMEMR1

  • ReMemR1 retains Paton’s biographical data and uses callback-enhanced memory to connect it with the later film-director evidence, yielding the correct answer, Is There Justice?.It synthesizes evidence from distant timesteps rather than relying on forward-only processing.Source: §B.5.2 · pp. 23–24

B.5.3 FAILURE ANALYSIS OF REMEMR1

The failure analysis identifies irrelevant recall queries and persistent false memories as two vulnerabilities in ReMemR1. The accompanying training-objective subsection states that the model uses a GRPO variant.

B.5.3 FAILURE ANALYSIS OF REMEMR1

  • Irrelevant callback queries can leave relevant historical context unretrieved, and one case ends with “People’s Party” instead of the ground-truth “Liberal Conservative.”The failure analysis says the query failure may partly explain why the answer focused on the party name rather than the ideology.Source: §B.5.3 · pp. 24–25
  • An early false memory can persist over later steps and override correct evidence; in the car example, the answer repeats Ford Falcon and Toyota Supra instead of sedan and station wagon.The agent read a later document listing the correct variants but appended that information rather than correcting its earlier claim.Source: §B.5.3 · p. 25

C.1 FULL EXPRESSION OF TRAINING OBJECTIVE

  • The training-objective subsection says ReMemR1 uses a variant of GRPO and introduces its full objective expression.The supplied notation fragments identify an importance-sampling ratio, token positions, a clipping ratio, and a KL term.Source: §C.1 · p. 26

C.2 TRAINING HYPERPARAMETERS

Training uses fixed document and chunk settings, while evaluation samples hard-difficulty examples from each benchmark. The section also reports model sizes, training hardware, and convergence times.

C.2 TRAINING HYPERPARAMETERS

  • Training uses 200 documents per sample (about 30K input tokens), chunks capped at 5000 tokens, and approximately six steps; rollouts use temperature 1 and at most 2048 tokens.The text says context-document counts varied at test time, whereas the training setup consistently used 200 documents per sample.Source: §C.2 · p. 26
  • The 3B model and variants converge after 100 hours on 16 H800 GPUs, while the 7B model converges after 80 hours on 32 H800 GPUs.Training uses the verl framework, sglang for trajectory generation, FSDP, and bfloat16 precision.Source: §C.2 · p. 26

C.3 EVALUATION SETTINGS

  • Evaluation training data uses only hard-difficulty samples, with 128 randomly seeded (seed 4) samples subsampled from each benchmark.The passage describes these questions as typically requiring at least two dependent pieces of evidence.Source: §C.3 · p. 27

D COMPLEXITY ANALYSIS

The complexity analysis concludes that ReMemR1 retains linear time and space complexity while adding retrieval over prior memories. The appendices also describe its prompts and disclose the limited use of an LLM in manuscript preparation.

D.1 BASELINE COMPLEXITY

  • The baseline processes T chunks with O(1) memory operations and a constant number of policy-network passes per update, giving O(T) time and O(T) space overall.The space total includes the document chunks and memory at each step.Source: §D.1 · p. 27

D.2 COMPLEXITY OF REMEMR1

  • ReMemR1 retains prior fixed-length memory vectors, adding O(T) space; exact maximum-similarity retrieval may cost O(t) per step, while practical heuristics or indexing keep expected total cost O(T).The passage describes the stored memories as having a smaller constant space term than the original text chunks.Source: §D.2 · p. 27
  • ReMemR1 preserves O(T) time and O(T) space complexity while enabling non-linear reasoning through retrieval.The section characterizes retrieval latency as negligible compared with the state-transition model, with total cost O(T) in expectation.Source: §D.2 · p. 27

E PROMPT TEMPLATE

  • Intermediate prompts ask the agent to update memory and optionally issue one recall query, while final prompts ask it to answer from the previous memory.The intermediate template supplies the question, document section, and previous memory; the final template requests a boxed answer.Source: E PROMPT TEMPLATE · pp. 27–28

F THE USE OF LARGE LANGUAGE MODELS

  • The authors report using an LLM only for proofreading, typo checking, and clarity edits, not for core intellectual contributions or original manuscript drafting.The authors take responsibility for the scientific claims, arguments, and final text.Source: F THE USE OF LARGE LANGUAGE MODELS · p. 28

Related papers

Loading 2509.23040v5…