Skip to content

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

🕒 Published (v1): 2026-01-01 · Source: ICLR · Venue: ICLR 2026 · link

Why this paper was selected

Revisitable memory with dynamic key-evidence retrieval for long-context LLM agents

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

ReMemR1 extends the "memorize while reading" paradigm for long-context LLM agents by adding a history-aware callback mechanism: at each step the agent generates a retrieval query over its entire memory history, enabling non-linear reasoning across millions of tokens. It is trained with GRPO using a multi-level reward that combines trajectory-level answer correctness with dense step-level information-gain signals. On multi-hop QA benchmarks it outperforms the strongest memory-agent baseline by up to 7.6% accuracy while adding negligible inference overhead (<0.2% latency).

Problem

Existing "memorize while reading" agents process documents in a single linear pass, compressing each chunk into a fixed-length memory buffer (\(m_t \leftarrow \pi_\theta(Q, c_t, m_t)\)). This creates three failure modes: (1) premature pruning — evidence encountered at step \(t\) may be discarded before its relevance is revealed at step \(t+k\); (2) progressive information loss — fixed-length buffers force overwriting of earlier content as context grows; (3) sparse supervision — RL training with only a final-answer reward gives weak gradient signal over the many intermediate memory-update steps.

Method

ReMemR1 augments the MDP state from \(s_t = m_t\) to \(s_t = (m_t, q_t)\), where \(q_t\) is a callback query enabling retrieval over the full memory history:

\[s_{t+1} = (m_{t+1}, q_{t+1}) = \pi_\theta\!\left(Q,\, c_t,\, m_t,\, \mathcal{E}\!\left(\{m_i\}_{i \leq t},\, q_t\right)\right)\]

The retrieval function \(\mathcal{E}(X, b) = \arg\max_{x \in X} \text{recall}(b, x)\) selects the historical memory with the highest word-overlap recall against the query — lightweight, no learned retriever needed. All intermediate memories are stored (compact model-generated summaries, not raw documents).

Training uses a multi-level GRPO objective combining: - Trajectory-level outcome reward \(R_\text{out}^{(g)}\): exact-match correctness of the final answer. - Step-level state rewards \(R_\text{state,t}^{(g)}\): sum of (a) memory information gain \(r_\text{memory,t} = \max_y \text{recall}(m_t, y) - \max_y \text{recall}(m_{t-1}, y)\), (b) callback retrieval bonus measuring additional recall of answer-relevant content beyond current memory + context, and (c) format reward for correct <callback>/<memory>/\box{} tags.

Advantages are normalized separately by trajectory group (for outcome) and by same-timestep group (for state rewards), then combined as \(\hat{A}_t^{(g)} = \alpha \hat{A}_\text{out}^{(g)} + (1-\alpha)\hat{A}_\text{state,t}^{(g)}\).

Key Contributions

  • History-augmented MDP state extending "memorize while reading" with a learned callback query component for non-linear evidence retrieval.
  • Word-overlap retrieval over stored intermediate memory summaries — no external vector index, negligible storage.
  • Multi-level GRPO reward design that provides dense, step-level supervision aligned with memory quality rather than only final-answer correctness.
  • Empirical demonstration that RL-trained callbacks outperform rule-based callback designs.

Results

  • HotpotQA (in-distribution, 7B): ReMemR1 82.8% vs. MemAgent 78.9% at 100 docs; 80.8% vs. 75.8% at 6400 docs. Gains are maintained across all context lengths (50–6400 documents).
  • 2WikiMultiHopQA (out-of-distribution, 7B): ReMemR1 63.9% vs. MemAgent 61.7% at 50 docs; 50.3% vs. 44.7% at 6400 docs — strongest OOD gap, indicating generalized retrieval skill.
  • 3B scale: up to +7.3% over MemAgent; +12.1% at 50 docs OOD.
  • Distant-evidence challenge (evidence in reverse reasoning order, inter-evidence gap >50% of total docs): ReMemR1 substantially outperforms MemAgent, which degrades sharply.
  • Error rate reduction: >20% relative reduction over MemAgent baseline.
  • Inference overhead: callback adds <2 s latency and <1 MB memory at 6400 documents (<0.2% of total inference time).
  • Training overhead: 1467.72 s/step vs. 1247.17 s/step (MemAgent); peak GPU memory 131.15 GB vs. 124.97 GB.

Limitations

  • Retrieval function is pure word-overlap recall — no semantic matching; may fail when paraphrasing or terminology shifts occur across memory states.
  • Training data is exclusively HotpotQA (padded to 200 docs / ~30K tokens); generalization to truly extreme lengths (millions of tokens) is not directly evaluated.
  • All intermediate memories must be stored; memory footprint grows linearly with document count, which may become non-trivial for very long corpora despite being compressed summaries.
  • The hyperparameter \(\alpha\) balancing outcome vs. state rewards requires tuning; sensitivity analysis is deferred to appendices.
  • Evaluation is confined to multi-hop QA; applicability to open-ended agentic tasks (tool use, code generation, long-horizon planning) is not demonstrated.

Relevance to Agentic AI / LLM Agents

ReMemR1 directly advances the design of memory-augmented LLM agents operating over long, unstructured document streams — a core challenge for agents in real-world workflows such as literature review, legal analysis, and multi-document reasoning. The callback mechanism operationalizes a form of episodic memory retrieval within an agent's own compressed state history, enabling non-Markovian reasoning without external retrieval infrastructure. The multi-level GRPO reward design is a transferable training recipe for any recurrent-state agent where intermediate behavior quality matters but is hard to supervise with sparse final rewards. This work connects to the broader line of RL-for-reasoning agents (DeepSeek-R1, MemAgent) and shows that selectively revisiting past internal states — rather than re-reading raw context — is both effective and scalable.