[논문 리뷰] Learning to Remember: End-to-End Training of Memory Agents for Long-Context Reasoning

Kehao Zhang, Shangtong Gui외 3명|arXiv (Cornell University)|2026. 02. 13.

Topic Modeling인용 수 0

한 줄 요약

본 논문은 Unified Memory Agent (UMA)를 제시한다. UMA는 길-context 과제에 대한 메모리 관리와 추론을 함께 학습하는 엔드-투-엔드 강화학습 프레임워크이며, 동적 상태 추적 벤치마크로서 Ledger-QA를 제시한다.

ABSTRACT

Long-context LLMs and Retrieval-Augmented Generation (RAG) systems process information passively, deferring state tracking, contradiction resolution, and evidence aggregation to query time, which becomes brittle under ultra long streams with frequent updates. We propose the Unified Memory Agent (UMA), an end-to-end reinforcement learning framework that unifies memory operations and question answering within a single policy. UMA maintains a dual memory representation: a compact core summary for global context and a structured Memory Bank that supports explicit CRUD (create, update, delete, reorganize) over key value entries, enabling proactive consolidation during streaming. To evaluate long-horizon memory behavior, we introduce Ledger-QA, a diagnostic benchmark for continuous state tracking where answers are latent values derived from accumulated updates rather than lo cal span retrieval. Across 13 datasets spanning Ledger-QA, Test-Time Learning, and Accurate Retrieval, UMA substantially outperforms long-context and RAG baselines on dynamic reasoning and learning tasks while remaining competitive on standard retrieval benchmarks, underscoring the importance of learned, end-to-end memory management.

연구 동기 및 목표

• 수동 검색을 넘어서 초장문 맥락에서 능동적이고 학습된 메모리 관리의 필요성을 동기부여한다.
• 단일 정책에서 메모리 연산(CRUD)과 질의 응답을 통합하는 UMA를 제안한다.
• 장기간에 걸친 연속 상태 추적을 위한 진단 벤치마크로 Ledger-QA를 소개한다.
• 엔드-투-엔드 메모리 최적화가 우수한 동적 추론과 경쟁력 있는 검색 성능을 산출함을 입증한다.

제안 방법

• 긴 컨텍스트 추론을 코어 요약과 구조화된 Memory Bank(키-값 항목에 대한 CRUD)를 포함하는 이중 메모리 상태를 가진 MDP로 형식화한다.
• 두 단계 아키텍처를 사용한다: Phase I은 청크에 걸친 순차적 메모리 유지, Phase II는 원문 텍스트와 구조화된 메모리에서의 검색을 결합한 하이브리드 QA.
• 메모리와 QA 이점을 추정하기 위해 중첩된 궤적 샘플링을 활용하는 Task-Stratified Group Relative Policy Optimization(GRPO)으로 학습한다.
• 도구 사용 성공과 최종 정답 정확성을 결합한 2단계 보상 설계를 사용하고, 메모리 및 QA 단계에 대해 계층화된 정규화를 통해 적절히 공로를 부여한다.
• Ledger-QA와 함께 13개 데이터셋에 걸친 표준 TTL/AR 벤치마크로 평가한다.

실험 결과

연구 질문

• RQ1 메모리 연산의 엔드-투-엔드 최적화가 초장문 맥락에서 긴 호라이즌 추론을 개선할 수 있는가?
• RQ2 통합 메모리+QA 정책이 동적 상태 추적 과제에서 검색 중심 벤치마크보다 우수한가?
• RQ3 메모리 유지 단계와 RL 학습이 전체 성능에 어떤 기여를 하는가?
• RQ4 Ledger-QA가 실제 상태 추적 능력을 로컬 범위 검색과 비교해 얼마나 잘 탐지하는가?

주요 결과

• UMA는 13개 데이터셋에 걸친 동적 추론 과제에서 긴 컨텍스트 및 RAG 기준선보다 상당히 우수하다.
• UMA는 표준 검색 벤치마크에서도 경쟁력을 유지하며, 학습된 메모리 관리가 일반화에 이점을 가져옴을 보인다.
• 타당성 분석은 메모리 유지와 RL 학습이 모두 최상의 성능에 결정적임을 보여준다.
• Task-Stratified GRPO는 이질적인 메모리와 QA 목표에 대한 효과적인 공로 할당을 제공한다.
• Ledger-QA의 도전은 지평이 커질수록 기준선의 취약성을 드러내는 반면, UMA는 견고한 정확도를 유지한다.