商汤力作！MaskGWM：自动驾驶世界模型再进化~

点击下方卡片，关注“自动驾驶之心”公众号

戳我-> 领取自动驾驶近15个方向学习路线

今天自动驾驶之心为大家分享商汤科技最新的工作！基于视频掩码重建的通用驾驶世界模型—MaskGWM！如果您有相关工作需要分享，请在文末联系我们！

自动驾驶课程学习与技术交流群事宜，也欢迎添加小助理微信AIDriver004做进一步咨询

>>点击进入→自动驾驶之心『世界模型』技术交流群

论文作者 | Jingcheng Ni等

编辑 | 自动驾驶之心

写在前面&出发点

预测环境变化的世界模型对于具有强大泛化能力的自动驾驶模型至关重要。现有的驾驶世界模型主要基于视频预测模型构建。尽管这些模型能够利用先进的基于扩散的生成器生成高保真视频序列，但它们的预测时长和整体泛化能力受到限制。为解决这一问题，本文探索将生成损失与MAE风格的特征级上下文学习相结合。具体而言，我们通过三个关键设计实现这一目标：（1）采用更具扩展性的扩散变换器（DiT）结构，并额外训练掩码构建任务；（2）设计与扩散相关的掩码令牌，以处理掩码重建与生成扩散过程之间的模糊关系；（3）通过使用行级别的掩码进行移位自注意力而非MAE中的掩码自注意力，将掩码构建任务扩展到时空领域。随后，我们采用行式跨视图模块以与这一掩码设计对齐。基于上述改进，我们提出了MaskGWM：一种融合视频掩码重建的泛化驾驶世界模型。我们的模型包含两个变体：专注于长期预测的MaskGWM-long和致力于多视图生成的MaskGWM-mview。在标准基准上的综合实验验证了所提出方法的有效性，包括Nuscene、OpenDV-2K和Waymo数据集的零样本验证。这些数据集上的定量指标显示，我们的方法显著提升了现有最先进驾驶世界模型的性能。

行业介绍

作为人工智能的关键应用之一，自动驾驶技术需要理解周围环境并执行正确动作。然而，对于现有方法而言，在复杂多变场景中泛化能力有限的问题仍未解决。例如，在天气变化、场景变异和运动模糊等情况中，感知性能可能会下降。解决这一问题的有效方法是使用世界模型，该模型可直接预测不同动作下的环境变化。这些模型有助于理解数据分布的复杂性，并像人类感知系统一样构建复杂的规律模式。

得益于基于扩散的图像生成系统的快速发展，最近一系列方法基于扩散模型构建世界模型。尽管能够生成高保真结果，这些方法仍面临长时预测和泛化能力的挑战。在这篇文章中，我们认为在构建自动驾驶通用世界模型方面仍存在两个问题。首先，大规模训练数据集与更具扩展性的变压器架构的结合仍在探索中。其次，一个基本问题仍未得到解答：基于扩散的生成损失是否足以构建一个通用的世界模型？由于扩散损失在像素级别操作，视觉语义的学习可能并不直接。另外，一些工作已表明扩散模型与自监督方法是互补的，有利于提高收敛速度和生成质量。基于上述分析，我们设计了MaskGWM，旨在提高现有方法的保真度、泛化能力和长时间序列预测能力。此外，通过集成多视图模块，我们的模型还能生成多视图案例。我们采用DiT作为主干网络，它更具扩展性，能够利用多种数据集的信息。此外，我们将掩码重建作为生成的补充任务。此外，通过自监督学习获得的特征在语义上更具意义，可以作为辅助监督进一步改善生成质量。然而，将现有的图像生成掩码重建集成到驾驶世界模型中并非直接可行，仍有两个问题需要解决：（1）如何增强扩散模型与掩码重建之间的协同作用。（2）我们应该使用哪种掩码策略来处理视频数据。

因此，MaskGWM引入如下策略来解决上述问题：（1）我们利用掩码标记来改善掩码重建与扩散模型之间的协同作用。（2）我们设计了一种针对时空域的双分支掩码重建策略。对于空间建模，我们使用在所有帧间共享的掩码，并通过空间变压器重建不可见标记。对于时间建模，我们引入各帧独立的掩码，并通过时间变压器恢复掩码标记。与空间分支不同，我们在掩码后直接连接时间维度上的未对齐标记，这相当于在行级别引入小范围的偏移。我们发现这个时间分支实现了在时间语境中预测掩码斑块和减少训练成本的双重目标。

网络模型结构&技术细节

在详细介绍了本文提出的算法模型技术细节之前，下图展示了我们提出的MaskGWM算法模型的整体网络结构，如下图所示。

掩码重建技术。 在训练环节，DiT网络被划分为编码器 和解码器 。编码时，假设时间步 的视频潜变量为 ，我们随机生成一个二元掩码 将 分为可见部分 和不可见部分 。在解码阶段，我们引入重建模块来处理不可见部分，对于不可见部分，我们统一使用掩码令牌 初始化其特征表示。具体来说，我们通过学习一个小型的transformer  得到与初始维度相同的特征，具体公式为 。为了强化生成过程与掩码重建之间的协同效应，我们提出了与扩散过程紧密相关的掩码策略。在扩散模型框架下，掩码重建可被视为 预测任务，而矫正流则采用 预测（即预测） 。为了平衡这两个预测目标，我们引入了额外的噪声编码网络 。由于 是明确给出的，因此更容易在扩散过程中恢复原始的掩码重建目标。此外，我们还考虑了时间步 的影响，将掩码令牌定义为 .

掩码策略。考虑到时空域的特性，我们设计了两种掩码： 和 ，前者关注空间上下文，掩码位置在不同帧共享；后者关注时序动态的建模，掩码位置在各帧独立；为了让两种不同的掩码关注各自的域，我们将 设计为双分支transformer，对 使用空间自注意力模块编码，对 使用时间自注意力模块编码：

对于时间域的掩码，我们进一步的修改了时序自注意力机制，对于掩码的位置，我们直接将其重新紧凑化，由于每一帧被掩码跳过的令牌数相同，重新紧凑化的特征图大小在不同帧是统一的。然而，引入这一策略会导致时序自注意力机制的不对齐，我们认为只要这种不对齐只要是发生在小范围内的，就对模型的性能有正向的帮助，因此我们引入行级别的掩码设计，具体来说对于 ，每行掩码掉的令牌数是相同的。这一掩码设计能够带来如下两个好处，首先被丢弃的令牌可以直接在计算中跳过，导致更高的训练效率，其次，我们在实验中发现这一设计的效果会更好。

多视角设计。为了验证模型的泛化能力，除了在关注于预测时长的单视角任务外，我们还额外衡量了模型在多视角任务上的性能，为了支持这一能力，我们引入行级别的跨视角注意力机制，可以与之前的掩码策略对齐，使得我们可以在多视角任务中继续使用提出的掩码策略。

实验分析

我们展示了一系列广泛的实验来评估模型的性能，我们使用SD3作为基础模型，使用FID和FVD作为评测指标，在Nuscenes，OpenDV和Waymo数据集上评估了模型的性能：

1）与最先进的方法的比较

表1将研究方法与单/多视角最先进的方法进行了比较，同时我们对单/多视角分别对齐了评估策略，研究方法在图像生成质量和视频生成质量均优于最先进方法。

2）消融实验

表3(a)比较了不同掩码令牌设计策略的效果，表3(b)比较了不同掩码策略的效果，实验结果显示，我们的模型在生成指标上有着显著的提升；表4(a)进一步地分析了行级别的掩码设计的影响，表4(b)验证了双分支使用独立的自注意力模块编码的影响，表4(c)验证了多视角模块的效果。

① 自动驾驶论文辅导来啦

② 国内首个自动驾驶学习社区

『自动驾驶之心知识星球』近4000人的交流社区，已得到大多数自动驾驶公司的认可！涉及30+自动驾驶技术栈学习路线，从0到一带你入门自动驾驶感知（端到端自动驾驶、世界模型、仿真闭环、2D/3D检测、语义分割、车道线、BEV感知、Occupancy、多传感器融合、多传感器标定、目标跟踪）、自动驾驶定位建图（SLAM、高精地图、局部在线地图）、自动驾驶规划控制/轨迹预测等领域技术方案、大模型，更有行业动态和岗位发布！欢迎扫描加入

 ③全网独家视频课程

端到端自动驾驶、仿真测试、自动驾驶C++、BEV感知、BEV模型部署、BEV目标跟踪、毫米波雷达视觉融合、多传感器标定、多传感器融合、多模态3D目标检测、车道线检测、轨迹预测、在线高精地图、世界模型、点云3D目标检测、目标跟踪、Occupancy、CUDA与TensorRT模型部署、大模型与自动驾驶、NeRF、语义分割、自动驾驶仿真、传感器部署、决策规划、轨迹预测等多个方向学习视频（扫码即可学习）

网页端官网：www.zdjszx.com

④【自动驾驶之心】全平台矩阵