Learning to Remember: End-to-End Training of Memory Agents for Long-Context Reasoning

2026-02-13

arXiv.org

AI Summary

为解决长上下文推理中状态跟踪不稳定的问题，本文提出统一记忆代理（UMA）通过一次构建查询无关的外部记忆并重复使用，结合任务分层GRPO方法进行训练。

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.

Problem

• Long-context Reasoning
• State Tracking
• Evidence Consolidation

Innovation

• Unified Memory Agent
• CRUD Operations
• Task-Stratified GRPO
• Ledger-QA