기술 요약: 확산 세계 모델을 위한 기억 전문가의 구성

문제 제기

세계 모델 (WM) 은 강화 학습 (RL) 및 계획에 필수적이며, 과거 관측과 일관된 plausible 미래 상태를 예측합니다. 그러나 기존 아키텍처는 기억 충실도와 계산 효율성 간의 근본적인 트레이드오프에 직면해 있습니다. 트랜스포머 기반 모델은 국소적 세부 사항을 보존하지만 이차적 어텐션 비용으로 인해 컨텍스트 길이에 제한을 받습니다. 반면, 순환 신경망과 상태 공간 모델 (SSM) 은 효율적으로 확장되지만 역사를 숨은 상태 (hidden states) 로 압축함으로써 시간이 지남에 따라 미세한 세부 사항을 필연적으로 손실합니다. 단일 아키텍처에서 컨텍스트 길이만 단순히 늘리는 것은 불안정한 학습과 prohibitive 추론 비용으로 이어집니다. 또한, 여러 모델을 단순하게 구성하면 "모드 붕괴 (mode collapse)"가 발생하여 결합된 분포가 공통 영역을 과도하게 예리하게 만들고 다양성을 상실하게 됩니다.

방법론

저자들은 **기억 전문가의 구성 (Composition of Memory Experts, CoME)**을 제안합니다. 이는 미래 - 과거 일관성을 단일 모놀리식 아키텍처에서 분리하는 확산 기반 프레임워크입니다. 대신 CoME 는 각각 특정 기억 역할을 처리하는 일련의 전문화된 전문가들을 활용하며, 추론 시 확률적으로 이를 융합합니다.

1. 대비 전문가의 곱 (Product of Contrastive Experts, PoCE)

다양성을 억제하고 가짜 모드를 증폭시킬 수 있는 표준 전문가의 곱 (PoE) 의 한계를 해결하기 위해, 저자들은 대비 전문가의 곱 (Product of Contrastive Experts) 공식을 도입합니다.

• 메커니즘: 이질적인 전문가 집합 $p_i(x)$에 대해, 구성은 $p(x) \propto \prod \tilde{p}_i(x)$로 정의되며, 여기서 $\tilde{p}_i(x) \propto p_i(x)^{\alpha_i} \bar{p}_i(x)^{1-\alpha_i}$입니다.
• 대비 기준선: $\bar{p}_i(x)$는 전문가의 무조건적 기준선 분포를 나타냅니다. 계수 $\alpha_i > 1$은 대비 인자로 작용합니다.
• 이론적 정당성: 커널 분산을 축소하고 국소 기하학을 왜곡하는 단순한 템퍼링 (tempering) 과 달리, 이 대비 혼합은 커널의 형태를 보존하면서 혼합 구성 요소를 재가중합니다. 이는 조건부 전문가가 유효 모드에 기준선보다 높은 확률을 할당하는 한, 가짜 모드 (전문가가 강력한 증거 없이 높은 확률을 할당하는 경우) 를 효과적으로 억제하면서 유효한 2 차 모드를 유지합니다.

2. 전문화된 기억 전문가

이 프레임워크는 세 가지 상호 보완적인 전문가를 구체화합니다.

• 단기 기억 (STM): 최근 컨텍스트 창 ($10\text{--}100$ 프레임) 을 조건으로 하는 확산 모델입니다. 확장된 컨텍스트 길이 (예: 슬라이딩 윈도우 어텐션) 를 가진 아키텍처를 사용하여 정교한 국소 역학과 즉각적인 시간적 의존성을 포착합니다.
• 장기 기억 (LTM): episodic 역사 ($100\text{--}1000$ 프레임) 를 가중치에 직접 저장하는 별도의 확산 모델입니다. 이는 장기 컨텍스트에 대한 경량 LoRA 어댑터의 **추론 시 미세 조정 (test-time finetuning)**을 통해 달성됩니다. 이 접근법은 회상 비용을 분산시켜, 모든 단계에서 전체 역사를 모델에 입력하지 않고도 과거 경험을 상수 시간으로 재사용할 수 있게 합니다. 백본을 고정하고 어댑터만 업데이트함으로써 catastrophic forgetting 을 방지합니다.
• 공간 장기 기억 (SLTM): 역사에서 추출된 공간 사전 지식 (예: 카메라 포즈, 포인트 맵) 을 통합하는 LTM 의 확장입니다. 이 전문가는 기하학적 및 공간적 일관성을 강제하여 반복적이거나 시각적으로 유사한 환경에서 위치를 명확히 합니다.

3. 통합 생성

최종 분포는 사전 학습된 사전 분포와 세 가지 대비 전문가 (STM, LTM, SLTM) 의 곱입니다. 샘플링 중에는 이러한 전문가들의 스코어 함수가 로그 공간에서 가산적으로 결합되어 디노이징 과정을 안내하며, 생성된 미래가 최근 역동성, 장기 역사, 공간적 일관성을 존중하도록 보장합니다.

주요 기여

1. 분리된 기억 아키텍처: 이 논문은 단일 백본이 아닌 전문화된 전문가들에 기억을 분배함으로써 기억 트레이드오프를 해결하는 것을 제안하며, 이를 전문가의 곱 문제로 공식화합니다.
2. 대비 전문가의 곱 (PoCE): 이질적인 전문가를 융합하기 위한 새로운 전략입니다. 저자들은 표준 PoE 와 달리 중복 모드를 분리함으로써 모드 붕괴와 과도한 확신을 방지한다는 이론적 및 실증적 증거를 제시합니다.
3. 외부 확산 기억: 가중치에 episodic 지식을 저장하기 위해 LoRA 어댑터의 추론 시 미세 조정을 사용하는 LTM 모듈의 도입입니다. 이는 SLTM 으로 확장되어 공간 구조를 통합하며, 내비게이션 작업의 정확도를 향상시킵니다.

실험 결과

이 방법은 합성 (Memory Maze, DMLab, Minecraft) 및 실세계 (RealEstate10K, RECON) 벤치마크에서 평가되었습니다.

• 재구성 품질: Memory Maze 데이터셋에서 CoME 는 LPIPS(0.097), SSIM(0.892), PSNR(23.07) 에서 최첨단 성능을 달성하여 풀 어텐션 트랜스포머, SSM, 슬라이딩 윈도우 모델을 포함한 기준선들을 크게 능가했습니다.
• 확장성: STM 기준선보다 $60\times$ 더 많은 연산이 필요한 풀 어텐션 모델과 달리, CoME 는 훈련 및 추론 비용의 일부로 우수한 일관성을 달성합니다.
• 컨텍스트 길이: 이 모델은 480 프레임의 컨텍스트까지 성능 일관성에서 포화 현상이 나타나지 않았으며, 이는 LTM 구성 요소의 효과를 입증합니다.
• 계획 및 내비게이션: 목표 조건부 내비게이션 작업 (RECON 벤치마크) 에서 CoME 는 절대 궤적 오차 (ATE) 를 0.96 으로 줄여 전문 내비게이션 세계 모델 (NWM) 및 기타 기준선들을 능가했습니다.
• 회상 능력: RealEstate10K 에 대한 역궤적 실험에서 CoME 는 순방향 롤아웃 후 초기 프레임을 성공적으로 회상한 반면, 기억이 없는 기준선들은 plausible 하지만 일관성 없는 장면을 생성했습니다.
• 절대 실험: 대비 공식화를 제거하면 성능이 붕괴되어 PoCE 메커니즘의 필요성을 확인시켰습니다. LTM 에 대한 LoRA 랭크와 컨텍스트 길이를 늘리는 것은 과적합 없이 일관적으로 결과를 개선했습니다.

중요성 및 주장

이 논문은 기억이 증강된 확산 세계 모델을 구축하기 위한 새로운 패러다임을 확립한다고 주장합니다. 기억을 전문화된 전문가들에 분배하고 대비 구성 전략을 사용함으로써 CoME 는 효율성과 고충실도 시간적 일관성 사이의 균형을 맞춥니다. 저자들은 이 접근법이 세계 모델이 이차적 비용을 치르거나 순환 압축에 내재된 세부 사항 손실을 겪지 않고도 긴 컨텍스트로 확장될 수 있게 한다고 주장합니다. 이 방법은 다양한 확산 기반 세계 모델에 적용 가능한 일반 프레임워크로 제시되며, 강화 학습 환경에서의 계획 및 의사 결정에 대한 신뢰성을 향상시킵니다. 이 연구는 기억이 단일 아키텍처의 제약이 아니라 분산 시스템의 속성이어야 함을 강조합니다.