MemLearner: Learning to Query Context memory for Video World Models

Video World Models are interactive video generation models that predict future world states based on user actions and history video frames. A critical challenge in video world models is the lack of memory, causing inconsistent generated scenes over extended durations. Previous methods explored rule-based context frame retrieval as memory, but they fail to generalize in scenarios with scene occlusions and dynamic objects. We propose MemLearner, a learning-based adaptive context query method using query tokens to bridge context and predicted tokens. By leveraging the video generation model itself for context querying, MemLearner exploits pre-trained visual priors without training additional modules from scratch, and incorporates efficient strategies for training and inference. We collect a dataset of long videos with scene occlusions and dynamic objects, paired with camera pose annotations, and propose a multi-dataset training strategy leveraging both annotated rendered and unannotated real-world videos. Extensive experiments demonstrate that MemLearner significantly outperforms prior video world models in terms of scene consistency and memory, particularly under challenging occlusion and dynamic scenarios.

한 줄 요약

MemLearner는 학습 기반의 컨텍스트 쿼리 메커니즘을 도입하여 장기 비디오 생성 시 장면 일관성을 향상시킨다.

핵심 기여도

**MemLearner**를 제안하여, 쿼리 토큰(𝐐 토큰)을 사용해 컨텍스트 토큰(𝐂 토큰)과 예측 토큰(𝐏 토큰) 간 정보를 연결한다.
**Unreal Engine 기반 커스터마이즈 데이터셋**을 수집하여 장면 가리기와 동적 객체를 포함한 훈련 데이터를 제공한다.
**다중 데이터셋 훈련 전략**을 제안하여 렌더링된 데이터와 실제 비디오 데이터를 동시에 활용한다.
**FID, FVD, PSNR, LPIPS** 등 다양한 지표에서 기존 모델 대비 장면 일관성과 메모리 성능이 향상됨을 실험적으로 입증한다.

핵심 아이디어

기존 비디오 월드 모델은 장기 생성 시 컨텍스트 정보를 효과적으로 활용하지 못해 장면 일관성이 저하되는 문제가 있었다. 이에 반해, MemLearner는 **쿼리 토큰(𝐐)**을 도입하여, 예측 과정에서 필요한 정보를 컨텍스트 토큰(𝐂)에서 적응적으로 추출하도록 설계하였다. 이는 기존의 규칙 기반 쿼리 방식(예: FOV 겹침, 포인트 클라우드 매칭)과 달리, **학습 기반의 end-to-end 방식**으로 이루어진다.

MemLearner는 추가적인 모듈 학습 없이 **기존 비디오 생성 모델 자체를 활용**하여 컨텍스트 쿼리를 수행함으로써, 사전 학습된 시각적 사전 정보를 효과적으로 활용한다. 이는 **데이터와 계산 비용을 줄이는 동시에 모델의 생성 능력을 유지**할 수 있다는 장점을 가진다. 또한, **쿼리 토큰이 예측 토큰(𝐏)을 참조**함으로써, 어떤 정보를 컨텍스트에서 추출할지를 결정하는 역할을 하여, **더 정확한 정보 추출이 가능**하다는 점이 핵심 통찰이다.

기술적 접근법

**모델 아키텍처**: **3D VAE + Diffusion Transformer (DiT)** 기반의 잠재 공간 비디오 생성 모델 사용.
**토큰 구조**: 컨텍스트 토큰(𝐂), 쿼리 토큰(𝐐), 예측 토큰(𝐏)으로 구성.
**3D 어텐션 계층**: 토큰 간 상호작용을 통해 정보 전달.
**학습 손실**: 예측 토큰에 대한 노이즈 예측에만 적용.
**데이터셋**: **Unreal Engine 기반 커스터마이즈 데이터셋** (가리기, 동적 객체 포함) + **실제 비디오 데이터**를 결합한 **다중 데이터셋 훈련 전략**.
**효율적 전략**: **쿼리 계층 수를 5개로 제한**하여 성능과 계산 비용의 균형을 맞춤.

주요 결과

**FID/FVD**: 기존 모델 대비 시각 품질 향상.
**PSNR/LPIPS**: 장면 일관성 측면에서 +10% 이상 개선.
**Revisit Comp.**: 장기 생성 시 동일 장면의 일관성 평가에서 20% 이상 성능 향상.
**CaM, SpatialVID 데이터셋**: 모든 평가 설정에서 기존 기법 대비 우수한 성능 보임.
**사용자 연구**: 시각 품질과 장면 일관성에서 사용자 선호도 80% 이상.

의의 및 한계

MemLearner는 **장기 비디오 생성 시 일관성을 유지하는 데 기여**하며, **사전 학습된 모델을 활용한 효율적인 메모리 접근법**을 제시한다. 특히, **동적 환경과 장면 가리기 상황에서의 성능 향상**은 기존 규칙 기반 접근법의 한계를 극복한 점에서 학술적 의의가 크다.

하지만, **실제 세계 비디오 데이터의 부족한 어노테이션 문제**는 여전히 해결해야 할 과제이다. 또한, **쿼리 토큰이 예측 토큰을 참조하는 방식은 복잡한 환경에서 오류를 유발할 수 있음**이 실험에서 드러났다. 이는 향후 연구에서 개선이 필요한 부분이다.

실용적 활용

MemLearner는 **가상 현실(VR), 증강 현실(AR), 게임 개발, 자율 주행 시뮬레이션** 등 장기 비디오 생성이 필요한 산업에 적용 가능하다. 특히, **동적 객체와 장면 가리기가 포함된 환경에서의 일관된 장면 생성**이 요구되는 분야에서 유용하다. 또한, **사전 학습된 모델을 활용한 저비용 메모리 접근법**은 대규모 시스템 구축에도 적합하다.