提出可控记忆更新与退出机制,提升长文本推理效率与准确率
When to Memorize and When to Stop: Gated Recurrent Memory for Long-Context Reasoning
- 引入双门控机制,仅在必要时更新记忆、及时退出循环
- 在多个任务上实现最高400%的推理速度提升,性能优于基线模型
- 适合需要高效处理长上下文的应用场景,如法律分析、科学推理
尽管长上下文推理对现实应用至关重要,但大语言模型(LLMs)在上下文变长时仍面临性能下降的问题。现有方法MemAgent通过类似RNN的循环处理上下文块,并更新文本记忆以生成答案,但存在两大缺陷:(i) 记忆可能因无差别更新而迅速膨胀;(ii) 循环缺乏退出机制,导致在已收集足够证据后仍进行无效计算。为此,我们提出GRU-Mem,引入两个受文本控制的门控机制,确保记忆仅在更新门开启时更新,且在退出门开启时立即终止循环。为训练该机制,我们在端到端强化学习中引入两个奖励信号 $r^{ ext{update}}$ 与 $r^{ ext{exit}}$,分别奖励正确更新与退出行为。实验表明,GRU-Mem在多个长上下文推理任务中均显著优于原始MemAgent,推理速度最高提升400%。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。