让大模型主动记忆并复用信息，提升长期推理能力。
Learning to Remember: End-to-End Training of Memory Agents for Long-Context Reasoning
- 构建一次性外部记忆库，支持多次问答复用。
- 在16k上下文预算下，多项指标表现最优。
- 无需任务微调即可跨场景迁移，适合复杂推理任务。
长上下文大模型与检索增强生成将状态追踪和证据整合推迟到查询时刻，当事实变化时易出错且依赖潜在状态。我们提出统一记忆代理（UMA），适用于一对多场景：从数据流中一次性构建与查询无关的外部记忆库，并在多个未来问答会话中复用。单一策略通过增删改查操作维护结构化记忆库，结合记忆库与原始上下文回答问题。任务分层GRPO利用从每个采样记忆状态分支出的问答轨迹平均奖励，分别对记忆维护及每题问答组进行监督与归一化。我们还引入Ledger-QA，一个用于评估长期状态追踪能力的诊断基准。在16k上下文预算下，UMA-Generalist在测试时学习与精确检索基准上均取得最高平均分，且无需任务特训即能迁移至Ledger-QA；UMA-Specialist经任务适配后进一步提升长期追踪性能。结果表明，学习式主动记忆管理对长上下文推理具有显著价值。
原文摘要 · Abstract (English)
Long-context LLMs and Retrieval-Augmented Generation defer state tracking and evidence consolidation to query time, which is brittle when facts evolve and answers depend on latent states. We introduce Unified Memory Agent (UMA) for a one-to-many setting: query-agnostic external memory is constructed once from a stream and reused across multiple future QA sessions. A single policy maintains a structured Memory Bank through CRUD operations and answers using both the Memory Bank and raw context. Task-Stratified GRPO uses the mean reward of QA trajectories branching from each sampled memory state to supervise memory maintenance, while normalizing memory and per-question QA groups separately. We also introduce Ledger-QA, a diagnostic benchmark for long-horizon state tracking over accumulated updates. At the 16k budget, UMA-Generalist achieves the highest average score among compared methods across the test-time-learning and accurate-retrieval benchmarks and transfers to Ledger-QA without task-specific training; UMA-Specialist further improves long-horizon tracking after task adaptation. These results support learned proactive memory management for long-context reasoning.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。