Abstract
While reasoning over long context is crucial for various real-world applications, it remains challenging for large language models (LLMs) as they suffer from performance degradation as the context length grows. Recent work MemAgent has tried to tackle this by processing context chunk-by-chunk in an RNN-like loop and updating a textual memory for final answering. However, this naive recurrent memory update faces two crucial drawbacks: (i) memory can quickly explode because it can update indiscriminately, even on evidence-free chunks; and (ii) the loop lacks an exit mechanism, leading to unnecessary computation after even sufficient evidence is collected. To address these issues, we propose GRU-Mem, which incorporates two text-controlled gates for more stable and efficient long-context reasoning. Specifically, in GRU-Mem, the memory only updates when the update gate is open and the recurrent loop will exit immediately once the exit gate is open. To endow the model with such capabilities, we introduce two reward signals $r^{\text{update}}$ and $r^{\text{exit}}$ within end-to-end RL, rewarding the correct updating and exiting behaviors respectively. Experiments on various long-context reasoning tasks demonstrate the effectiveness and efficiency of GRU-Mem, which generally outperforms the vanilla MemAgent with up to 400\% times inference speed acceleration.
Abstract (translated)
虽然在处理长文本上下文时进行推理对于许多现实世界的应用至关重要，但对于大型语言模型（LLMs）来说却是一个挑战。随着上下文长度的增加，这些模型的表现会逐渐下降。最近的工作MemAgent试图通过采用类似RNN的方式逐段处理上下文，并更新文本内存来最终提供答案，以此解决这一问题。然而，这种简单的递归记忆更新方式存在两个关键缺陷：（i）由于它可能在缺乏依据的情况下更新信息，因此记忆可能会迅速膨胀；（ii）循环中缺少退出机制，在收集到足够证据后仍进行不必要的计算。 为了解决这些问题，我们提出了GRU-Mem模型，该模型通过引入两种由文本控制的门来实现更稳定和高效的长上下文推理。具体来说，在GRU-Mem中，内存仅在更新门开启时才会被更新，并且一旦退出门开启，递归循环将立即结束。 为了赋予模型这些能力，我们提出了两个奖励信号$r^{\text{update}}$和$r^{\text{exit}}$，在端到端的强化学习框架内给予正确的更新行为和退出行为相应的奖励。实验表明，在各种长上下文推理任务上，GRU-Mem的有效性和效率均优于原始的MemAgent模型，并且推理速度提高了高达400%。
URL
https://arxiv.org/abs/2602.10560
