MemChain: 해석 가능한 메모리 트레이스를 학습하여 메모리 증강 LLM 에이전트 개발

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

메모리 증강 LLM 에이전트는 일반적으로 관련 정보를 검색하고 이를 직접 답변 모델에 제공하여 질문에 답합니다. 이러한 '검색을 증거로 활용'하는 방식은 검색된 정보가 이미 추론에 적합하다는 전제를 바탕으로 하며, 답변 모델은 중복, 충돌, 그리고 낮은 관련성을 해결해야 합니다. 또한, 장기 메모리 작업에서 상당한 컨텍스트 오버헤드가 발생합니다. 본 논문에서는 MemChain을 제안합니다. MemChain은 검색 후 학습 가능한 메모리 정책으로, 검색된 후보들을 답변에 적합한 활성 메모리로 변환하며, 이는 간결하고 명확한 증거 맥락으로 표현됩니다. 사용자 쿼리와 검색된 후보가 주어지면, MemChain은 먼저 질문에 기반한 증거 계획을 생성하고, 검색된 정보를 의미론적 역할과 의존성에 따라 정렬하여 순서화된 증거 트레이스를 구성하며, 마지막으로 명시적인 메모리 액션을 실행하여 답변 생성을 위한 간결한 증거 맥락을 생성합니다. 이러한 과정을 학습하기 위해, 우리는 두 단계로 구성된 학습 프레임워크를 도입했습니다. 지도 학습을 통해 정책은 구조적으로 유효한 계획, 트레이스, 액션 및 증거 컨텍스트를 생성하도록 훈련됩니다. 또한, Trace-Guided Memory Policy Optimization (TMPO)이라는 강화 학습 목표를 제안합니다. TMPO는 다운스트림 답변 품질을 최적화하고 동시에 트레이스의 정확성, 증거 지원, 구조적 유효성 및 여러 반복 과정에서의 답변 안정성을 유지하도록 메모리 정책을 조정합니다. LoCoMo와 LongMemEval-S 데이터셋에 대한 실험 결과, MemChain은 폐쇄형 및 개방형 가중치 답변 모델 모두에서 최첨단 성능을 지속적으로 달성하며, 동시에 답변 모델에 전달되는 메모리 컨텍스트를 크게 줄입니다.

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.