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

## Abstract

Abstract 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 corpus that is dynamically updated during a single-pass document scan, also known as the “memorize while reading” methods.
While this approach scales efficiently, it suffers from irreversible forward-only processing, information loss through overwriting, and sparse reinforcement learning signals.
To tackle these challenges, we present ReMemR1, a memory-augmented agent with callback-enhanced memory that allows selective retrieval from the entire memory history and allows non-linear reasoning and revisiting of early evidence.
To further strengthen training, we propose Reinforcement Learning with Multi-Level Rewards (RLMLR), 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 multi-hop memory utilizing.
Experiments on long-document QA show significant gains over existing memory-based approaches, which validates ReMemR1 as an effective solution for long-context reasoning agents.
Our code is available at https://github.com/syr-cn/ReMemR1 .

## 1 Introduction

Reasoning over vast, multi-document contexts remains a critical bottleneck for large language models (LLMs) (Hsieh et al., 2024; Team et al., 2024; Beltagy et al., 2020; Ding et al., 2023; Child et al., 2019).
This capability is crucial for real-world applications, such as synthesizing legal documents or reviewing scientific literature, where critical evidence for a single query can be scattered across millions of tokens (Beltagy et al., 2020; Ding et al., 2023; Child et al., 2019).
As illustrated in Figure [1](https://arxiv.org/html/2509.23040v1#S1.F1)(a), this extreme context length makes it difficult for LLMs to track long-range dependencies and faithfully synthesize disparate information into a coherent answer.

To mitigate this issue, recent research has explored the “memorize while reading” paradigm, which augments LLMs with an external memory module to process documents incrementally (Yu et al., 2025a; Li et al., 2025).
As shown in Figure [1](https://arxiv.org/html/2509.23040v1#S1.F1)(b), this framework employs a memory agent that digests the context documents chunk by chunk while maintaining its memory using an overwriting strategy.
This process can be formulated as a Markov decision process, where the agent’s state is defined solely by its fixed-length memory buffer, *i.e., *$s_{t}=m_{t}$.
At each step, the agent consumes a document chunk $c_{t}$ together with its previous memory $m_{t}$ and compresses them into a new memory $m_{t+1}$.
After a single linear pass through the entire document, the agent uses this final memory $m_{T}$ buffer to generate an answer for the given question.
By processing information in this manner, this paradigm reduces the quadratic complexity of long-context question answering to linear time.

Despite its efficiency, we identify the following intrinsic limitations in the existing “memorize while reading” paradigm:

- •
Irreversible Forward-Only Processing.
In multi-hop questions, evidence from different hops may be required for answering.
For example, in the scenario illustrated in Figure [1](https://arxiv.org/html/2509.23040v1#S1.F1), answering the question necessitates evidence from both the first-hop document (*e.g., *Doc 71) and the second-hop document (*e.g., *Doc 42).
During the search for first-hop evidence, the agent may encounter second-hop evidence but does not recognize its importance since the first-hop question remains unsolved.
As memory is updated, this second-hop evidence could be forgotten, making it inaccessible later.
This irreversible memory update limits the agent’s ability to revisit and integrate crucial information, especially in multi-hop reasoning tasks.
- •
Progressive Information Loss in Memory Overwriting.
The paradigm’s reliance on a fixed-length memory buffer necessitates constant information compression.
As illustrated in Figure [1](https://arxiv.org/html/2509.23040v1#S1.F1)(b), crucial early-stage details (*e.g., *“Dr Aris Thorne was a postdoc in Chicago" from Doc 42, step 5) can be inevitably lost after numerous overwrites.
This progressive degradation of memory makes it difficult to maintain the full context and impedes the ability to resolve complex queries that require synthesizing evidence spread across distant sections of the document.
- •
Sparse and Delayed Supervision.
Training these agents using reinforcement learning typically relies on a single reward signal, such as the correctness of the final answer.
This sparse reward, provided only at the end of the reasoning process, offers limited guidance for the long sequence of intermediate memory updates, leading to inefficient optimization and suboptimal memory management strategies, particularly in complex tasks where producing correct final answers is especially challenging.

Figure: Figure 1: Comparison of approaches for question answering in a long-context setting. (a) Full context input introduces substantial complexity and challenges the LLM to locate the correct information. (b) The “memorize while reading” paradigm processes documents by chunks to reduce context length at each step. Still, the irreversible and linear memory overwriting prevents the model from connecting distantly related information. (c) Our method augments the agent’s state to allow for callback of historical memories, enabling it to integrate relevant facts from earlier steps into its reasoning process.
Refer to caption: x1.png

To address these challenges, we introduce ReMemR1, a memory-augmented LLM agent that can callback historical memories when navigating long documents.
Our central contribution is to move beyond the restrictive state of the conventional MDP. Instead of passing only the memory $m_{t}$ during iteration, we augment the state to $s_{t}=(m_{t},q_{t})$, where $q_{t}$ is a callback query that enables retrieval over the agent’s entire memory history.
At each step, the agent not only updates its memory $m_{t}$ based on the new chunk $c_{t}$, but also generates a callback query $q_{t+1}$ to reach its past memories $\{m_{i}\}_{i\leq t}$ (Figure [2](https://arxiv.org/html/2509.23040v1#S1.F2)).
The retrieved information is then integrated into the context for the next state update. As depicted in Figure [1](https://arxiv.org/html/2509.23040v1#S1.F1)(c), this mechanism empowers the agent to construct non-linear reasoning paths, and selectively revisit critical facts from early stages to connect with new evidence.
This directly counters the progressive information loss and breaks the irreversible forward-only constraint.

We further develop an RL framework, Reinforcement Learning with Multi-Level Rewards (RLMLR). This framework addresses the sparse supervision problem by combining a trajectory-level outcome reward, which evaluates the correctness of the final answer, with dense, step-level state rewards. These step-level rewards provide fine-grained supervision by measuring the information gain in each memory update and promoting the effective use of the retrieval mechanism.
The advantages of these rewards are then calculated at the trajectory and step-level, correspondingly, for group relative policy optimization.
By shaping intermediate behavior, these rewards provide fine-grained supervision that alleviates the sparsity of traditional outcome-only RL.

Extensive experiments on both in-distribution and out-of-distribution benchmarks demonstrate that ReMemR1 consistently surpasses general-purpose LLMs and specialized memory agents.
Beyond overall performance, we further conduct systematic analyses of memory callback strategies and multi-level reward designs, confirming the superiority of our RL-driven framework.
Our results show that ReMemR1 not only alleviates progressive information loss in long contexts but also demonstrates a robust ability for retrieval and reasoning, leading to strong generalization across model scales and domains.

Figure: Figure 2: The comparison of state transition functions between previous work and our method. (left) Conventional memory agents use a restrictive state $s_{t}=m_{t}$, where the next memory $m_{t+1}$ only depends on the current context $c_{t}$ and memory $m_{t}$. (right) Our method introduces a history-augmented state $s_{t}=(m_{t},q_{t})$, where the agent generates a callback query $q_{t}$ to retrieve relevant information from its entire memory history $\{m_{i}\}_{i\leqslant t}$, enabling non-linear reasoning paths.
Refer to caption: x2.png

## 2 Method

In this section, we present ReMemR1, a memory-augmented agent that incorporates history-aware retrieval and reinforcement learning with multi-level rewards to enhance long-context reasoning.
We first review the formulation and limitations of conventional “memorize while reading” paradigm, where memory agents solve long-context QA through a single-pass scan that can be formulated as a Markov decision process (§[2.1](https://arxiv.org/html/2509.23040v1#S2.SS1)).
We then introduce our history-augmented state mechanism, which enriches the memory update process with a query component that enables retrieval over past memories and supports non-linear reasoning paths (§[2.2](https://arxiv.org/html/2509.23040v1#S2.SS2)).
Finally, we describe the proposed multi-level reward structure, which combines trajectory-level outcome rewards with step-level state rewards to provide more effective training supervision (§[2.3](https://arxiv.org/html/2509.23040v1#S2.SS3)).
Related work is discussed in Appendix [A](https://arxiv.org/html/2509.23040v1#A1).

### 2.1 Preliminaries: MDP Memory Agent for Long-Context QA

We consider the task of long-context question answering (QA), where each dataset sample is given as $(Q,Y)$. Here, $Q$ denotes a question and $Y$ is the set of all acceptable correct answers to that question (*i.e., *a candidate answer list, and answering with any element in $Y$ is regarded as correct). Each sample is further associated with a long document $C$, which is divided into small chunks ${c_{0},c_{1},\dots,c_{T-1}}$ and sequentially provided to the model.

Standard memory-augmented agents process long documents in a “memorize while reading” paradigm: the agent reads chunks one by one and continuously updates its memory to preserve important information.
This sequential procedure can be naturally cast as a Markov Decision Process (MDP), written as $(S,\mathcal{U},P,R)$. At each step $t$:

- •
The state $s_{t}\in S$ is defined by the agent’s memory $m_{t}$ (*i.e., *$s_{t}=m_{t}$), which serves as the sufficient statistic summarizing the past trajectory.
The agent also receives external inputs from the environment, consisting of the question $Q$ and the document chunk $c_{t}$.
- •
The action $u_{t}\in\mathcal{U}$ represents an update to the memory, which is determined by the policy $\pi_{\theta}$ given the current state and inputs.
- •
The transition $P(s_{t+1}\mid s_{t},u_{t})$ specifies how the next state is produced. In particular, the memory is updated as
$s_{t+1}=m_{t+1}=\pi_{\theta}(Q,c_{t},m_{t}),\text{\quad for }t\in[0,T-1]$
(1)
- •
The reward $R$ is defined based on the quality of the final answer after the entire document has been processed (§[2.3](https://arxiv.org/html/2509.23040v1#S2.SS3)).

The model begins with an empty memory, *i.e., *$m_{0}=\varnothing$. After all $T$ document chunks are processed, the agent produces a terminal output state by updating:

$$ $s_{T+1}=o=\pi_{\theta}(Q,\varnothing,m_{T}),$ (2) $$

where the empty input indicates that no document chunk is provided at this final step.

In this formulation, the memory $m_{t}$ is assumed to be a sufficient statistic of the entire history of previously processed chunks $\{c_{i}\}_{i<t}$.
However, this formulation is inherently restrictive.
First, in multi-hop reasoning, the agent may scan over evidence that is crucial for later hops but fail to recognize its importance at the time, since the preceding hop has not yet been resolved.
As the memory is updated, such overlooked evidence can be overwritten and thus lost for subsequent reasoning.
Second, because the memory is typically constrained to a fixed length to guarantee linear-time complexity, early evidence is progressively compressed and discarded as more chunks are processed.
Finally, the MDP structure itself prohibits the agent from revisiting past inputs once they are overwritten, further limiting its ability to integrate evidence scattered across distant parts of the document.

### 2.2 Memory Agent with History-Augmented State

To address these limitations, we extend the agent’s reasoning capability beyond a strictly forward trajectory by enabling it to revisit and incorporate past evidence on demand.
Specifically, the agent not only maintains the current memory $m_{t}$ but also generates a callback query $q_{t}$ to search over its history of memories $\{m_{i}\}_{i\leq t}$.
The retrieved content is then integrated into the state representation, yielding $s_{t}=(m_{t},q_{t})$.
This design allows the agent to selectively recall overlooked information and construct non-linear reasoning paths, rather than being confined to irreversible memory updates.

To realize this mechanism, at each step $t$ the agent receives the fixed question $Q$, the current document chunk $c_{t}$, and the current state $s_{t}$.
It is further equipped with a retrieval function $\mathcal{E}$, which selects relevant content from the previous memories $\{m_{i}\}_{i<t}$ on the overlap of words with the query $q_{t}$.
The state transition is then defined as

$$ $s_{t+1}=(m_{t+1},q_{t+1})=\pi_{\theta}\big(Q,c_{t},m_{t},\mathcal{E}(\{m_{i}\}_{i\leqslant t},q_{t})\big),$ (3) $$

where $\mathcal{E}(X,b)=\arg\max_{x\in X}\text{recall}(b,x)$, with $\text{recall}(a,b)$ denoting the proportion of words in $a$ that also appear in $b$.

The query component $q_{t+1}$ evolves alongside the memory, enabling the agent to iteratively refine its retrieval strategy over time.
This design frees the agent from a strictly linear trajectory through the document, allowing it to form non-linear reasoning paths by recalling earlier evidence and thereby mitigating the information loss inherent to fixed-length memory.

### 2.3 Reinforcement Learning with Multi-Level Reward Shaping

A primary challenge in training memory-augmented agents is the sparse and delayed nature of supervision.
For instance, a reward signal based solely on the final answer’s correctness provides weak guidance for the many intermediate steps leading to it.
To address this, we analyzed the agent’s reasoning process and made the key observations:
(1) In GRPO optimization, there are multiple rollouts for a single query $Q$ and document set $\{c_{t}\}_{t=0}^{T-1}$, yet they explore different reasoning paths leading to different answers.
(2) At each given step $t$, the agent across different trajectories sees the same external context $(Q,c_{t})$ but maintains a different internal state $s_{t}$. In this situation, the agent’s task is to integrate the current context with its evolving state to approach the correct answer.

Based on this insight, we introduce reinforcement learning with multi-level rewards (RLMLR), an RL algorithm tailored for the optimization of LLM memory agents.
As illustrated in Figure [3](https://arxiv.org/html/2509.23040v1#S2.F3)(b), this algorithm comprises two main components: a trajectory-level reward that evaluates the final outcome, and a dense, step-level state reward designed to shape the agent’s intermediate behaviors by measuring relative information gain.
These rewards are normalized across the corresponding trajectories and steps to acquire the overall advantage for group relative policy optimization (GRPO) (Shao et al., 2024) optimization.

#### 2.3.1 Trajectory-Level Outcome Rewards for Final Correctness

The ultimate measure of an agent’s success is its ability to answer the given question correctly. We capture this with a trajectory-level outcome reward, which is calculated based on the terminal state of each trajectory. Specifically, we first extract the predicted answer $\hat{y}^{(g)}$, enclosed in a \box{}, from the state $s_{T+1}^{(g)}$. The outcome reward is then computed using an exact match metric against the set of ground-truth answers $Y$:

$$ $R_{\text{out}}^{(g)}=\max_{y\in Y}{\mathbb{I}(\hat{y}^{(g)}=y)},$ (4) $$

where $\mathbb{I}(\cdot)$ is the indicator function that returns 1 if the condition is true and 0 otherwise.

Figure: Figure 3: Overview of RL with Multi-Level Rewards (RLMLR). (a) From the trajectories generated by the actor model, we compute outcome rewards at terminal states and state rewards at all states. (b) Each reward type is normalized at the corresponding level: state rewards across the states at the same step, and outcome rewards across all trajectories in the group.
Refer to caption: x3.png

#### 2.3.2 Step-Level Action Rewards for Behavior Shaping

To provide the dense, fine-grained supervision that outcome rewards lack, we introduce step-level state rewards. These rewards evaluate the quality of intermediate state updates within a trajectory, directly shaping the agent’s behavior toward greater efficiency and effectiveness.

- •
Information Gain in Memory Updates: To combat the progressive information loss discussed in the introduction, we use a rubric-based reward to measure the information gain in the agent’s memory.
After each update from $m_{t-1}$ to $m_{t}$, we assess the presence of crucial entities from the ground-truth answer.
If $m_{t}$ contains more information that are directly relevant to the ground truth $Y$ than $m_{t-1}$, we believe there’s a positive information gain achieved at time step $t$.
Building on such rationale, we use the change in recall as a reward:
$r_{\text{memory},t}^{(g)}=\max_{y\in Y}\text{recall}(m_{t}^{(g)},y)-\max_{y\in Y}\text{recall}(m_{t-1}^{(g)},y).$
(5)
- •
Bonus for Callback Retrievals: When the query component $q_{t}^{(g)}$ triggers a retrieval through $\mathcal{E}(\{m_{i}^{(g)}\}_{i\leq t},q_{t}^{(g)})$, the agent supplements its current memory with recalled information.
To encourage meaningful retrieval, we design a reward that measures the additional recall of critical information provided by the retrieved content beyond what is already available in the current memory $m_{t}^{(g)}$ and the immediate context $c_{t}$. Formally:
$r_{\text{callback},t}^{(g)}=\max_{y\in Y}\text{recall}\big(y,\mathcal{E}(\{m_{i}^{(g)}\}_{i\leq t},q_{t}^{(g)})\cup m_{t}^{(g)}\cup c_{t}\big)-\max_{y\in Y}\text{recall}(y,m_{t}^{(g)}\cup c_{t}).$
(6)
- •
Format Reward: To ensure that the agent’s outputs can be reliably parsed, we introduce a format reward $r_{\text{format},t}^{(g)}$ for all steps. For intermediate states, this reward checks for the correct usage of <callback> and <memory> tags. For the final step, it verifies the presence of the \box{} tag for the predicted answer.

The total step-level state reward at time $t$ for trajectory $g$ is the sum of these components:

$$ $R_{\text{state},t}^{(g)}=r_{\text{memory},t}^{(g)}+r_{\text{callback},t}^{(g)}+r_{\text{format},t}^{(g)}.$ (7) $$

#### 2.3.3 Training Objective

Given an actor model $\pi_{\theta}$ and a reference model $\pi_{\text{ref}}$, we sample a group of $G$ trajectories
$\{\tau^{(g)}\}_{g=1}^{G}$, where each trajectory $\tau^{(g)}=(s_{1}^{(g)},s_{2}^{(g)},\dots,s_{T+1}^{(g)})$ is generated according to the state-transition dynamics in §[2.2](https://arxiv.org/html/2509.23040v1#S2.SS2).
The optimization objective is a variant of GRPO (Shao et al., 2024) algorithm.
Refer to Appendix [C.1](https://arxiv.org/html/2509.23040v1#A3.SS1) for the full form of our training objective.

The normalized group advantage $\hat{A}_{t}^{(g)}$ is a composite of our multi-level rewards, with components calculated at different scales to reflect their distinct roles.
For the outcome reward, we compute a trajectory-level advantage $\hat{A}_{\text{out}}^{(g)}$ by comparing a trajectory’s outcome to the group average.
For the state rewards, we compute a step-level advantage $\hat{A}_{\text{state},t}^{(g)}$ by comparing a state’s reward to the average reward of states at the same step $t$ in the group.
Following (Liu et al., 2025b; c), we omit the standard deviation term during normalization to avoid introducing difficulty bias:

$$ $\hat{A}_{\text{out}}^{(g)}=R_{\text{out}}^{(g)}-\frac{1}{G}\sum_{k=1}^{G}{R_{\text{out}}^{(k)}},\qquad\hat{A}_{\text{state},t}^{(g)}=R_{\text{state},t}^{(g)}-\frac{1}{G}\sum_{k=1}^{G}{R_{\text{state},t}^{(k)}}.$ (8) $$

Finally, the overall advantage $\hat{A}_{t}^{(g)}$ in Eq. [10](https://arxiv.org/html/2509.23040v1#A3.E10) is a combination of these two components:

$$ $\hat{A}_{t}^{(g)}=\alpha\hat{A}_{\text{out}}^{(g)}+(1-\alpha)\hat{A}_{\text{state},t}^{(g)},$ (9) $$

where $\alpha$ is the hyperparameter that controls the importance of each term.

## 3 Experiments

In this paper, we conduct experiments to answer the following research questions (RQs):

RQ1: Does ReMemR1 outperform other memory agents or general-purpose LLMs on long-context tasks, and can it alleviate the progressive information loss?

RQ2: Does ReMemR1 achieve nonlinear document utilization through the callback mechanism?

RQ3: Does our proposed RLMLR help in shaping the intermediate behaviors of the memory agent?

RQ4: What’s the benefits of the RL-driven memory callback, comparing with rule-based design?

**Table 1: Long-context QA results on HotpotQA (Yang et al., 2018) and 2WikiMultiHopQA (Ho et al., 2020). Values are accuracy (%), rounded to 1 decimal. Bold denotes the best performances.**
|  |  | Number of Context Documents |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Scale | Method | 50 | 100 | 200 | 400 | 800 | 1600 | 3200 | 6400 |
| 3B | Qwen2.5 (Yang et al., 2024) | 59.4 | 57.0 | - | - | - | - | - | - |
| MemAgent (Yu et al., 2025a) | 70.3 | 69.4 | 60.9 | 68.8 | 60.9 | 60.2 | 59.4 | 58.8 |  |
| ReMemR1 (Ours) | 70.9 | 71.7 | 63.8 | 74.0 | 65.4 | 65.0 | 65.4 | 66.1 |  |
| 7B | Qwen2.5 (Yang et al., 2024) | 70.3 | 75.0 | - | - | - | - | - | - |
| R1-Distill (DeepSeek-AI et al., 2025) | 40.6 | 25.8 | 10.2 | 0.8 | 1.6 | 2.3 | 1.5 | 3.1 |  |
| Qwen2.5-1M (Yang et al., 2025b) | 75.8 | 71.9 | 68.0 | 67.2 | 69.5 | 54.7 | 22.7 | 0.0 |  |
| MemAgent (Yu et al., 2025a) | 81.8 | 78.9 | 78.9 | 77.0 | 79.7 | 72.1 | 74.0 | 75.8 |  |
| ReMemR1 (Ours) | 82.3 | 82.8 | 81.1 | 78.9 | 82.0 | 79.7 | 80.0 | 80.8 |  |

### 3.1 Experimental Setup

##### Datasets.

Our training data is sourced from HotpotQA (Yang et al., 2018).
We pad the context of each training sample with random documents to 200 (about 30K tokens) per sample.
For evaluation, we use the in-distribution (ID) HotpotQA and the out-of-distribution (OOD) 2WikiMultiHopQA (Ho et al., 2020) datasets.
The context documents of test data are also padded, ranging from 50 to 6400 documents per sample.
For more implementation and dataset details, refer to Appendix [C](https://arxiv.org/html/2509.23040v1#A3).

##### Baselines.

In our experiments, we compare our method against three categories of baselines:
(1) general LLMs, including Qwen2.5 models (Yang et al., 2024) and Qwen models distilled from DeepSeek-R1 (DeepSeek-AI et al., 2025).
(2) Long-context LLMs, including Qwen2.5-1M (Yang et al., 2025b);
(3) tailored memory agents, such as MemAgent (Yu et al., 2025a).
By default, we use the instruct version for all models.
For comparison with more baselines, refer to Appendix [B.1](https://arxiv.org/html/2509.23040v1#A2.SS1).

### 3.2 Main Results (RQ1)

As shown in Table [1(b)](https://arxiv.org/html/2509.23040v1#S3.T1.st2), our method consistently achieves the best accuracy across all model scales, datasets, and context lengths, surpassing both general-purpose LLMs and specialized memory agents.
Compared with MemAgent, it achieves up to 7.3% higher accuracy on 3B model and 7.6% on 7B model, underscoring the effectiveness of adaptive memory recall.
We further observe that as the number of context documents increases, the role of memory becomes increasingly critical.
Pure reasoning models and long-context models exhibit sharp performance degradation when facing very long contexts, while MemAgent mitigates this issue by adopting a “memorize while reading” strategy that stores salient information in a memory buffer.
Building upon this, our method equips the agent with an RL-driven memory callback mechanism that adaptively selects what and when to retrieve, thereby enhancing the quality of the maintained memory.
This advantage becomes increasingly evident as the document length grows, since in longer contexts important evidence is more likely to be overwritten or overlooked, amplifying the need for precise recall to preserve reasoning accuracy.
Notably, the gains are even more pronounced on the OOD 2WikiMultiHopQA dataset, indicating that our approach goes beyond memorizing dataset-specific patterns and instead acquires a genuine retrieval and reasoning ability, leading to stronger generalization across domains.

### 3.3 Distant Evidence Challenge (RQ2)

Figure: Figure 4: Accuracy on 2Wiki with distant evidences.
Refer to caption: x4.png

To rigorously test the effectiveness and accuracy of the proposed memory callback mechanism, we construct a more challenging evaluation setting.
Specifically, for each question, the supporting evidences are arranged in the reverse order of their required reasoning sequence, and the distance between successive evidence is enforced to exceed half of the total number of context documents.
This setup makes it infeasible for the model to rely on local context alone;
instead, it requires the model to identify and utilize interdependent evidences across long spans.

As shown in Figure [4](https://arxiv.org/html/2509.23040v1#S3.F4), our method surpasses MemAgent by large margin under this setting.
MemAgent suffers pronounced accuracy degradation due to its inherent inability to look back and reliably recall distant, scattered evidences.
In contrast, our RL-driven callback mechanism adaptively retrieves and maintains critical information, achieving far superior performance.
These results demonstrate that the proposed callback design is both effective and robust, particularly when reasoning requires nontrivial coordination of evidences over long contexts.

### 3.4 Ablation Studies

#### 3.4.1 Effectiveness of RLMLR (RQ3)

**Table 2: Accuracy on HotpotQA with different $\alpha$ values.**
|  |  | Number of Context Documents |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Method | $\alpha$ | 50 | 100 | 200 | 400 | 800 | 1600 | 3200 | 6400 |
| ReMemR1 | 1.0 | 70.3 | 73.4 | 61.5 | 59.6 | 60.9 | 64.1 | 62.5 | 63.3 |
| 0.8 | 70.9 | 71.7 | 63.8 | 74.0 | 65.4 | 65.0 | 65.4 | 66.1 |  |
| 0.5 | 71.7 | 68.5 | 62.2 | 66.1 | 63.0 | 58.3 | 59.6 | 65.4 |  |
| 0.2 | 68.8 | 68.5 | 55.9 | 62.5 | 53.5 | 45.7 | 49.6 | 52.0 |  |

In ReMemR1, we propose RLMLR to alleviate the sparse supervision problem by combining trajectory-level outcome rewards with step-level state rewards.
The balance between these two rewards is controlled by a hyperparameter $\alpha$, which determines how much weight is placed on final-answer correctness versus intermediate behavior shaping (Eq. [9](https://arxiv.org/html/2509.23040v1#S2.E9)).
We evaluate $\alpha\in\{1.0,0.8,0.5,0.2\}$ on Qwen2.5-3B Instruct to examine its impact.

Results in Table [2](https://arxiv.org/html/2509.23040v1#S3.T2) demonstrate that $\alpha=0.8$ consistently delivers the best accuracy across different context lengths.
A larger $\alpha$ (e.g., $1.0$) corresponds to using only outcome rewards, which neglects the benefits of dense step-level guidance and leads to weaker optimization.
Conversely, smaller values (e.g., $0.2$) overly emphasize step-level shaping, which distracts the model from optimizing for final correctness.
Based on these findings, we adopt $\alpha=0.8$ by default in all the other experiments, as it provides the best trade-off between global outcome rewards and local step-level supervision.

#### 3.4.2 RL-driven v.s. Rule-based Memory Callback (RQ4)

**Table 3: Comparison of accuracy (%) on HotpotQA and 2WikiMultiHopQA across different callback implementations. Bold denotes the best performance.**
|  |  | Number of Context Documents |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Benchmark | Method | 50 | 100 | 200 | 400 | 800 | 1600 | 3200 | 6400 |
| HotpotQA | MemAgent | 70.3 | 69.4 | 60.9 | 68.8 | 60.9 | 60.2 | 59.4 | 58.8 |
| MemAgent + rule-based callback | 69.5 | 66.4 | 57.0 | 60.9 | 61.4 | 53.9 | 61.7 | 60.9 |  |
| ReMemR1 (Ours) | 70.9 | 71.7 | 63.8 | 74.0 | 65.4 | 65.0 | 65.4 | 66.1 |  |
| 2WikiMultiHopQA | MemAgent | 41.4 | 45.3 | 42.2 | 41.4 | 38.3 | 28.9 | 26.7 | 25.9 |
| MemAgent + rule-based callback | 49.2 | 43.0 | 35.9 | 35.2 | 33.4 | 33.6 | 30.5 | 27.3 |  |
| ReMemR1 (Ours) | 53.5 | 50.4 | 42.5 | 41.7 | 37.0 | 36.2 | 35.4 | 37.8 |  |

A key component of ReMemR1 is the RL-driven memory callback, where the agent learns through reinforcement learning to generate informative queries that retrieve past evidence most relevant to the current step.
This mechanism allows the agent to dynamically determine when and what to recall during reasoning.
As an intuitive yet strong baseline, we design a rule-based memory callback, where the agent uses the question $Q$ itself as a fixed query for retrieval at every step.
This design is motivated by the fact that the question contains rich information about the target answer, and thus provides a natural heuristic for guiding memory recall without requiring additional training.

Table [3](https://arxiv.org/html/2509.23040v1#S3.T3) reports the results on HotpotQA and 2WikiMultiHopQA, with Qwen2.5-3B Instruct as the base model.
We observe that RL-driven memory callback consistently outperforms both the vanilla MemAgent and the rule-based callback on both datasets across all context lengths.
Notably, the rule-based callback does not always yield improvements and can even cause performance drops of up to $7.9\%$, highlighting that determining when and what to recall is non-trivial.
We also observe that the advantage of our method increases as the document length grows, indicating that effective memory recall becomes increasingly crucial in longer contexts.
These results confirm that learning adaptive recall strategies via RL is essential for robust and generalizable long-context reasoning.
Refer to Appendix [B.2](https://arxiv.org/html/2509.23040v1#A2.SS2) for extended discussion about the impact of RL training.

## 4 Conclusion

This work examined the inherent limitations of the prevailing “memorize while reading” paradigm for long-context question answering, including irreversible forward-only processing, progressive information loss from memory overwriting, and the sparsity of supervision signals.
To address these challenges, we proposed ReMemR1, a memory-augmented agent that enhances the state representation with callback queries, enabling retrieval from historical memories and facilitating non-linear reasoning paths.
To further improve training efficacy, we developed RLMLR, a reinforcement learning framework with multi-level rewards that combines trajectory-level outcome supervision with step-level state rewards.
Experiments across both in-distribution and out-of-distribution benchmarks show that ReMemR1 consistently surpasses general LLMs and prior memory agents, and remains robust under the challenging distant-evidence setting.
Ablation studies further confirm the necessity of the RLMLR training scheme and the RL-driven memory callback for enabling effective and generalizable long-context reasoning.
Looking ahead, we believe this work opens up new potential for future research on robust long-context understanding agents across diverse real-world domains.

## Ethics Statement

Our research is confined to computational experiments on publicly available benchmarks, specifically HotpotQA and 2WikiMultiHopQA. These datasets consist of publicly sourced text and do not contain personal information or other forms of sensitive data (Yang et al., 2018; Ho et al., 2020).
No human subjects were involved in any stage of our work, including data collection or model evaluation. The focus of this paper is on foundational research for long-context reasoning, and we do not develop or evaluate applications in high-stakes domains such as medicine, law, or finance.

We acknowledge the broader ethical challenges inherent in LLM-based systems, including the risk of perpetuating societal biases present in their training data. While our methodological focus is on reasoning capabilities, the introduction of a memory mechanism raises specific considerations regarding privacy and security. A system with the ability to store and recall information over long contexts could pose risks if deployed with private or proprietary data without robust safeguards. Any downstream application of this work should undergo evaluation for fairness, transparency, and potential discriminatory impacts.

## Reproducibility Statement

To ensure the reproducibility of our results, we provide an anonymous downloadable source code package in our abstract, as recommended by the conference guidelines.
This package includes:

- •
Complete code for generating our evaluation datasets from publicly available benchmarks (HotpotQA and 2WikiMultiHopQA) using fixed random seeds.
- •
Configuration files and instructions for setting up the experimental environment.
- •
The training procedure of ReMemR1, including the implementation of the callback mechanism, RLMLR, and runnable training scripts based on verl.
- •
Evaluation scripts for both baseline models and our proposed method.

In addition, detailed descriptions of the experimental setup and hyperparameters are reported in §[3.1](https://arxiv.org/html/2509.23040v1#S3.SS1) and Appendix [C](https://arxiv.org/html/2509.23040v1#A3).
We hope that these materials will enable researchers to fully replicate and further extend our work.

## References

- Ahmadian et al. (2024)
Arash Ahmadian, Chris Cremer, Matthias Gallé, Marzieh Fadaee, Julia Kreutzer, Olivier Pietquin, Ahmet Üstün, and Sara Hooker.
Back to basics: Revisiting reinforce style optimization for learning from human feedback in llms.
*arXiv preprint arXiv:2402.14740*, 2024.
- Beltagy et al. (2020)
Iz Beltagy, Matthew E Peters, and Arman Cohan.
Longformer: The long-document transformer.
*arXiv preprint arXiv:2004.05150*, 2020.
- Chen et al. (2025a)
Qiguang Chen, Libo Qin, Jinhao Liu, Dengyun Peng, Jiannan Guan, Peng Wang, Mengkang Hu, Yuhang Zhou, Te Gao, and Wangxiang Che.
Towards reasoning era: A survey of long chain-of-thought for reasoning large language models.
*arXiv preprint arXiv:2503.09567*, 2025a.
- Chen et al. (2025b)
Qiguang Chen, Libo Qin, Jinhao Liu, Dengyun Peng, Jiannan RGuan, Peng Wang, Mengkang Hu, Yuhang Zhou, Te Gao, and Wangxiang Che.
Towards reasoning era: A survey of long chain-of-thought for reasoning large language models.
*arXiv preprint arXiv:2503.09567*, 2025b.
- Chen et al. (2023)
Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian.
Extending context window of large language models via positional interpolation.
*arXiv preprint arXiv:2306.15595*, 2023.
- Child et al. (2019)
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever.
Generating long sequences with sparse transformers.
*arXiv preprint arXiv:1904.10509*, 2019.
- DeepSeek-AI et al. (2025)
DeepSeek-AI, Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, Xiaokang Zhang, Xingkai Yu, Yu Wu, Z. F. Wu, Zhibin Gou, Zhihong Shao, Zhuoshu Li, Ziyi Gao, Aixin Liu, Bing Xue, Bingxuan Wang, Bochao Wu, Bei Feng, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, Damai Dai, Deli Chen, Dongjie Ji, Erhang Li, Fangyun Lin, Fucong Dai, Fuli Luo, Guangbo Hao, Guanting Chen, Guowei Li, H. Zhang, Han Bao, Hanwei Xu, Haocheng Wang, Honghui Ding, Huajian Xin, Huazuo Gao, Hui Qu, Hui Li, Jianzhong Guo, Jiashi Li, Jiawei Wang, Jingchang Chen, Jingyang Yuan, Junjie Qiu, Junlong Li, J. L. Cai, Jiaqi Ni, Jian Liang, Jin Chen, Kai Dong, Kai Hu, Kaige Gao, Kang Guan, Kexin Huang, Kuai Yu, Lean Wang, Lecong Zhang, Liang Zhao, Litong Wang, Liyue Zhang, Lei Xu, Leyi Xia, Mingchuan Zhang, Minghua Zhang, Minghui Tang, Meng Li, Miaojun Wang, Mingming Li, Ning Tian, Panpan Huang, Peng Zhang, Qiancheng Wang, Qinyu Chen, Qiushi Du, Ruiqi Ge,
Ruisong Zhang, Ruizhe Pan, Runji Wang, R. J. Chen, R. L. Jin, Ruyi Chen, Shanghao Lu, Shangyan Zhou, Shanhuang Chen, Shengfeng Ye, Shiyu Wang, Shuiping Yu, Shunfeng Zhou, Shuting Pan, and S. S. Li.
Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning.
*CoRR*, abs/2501.12948, 2025.
doi: 10.48550/ARXIV.2501.12948.
- Ding et al. (2023)
Jiayu Ding, Shuming Ma, Li Dong, Xingxing Zhang, Shaohan Huang, Wenhui Wang, Nanning Zheng, and Furu Wei.
Longnet: Scaling transformers to 1,000,000,000 tokens.
*arXiv preprint arXiv:2307.02486*, 2023.
- Du et al. (2025)
Yiming Du, Wenyu Huang, Danna Zheng, Zhaowei Wang, Sebastien Montella, Mirella Lapata, Kam-Fai Wong, and Jeff Z Pan.
Rethinking memory in ai: Taxonomy, operations, topics, and future directions.
*arXiv preprint arXiv:2505.00675*, 2025.
- Duverger et al. (2024)
Romain Duverger, Alexis Bonnin, Romain Granier, Quentin Marolleau, Cédric Blanchard, Nassim Zahzam, Yannick Bidel, Malo Cadoret, Alexandre Bresson, and Sylvain Schwartz.
Metrology of microwave fields based on trap-loss spectroscopy with cold rydberg atoms.
*Physical Review Applied*, 22(4):044039, 2024.
- Gu & Dao (2023)
Albert Gu and Tri Dao.
Mamba: Linear-time sequence modeling with selective state spaces.
*arXiv preprint arXiv:2312.00752*, 2023.
- Gu et al. (2021)
Albert Gu, Karan Goel, and Christopher Ré.
Efficiently modeling long sequences with structured state spaces.
*arXiv preprint arXiv:2111.00396*, 2021.
- Ho et al. (2020)
Xanh Ho, Anh-Khoa Duong Nguyen, Saku Sugawara, and Akiko Aizawa.
Constructing a multi-hop qa dataset for comprehensive evaluation of reasoning steps.
In *Proceedings of the 28th International Conference on Computational Linguistics*, pp. 6609–6625, 2020.
- Hsieh et al. (2024)
Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg.
Ruler: What’s the real context size of your long-context language models?
*arXiv preprint arXiv:2404.06654*, 2024.
- Hu (2025)
Jian Hu.
Reinforce++: A simple and efficient approach for aligning large language models.
*arXiv preprint arXiv:2501.03262*, 2025.
- Jaech et al. (2024)
Aaron Jaech, Adam Kalai, Adam Lerer, Adam Richardson, Ahmed El-Kishky, Aiden Low, Alec Helyar, Aleksander Madry, Alex Beutel, Alex Carney, Alex Iftimie, Alex Karpenko, Alex Tachard Passos, Alexander Neitz, Alexander Prokofiev, Alexander Wei, Allison Tam, Ally Bennett, Ananya Kumar, Andre Saraiva, Andrea Vallone, Andrew Duberstein, Andrew Kondrich, Andrey Mishchenko, Andy Applebaum, Angela Jiang, Ashvin Nair, Barret Zoph, Behrooz Ghorbani, Ben Rossen, Benjamin Sokolowsky, Boaz Barak, Bob McGrew, Borys Minaiev, Botao Hao, Bowen Baker, Brandon Houghton, Brandon McKinzie, Brydon Eastman, Camillo Lugaresi, Cary Bassin, Cary Hudson, Chak Ming Li, Charles de Bourcy, Chelsea Voss, Chen Shen, Chong Zhang, Chris Koch, Chris Orsinger, Christopher Hesse, Claudia Fischer, Clive Chan, Dan Roberts, Daniel Kappler, Daniel Levy, Daniel Selsam, David Dohan, David Farhi, David Mely, David Robinson, Dimitris Tsipras, Doug Li, Dragos Oprica, Eben Freeman, Eddie Zhang, Edmund Wong, Elizabeth Proehl, Enoch Cheung, Eric Mitchell,
Eric Wallace, Erik Ritter, Evan Mays, Fan Wang, Felipe Petroski Such, Filippo Raso, Florencia Leoni, Francis Song, Fred von Lohmann, Freddie Sulit, Geoff Salmon, Giambattista Parascandolo, Gildas Chabot, Grace Zhao, Greg Brockman, Guillaume Leclerc, Hadi Salman, Haiming Bao, Hao Sheng, Hart Andrin, Hessam Bagherinezhad, Hongyu Ren, Hunter Lightman, Hyung Won Chung, Ian Kivlichan, Ian O’Connell, Ian Osband, Ignasi Clavera Gilaberte, and Ilge Akkaya.
Openai o1 system card.
*CoRR*, abs/2412.16720, 2024.
doi: 10.48550/ARXIV.2412.16720.
- Kaelbling et al. (1996)
Leslie Pack Kaelbling, Michael L Littman, and Andrew W Moore.
Reinforcement learning: A survey.
*Journal of artificial intelligence research*, 4:237–285, 1996.
- Katharopoulos et al. (2020)
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret.
Transformers are rnns: Fast autoregressive transformers with linear attention.
In *International conference on machine learning*, pp. 5156–5165. PMLR, 2020.
- Li & Liu (2024)
Yukun Li and Liping Liu.
Enhancing diffusion-based point cloud generation with smoothness constraint.
*arXiv preprint arXiv:2404.02396*, 2024.
- Li et al. (2025)
Zhiyu Li, Shichao Song, Chenyang Xi, Hanyu Wang, Chen Tang, Simin Niu, Ding Chen, Jiawei Yang, Chunyu Li, Qingchen Yu, et al.
Memos: A memory os for ai system.
*arXiv preprint arXiv:2507.03724*, 2025.
- Liu et al. (2025a)
Bang Liu, Xinfeng Li, Jiayi Zhang, Jinlin Wang, Tanjin He, Sirui Hong, Hongzhang Liu, Shaokun Zhang, Kaitao Song, Kunlun Zhu, et al.
Advances and challenges in foundation agents: From brain-inspired intelligence to evolutionary, collaborative, and safe systems.
*arXiv preprint arXiv:2504.01990*, 2025a.
- Liu et al. (2023)
Hao Liu, Matei Zaharia, and Pieter Abbeel.
Ring attention with blockwise transformers for near-infinite context.
*arXiv preprint arXiv:2310.01889*, 2023.
- Liu et al. (2025b)
Zichen Liu, Changyu Chen, Wenjun Li, Penghui Qi, Tianyu Pang, Chao Du, Wee Sun Lee, and Min Lin.
Understanding r1-zero-like training: A critical perspective.
*arXiv preprint arXiv:2503.20783*, 2025b.
- Liu et al. (2025c)
Zichen Liu, Changyu Chen, Wenjun Li, Penghui Qi, Tianyu Pang, Chao Du, Wee Sun Lee, and Min Lin.
Understanding r1-zero-like training: A critical perspective.
*arXiv preprint arXiv:2503.20783*, 2025c.
- Maharana et al. (2024)
Adyasha Maharana, Dong-Ho Lee, Sergey Tulyakov, Mohit Bansal, Francesco Barbieri, and Yuwei Fang.
Evaluating very long-term conversational memory of llm agents.
In *Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)*, pp. 13851–13870, 2024.
- OpenAI (2023)
OpenAI.
Gpt-4 technical report.
*arXiv preprint arXiv:2303.08774*, 2023.
- Ouyang et al. (2022)
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al.
Training language models to follow instructions with human feedback.
*Advances in neural information processing systems*, 35:27730–27744, 2022.
- Packer et al. (2023)
Charles Packer, Sarah Wooders, Kevin Lin, Vivian Fang, Shishir G Patil, Ion Stoica, and Joseph E Gonzalez.
Memgpt: Towards llms as operating systems.
*arXiv preprint arXiv:2310.08560*, 2023.
- Peng et al. (2023a)
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, et al.
Rwkv: Reinventing rnns for the transformer era.
*arXiv preprint arXiv:2305.13048*, 2023a.
- Peng et al. (2023b)
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole.
Yarn: Efficient context window extension of large language models.
In *The Twelfth International Conference on Learning Representations*, 2023b.
- Schulman et al. (2017)
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov.
Proximal policy optimization algorithms.
*arXiv preprint arXiv:1707.06347*, 2017.
- Shao et al. (2024)
Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, Y Wu, et al.
Deepseekmath: Pushing the limits of mathematical reasoning in open language models.
*arXiv preprint arXiv:2402.03300*, 2024.
- Su et al. (2024)
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu.
Roformer: Enhanced transformer with rotary position embedding.
*Neurocomputing*, 568:127063, 2024.
- Team et al. (2024)
Gemini Team, Petko Georgiev, Ving Ian Lei, Ryan Burnell, Libin Bai, Anmol Gulati, Garrett Tanzer, Damien Vincent, Zhufeng Pan, Shibo Wang, et al.
Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context.
*arXiv preprint arXiv:2403.05530*, 2024.
- Wan et al. (2025)
Fanqi Wan, Weizhou Shen, Shengyi Liao, Yingcheng Shi, Chenliang Li, Ziyi Yang, Ji Zhang, Fei Huang, Jingren Zhou, and Ming Yan.
Qwenlong-l1: Towards long-context large reasoning models with reinforcement learning.
*arXiv preprint arXiv:2505.17667*, 2025.
- Wang et al. (2025)
Kun Wang, Guibin Zhang, Zhenhong Zhou, Jiahao Wu, Miao Yu, Shiqian Zhao, Chenlong Yin, Jinhu Fu, Yibo Yan, Hanjun Luo, et al.
A comprehensive survey in llm (-agent) full stack safety: Data, training and deployment.
*arXiv preprint arXiv:2504.15585*, 2025.
- Xu et al. (2025)
Wujiang Xu, Kai Mei, Hang Gao, Juntao Tan, Zujie Liang, and Yongfeng Zhang.
A-mem: Agentic memory for llm agents.
*arXiv preprint arXiv:2502.12110*, 2025.
- Yan et al. (2025)
Sikuan Yan, Xiufeng Yang, Zuchao Huang, Ercong Nie, Zifeng Ding, Zonggen Li, Xiaowen Ma, Hinrich Schütze, Volker Tresp, and Yunpu Ma.
Memory-r1: Enhancing large language model agents to manage and utilize memories via reinforcement learning.
*arXiv preprint arXiv:2508.19828*, 2025.
- Yang et al. (2024)
An Yang, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, et al.
Qwen2. 5 technical report.
*arXiv preprint arXiv:2412.15115*, 2024.
- Yang et al. (2025a)
An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, et al.
Qwen3 technical report.
*arXiv preprint arXiv:2505.09388*, 2025a.
- Yang et al. (2025b)
An Yang, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoyan Huang, Jiandong Jiang, Jianhong Tu, Jianwei Zhang, Jingren Zhou, et al.
Qwen2. 5-1m technical report.
*arXiv preprint arXiv:2501.15383*, 2025b.
- Yang et al. (2018)
Zhilin Yang, Peng Qi, Saizheng Zhang, Yoshua Bengio, William Cohen, Ruslan Salakhutdinov, and Christopher D Manning.
Hotpotqa: A dataset for diverse, explainable multi-hop question answering.
In *Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing*, pp. 2369–2380, 2018.
- Yu et al. (2025a)
Hongli Yu, Tinghong Chen, Jiangtao Feng, Jiangjie Chen, Weinan Dai, Qiying Yu, Ya-Qin Zhang, Wei-Ying Ma, Jingjing Liu, Mingxuan Wang, et al.
Memagent: Reshaping long-context llm with multi-conv rl-based memory agent.
*arXiv preprint arXiv:2507.02259*, 2025a.
- Yu et al. (2025b)
Qiying Yu, Zheng Zhang, Ruofei Zhu, Yufeng Yuan, Xiaochen Zuo, Yu Yue, Tiantian Fan, Gaohong Liu, Lingjun Liu, Xin Liu, et al.
Dapo: An open-source llm reinforcement learning system at scale.
*arXiv preprint arXiv:2503.14476*, 2025b.

## Appendix A Related Work

We review three areas of prior research relevant to our long-context LLM agent: memory mechanisms for LLM-based agents, approaches for extending context length in language models, and reinforcement learning techniques for improving LLM reasoning abilities.

##### Memory Augmented LLM Agents.

The reasoning and planning capabilities of LLM agents are fundamentally limited by the fixed size of their context window (Hsieh et al., 2024; Maharana et al., 2024; Liu et al., 2025a).
To overcome this, researchers have built external memory systems to retain information across long interactions, enabling agents to recall past experiences and adapt their behavior (OpenAI, 2023; Wang et al., 2025; Du et al., 2025).
Early memory systems primarily focused on simple short-term memory (*e.g., *, prepending a conversation history to the prompt) and long-term memory (*e.g., *storing information in a vector database for retrieval) (Li & Liu, 2024; Duverger et al., 2024; Packer et al., 2023; Yan et al., 2025).
More recent approaches explore a “memorizing while reading" paradigm, where the LLM autonomously organizes its memory corpus during a single-pass scan through the documents (Xu et al., 2025; Li et al., 2025; Yu et al., 2025a).

##### Long-Context LLMs.

This long-context challenge in LLM has driven a variety of solutions, which can be broadly categorized into architectural modifications and context window extension techniques. Novel architectures, such as state space models (Gu et al., 2021; Gu & Dao, 2023; Peng et al., 2023a), achieve linear-time complexity and are highly efficient for long sequences. Other efforts focus on extending the context windows of attention-based LLMs. One approach involves developing more efficient attention mechanisms to reduce computational burden (Beltagy et al., 2020; Ding et al., 2023; Child et al., 2019; Katharopoulos et al., 2020; Liu et al., 2023). A complementary technical route modifies Rotary Position Embedding to enable models to extrapolate effectively beyond their original training length (Su et al., 2024; Chen et al., 2023; Peng et al., 2023b).

##### Reinforcement Learning in LLMs.

Reinforcement Learning (RL) (Kaelbling et al., 1996) has emerged as a powerful paradigm for post-training LLMs recently (Chen et al., 2025a; b; Jaech et al., 2024; DeepSeek-AI et al., 2025).
Early efforts focus on Reinforcement Learning from Human Feedback (RLHF) (Ouyang et al., 2022) using algorithms like Proximal Policy Optimization (PPO) to align the LLM with human preferences (Schulman et al., 2017).
More recent work has explored scaling this process by using outcome-based rewards.
These Techniques such as Group Relative Policy Optimization (GRPO) (Shao et al., 2024) and Reinforce (Hu, 2025) are central to this trend, which offer alternatives to traditional PPO that reduce the need for a separate value model or extensive human-annotated data (Ahmadian et al., 2024; Yu et al., 2025b).

## Appendix B Additional Results

### B.1 Comparison against more Baselines

We also conduct comparisons with a broader set of long-context models beyond 7B level.
The baselines include recent Qwen3 models (Yang et al., 2025a), 14B variant of Qwen2.5-1M (Yang et al., 2025b) and R1-Distill-Qwen (DeepSeek-AI et al., 2025), and the 32B long-context LLM QwenLong-L1-32B (Wan et al., 2025).

Table [4(b)](https://arxiv.org/html/2509.23040v1#A2.T4.st2) reports the extended comparison on both ID and OOD settings.
In the table, we observe:
(1) At high context lengths, ReMemR1 outperforms long-context LLMs that are four times larger.
On HotpotQA, ReMemR1 achieves 80.8% accuracy at 6400 documents, substantially higher than QwenLong-L1-32B (38.3%) and 14B-level R1-Distill-Qwen (31.3%). Similarly, on 2WikiMultiHopQA, ReMemR1 reaches 50.3% accuracy at 6400 documents, outperforming QwenLong-L1-32B (29.9%) and R1-Distill-Qwen-14B (32%).
This highlights ReMemR1’s robustness under extreme context scaling.
(2) At mid-range context lengths (200–800 documents), ReMemR1 remains highly competitive. For example, on HotpotQA at 400 documents, ReMemR1 (78.9%) surpasses QwenLong-L1-32B (73.4%) and all other baselines.

**Table 4: Extended long-context QA results on HotpotQA (Yang et al., 2018) and 2WikiMultiHopQA (Ho et al., 2020). Values are accuracy (%), rounded to 1 decimal.**
|  |  | Number of Context Documents |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Scale | Method | 50 | 100 | 200 | 400 | 800 | 1600 | 3200 | 6400 |
| $<$7B | Qwen3-4B (Yang et al., 2025a) | 75.0 | 75.8 | 69.5 | 63.3 | 60.2 | 21.9 | 18.8 | 18.8 |
| ReMemR1 (Qwen2.5-3B) | 70.9 | 71.7 | 63.8 | 74.0 | 65.4 | 65.0 | 65.4 | 66.1 |  |
| $\geqslant$7B | Qwen3-8B (Yang et al., 2025a) | 81.3 | 78.9 | 71.9 | 70.3 | 74.2 | 33.6 | 23.4 | 19.5 |
| R1-Distill-Qwen-7B (DeepSeek-AI et al., 2025) | 40.6 | 25.8 | 10.2 | 0.8 | 1.6 | 2.3 | 1.5 | 3.1 |  |
| R1-Distill-Qwen-14B (DeepSeek-AI et al., 2025) | 79.7 | 76.6 | 64.1 | 57.8 | 40.6 | 33.6 | 20.3 | 31.3 |  |
| Qwen2.5-1M-7B Yang et al. (2025b) | 75.8 | 71.9 | 68.0 | 67.2 | 69.5 | 54.7 | 22.7 | 0.0 |  |
| Qwen2.5-1M-14B Yang et al. (2025b) | 78.1 | 83.6 | 76.6 | 73.4 | 70.3 | 60.9 | 42.2 | 0.0 |  |
| QwenLong-L1-32B Wan et al. (2025) | 83.6 | 85.2 | 74.2 | 73.4 | 57.8 | 45.3 | 38.9 | 38.3 |  |
|  | ReMemR1 (Qwen2.5-7B) | 82.3 | 82.8 | 81.1 | 78.9 | 82.0 | 79.7 | 80.0 | 80.8 |

### B.2 Impact of RL Training

**Table 5: Ablation on RL training. We report accuracy (%) on HotpotQA and 2WikiMultiHopQA with and without RL. The based models are Qwen2.5-3B Instruct.**
|  |  |  | Number of Context Documents |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Benchmark | Method | Setting | 50 | 100 | 200 | 400 | 800 | 1600 | 3200 | 6400 |
| HotpotQA | MemAgent | w/o RL | 60.2 | 47.7 | 35.9 | 28.9 | 24.2 | 23.4 | 14.8 | 14.1 |
| ReMemR1 | w/o RL | 35.4 | 40.9 | 31.5 | 25.2 | 26.0 | 24.4 | 16.5 | 20.5 |  |
| MemAgent | w/ RL | 70.3 | 69.4 | 60.9 | 68.8 | 60.9 | 60.2 | 59.4 | 58.8 |  |
| ReMemR1 | w/ RL | 70.9 | 71.7 | 63.8 | 74.0 | 65.4 | 65.0 | 65.4 | 66.1 |  |
| 2WikiMultiHopQA | MemAgent | w/o RL | 37.5 | 30.5 | 32.0 | 22.7 | 16.4 | 16.4 | 16.4 | 15.6 |
| ReMemR1 | w/o RL | 26.0 | 25.2 | 26.8 | 18.9 | 16.5 | 17.3 | 22.8 | 22.0 |  |
| MemAgent | w/ RL | 41.4 | 45.3 | 42.2 | 41.4 | 38.3 | 28.9 | 26.7 | 25.9 |  |
| ReMemR1 | w/ RL | 53.5 | 50.4 | 42.5 | 41.7 | 37.0 | 36.2 | 35.4 | 37.8 |  |

We further examine the impact of reinforcement learning on long-context reasoning.
Table [5](https://arxiv.org/html/2509.23040v1#A2.T5) compares model performance with (w/) and without (w/o) RL across different numbers of context documents, where all methods use Qwen2.5-3B Instruct (Yang et al., 2024) as the foundational model.
Without RL, both our method and MemAgent suffer from sharp performance drops as the context length grows, indicating difficulties in optimizing with only supervised signals.
Introducing RL substantially improves accuracy on both HotpotQA and 2WikiMultiHopQA.
In particular, our method with RL consistently achieves the highest scores across most context lengths, outperforming MemAgent by a clear margin.

We also observe that without RL training, the two paradigms (MemAgent and ReMemR1) shows different behavior at different context length levels:

- •
$<$ 800 Documents. When the context length is relatively small, directly applying Qwen-3B on ReMemR1 without RL shows lower accuracies than MemAgent.
We find out this phenomenon is caused by the imperfect instruction-following in the untrained model.
As the callback mechanism provides an opportunity to include more information, it also introduces additional format requirements.
According to Figure [5](https://arxiv.org/html/2509.23040v1#A2.F5), the 3B-level LLM begins with around 0.6 average format reward, which means the LLM fail to extract the updated memory for 40% steps.
As the training processes and the format reward grows, ReMemR1 quickly learns the format requirements under the guidance of action-level rewards, resulting in quickly increasing early-stage rewards.
- •
$\geqslant$ 800 Documents. As the context length raises to more than 800 documents, ReMemR1 shows slower accuracy drop, resulting in about 6% improvements on both benchmarks.
This observation concurs with the findings in Section [3.4.2](https://arxiv.org/html/2509.23040v1#S3.SS4.SSS2), where rule-based callback yields better long-horizon performances, which validates the benefits of callback mechanism in preventing long-term information losses.
These results highlight the importance of reinforcement learning in stabilizing training and enabling effective reasoning under long-context settings.

Figure: Figure 5: Training dynamics of our method. ReMemR1 enables the LLM to generate both inner memory and callback queries, introducing additional formatting requirements. These constraints initially lead to a lower success rate due to frequent parsing errors, but performance rapidly improves after around 20 steps as the model learns to follow the required format.
Refer to caption: x5.png

## Appendix C Implementation Details

### C.1 Full Expression of Training Objective

Our model is optimized with a variant of GRPO objective.
The full expression of our training objective can be written as:

$$ $\displaystyle\underset{\theta}{\mathrm{argmax}}~J_{\text{GRPO}}(\theta)=$ $\displaystyle\mathbb{E}_{(Q,Y),\{\tau^{(g)}\}_{g=1}^{G}\sim\pi_{\theta_{\text{old}}}}\Bigg[\frac{1}{G(T+1)}\sum_{g=1}^{G}\sum_{t=1}^{T+1}\frac{1}{|s_{t}^{(g)}|}\sum_{i=1}^{|s_{t}^{(g)}|}\min\Bigg(\rho_{t,i}^{(g)}\,\hat{A}_{t}^{(g)},$ (10) $\displaystyle\quad\ \text{clip}\!\left(\rho_{t,i}^{(g)},\,1-\epsilon,\,1+\epsilon\right)\hat{A}_{t}^{(g)}\Bigg)-\beta\,\mathbb{D}_{\text{KL}}\!\left[\pi_{\theta}\,\|\,\pi_{\text{ref}}\right]\Bigg],$ $$

where $\rho_{t,i}^{(g)}$ is the importance sampling ratio:

$$ $\rho_{t,i}^{(g)}=\frac{\pi_{\theta}\!\left(s_{t,i}^{(g)}\mid s_{t,<i}^{(g)},\,s_{<t}^{(g)},\,Q,\,c_{t-1}\right)}{\pi_{\theta_{\text{old}}}\!\left(s_{t,i}^{(g)}\mid s_{t,<i}^{(g)},\,s_{<t}^{(g)},\,Q,\,c_{t-1}\right)}.$ (11) $$

Here, $s_{t,i}^{(g)}$ denotes the $i$-th token in the $t$-th state of trajectory $g$, $\epsilon$ is the clipping ratio, $\beta$ is the KL coefficient, and $\hat{A}_{t}^{(g)}$ is the normalized advantage. We assume $c_{T}=\varnothing$ for notational convenience.

### C.2 Training Hyperparameters

The training of ReMemR1 was built upon the verl(^1^11[https://github.com/volcengine/verl](https://github.com/volcengine/verl)) framework, with efficient trajectory generation powered by the sglang(^2^22[https://github.com/sgl-project/sglang](https://github.com/sgl-project/sglang)) engine.
We employed Fully Sharded Data Parallelism (FSDP) for distributed training, and used bfloat16 precision for both training and evaluation.
Table [6](https://arxiv.org/html/2509.23040v1#A3.T6) summarizes the primary hyperparameters used in our method.

Although we evaluated the model with varying numbers of context documents during testing, the training setup consistently used 200 documents per sample,
resulting in approximately 30K input tokens. Each document chunk $c_{t}$ was limited to a maximum length of 5000 tokens, yielding $T\approx 6$ during training.
At each timestep and at the final state, the model generated rollouts with a temperature of 1, up to a maximum of 2048 tokens.

The 3B version of ReMemR1 and its variants are trained on 16 H800 GPUs and converge after 100 hours.
The 7B model is trained on 32 H800 GPUs, reaching convergence after 80 hours.

**Table 6: Primary hyperparameters used in training.**
| Hyper-parameter | Value |
| --- | --- |
| Training Batch Size | 128 |
| Micro Training Batch Size | 8 |
| Total Converge Steps | 200$\sim$300 |
| Actor Model Learning Rate | $1\times 10^{-6}$ |
| Actor Model Warmup Steps | 20 |
| Rollout Temperature | 1 |
| Max Chunk Length | 5000 |
| Training Chunk Number T | 6 |
| Max Response Length | 2048 |
| KL Coefficient $\beta$ | 0.001 |
| Clip Ratio $\epsilon$ | 0.2 |
| Group Size $G$ | 16 |

### C.3 Evaluation Settings

To ensure the challenging nature of the samples, we only use samples from the hard difficulty level for training.
Questions in these datasets typically require at least two pieces of evidence to answer, and there exist dependencies between the evidence.
Due to the extraordinary computational cost of long-context QA, we subsample 128 samples from each benchmarks with a random seed of 4, following Yu et al. (2025a).

## Appendix D Complexity Analysis

In this section, we analyze the computational complexity of ReMemR1 and show that it preserves the linear complexity of conventional memory-agent approaches.

### D.1 Baseline Complexity

In the “memorize while reading” paradigm, the agent processes a sequence of $T$ document chunks $\{c_{1},c_{2},\dots,c_{T}\}$ in order.
At each step $t$, it updates the memory via:

$$ $m_{t+1}=\pi(Q,c_{t},m_{t}).$ (12) $$

Each update requires $O(1)$ memory operations and a constant number of forward passes through the policy network.
Thus, the overall time complexity is $O(T)$.
The space requirement is a summation of the document chunks and the memory at each step, which is $O(T+1)=O(T)$ in total.

### D.2 Complexity of ReMemR1

ReMemR1 augments the state by including a query component $q_{t}$ and a retrieval function $\mathcal{E}$ over past memories:

$$ $s_{t+1}=(m_{t+1},q_{t+1})=\pi(Q,c_{t},m_{t},\mathcal{E}(\{m_{i}\}_{i\leqslant t},q_{t})).$ (13) $$

This paradigm also performs the same number of state transition, which is $O(n)$ times of LLM generation.
Compared to Eq. [12](https://arxiv.org/html/2509.23040v1#A4.E12), our method includes two sources of computational overhead:

- •
Storage of previous memories. Although the state transition references $\{m_{i}\}_{i\leqslant t}$, each $m_{i}$ is itself a fixed-length vector (e.g., the hidden state of the model).
Maintaining this list across $T$ steps requires $O(T)$ additional space.
This is the same order as storing the original text chunks, but with a smaller constant term.
- •
Retrieval operation The retrieval function $\mathcal{E}$ computes similarity between $q_{t}$ and past memory states.
If implemented with exact maximum similarity search over $\{m_{i}\}_{i\leqslant t}$, the cost per step could be $O(t)$.
However, in practice, we use lightweight recall-based heuristics or an index that supports sublinear approximate nearest neighbor search.
This operation is negligible compared against the consumption of the state transition model $\pi_{\theta}$, which is often a 3B or 7B level LLM.
Thus, the total cost across $T$ steps remains $O(T)$ in expectation.

Therefore, ReMemR1 preserves the same asymptotic $O(T)$ time and $O(T)$ space complexity as the conventional memory-agent paradigm, while substantially enhancing the agent’s ability to perform non-linear reasoning through retrieval.

## Appendix E Prompt Template

We use separate prompt templates for the generation of intermediate states $s_{1\leqslant t\leqslant T}$ and the final states $s_{T+1}$. The prompts are listed below:

## Appendix F The Use of Large Language Models

In the preparation of this manuscript, we utilized an LLM as a writing assistance.
The use of the LLM was limited to proofreading for grammatical errors, checking for typos, and improving the clarity and readability of existing text.
The LLM was not used for any core intellectual contributions, including but not limited to research ideation, formulation of the methodology, analysis of results, or drafting of the original manuscript.
All scientific claims, arguments, and the final text are the sole work of the human authors, who pay full responsibility for all content.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: Please see also Weaknesses above.

1) What is the tradeoff introduced by varying alpha? I understand that a smaller value of alpha corresponds to more emphasis on step-based rewards, but conceptually, since the evaluation target is the final accuracy, the optimal value of alpha should be alpha = 1, which emphasizes only the final reward. Thus, the paper should better justify the need for step-based rewards. Given enough training iterations, will alpha=1 eventually converge to a better (i.e., higher accuracy) solution?

2) Why does ReMemR1 use the relatively simple word similarity as a reward in the step-based rewards? Would a more sophisticated information gain measure (e.g., one that evaluates the semantics) work better?

3) The experiment section also includes an ablation study of sorts that compares ReMemR1 to that of MemAgent with a rule-based callback, showing that ReMemR1’s learned callbacks outperform a rule-based one. It is not clear, however, whether this is a true ablation study on ReMemR1; is MemAgent equivalent to ReMemR1 without revisitable memory? It’s also not clear why MemAgent + rule-based callback often performs worse than MemAgent itself in Table 3.

4) In Figure 5 of Appendix B, it’s not clear why the paper states that performance rapidly improves around epoch 20. There seems instead to be a smoothly decreasing increase in performance across the training duration.
soundness: 3
strengths: + The experiment results show a clear gain in performance relative to the baselines for both in-distribution and out-of-distribution test datasets, for most document context lengths. The accuracy of ReMemR1 is also fairly consistent as the number of context documents increases.

+ The proposed idea of using learned queries to revisit past context information, while simple to conceptualize, is intuitively likely to work. It also appears to use about the same amount of memory as the usual single-pass approaches (though I’d appreciate if the authors could confirm this explicitly).
summary: This paper addresses the challenge of question-answering over long contexts, in which large language models (LLMs) typically utilize a memory corpus to summarize a single pass over the context in order to answer the question. This approach, which uses only forward processing, is storage-efficient but often suboptimal, as the importance of some parts of the context may only be apparent after scanning the initial context. This paper proposes to mitigate this challenge by appending a query and retrieval function to the agent memory for each document chunk. This query can retrieve information from prior chunks, allowing them to be “revisited.” The paper further introduces chunk-level rewards that measure the information gain in memory updates to encourage retention of new information relevant to the ground truth, using reinforcement learning to learn the memories and retrievals for each document chunk. Experiments on multiple short and long context datasets show that the proposed ReMemR1 technique outperforms MemAgent as well as existing standard and long-context LLMs.
weaknesses: --It seems like the presence of learned queries in ReMemR1 would increase the LLM inference time, as the model would need to query past context while making an inference. How much more time would this take in practice (and if I am incorrect and the inference time is comparable to that of single-pass memory agents, can the authors explain why?)

-- The experiments pad the contexts of both HotPotQA and 2WikiMultiHopQA with random documents. This would seem to give an unrealistic picture of true long-context documents, which would very likely have different patterns from random. Indeed, it seems that the random padding would simply lead to the model learning not to query these random documents.

---

confidence: 4
contribution: 3
initial_rating: 4
presentation: 2
questions: 1. On the RL training level, is each state treated as one training instance in RLMLR?

2. What is the precise format of the memory? Can you provide more qualitative examples of the memories and the generated queries?
soundness: 3
strengths: 1. The proposed method allows memory agents to selectively retrieve information from its previous memories. This counters the potential information loss in standard memory agent methods and enables non-linear reasoning paths.

2. The propsoed RLMLR introduces a RL-learning framework with multi-level reward design. This provides more dense and step-level rewards for each states in the reasoning trajectory, which offers better guidance than RL methods that solely rely on the final outcome.

3. The authors perform extensive experiments to show that ReMemR1 can outperfor memory agents and lon-context LLMs in open-domain QA benchmarks. They also highlight the performance gains on out-of-distribution datasets.
summary: This paper introduces ReMemR1, an LLM training paradigm designed to overcome memory-related challenges in long-context QA. Existing long-context methods process documents in a single forward pass, while memory agents may result in irreversible information loss due to the linear processing of memories. In contrast, ReMemR1 addresses this by using callback-enhanced memory, in which the agent generates a callback query each step to selectively retrieve and revisit information from its previous memories and thereby enabling non-linear reasoning. To train this agent effectively, the authors propose RLMLR and combines sparse outcome-based rewards with dense, step-level rewards based on the "information gain" in memory updates and callback retrievals. Experiments show that the proposed ReMemR1 can outperforms existing memory-based approaches and long-context alternatives.
weaknesses: 1. Limited datasets / baselines. For open-domain QA benchmarks, the authors should also consider one-hop datasets like triviaQA, natural questions and further multi-hop questions like bamboogle, musique. Further long-context baselines like IterDRAG and search agents like Search-R1 should also be considered as baselines.

2. The method requires storing the entire history of all previous memory states over many retrieved document chunks. For each step, the LLM then perofrms a retrieval search across this history and generate new states / memory or final answer. This adds additional compute, memory overhead and latency compared to standard single-pass methods or search agents.

3. The authors adopt a very simple retrieval function based on word overlapping, other alternative methods like embedding-based retrieval is not discussed. Furthermore, the proposed method is designed to return one single past memory with the maximum word overlapping. This could be a limitation preventing the agent from synthesizing information across multiple documents / reasoning steps.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 2
questions: 1. How does the memory overhead scale with task complexity and document length compared to the "memorize while reading" paradigm?
2. What are the typical failure modes of the proposed retrieval mechanism?
soundness: 3
strengths: 1. The paper provides a clear motivation for addressing the limitations of "memorize while reading" paradigm. 
2. The experimental results demonstrate consistent improvements across multiple model sizes and varying context lengths. 
3. The proposed RLMLR approach is clearly described.
summary: The paper presents a novel approach for long-text QA comprised of two components. The first component is a memory-augmented agent that stores multiple historical memories while processing the long text rather than rewriting a single memory state. It allows the agent to revisit past memories via query-based retrieval, enabling non-linear reasoning paths.  The second component is a modification of GRPO with a detailed reward function that combines traditional outcome rewards and step-level state rewards to improve the relevance of the generated reasoning sequences. The proposed method is evaluated on two MHQA datasets and shows improvement over vanilla model and a MemAgent baseline.
weaknesses: 1. Poor figure presentation. Throughout the paper, figures are consistently misaligned with their textual references, which significantly disrupts the reading flow and makes it difficult to follow the arguments.
2. The memory efficiency of the proposed method compared to the "memorize while reading" baseline is not addressed. The proposed approach requires storing multiple historical memory states, and it is not discussed how such large memory storage affects inference time and cost.
3. The statement about per-step reward innovation should be clarified. Per-step rewards are a common practice in RL. 
4. The paper does not provide the callback mechanism failures analysis.

---

confidence: 4
contribution: 3
initial_rating: 4
presentation: 3
questions: 1. How does ReMemR1 perform when the memory corpus exceeds the retrieval limit or available storage capacity? Or how does other methods enlarge their memory storage?

2. Does the system employ any pruning or compression strategy to prevent unbounded memory growth?
soundness: 3
strengths: - The paper is clearly written, well-organized, and easy to follow, which facilitates understanding of the technical contributions.
    
- The proposed method achieves consistent gains on both in-distribution and out-of-distribution datasets, with notable advantages in challenging long-context scenarios.
    
- The idea of non-linear reasoning via callback-based memory is conceptually interesting and potentially useful for improving long-context reasoning.
summary: This paper presents ReMemR1, a memory-augmented large language model (LLM) agent for long-context question answering (QA). The key innovation is a callback-enhanced memory mechanism that enables the model to revisit and retrieve earlier memory states instead of relying solely on a forward-only “memorize while reading” paradigm. To address the sparse reward issue commonly faced in reinforcement learning (RL)-based memory agents, the authors propose Reinforcement Learning with Multi-Level Rewards (RLMLR), which combines trajectory-level (final-answer) and step-level (intermediate) rewards. Experiments on two datasets demonstrate consistent improvements over strong baselines such as MemAgent and Qwen2.5, particularly under long-context and distant-evidence settings.
weaknesses: - While the callback mechanism is intuitively appealing, it can be viewed as an incremental extension of existing memory-retrieval or recurrent-state models. The distinction from prior works (e.g., MemAgent) is not sufficiently emphasized.
    
- The paper provides limited formal or analytical insight into *why* the proposed retrieval and multi-level reward structure lead to improved reasoning or generalization.
    
- The authors argue that existing memory-based agents suffer from irreversible forward-only processing and information loss because they cannot revisit previous states. However, this argument is not entirely convincing. When the policy updates its memory at each step, it already considers both the previous memory and the current input (query or document chunk), which inherently allows recursive integration of prior information. It is unclear why an additional “memory recall” mechanism is strictly necessary, rather than improving the existing update function to retain or reweight key information.
    
- The proposed design retrieves relevant content from all previous memory states, which implies that the model must maintain access to the entire memory history. This could be computationally expensive and may not scale to extremely long contexts. The paper should clarify whether ReMemR1 explicitly stores all intermediate memories or uses a compressed or differentiable representation to manage memory growth.
    
- Important implementation aspects, such as retrieval efficiency, computational overhead, and inference latency, are under-discussed. It remains unclear how the proposed approach scales in both training and inference compared to other memory-based methods.

- Additional baselines, especially recent long-context or memory-augmented LLMs, would make the evaluation more comprehensive and convincing.


Decision: Accept (Poster)