PaperlayerView exact revision on arXiv

Source-linked AI summary

MemLearner: Learning to Query Context memory for Video World Models

Jiwen Yu, Jianxiong Gao, Jianhong Bai, Yiran Qin, Kaiyi Huang, Quande Liu, Xintao Wang, Pengfei Wan, Kun Gai, Xihui Liu

arXiv:2606.31734v1cs.CV

TL;DR

Video World Models struggle to maintain scene consistency over long durations because existing rule-based memory retrieval fails with occlusions and dynamic objects. MemLearner learns to query historical context, and experiments show improved scene consistency and memory, especially in these challenging scenarios.

  • Problem

    Video World Models lack effective memory for maintaining scene consistency over extended durations, while rule-based retrieval struggles with occlusions and dynamic objects.

  • Method

    MemLearner replaces rule-based retrieval with learned adaptive context querying and combines occlusion-rich rendered data with real-world videos through multi-dataset training.

  • Results

    MemLearner significantly outperforms prior Video World Models in scene consistency and memory, particularly under occlusion and dynamic-object scenarios.

  • Takeaways & Limitations

    The paper supports learning-based context query as a more adaptive memory mechanism for Video World Models in complex dynamic scenes.

  • Takeaways & Limitations

    At the current 1B scale, generation degrades when more than five characters interact in a scene, causing inconsistent appearances or missing objects upon revisit.

Abstract

from arXiv · show

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.

1 Introduction

Video World Models offer photorealistic, scalable world simulation but suffer from inconsistent scenes when limited context windows impair memory. MemLearner addresses this with adaptive context querying, mixed rendered and real-video training, and demonstrated gains in consistency and memory under occlusion and dynamic objects.

  • Video World Models are promising because videos capture real-world dynamics photorealistically and abundant internet video provides scaling potential.
  • Limited context windows cause later generated scenes to become inconsistent with earlier ones, motivating memory representations based on reconstruction, compression, or context-frame retrieval.
  • Rule-based retrieval using FOV overlap or point-cloud and surfel matching fails in complex scenes with occlusions, moving objects, and changing dynamics.
  • MemLearner uses query tokens to adaptively extract information from context tokens and provide it to predicted tokens, leveraging the video model’s pretrained visual priors.The method avoids additional scratch-trained modules and introduces efficient training and inference strategies for long video sequences.
  • MemLearner combines an Unreal Engine dataset with occlusions and dynamic objects and a multi-dataset strategy that leverages rendered and real videos.The training setup addresses the complementary strengths of rendered videos with accurate poses and real videos with diverse, dynamic content.
  • MemLearner significantly outperforms previous Video World Models in scene consistency and memory, particularly under occlusion and dynamic-object scenarios.

2 Related Work

Related work defines video world models, reviews video generation and interactive control, and identifies memory as a central limitation caused by finite context windows. Existing retrieval methods reduce computation but rely on non-generalizing rules, motivating learnable context querying as a distinct approach.

  • World Model: World models predict future world states from historical states and current actions, with world states represented through language, semantics, 3D/4D structures, or videos.Video world models have recently benefited from video photorealism, large datasets, and advances in video generation.
  • Video Generation Models: Diffusion Transformer architectures dominate video generation because they produce highly realistic videos, while next-token and hybrid architectures generally achieve lower generation quality.Interactive control includes camera pose, trajectory, and action control.
  • Memory in Video World Models: Memory enables consistent long-video generation and supports planning, interaction, and state modeling, but limited context windows often cannot retain enough historical information.Additional conditional information is therefore needed to preserve relevant past visual memories.
  • Context Retrieval: Context retrieval reduces computational overhead by conditioning generation on selected historical frames, but rule-based retrieval lacks generalization across videos and scenes.The proposed direction is a learning-based context query method that learns which historical information is useful.
  • Learnable Query Tokens: Learnable query tokens have been used to aggregate visual features in multimodal language models and select informative frames in video understanding, whereas this work targets memory in video world models.The paper frames its problem as fundamentally different from visual-context compression and recognition-oriented frame selection.

3 MemLearner

MemLearner uses unperturbed context and query tokens alongside noised predicted tokens in a latent video diffusion model. It adaptively queries context through the generative model while reducing computation by restricting context access and attention patterns.

  • Introducing Adaptive Query Tokens: Different predicted frames and diffusion timesteps require different context guidance, motivating adaptive query tokens with varying attention distributions.Attention similarity analysis reports markedly different query-token attention distributions across predicted frames and diffusion timesteps.
  • Architecture Design Insights: A jointly trained from-scratch context query module fails to establish meaningful context modeling, producing near-zero query-context similarity and hindering gradient propagation.The reported failure motivates using the video generation model itself for context querying rather than adding a separately initialized module.
  • Architecture Details: MemLearner represents video latents as context, query, and predicted tokens, leaving context and query unperturbed while applying Gaussian noise only to predicted tokens.The diffusion objective is applied to the predicted-token noise.
  • Efficiency Strategies: MemLearner improves efficiency by querying context only in shallow layers and retaining three attention patterns: Q→P, Q→C, and P→{P,Q}.The three patterns let Q identify prediction-relevant context information, extract it from C, and let P use both P and Q; large context lengths make full attention expensive.

4 Dataset

The paper introduces a rendered Unreal Engine dataset designed for long-video world-model evaluation with precise camera poses, occlusions, dynamic objects, and diverse scenes. It further combines rendered, estimated-pose, and pose-unannotated real-world data through dedicated camera encoders.

  • Dataset requirements: The dataset targets long videos with accurate per-frame camera annotations, occlusion relationships, dynamic objects, and sufficient diversity.Existing long-video datasets do not satisfy these requirements simultaneously.
  • Rendered dataset: The authors collect a rendered Unreal Engine dataset featuring customized diverse scenes and dynamic objects with occlusion relations.The collection is designed to address limitations of existing long-video datasets.
  • Data sources: Rendered, estimated-pose, and pose-unannotated real-world datasets provide complementary trade-offs in pose accuracy, visual realism, diversity, and dynamics.Rendered data has precise poses but limited realism and diversity; estimated-pose data is photorealistic but less precise and less dynamic; pose-unannotated data is diverse and dynamic but lacks accurate camera annotations.
  • Multi-dataset training: The multi-dataset training strategy assigns each dataset type a dedicated camera encoder according to its camera-annotation quality.Rendered and estimated-pose data use their respective annotations, while pose-unreliable data receives zero camera parameters, R = 0 and t = 0, through a third encoder.

5 Experiments

Experiments evaluate MemLearner across challenging occlusion/dynamic scenes, non-occluded benchmarks, and real-world videos using memory and visual-quality metrics. MemLearner consistently achieves strong scene consistency, memory, and visual quality, with ablations identifying essential query design choices.

  • Evaluation Protocol: Evaluation spans occluded and dynamic scenes, the CaM dataset without such challenges, and real-world SpatialVID videos, using PSNR/LPIPS for memory and FID/FVD for quality.A 5% held-out split has no video overlap between training and testing.
  • Main Results: MemLearner consistently outperforms baselines across all three evaluation settings and is preferred in user studies for visual quality and scene consistency.Quantitative comparisons report the best performance across all evaluated metrics and benchmarks.
  • Qualitative Results: Qualitative results show robust handling of indoor and outdoor scenes with occlusions and dynamic objects, while mixed real-world training enables generalization without adaptation.The model preserves scene layout and object appearance across revisits in real-world scenes.
  • Cross-Dataset Evaluation: On CaM, CaM and MemLearner perform comparably, whereas CaM degrades sharply on the occlusion/dynamic dataset and MemLearner remains robust.This comparison supports MemLearner’s effectiveness under occlusions and dynamic objects.
  • Ablation Studies: Ablations show that Q queries attending to P keys/values is essential, five query layers balance cost and memory performance, and explicit camera poses are unnecessary.Adding Q-to-Q attention does not materially improve performance, while removing Q-to-P attention causes significant degradation.

6 Conclusion

MemLearner shifts video-world-model memory from rule-based retrieval to learned context querying, supported by annotated rendered data and multi-dataset training. The authors identify model capacity, data scale, and memory compression as key remaining challenges and future directions.

  • Contributions: MemLearner introduces learning-based context querying as a paradigm shift from rule-based context retrieval for video world-model memory.The method is validated using rendered data with accurate camera annotations, customized occlusions and dynamic objects, and sufficient revisit patterns.
  • Contributions: The proposed multi-dataset training strategy leverages both rendered and real-world videos to support context-memory learning.Rendered data provide accurate camera annotations and controlled occlusions, dynamic objects, and revisit patterns.
  • Limitations and Future Work: At the current 1B model scale, scenes with many simultaneously interacting characters in large environments remain challenging.The authors identify larger models and long videos with richer dynamic interactions as necessary future resources.
  • Limitations and Future Work: Memory should not scale linearly with generation time, motivating future work on context compression, summarization, updating, editing, and selective forgetting.Context compression is described as orthogonal to MemLearner and can precede learned querying over compressed tokens.

A Details of Collected Dataset

The collected dataset combines diverse photorealistic and stylized 3D scenes with dynamic objects, automated camera trajectories, and captioned video segments for training.

  • 3D Scenes and Dynamic Objects: The dataset contains 13 diverse 3D scene assets spanning streets, shopping malls, rural areas, and indoor and outdoor environments.Photorealistic assets minimize the domain gap with real-world videos, while stylized scenes increase diversity.
  • 3D Scenes and Dynamic Objects: Dynamic objects, including human characters with diverse appearances, are added to increase scene dynamics.
  • Automated Camera Trajectory Generation: An Unreal Engine Blueprint script generates randomized camera motion, performs obstacle avoidance and exploration, and logs camera poses at each timestep.The logged poses are parsed into sequence files.
  • Data Preprocessing and Text Annotation: A pre-trained multimodal large language model captions rendered frames every 77 frames, with nearest-segment captions conditioning randomly sampled 77-frame training clips.

B Details of Internal Model Architecture · C Supplementary Experimental Results

The appendix describes a latent diffusion-transformer text-to-video architecture using 3D-VAE representations and 3D attention over context, query, and predicted tokens. It also provides implementation pseudocode for this attention mechanism, while no supplementary experimental findings are included in the supplied passages.

  • B Details of Internal Model Architecture: The base text-to-video model uses a latent diffusion transformer architecture.Its overview is shown in Figure 9.
  • B Details of Internal Model Architecture: Videos are encoded into latent representations through a 3D-VAE before entering the diffusion transformer.
  • B Details of Internal Model Architecture: Earlier UNet- and transformer-based approaches often append a separate 1D temporal-attention module for temporal modeling.
  • B Details of Internal Model Architecture: The appendix characterizes spatially and temporally decoupled attention as suboptimal.
  • B Details of Internal Model Architecture: The designed 3D attention processes context, query, and predicted tokens efficiently.
  • B Details of Internal Model Architecture: Corresponding pseudocode for the 3D-attention implementation is provided in Figure 10.

C.1 Results on Open-Source Video Models

MemLearner generalizes to the open-source Wan 2.1 (T2V-1.3B) model, where it consistently outperforms CaM and VMem. All methods perform worse than on the internal model, likely because Wan 2.1 has weaker generation capability at the 1.3B scale.

  • Evaluation Setup: MemLearner is evaluated on Wan 2.1 (T2V-1.3B), an open-source video Diffusion Transformer compatible with its context query mechanism.Wan 2.1 uses Transformer blocks containing 3D attention modules, and the same training protocol as Sec. 5.1 is followed.
  • Results: MemLearner consistently outperforms CaM and VMem on Wan 2.1, supporting the architecture-agnostic nature of its learning-based context query approach.The comparison is reported in Table 8.
  • Results: All methods achieve lower absolute performance on Wan 2.1 than on the internal model, which is attributed to weaker generation capability at the 1.3B scale.The internal-model comparison is referenced as Table 2 in the main text.

C.2 Attention Visualization Results

Attention visualizations show that context information is queried adaptively across predicted frames and diffusion timesteps, while a separate context query module fails to establish meaningful context modeling. The resulting information-flow failure causes the generation model to ignore that module, consistent with the main-text quantitative results.

  • Adaptive Attention Across Frames and Timesteps: Query tokens exhibit markedly different attention distributions across predicted frames and diffusion timesteps, supporting adaptive context querying.Early diffusion timesteps attend broadly to context tokens for global scene layout, whereas later timesteps focus on fine-grained local correspondences.
  • Failure of the Separate Module: The separate context query module produces near-zero attention similarity between query and context tokens, indicating ineffective context modeling.The visualization shows a failure to establish meaningful information flow between the query and context tokens.
  • Failure of the Separate Module: This ineffective information flow prevents gradient propagation to the video generation model, which ignores the separate module’s output and degrades to a text-to-video model.This behavior is consistent with the quantitative results reported in Tab. 2 of the main text.

C.3 Evaluation on Real-World SpatialVID Dataset … C.8 Comparison with a Separately Trained Query Module

Across real-world evaluation and supplementary analyses, MemLearner generalizes to SpatialVID, is preferred by users, and improves realism while avoiding costly attention designs and outperforming a separately trained query module.

  • C.3 Evaluation on Real-World SpatialVID Dataset: MemLearner is evaluated on 1,000 held-out SpatialVID sequences, which extend to 900 frames and include per-frame camera pose annotations.SpatialVID was selected for its long videos and pose annotations, enabling controlled real-world evaluation.
  • C.3 Evaluation on Real-World SpatialVID Dataset: MemLearner consistently outperforms baselines across all SpatialVID metrics, although its margins are smaller because the benchmark filters large camera motions and lacks frequent revisits.These characteristics make SpatialVID a relatively easy benchmark for evaluation.
  • C.5 Additional Attention Computational Comparison: Using context tokens as queries in 3D attention adds substantial computational overhead while providing negligible or degraded performance.The comparison covers alternative attention settings across the model’s 28 layers.
  • C.6 Additional Dataset Ablation Results: Additional qualitative results show that incorporating real-world training data improves the realism of generated videos.This evidence is presented in Figure 13.
  • C.7 Additional Comparison Results: Additional qualitative comparisons across different methods are provided to further assess their generated-video behavior.These comparisons are presented in Figure 14.
  • C.8 Comparison with a Separately Trained Query Module: A separately pre-trained query module with auxiliary L1 supervision and progressive fine-tuning remains far below MemLearner.The alternative is pre-trained using an L1 loss between query and predicted tokens before joint fine-tuning with the video generation model.

C.9 Ablation on Query Token Initialization … C.12 Additional Video Quality Evaluation: VBench and CLIP Similarity

Across four evaluations, MemLearner remains robust to query-token initialization, outperforms geometry-based and prior learned retrieval methods, transfers zero-shot to Epic-Kitchens, and leads on VBench and CLIP similarity.

  • C.9 Ablation on Query Token Initialization: The comparable initialization results support using randomly sampled noise by default, matching diffusion-model input distributions.MemLearner’s default Q-token initialization is randomly sampled noise.
  • C.9 Ablation on Query Token Initialization: Noise-initialized and noisy-P-token-initialized query tokens perform comparably in MemLearner.The Q-to-P attention enables query tokens to learn the predicted tokens’ needs at each forward pass.
  • C.10 Comparison with Geometry-Based Retrieval: MemLearner outperforms VRAG across all metrics, demonstrating an advantage over geometry-based context retrieval.VRAG uses geometric position cues to alleviate wall occlusion, whereas MemLearner learns context queries.
  • C.11 Zero-Shot Transfer to Real-World Epic-Kitchens: On zero-shot Epic-Kitchens transfer, MemLearner consistently outperforms CaM and VMem.Epic-Kitchens differs substantially from the training data, indicating generalization beyond the rendered occlusion and dynamic-object patterns.
  • C.12 Additional Video Quality Evaluation: VBench and CLIP Similarity: The additional evaluation measures five VBench dimensions and CLIP similarity to assess video quality and semantic consistency.The VBench dimensions are Background Consistency, Temporal Flickering, Motion Smoothness, Aesthetic Quality, and Imaging Quality.
  • C.12 Additional Video Quality Evaluation: VBench and CLIP Similarity: MemLearner consistently outperforms CaM and VMem across all reported VBench and CLIP metrics.These results confirm stronger video quality and semantic consistency than the compared methods.
  • C.12 Additional Video Quality Evaluation: VBench and CLIP Similarity: The additional qualitative comparison reports results for GT, Ours, CaM, DFoT, and FramePack.The figure is identified as additional qualitative comparison results.
Loading 2606.31734v1…