MemTrain: 자기지도 학습 기반의 컨텍스트 메모리 학습

메모리(Memory)는 장기적인 목표를 수행하는 LLM 에이전트(LLM agents)에게 필수적인 능력으로, 에이전트가 확장된 상호작용 전반에 걸쳐 축적된 정보를 보존하고 활용할 수 있게 합니다. 기존의 메모리-에이전트 접근 방식은 일반적으로 다운스트림 태스크(downstream tasks)에 대해 강화학습 (Reinforcement Learning)을 사용하여 엔드투엔드 (end-to-end) 방식으로 학습됩니다. 그러나 메모리 집약적인 시나리오를 위해 고품질의 주석이 달린 문제(annotated problems)를 수집하는 것은 비용이 많이 들며, 결과적으로 생성된 학습 데이터는 일반적인 메모리 동작을 포괄할 만큼 충분한 다양성을 갖추지 못하는 경우가 많습니다.

본 연구에서는 보다 효과적인 다운스트림 사후 학습 (post-training)을 위해 LLM 에이전트의 컨텍스트 메모리 (context-memory) 능력을 전반적으로 향상시키는 자기지도 학습 (self-supervised) 프레임워크인 MemTrain을 제안합니다. MemTrain은 레이블이 없는 Wikipedia 코퍼스(corpora)를 대상으로 두 가지 결합된 프록시 태스크 (proxy tasks)를 도입합니다: (1) 엔드투엔드 마스크 재구성 (masked reconstruction) 목적 함수로, 모델이 여러 차례의 메모리 업데이트 이후에 마스킹된 엔티티 (masked entities)를 복구하도록 요구하며, 이를 통해 최종 결과의 관점에서 메모리 유지 (memory maintenance)를 장려합니다. (2) 중간 메모리 회상 (intermediate memory recall) 목적 함수로, 모델이 중간 메모리 상태를 사용하여 마스킹된 과거 정보를 재구성하도록 요구하며, 상호작용 과정 전반에 걸쳐 충실한 압축 (faithful compression)과 메모리 완전성 (memory completeness)을 장려합니다.

두 목적 함수는 GRPO를 사용하여 공동으로 최적화됩니다. 긴 텍스트 QA 및 검색 기반 QA 벤치마크에 대한 광범위한 실험 결과, MemTrain은 다양한 모델에 걸쳐 다운스트림 메모리 집약적 추론 성능을 일관되게 향상시키며, 직접적인 태스크별 사후 학습 (task-specific post-training) 대비 최대 17.67포인트의 성능 향상을 달성함을 입증했습니다.