跳转至

学会记忆：面向长文本推理的记忆智能体端到端训练

文章背景与核心概要

长文本大语言模型（LLM）和检索增强生成（RAG）在处理随时间演变的事实或依赖潜在状态的答案时，往往在状态追踪和证据整合方面表现不佳。为了克服这一局限性，作者推出了统一记忆智能体（Unified Memory Agent, UMA），专为“一对多”场景设计：从数据流中构建一次与查询无关的外部记忆，并在未来的多个问答（QA）会话中重复使用。

该研究的主要贡献包括：1）任务分层群组相对策略优化（Task-Stratified GRPO）：一种监督方法，利用从每个采样记忆状态分叉出来的QA轨迹的平均奖励来指导记忆维护，并分别对记忆和每个问题的QA组进行归一化；2）Ledger-QA：一个新的诊断基准，专为经过累积更新的长周期状态追踪而设计。实验表明，在16k上下文预算下，UMA-Generalist在测试时学习和检索基准测试中取得了最高的平均分（零样本迁移至Ledger-QA），而UMA-Specialist则在任务适配后进一步优化了长周期追踪能力。

摘要 (Abstract)

长文本大语言模型（LLMs）和检索增强生成（RAG）将状态追踪和证据整合推迟到查询阶段进行，当事实发生演变且答案依赖于潜在状态时，这种做法往往很脆弱。我们引入了统一记忆智能体（UMA），用于一对多设置：从流中构建一次与查询无关的外部记忆，并在多个未来的问答（QA）会话中重复使用。单一策略通过CRUD（增删改查）操作维护结构化的记忆库，并同时使用记忆库和原始上下文进行回答。任务分层GRPO使用从每个采样的记忆状态分叉出的QA轨迹的平均奖励来监督记忆维护，同时分别规范记忆和每个问题的QA组。我们还引入了Ledger-QA，这是一个用于跨累积更新的长视界状态追踪的诊断基准。在16k预算下，UMA-Generalist在测试时学习和精确检索基准中实现了比较方法中的最高平均分，并在无需特定任务训练的情况下迁移到Ledger-QA；UMA-Specialist在任务自适应后进一步改善了长视界追踪。这些结果支持了用于长文本推理的学得式主动记忆管理。

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.

访问与资源 (Access & Resources)