MemLearner: Learning to Query Context memory for Video World Models
Abstract Overview
This paper presents MemLearner, a memory mechanism for video world models that replaces rule-based context retrieval with a learned context-query process. The method introduces query tokens that bridge historical context tokens and future prediction tokens, and performs this querying inside the pretrained video generation model rather than with a separately trained module. The paper also contributes a rendered long-video dataset designed for revisits, occlusions, dynamic objects, and accurate camera poses, together with a multi-dataset training strategy that mixes rendered and real videos through separate camera encoders. Experiments evaluate scene consistency, memory, visual quality, efficiency, and robustness across synthetic and real-world settings.
Novelty
The main novelty is a learning-based adaptive context query mechanism for video world models, using query tokens inside the pretrained diffusion transformer itself instead of hand-crafted retrieval rules or a separately trained memory module. The work is also distinctive in pairing this mechanism with a new occlusion- and dynamics-focused long-video dataset and a multi-dataset training strategy for heterogeneous camera annotation quality.
Results
MemLearner achieves the best reported results in the paper’s main comparison table, improving both memory-related metrics and visual-quality metrics over prior methods on the authors’ challenging dataset. It is especially stronger than rule-based retrieval baselines under occlusion and dynamic-object scenarios, while on the simpler CaM dataset without those challenges the gains are smaller but still positive. Ablations further show that using the pretrained video model itself for querying is effective, whereas a separate scratch-trained query module fails to learn useful context conditioning.
Key Points
- MemLearner uses query tokens to adaptively extract useful historical information for future video generation, with querying performed within the pretrained video diffusion model.
- The paper introduces efficiency strategies that restrict context querying to early layers and remove unnecessary attention patterns to reduce training and inference cost.
- The proposed method is supported by a new long-video rendered dataset with occlusions, dynamic objects, revisits, and camera poses, and shows robust improvements across synthetic and real-world evaluations.
References
- arXiv: https://arxiv.org/abs/2606.31734v1
- Fugu-MT: https://fugumt.com/fugumt/paper_check/2606.31734v1
- Hugging Face Papers: https://huggingface.co/papers/2606.31734
- Project: https://yujiwen.github.io/memlearner/