DiffusionDrive论文解读（快速理解版）

原文：DiffusionDrive: Truncated Diffusion Model for End-to-End Autonomous Driving

作者单位：华中科技大学 & 地平线

发表时间：2024.11

项目地址：GitHub - hustvl/DiffusionDrive: Truncated Diffusion Model for Real-Time End-to-End Autonomous Driving

Disclamer

这篇总结仅仅列出了论文中的精华部分：作者的动机，发现的问题，以及如何解决这些问题。其中有一些个人的理解一般会放在括号或者引用字体中，如有错误还望和大家一起讨论更正!

简介

本文将diffusion的网络结构的思想引入到端到端自动驾驶的任务中，并着重解决了

<1> 多步降噪带来的计算量问题

<2>模态坍缩的问题

方案亮点：

• 加入了prior multi-mode anchors来减少diffusion中的计算量，让模型能够从anchored Gaussian Distribution里面学习到denoising从而得到最终的多模态轨迹规划。

解读：这里的改进实际上是把random Gaussian noise改成了基于“预先设计好的多条可能的轨迹线”（也就是文中多次提到的anchors）来采样gaussian noise。相当于从源头上输出上就缩小了模型的搜索空间。

• 级联(cascade)的diffusion decoder，加强模型与场景上下文的互动

效果：

• 在降噪阶段有10倍的速度提升，只需要2步去噪
• 在planning-oriented NAVSIM数据集上，达到88.1PDMS
• 在4090显卡上的推理速度为45 FPS

Introduction：

之前的一些work(Transfuser, UniAD, VAD)直接去回归 一条确定的自车轨迹 ，没有考虑到驾驶行为的不确定性和多模态性。直到VADv2中，把连续的动作空间 离散化 成 一系列固定的anchor轨迹 （4096个anchors），然后在预测出的置信度上去采样这些anchors，来得到多模态的轨迹预测。然而这样做的缺点是，“预先定义好的anchor轨迹”这本身就限制了预测轨迹的数量和质量——简而言之，无法覆盖全部场景。

与之相比，diffusion模型可以通过去噪过程，基于高斯分布在连续的动作空间来采样多模态的驾驶动作。

一些早期的尝试：直接在Transfuser的模型基础上，把单模态轨迹回归的方式替换成了简单的diffusion策略，（确定的mlp-regression替换成conditional diffusion model）但是发现了两个问题：(1) 20步的去噪步骤带来了极大的计算量（延时从0.2ms --> 6.5ms） (2) 从高斯噪声采样得来的轨迹之间有非常严重的overlap，甚至坍缩成了单模态输出。

由此引出了改进点：将高斯分布分割成多个以prior anchors为中心的sub-Gaussian distributions——叫做anchored Gaussian distribution。利用truncated diffusion policy，把去噪的步骤从20降低到了2

为了加强conditional scene context，使用了基于transformer的diffusion decoder，会与感知模块的structured queries，BEV，PV features一起交互，通过一个sparse deformable attention机制。另外，作者引入了一个cascade的机制来在每一步去噪步骤对轨迹重建进行refine。

主要贡献点：

1. 第一次在端到端自动驾驶领域引入了diffusion模型，并提出了一个truncated diffusion策略来解决原始diffusion模型中的“模态坍缩”和大计算量的问题。
2. 设计了一个高效的diffusion decoder与conditional information进行交互，达到了更好的重建效果。
3. 在NAVSIM数据集上达到了88.1 PDMS的跑分，在4090显卡上的推理速度为45 FPS
4. 展示了DiffusionDrive可以生成多模态且可行的轨迹

Method

3.1 Diffusion简介

Conditional diffusion model

一种forward diffusion的过程，慢慢地在data sample中加入噪声：

3.2 Vanilla Diffusion Policy存在的问题

3.2.1 将Transfuser变成conditional diffusion model

将模型中的MLP回归层替换成基于UNet架构的[1]。在评估过程中，采样了一个随机噪声，然后逐步进行refine（20步），相比于原始的Transfuser，这个改动给planning跑分带来了小幅度的提升（PDMS↑ +0.6）。但同时带来了两个巨大的问题：

问题1：模态坍缩

作者从高斯分布中随机选取了20个噪声作为初始状态，并用20步对其进行去噪过程。如下图所示，20步之后所有不同的输出模态都坍缩成了几乎是同一条轨迹。

问题2：计算量大

20步的去噪过程来将随机噪声变成一个可执行的轨迹，计算量太大，将原本的FPS从60降低到了7.

3.3 解决方案 - Truncated Diffusion

因为人类驾驶员是有固定的驾驶模式的，所以决定不再从标准的高斯分布中随机初始化，而是设计了一个 anchored Gaussian distribution 。

一、训练阶段

二、推理阶段

1. 输入的初始状态还是从anchored Gaussion distribution中采样得到“噪声轨迹”
2. 在每一步的去噪过程中，上一步得到的预估轨迹会传递到下一步，最终预测出 条轨迹和其对应的置信度
3. 更新规则 ：在每次得到了当前timestep的预测之后，作者使用了DDIM[2]的更新规则来为下一个timestep去采样轨迹
4. 优点 ：训练时使用的预测轨迹数 不会影响到推理时的预测轨迹数 。也就是, 可以是任意数字，可以根据不同部署平台的计算资源任意调整。

3.4 模型结构

Diffusion Decoder

1. 输入从anchored Gaussian distribution采样来的噪声轨迹
2. 使用deformable spatial cross-attention来与BEV或者PV特征进行交互，这个交互是基于轨迹坐标系的（想必是根据轨迹的坐标点进行reference point的选择）
3. 在轨迹features和agent/map信息之间做cross attention，然后接上一个FFN
4. Timestep Modulation layer，用来对diffusion timestep信息进行编码，而后接上一个MLP来预测置信度以及offset(相对于初始的噪声轨迹坐标）