AiO-Mem, 검색·기억 저장을 한 보상으로 학습하는 메모리 에이전트 공개

연구진이 2026년 10월 15일 arXiv에 All-in-One Memory(AiO-Mem)를 공개했다. 대규모 언어 모델이 쓰는 외부 메모리에서 검색, 기억 쓰기, 검색 중단 결정을 하나의 보상 공간에서 평가해 강화 학습으로 함께 최적화하는 방법이다.

논문은 여러 세션에 걸친 대화나 긴 문서처럼 오래된 기록에 근거해 답해야 하는 질의를 대상으로 한다. 연구진은 여러 장문 맥락 메모리 벤치마크에서 AiO-Mem이 정적 메모리 시스템과 강화학습으로 훈련한 기존 메모리 에이전트보다 답변 정확도를 꾸준히 높였다고 밝혔다.

작동 방식

연구진에 따르면 기존 메모리 방법은 검색과 기억 구축을 따로 최적화하는 경우가 많았다. 검색, 쓰기, 중단 결정을 함께 평가할 공통 보상 공간도 없었다. AiO-Mem은 이 동작들을 공유하는 증거 상태 위의 전이로 정리하고, 각 행동의 의미적 품질을 통합 보상 공간에서 직접 평가한다.

먼저 서로 다른 형태의 기록을 출처를 추적할 수 있는 다중 관점 이벤트로 변환한다. 검색과 기억 쓰기가 이 이벤트를 공통 기반으로 쓴다. 통합 메모리 정책 액터는 이 증거 상태 위에서 상호작용 전체 궤적을 끝까지 실행한다.

궤적에 점수를 매기는 일은 별도의 궤적 조건부 프로세스 보상 모델이 맡는다. 이 모델은 증거에 근거한 지도 미세 조정과 궤적 수준 직접 선호 최적화(DPO)로 훈련한다. 모델이 준 점수는 행동별 균형을 맞춘 궤적 보상으로 합산되고, Process-Aligned DAPO(PA-DAPO)가 이 보상으로 메모리 정책을 최적화한다.

검색과 쓰기, 중단을 한 보상으로 묶으면 마지막 답의 정답 여부만으로는 알기 어려운 중간 행동의 질도 훈련 신호에 들어간다. 다만 이는 초록에 적힌 설계에서 끌어낸 해석이며, 어느 구성 요소가 성능 향상에 얼마나 기여했는지는 초록만으로는 알 수 없다.