[论文解读] Learning to Remember: End-to-End Training of Memory Agents for Long-Context Reasoning

Kehao Zhang, Shangtong Gui+3 more|arXiv (Cornell University)|Feb 13, 2026

Topic Modeling被引用 0

一句话总结

论文提出统一内存代理（UMA），一个端到端强化学习框架，联合学习内存管理与推理以处理长上下文任务，并提出 Ledger-QA 作为动态状态跟踪基准。

ABSTRACT

Long-context LLMs and Retrieval-Augmented Generation (RAG) systems process information passively, deferring state tracking, contradiction resolution, and evidence aggregation to query time, which becomes brittle under ultra long streams with frequent updates. We propose the Unified Memory Agent (UMA), an end-to-end reinforcement learning framework that unifies memory operations and question answering within a single policy. UMA maintains a dual memory representation: a compact core summary for global context and a structured Memory Bank that supports explicit CRUD (create, update, delete, reorganize) over key value entries, enabling proactive consolidation during streaming. To evaluate long-horizon memory behavior, we introduce Ledger-QA, a diagnostic benchmark for continuous state tracking where answers are latent values derived from accumulated updates rather than lo cal span retrieval. Across 13 datasets spanning Ledger-QA, Test-Time Learning, and Accurate Retrieval, UMA substantially outperforms long-context and RAG baselines on dynamic reasoning and learning tasks while remaining competitive on standard retrieval benchmarks, underscoring the importance of learned, end-to-end memory management.

研究动机与目标

• 在超长上下文中超越被动检索，强调对主动、学习型内存管理的需求。
• 提出将内存操作（CRUD）与问答统一在单一策略中的 UMA。
• 引入 Ledger-QA 作为对长时程连续状态跟踪的诊断性基准。
• 证明端到端的内存优化能够带来更出色的动态推理和具有竞争力的检索性能。

提出的方法

• 将长上下文推理公式化为一个二重内存状态的MDP：核心摘要与结构化内存库（对键-值条目的CRUD）。
• 采用两阶段架构：阶段I对分块进行序列化的记忆维护，阶段II进行混合问答，结合原始文本检索与结构化内存。
• 通过任务分层的组相对策略优化（GRPO）进行训练，利用嵌套轨迹采样来估计记忆与问答的优势。
• 采用两阶段奖励设计，结合工具使用成功与最终答案正确性，并通过分层归一化在记忆与问答步骤之间进行适当信誉分配。
• 在 Ledger 风格的动态状态跟踪（Ledger-QA）和标准的 TTL/AR 基准上对13个数据集进行评估。

实验结果

研究问题

• RQ1 端到端优化内存操作是否能在超长情境下提升长远推理？
• RQ2 统一的内存+问答策略是否在动态状态跟踪任务上优于以检索为中心的基线？
• RQ3 记忆维护阶段与强化学习训练对整体性能的贡献是什么？
• RQ4 Ledger-QA 如何检验真实状态跟踪能力与局部跨度检索之间的差异？

主要发现

• UMA 在13个数据集的动态推理任务中显著超越了长上下文与基于RAG的基线。
• UMA 在标准检索基准上仍具竞争力，表明学习型内存管理对泛化有益。
• 消融研究显示记忆维护与强化学习训练对达到峰值性能都至关重要。
• 任务分层的GRPO 为异质性内存和问答目标提供了有效的信誉分配。
• Ledger-QA 的挑战性揭示了基线在 horizon 增长时的脆弱性，而 UMA 能维持稳健的准确性。