@inproceedings{zhang-etal-2026-learning-remember,
title = "Learning to Remember: End-to-End Training of Memory Agents for Long-Context Reasoning",
author = "Zhang, Kehao and
Gui, Shangtong and
Yang, Sheng and
Wei, Chen and
Feng, Yang",
editor = "Lee, Hung-Yi and
Dev, Sunipa and
Li, Jessy and
Elliott, Desmond",
booktitle = "Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing",
month = oct,
year = "2026",
address = "Budapest, Hungary",
publisher = "Association for Computational Linguistics",
url = "https://aclanthology.org/2026.emnlp-main.1299/",
pages = "29575--29599",
ISBN = "979-8-89176-459-0",
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."
}<?xml version="1.0" encoding="UTF-8"?>
<modsCollection xmlns="http://www.loc.gov/mods/v3">
<mods ID="zhang-etal-2026-learning-remember">
<titleInfo>
<title>Learning to Remember: End-to-End Training of Memory Agents for Long-Context Reasoning</title>
</titleInfo>
<name type="personal">
<namePart type="given">Kehao</namePart>
<namePart type="family">Zhang</namePart>
<role>
<roleTerm authority="marcrelator" type="text">author</roleTerm>
</role>
</name>
<name type="personal">
<namePart type="given">Shangtong</namePart>
<namePart type="family">Gui</namePart>
<role>
<roleTerm authority="marcrelator" type="text">author</roleTerm>
</role>
</name>
<name type="personal">
<namePart type="given">Sheng</namePart>
<namePart type="family">Yang</namePart>
<role>
<roleTerm authority="marcrelator" type="text">author</roleTerm>
</role>
</name>
<name type="personal">
<namePart type="given">Chen</namePart>
<namePart type="family">Wei</namePart>
<role>
<roleTerm authority="marcrelator" type="text">author</roleTerm>
</role>
</name>
<name type="personal">
<namePart type="given">Yang</namePart>
<namePart type="family">Feng</namePart>
<role>
<roleTerm authority="marcrelator" type="text">author</roleTerm>
</role>
</name>
<originInfo>
<dateIssued>2026-10</dateIssued>
</originInfo>
<typeOfResource>text</typeOfResource>
<relatedItem type="host">
<titleInfo>
<title>Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing</title>
</titleInfo>
<name type="personal">
<namePart type="given">Hung-Yi</namePart>
<namePart type="family">Lee</namePart>
<role>
<roleTerm authority="marcrelator" type="text">editor</roleTerm>
</role>
</name>
<name type="personal">
<namePart type="given">Sunipa</namePart>
<namePart type="family">Dev</namePart>
<role>
<roleTerm authority="marcrelator" type="text">editor</roleTerm>
</role>
</name>
<name type="personal">
<namePart type="given">Jessy</namePart>
<namePart type="family">Li</namePart>
<role>
<roleTerm authority="marcrelator" type="text">editor</roleTerm>
</role>
</name>
<name type="personal">
<namePart type="given">Desmond</namePart>
<namePart type="family">Elliott</namePart>
<role>
<roleTerm authority="marcrelator" type="text">editor</roleTerm>
</role>
</name>
<originInfo>
<publisher>Association for Computational Linguistics</publisher>
<place>
<placeTerm type="text">Budapest, Hungary</placeTerm>
</place>
</originInfo>
<genre authority="marcgt">conference publication</genre>
<identifier type="isbn">979-8-89176-459-0</identifier>
</relatedItem>
<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.</abstract>
<identifier type="citekey">zhang-etal-2026-learning-remember</identifier>
<location>
<url>https://aclanthology.org/2026.emnlp-main.1299/</url>
</location>
<part>
<date>2026-10</date>
<extent unit="page">
<start>29575</start>
<end>29599</end>
</extent>
</part>
</mods>
</modsCollection>
%0 Conference Proceedings
%T Learning to Remember: End-to-End Training of Memory Agents for Long-Context Reasoning
%A Zhang, Kehao
%A Gui, Shangtong
%A Yang, Sheng
%A Wei, Chen
%A Feng, Yang
%Y Lee, Hung-Yi
%Y Dev, Sunipa
%Y Li, Jessy
%Y Elliott, Desmond
%S Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing
%D 2026
%8 October
%I Association for Computational Linguistics
%C Budapest, Hungary
%@ 979-8-89176-459-0
%F zhang-etal-2026-learning-remember
%X 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.
%U https://aclanthology.org/2026.emnlp-main.1299/
%P 29575-29599
Markdown (Informal)
[Learning to Remember: End-to-End Training of Memory Agents for Long-Context Reasoning](https://aclanthology.org/2026.emnlp-main.1299/) (Zhang et al., EMNLP 2026)
ACL