2025-07-17 超越SOTA近40%!西交I2-World:超强OCC世界模型实现3G训练显存37 FPS推理

原文链接:https://mp.weixin.qq.com/s/AKpQOf3vMe1p4zv8pCAG2Q

基于OCC的世界模型通过预测3D场景的变化并生成新场景,在解决自动驾驶系统中的corner case方面具有巨大潜力。尽管tokenization已经彻底改变了图像和视频生成领域,但如何高效地对复杂的3D场景进行分词化仍然是3D世界模型的一个关键挑战。为此西交的团队提出了I2-World,一种高效的4D OCC预测框架。本文的方法将场景分词化解耦为场景内分词器和场景间分词器。场景内分词器采用多尺度残差量化策略,以层次结构压缩3D场景同时保留空间细节。场景间分词器则残差聚合timestep之间的依赖关系。这种双重设计在保持3D分词器紧凑性的同时保留了4D分词器的动态表达能力。与decoder-only的GPT式自回归模型不同,I2-World采用了编码器-解码器架构。编码器从当前场景中聚合空间上下文,并预测一个变换矩阵以实现对场景生成的高层次控制。解码器在此矩阵和历史token的基础上,确保生成过程中的时间一致性。实验表明I2-World在4D OCC预测任务中达到了最先进的性能,在mIoU和IoU方面分别比现有方法高出25.1%和36.9%,同时表现出卓越的计算效率:只需要2.9 GB的训练显存,并且实现了37 FPS的实时推理速度。

  • 论文链接:https://arxiv.org/abs/2507.09144
  • 开源代码:https://github.com/lzzzzzm/II-World

引言

3D OCC 提供了更多关于 3D 场景的几何和细节信息,使其比传统的 3D 边界框和点云更适合自动驾驶系统。随着生成式人工智能的发展,基于占据空间的世界模型展示了其作为世界仿真器的潜力,并能够在复杂交通场景中解决corner case。

受图像和视频生成中分词化技术进步的启发,将 3D 场景进行分词化已经成为开发 3D 世界模型的基础步骤。然而,在保留时空动态的同时高效地将 3D 场景压缩为紧凑的标记仍然是一个关键挑战。如图 1a 所示,现有的分词器分为两类:3D 场景分词器和 4D 场景分词器。3D 场景分词器将单个 3D 场景压缩成高度紧凑的潜在表示形式。虽然这些分词器能够实现准确的场景重建,但它们无法建模时间动态,严重限制了它们的预测性能。相比之下,4D 场景分词器——受视频生成框架启发,直接处理 4D 时空标记,将动态场景演变嵌入到标记空间中。尽管这种方法提高了预测保真度,但由此产生的高维标记给下游自回归模型或扩散框架带来了巨大的计算开销,使其在对耗时敏感的应用(如自动驾驶)中不切实际。

本文提出了 I2-World,一种高效的 4D OCC预测世界模型。它主要包括两个关键组件:I2-Scene Tokenizer,一种高压缩但动态保持的分词方法,解决了之前场景分词器的局限性;以及 I2-Former,一种变换矩阵引导的自回归框架。

I2-Scene Tokenizer 将分词过程解耦为两个互补的组件:场景内分词器和场景间分词器。场景内分词器采用多尺度分词策略,专注于捕捉当前场景中的细粒度细节和静态上下文。相比之下,场景间分词器维护一个内存队列来存储历史场景标记,并沿时间戳使用时间量化来建模场景内的动态运动。通过集成这些组件,本文的 I2-Scene Tokenizer 产生了与 3D 分词器相当的高压缩标记,同时保留了 4D 分词器的时间建模能力。

为了结合提出的 I2-Scene Tokenizer,本文引入了 I2-Former,这是一种不同于仅解码器的 GPT 类自回归模型的混合架构。相反,I2-Former 包括两个核心组件:场景内编码器和场景间解码器。场景内编码器层次化聚合来自当前场景标记的空间上下文。同时,它利用多种自车动作的嵌入来预测一个变换矩阵,该矩阵将当前场景表示映射到下一个时间步。然后,这个矩阵作为条件指导信号,用于场景间解码器,后者动态整合历史和当前场景标记以预测下一个时间步的场景标记。通过解耦静态上下文聚合和动态时间推理,I2-Former 实现了计算效率和稳健的预测准确性。

I2-Former 的解耦设计实现了精细的可控生成。用户可以通过操纵变换矩阵或调整各种控制动作(例如,速度、转向命令)直观地引导场景预测,从而灵活适应不同的驾驶场景。

总结来说,本文的主要贡献可以总结如下:

  • 提出了 I2-Scene Tokenizer,这是一种新颖的 3D 场景分词框架,在保留 3D 分词器计算效率的同时实现了 4D 方法的时间表达能力。
  • 介绍了 I2-Former,这是一种遵循编码器-解码器范式的自回归架构。解码器的生成依赖于一个变换矩阵,从而实现更高保真度且更详细的场景生成。
  • I2-World 在 Occ3D 基准测试中建立了新的最先进水平,在 mIoU 上比之前表现最好的方法提升了 25.1%,在 IoU 上提升了 42.9%。此外,它达到了前所未有的计算效率,训练内存仅需 2.9 GB,推理速度达到每秒 37 FPS。

相关工作

世界模型利用智能体的动作和历史观测来预测未来3D场景的演变,使智能体能够推理动态环境。基于占据空间的世界模型的最新进展探索了多种3D场景分词化和生成策略。像OccWorld这样的方法采用基于VQ-VAE的分词化技术结合自回归变换器来预测未来场景,而RenderWorld引入了一种AM-VAE,将空体素和非空体素分离以进行结构化生成。语言增强的方法如Oc-cLLaMA和Occ-LLM将3D占据标记与文本场景描述相结合,通过大语言模型引导生成。其他工作如DriveWorld和UniWorld利用4D占据重建进行预训练,而Drive-OccWorld则加入了动作条件以实现可控的场景合成。为了提高效率,DFIT-OccWorld采用了非自回归流水线,而OccSora和DOME则使用4D分词器与扩散模型进行长时程生成。UniScene提出了一种时间感知的VAE,大幅提升了重建性能。尽管有这些进展,现有方法在优化分词化效率和动态保真度方面仍面临挑战。

算法详解

I2-World整体架构

I2-Scene Tokenizer

I2-Former

实验结果分析

实现细节

与现有工作一致,本文利用2秒的历史数据来预测未来3秒的占据空间。本文提出的I2-Scene Tokenizer使用一个包含512个 entries和128维潜在特征的共享码本,以编码紧凑的时空表示。训练采用AdamW优化器,基础学习率为1 × 10^-3,全局batch大小为128,并为每个组件制定了不同的训练计划:I2-Scene Tokenizer训练24个周期,而I2-Former则训练48个周期。在评估中,本文使用真实变换矩阵来引导生成,并采用mIoU和IoU作为4D占据空间预测任务的评估指标。FPS的计算方法遵循OccWorld,并在RTX 4090 GPU上进行测量。

主要结果

4D占据空间预测

如表1所示,I2-World-O在当前SOTA的方法基础上取得了显著的提升,在mIoU方面提升了25.1%(39.73 vs. 31.75),在IoU方面提升了36.9%(49.80 vs. 36.36)。值得注意的是,尽管现有解决方案依赖于资源密集型架构(如大语言模型或扩散模型),本文的系统仍保持了极高的效率——仅需2.9 GB的训练内存,并在37.04 FPS下实现实时性能。端到端的I2-World-STC变体也表现出特别有希望的结果,在mIoU指标上比之前的方法高出50.9%(18.97 vs. 12.57),在IoU指标上高出40.9%(28.77 vs. 20.41)。

泛化能力

为了评估本文方法的泛化能力,本文在Occ3D-Waymo数据集上进行了零样本的4D占据空间预测。本文的复制粘贴方法利用当前帧的重建进行预测。如表2所示,本文在不同采样率下评估了本文的方法。结果表明,本文的方法在Waymo数据集上表现优异,显示出其作为自动标注方法的潜力。

消融研究

I2 Scene Tokenizer的有效性

如表3所示,本文对I2 Scene Tokenizer进行了消融研究,以评估每个设计组件的贡献。本文的基线采用来自OccWorld的单尺度Scene Tokenizer,并通过翻转BEV平面来增强分词器的泛化能力。Inter-Scene分词模块改善了时间建模,带来了5.7%的mIoU提升和1.8%的IoU提升,而对齐策略进一步增强了时间信息建模,带来了15.9%的mIoU提升和5.1%的IoU提升。值得注意的是,Inter-Scene分词作为一个即插即用模块,可以与其他分词器兼容。通过将Inter-Scene分词与多尺度Intra-Scene分词结合,本文实现了最佳性能。结果表明,解耦的时空分词能够有效应对3D场景理解的不同方面,同时实现高效、协同的整体建模。

I2 Former的有效性

如表4所示,本文分析了Inter-Decoder和Intra-Encoder中各个组件的贡献。首先,本文评估了一个基线配置,其中普通的Inter-Decoder在无条件的情况下预测未来的占据空间,取得了17.1%的mIoU和27.2%的IoU。这一结果突出了潜在标记在编码动态场景信息方面的固有能力,尽管预测准确性有限。接下来,本文引入平移作为条件信号,而旋转固定为起始帧。这一修改达到了67.8%的mIoU和31.4%的IoU,比基线提升了50.7%的mIoU和4.2%的IoU。相比之下,仅引入旋转作为条件信号仅取得18.7%的mIoU和5.4%的IoU,表明单独的旋转对性能的贡献微乎其微。最后,结合平移和旋转条件带来了显著的提升,达到了85.3%的mIoU和38.9%的IoU,这强调了这些空间变换在建模场景动态中的互补作用。Intra-Encoder、轻量级时间融合和多尺度策略仅带来GPU内存的有限增加,却带来了16.0%的mIoU提升和16.6%的IoU提升,证明了本文设计的效率和有效性。

可视化

4D占据空间预测

在图3中,本文可视化了Occ3D-nuScenes验证集上的4D占据空间预测结果。与DOME相比,本文的方法在3D场景中更好地保持了空间一致性,并更准确地仿真了移动目标的运动,从而展示了本文方法在时间动态建模方面的优越能力。

可控生成

图4通过高层命令和细粒度变换展示了本文框架的可控性。对于高层控制,不同的命令(例如“左转/右转”)生成6秒的场景,其中自我车辆的动作动态影响周围代理。例如,“右转”命令导致卡车碰撞,展示了本文模型仿真复杂交互和边缘情况的能力。细粒度控制通过直接操作变换矩阵实现,允许精确的场景生成(米/弧度级精度)。底行展示了一个失败案例:由变换矩阵(训练期间未见过)表示的倒车场景导致不现实的静态代理行为。本文认为这个问题可以通过在数据集中多样化变换矩阵的分布来缓解。

结论

本文提出了一种新的3D场景分词框架用于3D场景生成。本文的I2-Scene Tokenizer实现了高效的3D场景压缩,生成紧凑的潜在表示,保留了多尺度空间细节和时间动态,从而在降低计算成本的同时实现高保真生成。互补的I2-Former框架引入了基于变换矩阵的生成,实现了细粒度控制,同时保持生成质量。实验验证了本文框架的有效性,并展示了其作为自动化场景标注解决方案的潜力。

参考

[1]-World: Intra-Inter Tokenization for Efficient Dynamic 4D Scene Forecasting

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注


往期评论