何时记忆何时停止：用于长上下文推理的门控循环记忆

尽管在长上下文中进行推理对许多现实应用至关重要，但随着上下文长度的增长，大型语言模型（LLMs）的性能会下降。近期的工作MemAgent尝试通过类似RNN的循环逐块处理上下文，并更新文本记忆以最终回答问题。然而，这种简单的递归记忆更新面临两个关键缺点：（i）记忆可能迅速膨胀，因为其可以无差别地更新，甚至在没有证据的块上；（ii）循环缺乏退出机制，导致即使收集到足够证据后仍进行不必要的计算。为了解决这些问题，我们提出了GRU-Mem，该方法引入了两个由文本控制的门控机制，以实现更稳定和高效的长上下文推理。具体而言，在GRU-Mem中，只有当更新门开启时才会更新记忆，而一旦退出门开启，循环将立即终止。为了赋予模型这些能力，我们在端到端强化学习中引入了两个奖励信号 $r^{ ext{update}}$ 和 $r^{ ext{exit}}$，分别奖励正确的更新和退出行为。在各种长上下文推理任务上的实验表明，GRU-Mem在效果和效率上均优于原始的MemAgent，推理速度加快高达400%。