MemLearner: 쿼리 토큰 기반의 학습형 컨텍스트 질의로 장기 일관성을 확보한 Video World Models

Video World Models는 상호작용 환경에서 과거 장면을 기억하지 못해 장기 생성에서 장면 불일치가 발생하는 문제가 핵심적이다. 본 논문은 규칙 기반의 컨텍스트 검색이 차폐나 동적 객체 상황에서 실패하는 점을 지적하고, pretrained 생성 모델의 시각적 priors를 활용하여 학습 가능한 쿼리 메커니즘으로 이를 보완했다. 이 방식은 추가적인 3D 재구성이나 별도 모듈 없이 컨텍스트 활용 능력을 향상시켜 실세계 도메인으로의 일반화 가능성을 확보했다.

문제 출발점은 과거 프레임들을 모두 보존한 채로 장시간 비디오를 생성하면 연산과 메모리 비용이 급증하고, 규칙 기반 검색은 occlusion이나 동적 객체에서 실패한다는 점이다. Transformer 기반의 Diffusion Generator는 이미 개별 프레임의 시각적 prior를 학습하고 있으므로 이 내부를 질의하는 형태로 컨텍스트 접근을 학습시키면 별도 재구성 모듈 없이도 과거 정보를 효과적으로 활용할 수 있다. 이 접근은 C, Q, P 토큰을 하나의 입력으로 병합하여 3D attention 내에서 상호작용시키는 방식으로 구현되어, 모델이 언제 어떤 과거 정보를 참조해야 하는지를 diffusion loss를 통해 간접적으로 학습한다.

쿼리 토큰의 역할은 두 단계로 이해할 수 있다. 첫째, P 토큰(예측 대상)을 Q 토큰에 쿼리로 주어 P가 어떤 스펙트럼의 정보를 필요로 하는지 Q가 판단한다. 둘째, Q 토큰이 해당 요구에 따라 C 토큰들에서 키·값 조회를 수행하여 필요한 세부 정보를 추출하고 이 결과를 다시 P 토큰 쪽으로 전달한다. 이 흐름은 입력 단계의 noisy P_t와 unperturbed C, Q 토큰이 동시에 DiT에 들어가는 설계에서 자연스럽게 작동하며 중간 층의 attention 유사도 분석에서 Q의 attention 분포가 시간·프레임별로 달라지는 것이 확인되었다.

이 접근의 핵심적 차이는 별도 모듈을 처음부터 학습하는 것이 아니라 pretrained 생성 모델의 내부 표현을 활용하여 질의 능력을 획득한다는 점이다. 그 결과로 모델은 occlusion과 동적 객체가 존재하는 장면에서 규칙 기반 방법보다 더 강건한 컨텍스트 선택을 수행했고, 실험에서 PSNR과 LPIPS 등 메모리 관련 지표에서 유의한 개선을 나타냈다.

전체 접근 방식은 latent video diffusion 모델과 causal 3D VAE, 그리고 Diffusion Transformer(DiT)를 기반으로 한다. 입력으로는 과거 프레임의 latent인 C 토큰, 학습 가능한 초기화된 Q 토큰(노이즈로 샘플링), 그리고 예측할 P 토큰(입력에서 노이즈가 더해진 형태)을 프레임 차원으로 연결하여 DiT에 공급한다. 학습 손실은 예측된 노이즈와 실제 가우시안 노이즈의 차이에 대해 계산되므로 Q 토큰은 직접적인 지도 없이도 diffusion loss를 통해 유용한 컨텍스트 추출 기능을 획득한다.

핵심 메커니즘은 attention 패턴의 제한과 Query Layer 분할에 있다. 전체 DiT를 n+m 층으로 보고 입력 근처의 n개 Query Layers에서만 C, Q, P가 상호작용하게 하며, 이후의 m개 Generative Layers에서는 C를 제거하고 Q, P만으로 생성 연산을 수행한다. 또한 유지되는 attention 패턴은 Q→P, Q→C, P→{P,Q} 세 가지로 한정되어 C가 직접 쿼리로 동작하는 불필요한 연산을 제거하여 복잡도를 낮췄다.

데이터와 학습 전략 측면에서는 렌더링 데이터의 정확한 카메라 주석을 활용하되 실세계 데이터의 사실감을 얻기 위해 멀티데이터셋 학습을 수행했다. 각 데이터 타입별로 전용 카메라 인코더를 두어 포즈 품질 차이가 서로 간섭하지 않게 했고, 포즈가 없는 데이터는 제로 포즈로 처리했다. 구현 세부에서는 P 토큰이 77 프레임을 표현하고 latent temporal compression ratio가 4여서 P 토큰은 20 latent frames를 대표하도록 구성했으며, Query Layer 수는 실험적으로 5층이 비용과 성능의 균형을 이루었다.

메인 벤치마크에서 MemLearner는 수치 상으로 기존 방법들보다 우수한 성능을 보였다. 논문에 보고된 수치에 따르면 GT 비교에서 MemLearner의 PSNR은 21.23, LPIPS는 0.2904였고 FID는 112.75, FVD는 835.98로 측정되었다. 재방문 비교에서는 PSNR 18.57과 LPIPS 0.3230을 기록했으며, 비교군 가운데 Separate Module 설계는 GT Comp에서 PSNR 9.16, LPIPS 0.6567로 크게 저하되어 별도 from-scratch 모듈의 한계를 숫자로 확인했다.

Ablation 결과는 설계의 핵심 요소를 검증했다. Query Layers 수를 조절한 시험에서 5층이 성능과 FPS의 균형을 제공했으며, Q가 P를 쿼리로 삼아 무엇을 뽑아야 하는지 결정하는 attention 패턴이 제거될 경우 성능이 크게 저하되었다. 카메라 임베딩 실험에서는 P 토큰에만 카메라 정보를 넣어도 성능 저하가 거의 없었고, 이는 모델이 명시적 포즈 없이도 기하학적 대응을 일부 학습할 수 있음을 시사했다.

데이터 혼합 관련 실험에서 렌더링 전용 학습은 메모리 성능을 학습시키는 데 유리했으나 실세계 분포와의 격차가 존재했다. 실세계 데이터만으로는 재방문 패턴이 부족하여 메모리 능력을 학습하기 어려웠고, 렌더링 데이터와 실세계 데이터를 적절히 혼합하는 전략이 사실감과 일관성 사이에서 균형을 제공했다.

전체 아키텍처는 causal 3D VAE로 입력 비디오를 시공간적으로 압축한 뒤 Diffusion Transformer(DiT)에 C, Q, P 토큰을 프레임 차원으로 연결하여 입력하는 구조로 이루어져 있다. 각 DiT 블록은 2D spatial attention, 3D spatial-temporal attention, cross-attention, FFN 순서로 구성되며 3D attention의 입력은 F_in={C,Q,P}로 정의된다. 학습 손실은 noisy P_t에 대한 노이즈 예측 오차로 계산되므로 C와 Q에는 직접 노이즈를 추가하지 않고 P에만 입력 단계에서 가우시안 노이즈를 더했다.

효율성 관점에서 두 가지 기법이 수학적으로 반영되었다. 첫째, Query Layers에서만 C,Q,P가 상호작용하도록 하여 계산 복잡도를 낮췄다. 이 경우 Query Layers의 출력은 C_out=C로 고정하고 Q_out=Q+g(Q,{C,P},{C,P}), P_out=P+g(P,{P,Q},{P,Q})로 정리된다. 둘째, Generative Layers에서는 C를 제거하여 Q_out=Q+g(Q,P,P)와 P_out=P+g(P,{P,Q},{P,Q})만 연산하도록 하여 불필요한 q·k·v 계산을 줄였다.

학습 및 구현 세부에서는 내부 1B 파라미터 텍스트-투-비디오 DiT를 주 모델로 사용했고 레이어 총 28층 중 5층을 Query Layers로 구성했다. P 토큰은 생성당 77 프레임을 나타내며 latent temporal compression ratio는 4이므로 P 토큰은 20개의 latent video frame을 대표하도록 설정되었다. 훈련은 20,000회 이상의 이터레이션, 배치 사이즈 8, 학습률 5×10^-5로 수행되었고 Classifier-Free Guidance에 50 스텝을 사용하여 텍스트 조건을 적용했다.

데이터 파이프라인은 세 가지 구성으로 운영되었고 렌더링 전용 구성과 mixed 구성(렌더링 75%+실세계 25%)을 비교했다. 멀티데이터 전략에서는 각 데이터 타입별 카메라 인코더를 분리하여 포즈 주석의 품질 차이가 공동 학습에 미치는 영향을 억제했다. 하드웨어 및 학습 시간에 대한 구체 수치는 본문에 명시되지 않았으나 모델 크기와 레이어 수 설계가 대규모 대안 대비 실험 비용을 절감하는 방향으로 선택되었다.

논문에 명시된 주된 한계는 모델 용량과 학습 데이터 규모의 제약이다. 현재 1B 규모에서는 다수의 동적 객체가 동시에 상호작용하는 대규모 장면에서 객체 추적과 외형 일관성 유지에 실패하는 사례가 관찰되며 특히 다섯 명 이상의 캐릭터가 동시에 상호작용할 때 품질 저하가 두드러졌다. 또한 본 접근은 전체 컨텍스트 저장 전략에 기반하므로 메모리가 생성 시간에 선형적으로 증가하는 문제는 여전히 남아 있으며 컨텍스트 압축·요약·갱신·선택적 망각 등 실용적 메모리 관리 기술과의 결합이 필요하다고 명시되었다.