Technical Summary: MemLearner: Learning to Query Context Memory for Video World Models

Problem Statement

Video World Models aim to predict future world states based on historical states and user actions, offering a promising path for interactive video generation. However, a critical bottleneck remains: scene inconsistency over extended durations. Current models suffer from limited context windows, leading to memory loss where later generated frames contradict earlier ones.

Existing solutions attempt to address this via memory mechanisms such as 3D reconstruction, compressed feature representations, or context retrieval. While context retrieval is computationally efficient, prior methods rely on rule-based heuristics (e.g., Field of View overlap or point cloud estimation). These rigid rules fail to generalize in complex scenarios involving scene occlusions (e.g., walls blocking views between camera positions) and dynamic objects (e.g., moving characters), resulting in inconsistent reconstructions when the camera revisits a scene.

Methodology: MemLearner

The authors propose MemLearner, a paradigm shift from rule-based retrieval to a learning-based adaptive context query mechanism. Instead of hard-coding retrieval rules, MemLearner enables the network to learn how to query relevant information from historical frames end-to-end.

Core Architecture and Design

1. Query Tokens (Q Tokens): The method introduces learnable query tokens that act as an information bridge between Context Tokens (C tokens) (historical frames) and Predicted Tokens (P tokens) (future frames).

   • Mechanism: Q tokens attend to C tokens to adaptively extract relevant context information. P tokens then attend to Q tokens as the generation condition.
   • Integration: Crucially, MemLearner leverages the video generation model itself for context querying. Rather than introducing a separate, scratch-trained query module (which the authors found fails to learn useful information due to gradient propagation issues), the C, Q, and P tokens are fed jointly into the pre-trained Diffusion Transformer (DiT). This exploits the model's existing visual priors without requiring additional modules trained from scratch.

2. Efficient Strategies: To handle the computational cost of long video sequences, two efficiency strategies are employed:

   • Layer-Split: Context querying (interaction between C, Q, and P) occurs only in the shallow Query Layers (e.g., the first 5 layers). The deeper Generative Layers process only Q and P tokens, removing the expensive context tokens from the computation graph after the initial extraction.
   • Attention Pruning: Unnecessary attention patterns are removed. Specifically, Context tokens (C) are never used as queries; they only serve as keys/values. This significantly reduces overhead while preserving the ability of Q tokens to extract information from C and P tokens.

Dataset and Training Strategy

Recognizing the lack of existing datasets that simultaneously offer long durations, precise camera poses, occlusions, and dynamic objects, the authors:

• Collected a Custom Rendered Dataset: Using Unreal Engine, they generated 16.7 hours of video across 13 scenes with customized occlusions, dynamic objects (humans, animals), and automated camera trajectories with collision avoidance.
• Multi-Dataset Training Strategy: To leverage the realism of real-world data and the precision of rendered data, they propose a strategy using dedicated camera encoders for different data types. Rendered data (precise poses), estimated-pose real data, and unannotated real data are processed by separate encoders. Unannotated data is fed with zero camera parameters. This isolates annotation quality differences, allowing the model to benefit from diverse sources without interference.

Key Contributions

1. MemLearner Framework: A novel, learning-based adaptive context query method that utilizes query tokens to extract valid memory from context, significantly improving scene consistency and generalization in Video World Models.
2. Custom Dataset: The creation of a specialized video world model dataset based on Unreal Engine, explicitly designed with occlusions, dynamic objects, and revisit scenarios, paired with a multi-dataset training strategy to combine rendered and real-world data.
3. Empirical Validation: Extensive experiments demonstrating that MemLearner outperforms state-of-the-art baselines (including rule-based retrieval and separate module designs) in scene consistency, particularly under challenging occlusion and dynamic object scenarios.

Results

• Quantitative Performance: MemLearner achieves superior performance across all metrics (PSNR, LPIPS, FID, FVD) compared to baselines like CaM (rule-based FOV retrieval), VMem (point-cloud retrieval), and FramePack.
  • On the authors' custom dataset (with occlusions/dynamics), MemLearner significantly outperforms CaM (e.g., PSNR 21.23 vs. 19.85 in Ground Truth Comparison).
  • On the CaM dataset (without occlusions), performance is comparable, indicating the method does not degrade in simpler scenarios.
• Qualitative Results: Visual comparisons show MemLearner maintains consistent object appearance and scene layout during camera revisits, whereas baselines often fail to reconstruct occluded areas or dynamic objects correctly.
• Ablation Studies:
  • Separate Module Failure: A design using a separate context query module (Fig. 2b) failed to learn, performing worse than a standard text-to-video model, validating the necessity of leveraging the pre-trained DiT for querying.
  • Attention Patterns: Removing the pattern where Q tokens attend to P tokens (to determine what to query) caused significant performance degradation, confirming the importance of this interaction.
  • Camera Embeddings: The model learns to query useful information even without explicit camera pose inputs for context tokens, suggesting implicit learning of geometric correspondences.
  • Generalization: The method generalizes to open-source models (Wan 2.1) and zero-shot transfers to real-world datasets (Epic-Kitchens, SpatialVID) without retraining.

Significance and Claims

The paper claims that MemLearner represents a paradigm shift in context memory for Video World Models. By moving from hand-crafted, rule-based retrieval to end-to-end learned adaptive querying, the approach overcomes the fundamental limitations of handling occlusions and dynamic objects.

The authors emphasize that their method:

• Exploits Pre-trained Priors: It utilizes the visual knowledge already embedded in the video generation model, avoiding the need for training additional modules from scratch.
• Handles Complex Scenarios: It is the first method to effectively manage scene consistency in the presence of occlusions and dynamic objects, a critical requirement for realistic interactive world simulation.
• Scalability: The efficient strategies (layer-splitting and attention pruning) make the approach viable for long video generation without prohibitive computational costs.

The paper concludes that while MemLearner advances the field, challenges remain in scaling to larger models for complex multi-entity interactions and exploring context compression and selective forgetting for truly intelligent memory systems.