InfoMem：以答案条件化信息增益训练长上下文记忆智能体

摘要

长上下文任务要求大语言模型从大上下文中识别并保留与答案相关的信息。块式内存代理通过顺序读取文档块、更新紧凑内存并从累积内存生成最终答案来解决此问题。然而，现有的基于强化学习的分块智能体要么依赖于稀疏的最终答案奖励，要么使用词汇中间奖励来进行记忆和检索操作。这些信号监督任务成功或局部重叠，但不直接评估最终记忆是否支持真实答案。我们提出了 InfoMem，一种用于训练分块记忆代理的奖励机制，它使用答案条件信息来评估最终记忆效用。 InfoMem 测量最终记忆使模型的每个标记对真实答案的对数似然度增加了多少。为了稳定 RL 优化，InfoMem 仅将此信号应用于成功的轨迹，并在奖励组合之前将其标准化。在相同的 GRPO 框架和训练预算下，InfoMem 比可比较的记忆代理 RL 基线提高了长上下文记忆代理的性能。分析表明，有效的最终记忆奖励应该在成功的轨迹上运行，在奖励构成之前进行标准化，并以答案而不是查询为条件。我们的代码可在 https://github.com/GenSouKa1/InfoMem 获取。

InfoMem：以答案条件化信息增益训练长上下文记忆智能体：论文配图 — 图 1：用于分块长上下文 RL 的 InfoMem 概述。 InfoMem 通过比较有或没有最终记忆的真实答案 y*y^{*} 的教师强制每个标记的平均对数似然，使用答案条件信息增益来测量最终记忆效用。在 GRPO 训练期间，信息增益监督仅应用于成功的轨迹，在成功的推广中标准化，并与策略优化的稀疏结果奖励相结合。