MaskGWM: Een generaliseerbaar rijwereldmodel met videomaskerreconstructie

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

Samenvatting

Wereldmodellen die omgevingsveranderingen voorspellen op basis van acties zijn cruciaal voor autonome rijsystemen met sterke generalisatie. De heersende rijsimulatiemodellen zijn voornamelijk gebaseerd op videovoorspellingsmodellen. Hoewel deze modellen hoogwaardige videosequenties kunnen genereren met geavanceerde op diffusie gebaseerde generatoren, worden ze beperkt door hun voorspellingsduur en algehele generalisatievermogen. In dit artikel onderzoeken we een oplossing voor dit probleem door het combineren van generatieverlies met MAE-stijl contextleren op featureniveau. Concreet realiseren we dit doel met drie belangrijke ontwerpen: (1) Een schaalbaarder Diffusion Transformer (DiT)-structuur getraind met een extra maskerconstructietaak. (2) We introduceren diffusiegerelateerde maskertokens om de vage relaties tussen maskerreconstructie en het generatieve diffusieproces aan te pakken. (3) We breiden de maskerconstructietaak uit naar het ruimtelijk-temporele domein door rijgewijze maskers te gebruiken voor verschoven zelf-attentie in plaats van gemaskeerde zelf-attentie zoals in MAE. Vervolgens passen we een rijgewijze cross-view module aan om dit maskerontwerp te ondersteunen. Op basis van deze verbeteringen stellen we MaskGWM voor: een generaliseerbaar rijsimulatiemodel geïmplementeerd met videomaskerreconstructie. Ons model bevat twee varianten: MaskGWM-long, gericht op langetermijnvoorspelling, en MaskGWM-mview, toegewijd aan multiview-generatie. Uitgebreide experimenten op standaardbenchmarks valideren de effectiviteit van de voorgestelde methode, waaronder normale validatie van de Nuscene-dataset, langetermijnrollouts van de OpenDV-2K-dataset en zero-shot-validatie van de Waymo-dataset. Kwantitatieve metingen op deze datasets tonen aan dat onze methode de state-of-the-art rijsimulatiemodellen aanzienlijk verbetert.

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.