LLM 장기 대화 기억을 강화하는 MemBuilder: 밀도 보상 기반 강화학습 프레임워크

초록

MemBuilder는 4B 파라미터 LLM이 다차원 외부 메모리를 스스로 구축하도록 훈련하는 강화학습 시스템이다. 세션‑단위 합성 질문 생성으로 밀도 보상을 제공하고, 메모리 유형별 사용 빈도에 따라 기여도 가중치를 부여해 다차원 메모리의 기여를 명시적으로 반영한다. 실험에서 LoCoMo 등 장기 대화 벤치마크에서 Claude 4.5 Sonnet을 능가하는 성능을 기록한다.

상세 분석

본 논문은 장기 대화에서 일관성을 유지하기 위해 외부 메모리를 활용하는 기존 접근법의 두 가지 근본적인 한계를 짚는다. 첫째, 대화 전체를 하나의 트래젝터리로 보고 최종 QA 결과에만 보상을 주는 ‘희소 트래젝터리 보상’은 어느 세션의 메모리 조작이 최종 성과에 기여했는지 식별하기 어렵다. 둘째, 다차원 메모리(코어, 에피소드, 의미, 절차) 각각에 동일한 전역 보상을 적용하면, 실제 downstream 작업에서 어느 메모리 유형이 더 많이 활용되는지에 대한 신호가 사라진다.

이를 해결하기 위해 MemBuilder는 두 가지 핵심 기술을 도입한다. ① 밀도 세션‑레벨 보상: 각 대화 세션 τ마다, 현재 메모리 뱅크 Mτ‑1과 세션 텍스트 sτ를 입력으로 전문가 LLM(Claude 4.5 Sonnet)이 합성 질문‑답 쌍을 J개 생성한다. 질문은 단일 세션, 다세션, 시간 추론 등 세 가지 유형으로 구성돼 메모리의 다양한 측면을 평가한다. 이후 정책 모델이 수행한 메모리 조작을 기반으로 생성된 메모리 M(i)τ에 대해 동일한 질문들을 재질문하고, 정답 일치율을 r_task으로 계산한다. 이렇게 하면 매 세션마다 즉각적인 피드백이 제공돼, 메모리 삽입·업데이트·병합 등의 행동이 어느 정도 품질을 갖추었는지 즉시 알 수 있다.

② 기여도 인식 그래디언트 가중치: 세션‑레벨 보상 r은 전역 보상이지만, 실제 QA 과정에서 각 메모리 유형이 얼마나 자주 검색되는지를 h(m)으로 기록한다. 가장 많이 사용된 유형 d에 대해 가중치 α(>1)를 부여하고, 나머지는 1로 둔다. 코어 메모리는 항상 컨텍스트에 포함되므로 고정 가중치 1을 적용한다. 이 가중치는 PPO‑계열의 정책 업데이트 시 중요도 비율 ρ(m)와 결합돼, 영향력이 큰 메모리 유형에 대해 더 큰 정책 변화가 일어나도록 설계되었다.

모델 아키텍처는 4B 규모의 Qwen‑3 기반 LLM 하나가 네 개의 메모리 유형에 대해 각각 별도의 프롬프트를 받아 행동을 선택하도록 한다. 행동 공간은 Core(append, replace, rewrite), Episodic(add, update, merge), Semantic(add, update, skip), Procedural(add, update) 로 정의돼, 기존 시스템이 ‘삭제·덮어쓰기’만 하는 것과 달리 업데이트 시 이전 엔트리를 보존하고, merge 연산을 통해 시간 구간을 아우르는 요약을 미리 생성한다.

학습 파이프라인은 먼저 Claude 4.5 Sonnet으로부터 전문가 궤적을 수집해 SFT( supervised fine‑tuning) 단계에서 형식 오류를 최소화한다. 이후 ADRPO(Attributed Dense Rewards Policy Optimization) 단계에서 N개의 롤아웃을 샘플링하고, 각 롤아웃에 대해 위에서 정의한 밀도 보상과 기여도 가중치를 적용해 GRPO 기반의 클리핑 손실을 최적화한다. 정규화 항 λ·ℓ 은 메모리 길이 과다를 억제하고, KL 다이버전스 항 β·D_KL 은 정책이 레퍼런스 정책(전문가 모델)에서 크게 벗어나지 않도록 한다.

실험에서는 LoCoMo, Long‑MemEval, PerL‑TQA 등 장기 대화 QA 벤치마크를 사용했다. MemBuilder는 84.23%의 LoCoMo 정확도를 달성했으며, 이는 Claude 4.5 Sonnet(≈81%) 및 최신 공개 모델들을 모두 앞선 결과다. Ablation study에서 (1) 밀도 보상 없이 전역 보상만 사용했을 때 성능이 4~6% 감소하고, (2) 기여도 가중치를 제거했을 때 특정 메모리 유형(특히 Episodic)의 활용도가 급감함을 확인했다. 또한 메모리 크기와 토큰 제한을 변동시켜도 성능 저하가 미미해, 제안된 길이 페널티와 메모리 압축 메커니즘이 실제 서비스 환경에서도 안정적임을 보여준다.

결론적으로 MemBuilder는 (1) 세션‑단위 합성 QA를 통한 밀도 보상으로 장기 대화의 단계적 학습 신호를 제공하고, (2) 메모리 사용량 기반 기여도 가중치로 다차원 메모리의 효과적 학습을 가능하게 함으로써, 경량 LLM이 고가의 폐쇄형 모델을 대체할 수 있는 실용적인 메모리‑증강 프레임워크를 제시한다.