MemChain: Learning Interpretable Memory Traces for Memory-Augmented LLM Agents

2026-07-27

AI Summary

This work addresses the limitations of conventional memory-augmented LLM agents, which directly inject retrieved memories and often introduce redundant, conflicting, or weakly relevant information, leading to substantial context overhead in long-horizon tasks. To overcome this, the authors propose MemChain, a trainable post-retrieval memory strategy that transforms raw memories into compact, interpretable, answer-oriented active memories by generating question-conditioned evidence plans, constructing ordered semantic dependency chains, and performing explicit memory operations. MemChain introduces, for the first time, an interpretable memory trace structure and employs a two-stage training framework combining supervised trace learning with TMPO—a reinforcement learning objective based on answer quality—to optimize memory strategies. Evaluated on LoCoMo and LongMemEval-S benchmarks, MemChain achieves state-of-the-art performance on both frozen closed-source and open-source answer models while significantly compressing the memory length passed to the answer model.

Abstract

Memory-augmented LLM agents typically answer queries by retrieving relevant memories and feeding them directly to an answer model. This retrieval-as-evidence paradigm assumes retrieved memories are already suitable for reasoning, leaving the answer model to resolve redundancy, conflicts, and weak relevance while incurring substantial context overhead in long-term memory tasks. We propose MemChain, a trainable post-retrieval memory policy that transforms retrieved candidates into answer-facing active memory, represented as a compact and grounded evidence context. Given a user query and retrieved candidates, MemChain first generates a question-conditioned evidence plan, then constructs an ordered grounded evidence trace that organizes retrieved memories according to their semantic roles and dependencies, and finally executes explicit memory actions to produce a concise evidence context for answer generation. To train the mediator, we introduce a two-stage learning framework. Supervised trace learning first teaches the policy to generate structurally valid plans, traces, actions, and evidence contexts. We then propose Trace-Guided Memory Policy Optimization (TMPO), a reinforcement learning objective that optimizes the memory policy using downstream answer quality while jointly encouraging trace grounding, evidence support, structural validity, and answer stability across multiple rollouts. Experiments on LoCoMo and LongMemEval-S demonstrate that MemChain consistently achieves state-of-the-art performance across both closed-source and open-weight frozen answer models while substantially reducing the memory context passed to the answer model.

Problem

• memory-augmented LLM agents
• retrieval-as-evidence
• context overhead
• memory redundancy
• evidence relevance

Innovation

• MemChain
• memory-augmented LLM agents
• evidence trace
• trace-guided policy optimization
• compact memory context