Abstract
World models that forecast environmental changes from actions are vital for autonomous driving models with strong generalization. The prevailing driving world model mainly build on video prediction model. Although these models can produce high-fidelity video sequences with advanced diffusion-based generator, they are constrained by their predictive duration and overall generalization capabilities. In this paper, we explore to solve this problem by combining generation loss with MAE-style feature-level context learning. In particular, we instantiate this target with three key design: (1) A more scalable Diffusion Transformer (DiT) structure trained with extra mask construction task. (2) we devise diffusion-related mask tokens to deal with the fuzzy relations between mask reconstruction and generative diffusion process. (3) we extend mask construction task to spatial-temporal domain by utilizing row-wise mask for shifted self-attention rather than masked self-attention in MAE. Then, we adopt a row-wise cross-view module to align with this mask design. Based on above improvement, we propose MaskGWM: a Generalizable driving World Model embodied with Video Mask reconstruction. Our model contains two variants: MaskGWM-long, focusing on long-horizon prediction, and MaskGWM-mview, dedicated to multi-view generation. Comprehensive experiments on standard benchmarks validate the effectiveness of the proposed method, which contain normal validation of Nuscene dataset, long-horizon rollout of OpenDV-2K dataset and zero-shot validation of Waymo dataset. Quantitative metrics on these datasets show our method notably improving state-of-the-art driving world model.
Abstract (translated)
世界模型，用于从行动预测环境变化，在具有强泛化能力的自动驾驶模型中至关重要。当前主流的驾驶世界模型主要建立在视频预测模型基础上。尽管这些模型能够利用先进的扩散生成器产生高保真度的视频序列，但它们受制于预测时长和整体泛化的局限性。本文探索通过结合生成损失与类似MAE（掩码图像建模）特征级上下文学习的方法来解决这一问题。 具体而言，我们采用以下三个关键设计来实现这一目标： 1. 一个更可扩展的Diffusion Transformer (DiT) 结构，并且在训练过程中加入了额外的掩码构建任务。 2. 设计了与扩散过程相关的掩码令牌以处理掩码重建和生成式扩散过程之间的模糊关系。 3. 将掩码构建任务扩展到时空域，通过使用行级掩码进行移位自我注意机制（而不是MAE中的屏蔽自我注意），并采用行级跨视图模块来与此设计对齐。 基于上述改进，我们提出了MaskGWM：一个具备视频掩码重建的可泛化驾驶世界模型。该模型包含两种变体： - MaskGWM-long: 专注于长时预测。 - MaskGWM-mview: 面向多视角生成。 在标准基准上的全面实验验证了所提方法的有效性，包括Nuscene数据集的标准验证、OpenDV-2K数据集的长时滚动和Waymo数据集的零样本验证。这些数据集上的定量指标显示我们的方法显著提升了最先进的驾驶世界模型的表现。
URL
https://arxiv.org/abs/2502.11663
