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

Kehao Zhang, Shangtong Gui, Sheng Yang, Chen Wei, Yang Feng


Abstract
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.
Anthology ID:
2026.emnlp-main.1299
Volume:
Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing
Month:
October
Year:
2026
Address:
Budapest, Hungary
Editors:
Hung-Yi Lee, Sunipa Dev, Jessy Li, Desmond Elliott
Venue:
EMNLP
SIG:
SIGDAT
Publisher:
Association for Computational Linguistics
Note:
Pages:
29575–29599
Language:
URL:
https://aclanthology.org/2026.emnlp-main.1299/
DOI:
Bibkey:
Cite (ACL):
Kehao Zhang, Shangtong Gui, Sheng Yang, Chen Wei, and Yang Feng. 2026. Learning to Remember: End-to-End Training of Memory Agents for Long-Context Reasoning. In Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing, pages 29575–29599, Budapest, Hungary. Association for Computational Linguistics.
Cite (Informal):
Learning to Remember: End-to-End Training of Memory Agents for Long-Context Reasoning (Zhang et al., EMNLP 2026)
Copy Citation:
PDF:
https://aclanthology.org/2026.emnlp-main.1299.pdf