MaskGWM : Un modèle généralisable du monde de la conduite avec reconstruction de masque vidéo

MaskGWM: A Generalizable Driving World Model with Video Mask Reconstruction

February 17, 2025

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

cs.AI

Résumé

Les modèles du monde capables de prédire les changements environnementaux à partir des actions sont essentiels pour les modèles de conduite autonome dotés d'une forte généralisation. Les modèles du monde de conduite prédominants s'appuient principalement sur des modèles de prédiction vidéo. Bien que ces modèles puissent produire des séquences vidéo de haute fidélité grâce à des générateurs avancés basés sur la diffusion, ils sont limités par leur durée de prédiction et leurs capacités de généralisation globales. Dans cet article, nous explorons la résolution de ce problème en combinant une perte de génération avec un apprentissage contextuel au niveau des caractéristiques de style MAE. Plus précisément, nous concrétisons cet objectif avec trois conceptions clés : (1) Une structure de Diffusion Transformer (DiT) plus évolutive entraînée avec une tâche supplémentaire de construction de masque. (2) Nous concevons des tokens de masque liés à la diffusion pour gérer les relations floues entre la reconstruction de masque et le processus de diffusion générative. (3) Nous étendons la tâche de construction de masque au domaine spatio-temporel en utilisant un masque par ligne pour l'auto-attention décalée plutôt que l'auto-attention masquée dans MAE. Ensuite, nous adoptons un module de vue croisée par ligne pour s'aligner sur cette conception de masque. Sur la base de ces améliorations, nous proposons MaskGWM : un modèle du monde de conduite généralisable intégrant une reconstruction vidéo par masque. Notre modèle comprend deux variantes : MaskGWM-long, axé sur la prédiction à long terme, et MaskGWM-mview, dédié à la génération multi-vues. Des expériences approfondies sur des benchmarks standards valident l'efficacité de la méthode proposée, comprenant une validation normale sur le dataset Nuscene, un déploiement à long terme sur le dataset OpenDV-2K et une validation en zero-shot sur le dataset Waymo. Les métriques quantitatives sur ces datasets montrent que notre méthode améliore notablement l'état de l'art des modèles du monde de conduite.

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.