Kehao Zhang
Author directory2026
Learning to Remember: End-to-End Training of Memory Agents for Long-Context Reasoning
Kehao Zhang | Shangtong Gui | Sheng Yang | Chen Wei | Yang Feng
Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing
Kehao Zhang | Shangtong Gui | Sheng Yang | Chen Wei | Yang Feng
Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing
Long-context LLMs and Retrieval-Augmented Generation (RAG) systems expand accessible context, but still defer state tracking and evidence consolidation to query time. This passive paradigm is brittle in streaming settings where facts are updated and answers depend on latent states rather than explicit spans. We introduce Unified Memory Agent (UMA), an end-to-end reinforcement-learning framework that trains a single policy to maintain external memory and answer questions. UMA updates a structured memory bank via CRUD operations and combines it with raw-context retrieval for QA. To jointly optimize memory maintenance and question answering, we propose Task-Stratified GRPO, which normalizes heterogeneous trajectories separately for more effective credit assignment. We also introduce Ledger-QA, a controlled benchmark for long-horizon state tracking over accumulated updates. Experiments on Ledger-QA and 13 TTL/AR benchmarks show that UMA improves long-horizon tracking, achieves the best average TTL/AR score among compared methods, and remains competitive on retrieval-centric benchmarks. Our results highlight learned proactive memory management as a complement to longer context windows and query-time retrieval.