When to Memorize and When to Stop: Gated Recurrent Memory for Long-Context Reasoning

📄 arXiv: 2602.10560v1 📥 PDF

作者: Leheng Sheng, Yongtao Zhang, Wenchang Ma, Yaorui Shi, Ting Huang, Xiang Wang, An Zhang, Ke Shen, Tat-Seng Chua

分类: cs.CL, cs.AI

发布日期: 2026-02-11

备注: 26 pages

💡 一句话要点

提出GRU-Mem，通过门控循环记忆网络解决长文本推理中效率和稳定性问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 长文本推理 门控循环单元 强化学习 记忆网络 效率优化

📋 核心要点

  1. 大型语言模型在长文本推理中面临性能下降，原因是无法有效区分关键信息和冗余信息。
  2. GRU-Mem通过引入更新门和退出门，控制记忆的更新和循环的退出，从而提高效率和稳定性。
  3. 实验表明，GRU-Mem在长文本推理任务中优于MemAgent，推理速度提升显著。

📝 摘要（中文）

长文本推理在现实应用中至关重要，但大型语言模型(LLMs)的性能会随着上下文长度的增加而下降。MemAgent尝试通过RNN式循环逐块处理上下文并更新文本记忆来解决这个问题。然而，这种朴素的循环记忆更新存在两个关键缺陷：(i)记忆可能迅速爆炸，因为它会无差别地更新，即使在没有证据的块上也是如此；(ii)循环缺乏退出机制，导致在收集到充分证据后仍进行不必要的计算。为了解决这些问题，我们提出了GRU-Mem，它结合了两个文本控制的门，以实现更稳定和高效的长文本推理。具体来说，在GRU-Mem中，记忆仅在更新门打开时更新，并且循环一旦打开退出门就会立即退出。为了赋予模型这种能力，我们在端到端强化学习中引入了两个奖励信号$r^{ ext{update}}$和$r^{ ext{exit}}$，分别奖励正确的更新和退出行为。在各种长文本推理任务上的实验表明了GRU-Mem的有效性和效率，它通常优于原始MemAgent，推理速度提高了高达400%。

🔬 方法详解

问题定义：现有的大型语言模型在处理长文本推理任务时，随着上下文长度的增加，性能会显著下降。MemAgent虽然尝试使用循环记忆机制来处理长文本，但其无差别更新记忆的方式容易导致记忆爆炸，且缺乏退出机制导致计算冗余。

核心思路：GRU-Mem的核心思路是通过引入门控机制来控制记忆的更新和循环的退出。具体来说，模型学习何时应该更新记忆（更新门）以及何时应该停止循环（退出门），从而避免不必要的计算和记忆爆炸。这种设计使得模型能够更有效地利用长文本中的信息。

技术框架：GRU-Mem的整体框架类似于MemAgent，采用RNN式的循环结构逐块处理上下文。不同之处在于，GRU-Mem在每个循环步骤中引入了两个门：更新门和退出门。更新门决定是否将当前块的信息更新到记忆中，而退出门决定是否结束循环。模型的输入是上下文块和当前的记忆状态，输出是更新门和退出门的概率以及更新后的记忆状态。

关键创新：GRU-Mem的关键创新在于引入了文本控制的门控机制，通过强化学习来学习更新门和退出门的策略。与传统的RNN或LSTM不同，GRU-Mem的门控信号不是基于隐藏状态，而是基于文本内容，这使得模型能够更好地理解上下文并做出更明智的决策。

关键设计：GRU-Mem使用强化学习来训练更新门和退出门。具体来说，模型的目标是最大化奖励函数，该奖励函数包括两个部分：更新奖励和退出奖励。更新奖励鼓励模型在需要更新记忆时打开更新门，而退出奖励鼓励模型在收集到足够信息时打开退出门。奖励信号的设计至关重要，需要仔细调整以平衡更新和退出的需求。模型使用策略梯度方法来优化策略网络，该策略网络输出更新门和退出门的概率。

📊 实验亮点

实验结果表明，GRU-Mem在多个长文本推理任务上优于基线模型MemAgent。在推理速度方面，GRU-Mem实现了高达400%的加速。这表明GRU-Mem能够更有效地利用长文本中的信息，并在保证性能的同时显著降低计算成本。这些结果验证了GRU-Mem的有效性和效率。

🎯 应用场景

GRU-Mem适用于需要处理长文本上下文的各种应用场景，例如文档摘要、问答系统、信息检索和对话系统。通过提高长文本推理的效率和稳定性，GRU-Mem可以帮助这些应用更好地理解和利用长文本中的信息，从而提升用户体验和系统性能。该研究对于提升LLM在实际应用中的能力具有重要价值。

📄 摘要（原文）

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.