MaskGWM: Универсальная модель вождения с реконструкцией видео через маскирование

MaskGWM: A Generalizable Driving World Model with Video Mask Reconstruction

February 17, 2025

Авторы: Jingcheng Ni, Yuxin Guo, Yichen Liu, Rui Chen, Lewei Lu, Zehuan Wu

cs.AI

Аннотация

Модели мира, прогнозирующие изменения окружающей среды на основе действий, имеют ключевое значение для моделей автономного вождения с высокой способностью к обобщению. Преобладающие модели мира для вождения в основном строятся на основе моделей прогнозирования видео. Хотя эти модели могут создавать высококачественные видеопоследовательности с использованием современных генераторов на основе диффузии, они ограничены продолжительностью прогнозирования и общими способностями к обобщению. В данной работе мы исследуем решение этой проблемы путем объединения потерь генерации с обучением контекста на уровне признаков в стиле MAE. В частности, мы реализуем эту цель с помощью трех ключевых элементов: (1) Более масштабируемая структура Diffusion Transformer (DiT), обученная с дополнительной задачей маскирования. (2) Мы разрабатываем токены маски, связанные с диффузией, для работы с нечеткими отношениями между реконструкцией маски и процессом генеративной диффузии. (3) Мы расширяем задачу маскирования на пространственно-временную область, используя построчные маски для сдвинутого самовнимания вместо маскированного самовнимания в MAE. Затем мы применяем построчный кросс-видовой модуль для согласования с этим дизайном маски. На основе указанных улучшений мы предлагаем MaskGWM: обобщаемую модель мира для вождения, воплощенную с реконструкцией видео-маски. Наша модель включает два варианта: MaskGWM-long, ориентированный на долгосрочное прогнозирование, и MaskGWM-mview, предназначенный для генерации с нескольких видов. Комплексные эксперименты на стандартных бенчмарках подтверждают эффективность предложенного метода, включая обычную валидацию на наборе данных Nuscene, долгосрочное прогнозирование на наборе данных OpenDV-2K и валидацию с нулевым обучением на наборе данных Waymo. Количественные метрики на этих наборах данных показывают, что наш метод значительно улучшает современные модели мира для вождения.

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.