MaskGWM: Um Modelo Generalizável de Mundo de Condução com Reconstrução de Máscara de Vídeo

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

Resumo

Modelos de mundo que preveem mudanças ambientais a partir de ações são essenciais para modelos de direção autônoma com forte generalização. O modelo de mundo de direção predominante é principalmente baseado em modelos de previsão de vídeo. Embora esses modelos possam produzir sequências de vídeo de alta fidelidade com geradores avançados baseados em difusão, eles são limitados por sua duração preditiva e capacidades gerais de generalização. Neste artigo, exploramos resolver esse problema combinando perda de geração com aprendizado de contexto em nível de características no estilo MAE. Em particular, instanciamos esse objetivo com três projetos principais: (1) Uma estrutura mais escalável de Transformador de Difusão (DiT) treinada com uma tarefa adicional de construção de máscaras. (2) Criamos tokens de máscara relacionados à difusão para lidar com as relações difusas entre a reconstrução de máscaras e o processo de difusão generativa. (3) Estendemos a tarefa de construção de máscaras para o domínio espaço-temporal utilizando máscaras em linha para auto-atenção deslocada, em vez de auto-atenção mascarada no MAE. Em seguida, adotamos um módulo de visão cruzada em linha para alinhar com esse design de máscara. Com base nas melhorias acima, propomos o MaskGWM: um Modelo de Mundo de Direção Generalizável incorporado com Reconstrução de Máscara de Vídeo. Nosso modelo contém duas variantes: MaskGWM-long, focada em previsão de longo prazo, e MaskGWM-mview, dedicada à geração multi-visão. Experimentos abrangentes em benchmarks padrão validam a eficácia do método proposto, que inclui validação normal do conjunto de dados Nuscene, rollouts de longo prazo do conjunto de dados OpenDV-2K e validação zero-shot do conjunto de dados Waymo. Métricas quantitativas nesses conjuntos de dados mostram que nosso método melhora significativamente o estado da arte em modelos de mundo de direção.

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.