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

arXiv: https://arxiv.org/abs/2602.18493 · HTML: https://mlsanity.com/paper/2602.18493/

## Key facts

- Cited by 8 papers in the map (data through 14 Oct 2026).
- Placed in Reasoning for agents, planning and code, one of 17 directions in LLM reasoning.
- Builds on 30 of its 54 references that are in the map.
- First posted on arXiv 13 Feb 2026.

## Placed in

- LLM reasoning / [Reasoning for agents, planning and code](https://mlsanity.com/lr/reasoning-for-agents-planning-and-code/)
- Retrieval-augmented generation / [Agentic RAG, agents and memory](https://mlsanity.com/rag/agentic-rag-agents-and-memory/)
- LLM agents and tool use / [Agent memory systems](https://mlsanity.com/agents/agent-memory-systems/)

## Builds on (most cited first, 30 of 30)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [SQuAD: 100,000+ Questions for Machine Comprehension of Text](https://mlsanity.com/paper/1606.05250/) (2016)
- [HotpotQA: A Dataset for Diverse, Explainable Multi-hop Question Answering](https://mlsanity.com/paper/1809.09600/) (2018)
- [Lost in the Middle: How Language Models Use Long Contexts](https://mlsanity.com/paper/2307.03172/) (2023)
- [HybridFlow: A Flexible and Efficient RLHF Framework](https://mlsanity.com/paper/2409.19256/) (2024)
- [MemGPT: Towards LLMs as Operating Systems](https://mlsanity.com/paper/2310.08560/) (2023)
- [A-MEM: Agentic Memory for LLM Agents](https://mlsanity.com/paper/2502.12110/) (2025)
- [SnapKV: LLM Knows What You are Looking for Before Generation](https://mlsanity.com/paper/2404.14469/) (2024)
- [Evaluating Very Long-Term Conversational Memory of LLM Agents](https://mlsanity.com/paper/2402.17753/) (2024)
- [Mem0: Building Production-Ready AI Agents with Scalable Long-Term Memory](https://mlsanity.com/paper/2504.19413/) (2025)
- [LongMemEval: Benchmarking Chat Assistants on Long-Term Interactive Memory](https://mlsanity.com/paper/2410.10813/) (2024)
- [A Survey on the Memory Mechanism of Large Language Model based Agents](https://mlsanity.com/paper/2404.13501/) (2024)
- [Zep: A Temporal Knowledge Graph Architecture for Agent Memory](https://mlsanity.com/paper/2501.13956/) (2025)
- [Memory in the Age of AI Agents](https://mlsanity.com/paper/2512.13564/) (2025)
- [Beyond Goldfish Memory: Long-Term Open-Domain Conversation](https://mlsanity.com/paper/2107.07567/) (2021)
- [Evaluating Memory in LLM Agents via Incremental Multi-Turn Interactions](https://mlsanity.com/paper/2507.05257/) (2025)
- [MemAgent: Reshaping Long-Context LLM with Multi-Conv RL-based Memory Agent](https://mlsanity.com/paper/2507.02259/) (2025)
- [MEM1: Learning to Synergize Memory and Reasoning for Efficient Long-Horizon Agents](https://mlsanity.com/paper/2506.15841/) (2025)
- [Memory-R1: Enhancing Large Language Model Agents to Manage and Utilize Memories via Reinforcement Learning](https://mlsanity.com/paper/2508.19828/) (2025)
- [MIRIX: Multi-Agent Memory System for LLM-Based Agents](https://mlsanity.com/paper/2507.07957/) (2025)
- [MemoRAG: Boosting Long Context Processing with Global Memory-Enhanced Retrieval Augmentation](https://mlsanity.com/paper/2409.05591/) (2024)
- [MemGen: Weaving Generative Latent Memory for Self-Evolving Agents](https://mlsanity.com/paper/2509.24704/) (2025)
- [SCM: Enhancing Large Language Model with Self-Controlled Memory Framework](https://mlsanity.com/paper/2304.13343/) (2023)
- [Memory Transformer](https://mlsanity.com/paper/2006.11527/) (2020)
- [Look Back to Reason Forward: Revisitable Memory for Long-Context LLM Agents](https://mlsanity.com/paper/2509.23040/) (2025)
- [PerLTQA: A Personal Long-Term Memory Dataset for Memory Classification, Retrieval, and Synthesis in Question Answering](https://mlsanity.com/paper/2402.16288/) (2024)
- [QwenLong-L1.5: Post-Training Recipe for Long-Context Reasoning and Memory Management](https://mlsanity.com/paper/2512.12967/) (2025)
- [Self-Updatable Large Language Models by Integrating Context into Model Parameters](https://mlsanity.com/paper/2410.00487/) (2024)
- [Mem-T: Densifying Rewards for Long-Horizon Memory Agents](https://mlsanity.com/paper/2601.23014/) (2026)
- [Convomem Benchmark: Why Your First 150 Conversations Don't Need RAG](https://mlsanity.com/paper/2511.10523/) (2025)

## Built on by (most cited first, 8 of 8)

- [TOKI: A Bitemporal Operator Algebra for Contradiction Resolution in LLM-Agent Persistent Memory](https://mlsanity.com/paper/2606.06240/) (2026)
- [PM-Bench: Evaluating Prospective Memory in LLM Agents](https://mlsanity.com/paper/2607.12385/) (2026)
- [Mem-π: Adaptive Memory through Learning When and What to Generate](https://mlsanity.com/paper/2605.21463/) (2026)
- [DimMem: Dimensional Structuring for Efficient Long-Term Agent Memory](https://mlsanity.com/paper/2605.15759/) (2026)
- [MMLA: Memory-Mediated Learning Architecture for Predictive Dual-State Adaptation](https://mlsanity.com/paper/2606.28876/) (2026)
- [D2ACCI: A Dual-Loop Diagnostic Protocol for Evidence-Preserving Agent Memory](https://mlsanity.com/paper/2608.17756/) (2026)
- [MoM: Memory of Memory](https://mlsanity.com/paper/2609.25054/) (2026)
- [LongHarness: Advancing Agentic Reasoning over 100M-Token Corpora via Stateful Evidence Management](https://mlsanity.com/paper/2610.13075/) (2026)

## Abstract (arXiv)

Long-context LLMs and Retrieval-Augmented Generation defer state tracking and evidence consolidation to query time, which is brittle when facts evolve and answers depend on latent states. We introduce Unified Memory Agent (UMA) for a one-to-many setting: query-agnostic external memory is constructed once from a stream and reused across multiple future QA sessions. A single policy maintains a structured Memory Bank through CRUD operations and answers using both the Memory Bank and raw context. Task-Stratified GRPO uses the mean reward of QA trajectories branching from each sampled memory state to supervise memory maintenance, while normalizing memory and per-question QA groups separately. We also introduce Ledger-QA, a diagnostic benchmark for long-horizon state tracking over accumulated updates. At the 16k budget, UMA-Generalist achieves the highest average score among compared methods across the test-time-learning and accurate-retrieval benchmarks and transfers to Ledger-QA without task-specific training; UMA-Specialist further improves long-horizon tracking after task adaptation. These results support learned proactive memory management for long-context reasoning.