MaskGWM: Un Modelo Generalizable del Mundo de Conducción con Reconstrucción de Máscaras de Video

MaskGWM: A Generalizable Driving World Model with Video Mask Reconstruction

February 17, 2025

Autores: Jingcheng Ni, Yuxin Guo, Yichen Liu, Rui Chen, Lewei Lu, Zehuan Wu

cs.AI

Resumen

Los modelos del mundo que predicen cambios ambientales a partir de acciones son fundamentales para los modelos de conducción autónoma con una fuerte generalización. El modelo predominante del mundo de conducción se basa principalmente en modelos de predicción de video. Aunque estos modelos pueden producir secuencias de video de alta fidelidad con generadores avanzados basados en difusión, están limitados por su duración predictiva y capacidades generales de generalización. En este artículo, exploramos resolver este problema combinando la pérdida de generación con el aprendizaje de contexto a nivel de características al estilo MAE. En particular, instanciamos este objetivo con tres diseños clave: (1) Una estructura más escalable de Transformador de Difusión (DiT) entrenada con una tarea adicional de construcción de máscaras. (2) Diseñamos tokens de máscara relacionados con la difusión para manejar las relaciones difusas entre la reconstrucción de máscaras y el proceso de difusión generativa. (3) Extendemos la tarea de construcción de máscaras al dominio espacio-temporal utilizando máscaras por filas para la autoatención desplazada en lugar de la autoatención enmascarada en MAE. Luego, adoptamos un módulo de vista cruzada por filas para alinearnos con este diseño de máscara. Basándonos en las mejoras anteriores, proponemos MaskGWM: un Modelo del Mundo de Conducción Generalizable que incorpora la Reconstrucción de Máscaras de Video. Nuestro modelo contiene dos variantes: MaskGWM-long, enfocado en la predicción a largo plazo, y MaskGWM-mview, dedicado a la generación multivista. Experimentos exhaustivos en puntos de referencia estándar validan la efectividad del método propuesto, que incluye la validación normal del conjunto de datos Nuscene, el despliegue a largo plazo del conjunto de datos OpenDV-2K y la validación de cero disparos del conjunto de datos Waymo. Las métricas cuantitativas en estos conjuntos de datos muestran que nuestro método mejora notablemente el estado del arte en modelos del mundo de conducción.

English

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.