用答案条件信息增益提升长文本记忆代理的准确性
InfoMem: Training Long-Context Memory Agents with Answer-Conditioned Information Gain

- 基于答案条件信息增益设计新型奖励机制
- 在相同训练条件下性能超越现有强化学习基线
- 适合研究长上下文理解与记忆增强的模型
长上下文任务要求大模型从大量文本中识别并保留与答案相关的信息。分块记忆代理通过依次读取文档块、更新紧凑记忆,并从累积记忆生成最终答案来应对这一挑战。然而,现有基于强化学习的分块代理要么依赖稀疏的最终答案奖励,要么使用词法中间奖励来监督记忆和检索动作。这些信号仅指导任务成功或局部重叠,无法直接评估最终记忆是否支持真实答案。我们提出InfoMem,一种用于训练分块记忆代理的奖励机制,通过答案条件信息增益评估最终记忆的效用。InfoMem测量最终记忆使模型对真实答案的逐标记对数似然提高的程度。为稳定强化学习优化,InfoMem仅在成功轨迹上应用该信号,并在奖励组合前进行归一化。在相同的GRPO框架和训练预算下,InfoMem显著优于同类强化学习基线。分析表明,有效的最终记忆奖励应作用于成功轨迹,归一化后再组合,且需以答案而非查询为条件。代码已公开于https://github.com/GenSouKa1/InfoMem。
原文摘要 · Abstract (English)
Long-context tasks require LLMs to identify and preserve answer-relevant information from large contexts. Chunk-wise memory agents address this issue by sequentially reading document chunks, updating a compact memory, and generating the final answer from the accumulated memory. However, existing RL-based chunk-wise agents either rely on sparse final-answer rewards or use lexical intermediate rewards for memory and retrieval actions. These signals supervise task success or local overlap, but do not directly evaluate whether the final memory supports the ground-truth answer. We propose InfoMem, a reward mechanism for training chunk-wise memory agents that evaluates final-memory utility using answer-conditioned information. InfoMem measures how much the final memory increases the model's per-token log-likelihood of the ground-truth answer. To stabilize RL optimization, InfoMem applies this signal only to successful trajectories and normalizes it before reward composition. Under the same GRPO framework and training budget, InfoMem improves long-context memory-agent performance over comparable memory-agent RL baselines. Analyses show that effective final-memory rewards should operate on successful trajectories, be normalized before reward composition, and be conditioned on the answer rather than the query. Our code is available at https://github.com/GenSouKa1/InfoMem.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。