MaskGWM: Un Modello Generale del Mondo di Guida con Ricostruzione di Maschere Video

MaskGWM: A Generalizable Driving World Model with Video Mask Reconstruction

February 17, 2025

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

cs.AI

Abstract

I modelli del mondo che prevedono i cambiamenti ambientali a partire dalle azioni sono fondamentali per i modelli di guida autonoma con una forte capacità di generalizzazione. I modelli del mondo di guida predominanti si basano principalmente su modelli di previsione video. Sebbene questi modelli possano produrre sequenze video ad alta fedeltà grazie a generatori avanzati basati su diffusione, sono limitati dalla durata predittiva e dalle capacità complessive di generalizzazione. In questo articolo, esploriamo la risoluzione di questo problema combinando la perdita di generazione con l'apprendimento contestuale a livello di feature in stile MAE. In particolare, concretizziamo questo obiettivo con tre elementi chiave: (1) Una struttura più scalabile di Diffusion Transformer (DiT) addestrata con un'ulteriore attività di costruzione di maschere. (2) Progettiamo token di maschera correlati alla diffusione per gestire le relazioni sfumate tra la ricostruzione delle maschere e il processo di diffusione generativa. (3) Estendiamo l'attività di costruzione delle maschere al dominio spazio-temporale utilizzando maschere per righe per l'attenzione auto-attenzionale spostata, piuttosto che l'attenzione auto-attenzionale mascherata come in MAE. Successivamente, adottiamo un modulo cross-view per righe per allinearsi a questo design di maschera. Sulla base di questi miglioramenti, proponiamo MaskGWM: un modello del mondo di guida generalizzabile che incorpora la ricostruzione video con maschere. Il nostro modello include due varianti: MaskGWM-long, focalizzata sulla previsione a lungo termine, e MaskGWM-mview, dedicata alla generazione multi-vista. Esperimenti completi su benchmark standard convalidano l'efficacia del metodo proposto, che include la validazione normale del dataset Nuscene, il rollout a lungo termine del dataset OpenDV-2K e la validazione zero-shot del dataset Waymo. Le metriche quantitative su questi dataset dimostrano che il nostro metodo migliora significativamente lo stato dell'arte dei modelli del mondo di guida.

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.