3D 视觉感知任务,例如基于多摄像头图像的3D检测,是自动驾驶和辅助系统的必要组成部分。然而,设计高效的计算方法仍然是一个重大挑战。
本文提出了一种基于 Mamba 的框架 MamBEV,该框架利用线性时空 SSM 基准注意机制学习统一的鸟瞰图 (BEV) 表示。这种方法可以显著提高多种3D感知任务的计算和内存效率。
此外,作者引入了基于 SSM 的跨注意力机制,类似于标准的跨注意力机制,其中BEV Query 表示可以与相关图像特征进行交互。
广泛的实验表明,MamBEV 在多种视觉感知指标上表现出色,特别是在输入缩放效率方面优于现有基准模型。
代码可以在 https://github.com/amaigsu/MamBEV 查看。
1 INTRODUCTION
自动构建物体周围环境的鸟瞰视图(BEV)对于自动驾驶和驾驶辅助系统等任务非常有益(Wang等,2023a)。这些方法通常会整合多视角摄像头接收到的信号,并将其转换为周围环境的俯视图。此外,由于这些系统在移动边缘环境中运行,因此在考虑构建准确性的同时还需关注计算成本(Ke等,2024)。
特斯拉汽车(Tesla, 2021)中已部署的BEV系统实例可以供参考。这些详细的构建方法可用于下游感知、预测和规划任务(Casas等, 2021;Hu等, 2023)。构建BEV模型主要有两种方法:推法和拉法。推法通过像素级深度预测将二维图像特征投影到三维空间,从而将扁平图像转换为具有空间 Aware 的三维表示。而拉法则从三维先验中抽取点,将其投影到二维图像平面,从而使模型能够在没有显式深度预测的情况下提取三维信息。然而,许多此类方法依赖于使用 Transformer 昂贵的注意力机制来学习准确的表示。因此,作者受到启发,研究更高效的替代方案以取代基于 Transformer 的架构。
近年来,Gu等人(2021)和Gu & Dao(2023)展示了状态空间模型(SSMs)可以通过线性近似替代 Transformer 注意力的二次计算复杂度。这些模型已经在大规模语言模型(例如Dao & Gu,2024的Codestrals Mamba)中显示出与 Transformer 相当的效果。此外,也有证据表明,SSM的注意力可以替代 Transformer 的注意力来学习有效的视觉表示(Zhu等人,2024;Liu等人,2024;Patro & Agneeswaran,2024)。尽管Mamba与Self-Attention一样,在捕捉序列内部依赖关系方面有效,但在涉及多种输入模态的BEV场景中,它在动态信息交换方面存在困难。
受这些发现的启发,作者研究如何使用SSMs生成BEV表示。这很重要,因为在多视图视频中捕捉时空关系是非常昂贵的。作者探讨了SSMs内部的线性注意力如何用于解决这些问题。然而,将SSMs纳入三维表示学习任务尚不明确。此外,如何融合不同的视觉表示仍不清楚,因为这是学习BEV表示的一个关键步骤。
作者的论文通过以下贡献来解决这些问题:
作者提出了一种基于SSM的方法MamBEV,其性能可以超越现有的Transformer架构。
作者提出了一种新颖的空间交叉Mamba方法,类似于标准的交叉注意力机制,其中集合映射机制能够实现两种不同模态之间的关联和融合。在作者的情况下,通过匹配BEV Query 表示与相应的图像特征来促进两个模态信息的直接集成。
提供了详尽的消融研究以展示模型扩展和其他特性。作者开源了代码1,并为未来的实验提供了强大的 Baseline 和评估框架。
2BACKGROUND
2.1 LEARNING BEV REPRESENTATIONS
BEV 表示在自主驾驶中被广泛应用,因为它们能够提供一种密集统一的场景表示,适用于多种任务。其中一种方法基于 LSS(Philion & Fidler, 2020),通过利用相机内参将二维图像投影到三维空间。例如,VideoBEV(Han 等人, 2024)采用了一种基于 LSS 的 3D BEV 检测方法,并引入了定制化的递归式时序融合块和时序嵌入模块,以更有效地利用长期信息。然而,这种方法对深度估计的准确性有很高的依赖性,以确保特征投影的精度。
另一方面,BEVFormer(Li 等人, 2022a)通过直接从二维图像空间提取特征并将其投射到 BEV 空间来简化这一过程,从而将所有的高度信息压缩到一个单一的 BEV Query 中。该方法采用变形注意力机制(Zhu 等人, 2020),这是一种稀疏采样方法,作为 BEV 编码器的核心,用于将多视图图像特征映射到 BEV 特征上。BEVFormerV2(Yang 等人, 2023)在此基础上增加了额外的 Head 和图像增强技术以改进对模型主干的监督。此外,它还使用了基于 CNN 的时序编码器来替换 BEVFormer 中使用的递归式变形时序注意力机制。
Deformable Attention 通过减少内存消耗提供了一定优势,但也存在明显局限性。例如,Deformable Attention 只对 Backbone 提供稀疏且有限的监督,因为它并没有与所有图像特征进行交互(Yang et al., 2023)。在特征精炼方面,Deformable Attention 强烈依赖于 Backbone 来生成强大的特征表示。
例如,在Deformable Attention Transformer(DAT)中,Xia等人(2022),多头全注意力依然必要,以有效细化特征,这突显了可变形注意力在细化空间表示方面的局限性。从硬件角度来看,由于采样过程中依赖随机内存访问,可变形注意力对于现代GPU架构来说不够优化,从而对模型吞吐量产生负面影响(Zhu等人,2020)。
2.2 LINEAR ATTENTION WITH STATE SPACE MODELS
状态空间模型(SSMs),尤其是结构化SSMs,提供了与主流 Transformer 架构相比具有有效线性复杂度的选择方案,适用于序列建模。其核心在于,SSMs独立地处理每个通道的数据,并更新一个在输入序列中随时间演化的隐状态
。这一隐状态的演变由一组可学习参数控制,分别表示为
。Mamba Gu & Dao (2023) 引入了选择性SSMs(S6),该模型能够根据输入动态调整
参数矩阵的值,从而有效地充当了一个筛选或门控机制。这一机制使得SSMs能够更好地建模信息密度差异较大的输入序列,例如文本,因为它们可以忽略信息稀疏的Tokens并记忆信息密集的Tokens。Mamba 还包含了大量的硬件优化,这使得状态更新可以通过关联扫描实现并行计算。
Mamba-2(道 & 库,2024)通过一种对状态空间模型(SSM)与注意力机制之间联系的新颖理解,即状态空间二重性(SSD),克服了其前身Mamba-1的一些关键弱点。SSD表明,SSM计算可以通过涉及结构化矩阵乘法的双形式来表达。这使得Mamba-2能够利用GPU上的高度优化的矩阵乘法单元,从而相较于Mamba-1基于扫描的实现方式显著提高了速度。
虽然Mamba-2在自回归任务中表现出色,但其底层的SSM框架本质上是以因果方式运行的,这限制了它在非因果场景中的应用。Hydra(Hwang等人,2024)通过利用拟可分离矩阵混合器来解决这一问题,这些混合器将Mamba-2中发现的半可分离矩阵混合器推广到同时包括下三角和上三角部分。这种变化使得信息双向 Stream 能够在计算和参数量略有增加的情况下实现。为了简化讨论,在本文后续部分提及Hydra块时,作者使用Mamba来指代,因为Hydra采用了与Mamba-2相同的块结构和SSM公式。
SSM 内部函数。状态空间模型可以被视为一种系统,该系统通过 ( h(t) \in \tilde{\mathbb{R}}^N ) 将信号 ( x(t) \in \mathbb{R} ) 映射到 ( y(t) \in \mathbb{R} ),其形式可以表示为:
其中,
是演化参数,
,
和
是投影参数。
离散化内部函数。在机器学习应用中,大多数输入都不是连续信号,因此为了使这些系统适应离散输入序列,系统本身必须进行离散化处理。离散化的版本可以表示为
其中使用了时间尺度参数
将连续参数
和
转换为离散参数
和
。
该操作显式地将隐藏状态更新参数
和
联系在一起。在实践中,这种操作并不是必需的,因为模型可以直接学习离散化系统,如 (Gu & Dao, 2023) 所指出的那样。然而,引入一个离散化函数为解释特定输入与隐藏状态之间的相互作用以及直接控制隐藏状态是否由某个特定的 token 或一类 token 更新提供了有用的机制。Ali 等人 (2024) 指出,每个通道具有独立的离散化参数等同于注意力中的多个头,其中每个头对应于其自己的通道。
在 Mamba-2 中,作者发现通过在多个通道之间共享一个离散化因子来捆绑它们可以降低计算复杂度,同时保留类似的表达能力。当将多个通道分组为一个头时,可以共享
矩阵,并且需要根据输入计算的
数量也会减少。此外,他们还发现,将
矩阵从
中的对角矩阵简化为标量可以显著减少计算成本,而对表达能力的影响很小。图中的 SSM 组件指的是方程 2 和方程 3 中的操作。
3 METHODS
BEV构建问题是监督学习算法,它将特定时间步
处
个相机视角的相机图像
映射到一个二维平面鸟瞰图上,表示物体的位置。给定包含
个物体的一个BEV场景,包含了边界框
、类别
以及轨迹信息
。该方法首先学习一种编码转换
,以获取隐含的BEV Query 表示
,其中
和
表示BEV网格的空间形状,而
则表示隐含维度。随后,学习一种解码转换来预测物体信息
,并通过采样BEV网格来实现。
BEV 方法概览。图2提供了作者学习BEV构建的整体方法的可视化。许多现有的组件遵循了Li等人(2022a)、Li等人(2023b)和Liu等人(2023)的方法学。编码函数
由ResNet101主干、特征金字塔网络(FPN)、BEV编码器和时间融合模块参数化。该Pipeline分为三个主要阶段。首先,不同尺度的图像特征图
作为中间主干输出生成。然后,通过FPN中的水平卷积将这些特征图的通道数减少到统一大小
(Lin等人,2017)。其次,使用基于SSM的新颖Pipeline来建模特征表示和交互(具体内容将在下面讨论)。最后,解码函数
作用于此表示进行预测。与Zhu等人(2020)、Li等人(2022a)和Yang等人(2023)的方法类似,解码器使用交替层的分组目标 Query 自注意力和目标 Query 到BEV特征的可变形交叉注意力。每层的解码头遵循Li等人(2022b)中提出的掩蔽解码头,预测边界框、其属性以及物体类别。作者没有使用地图分割头,仅使用了检测Head,并且该头是通过3D匈牙利集合匹配优化的,使用平滑
损失对边界框进行优化,使用Focal Loss对类别预测进行优化。
接下来,作者讨论图2中所示的新型编码器。
3.1 SPATIAL CROSS MAMBA
使用Mamba作为BEV(Bird's Eye View)的中心挑战是如何以高效的方式进行跨注意力操作。对于SSMs(State Space Models),存在一些简单的跨注意力改编方法,但这些方法效率低下或不适用于该问题。例如,通过将状态大小N扩展到等于图像特征的数量T,Mamba-2能够存储所有关于先前Token的信息,但会回到Transformer的复杂度(Dao & Gu, 2024)。为了实现这种简单的Mamba跨注意力操作,首先计算最终的隐藏状态hT ∈ RTxD',然后找到y_Q。
其中 ( C_{i} \in \mathbb{R}^{1 \times T} ) 是相应 Query ( q_{i} \in Q ) 的函数。计算 ( h_{T} ) 需要 (\hat{T}NP = \hat{T}^2P) FLOPs,其中 ( P ) 是每个头的潜在维度,( TN = T^2 ) 内存,在这种情况下相当于 Transformer 中计算自注意力时使用的 FLOPs 和内存,即 ( T^2N ) 和 ( T^2 )。计算 ( y_{i} ) 需要 ( T|Q|\alpha D ) FLOPs 和 ( T|Q| ) 内存,给出的总体计算和内存复杂度类似于 Transformer 的交叉注意力。为了降低朴素实现的计算成本,作者提出如图3 所示的任务特异性适应方法,允许作者将隐藏维度的大小减少到 ( N << T )。
3.1.1 REDUCING STATE SIZE
简单地减少 (N) ,如果图像的信息量稀疏,仍然可能获得可接受的结果,然而这一点不能保证。(h_{T}) 的结果可能会损失序列初始部分图像特征的信息。这源于衰减参数 (\exp(\Delta A) \in [0,1], A < 0, \Delta \ge \bar{0}) 的定义,该定义确保除了 (\Delta_{i}=0, i=1,\ldots,T) 以外,(x_{0}) 对 (y_{T}) 的影响会随着 (T) 增加而减弱。作者提出的解决方案是使用 (C_{i}h_{k}) 计算 (y_{i}) ,其中图像特征向量 (k) 很可能与 (q_{i}) 相关。由于多个图像区域可能与单一BEV Query 相关,作者利用 (Z) 个 (C_{i}) 来关注图像特征图上的 (Z) 个位置,然后
为了为 Query
选择
个特征图位置,作者首先将二维BEV位置
提升至三维 Pillar (x,y,z),然后将均匀分布的
个 Pillar 点投影到每张图像上,参考Lang等人的工作(2018)。由此得到的位置称为参考点
,这些位置基于ego车辆的相机校准假设并无遮挡情况下,在图像中相应于BEV位置处物体出现的位置。落在图像边界内的位置数量
远远小于
。作者用
的数量M=f(|Q|, Z, Orov)近似表示相机的数量,因为大多数
在单个视角下可见。
3.1.2 BEV POSITION AWARE MERGE
和
被展平为一维序列,以便可以被 Mamba 处理。
根据遍历顺序
进行展平。然后,
通过一维参考点
与图像特征
的展平序列进行合并。
其中,
分别表示参考位置在图像上的归一化 x 坐标和 y 坐标。 argsort
函数返回 $R_{1\mathrm{D}}
排序后每个元素的索引,并将其添加回 $R_{\mathrm{1D}}
以重新排序索引。这种形式的时间复杂度为
,尽管实际上该操作的成本可以忽略不计。结果是一个一维索引列表,假设图像特征已使用行优先遍历方式展平,则该列表对应于目标位置左侧一个像素位置及上方一个像素位置。在其他遍历方法中,通过
进行更新索引。从 $R_{\mathrm{1D}}$
开始,作者生成一个 Mask
,用于 Token 每个 Query
在命中相机视图中的插入点。为了完成合并操作,作者将 Query
复制到其相应的参考位置,并将值
复制到 Mask 输出张量中。此操作对每个相机及其遍历方法重复执行。图4 显示了输入序列的一个示例。此外,在表9 中还进行了消融研究,分析了不同遍历顺序的影响。
3.1.3 CROSS QUASI-SEPARABLE STATE SPACE MODEL
Mamba-2 为选择性状态空间模型提供了两种计算方法:作为一种矩阵混合器,SSD 和关联扫描器 SSM。将 SSM 重新构想为结构化的矩阵混合器
,使得可以通过批量矩阵乘法和缩短的关联扫描实现快速并行计算。实际上,在训练和推理过程中作者都使用了 SSD,因为它更适合硬件优化,但作者并未调整 Kernel 以反映交叉准可分离状态空间模型(XQSSM)模块的真实计算复杂度。
算法2(附录中)的简单顺序实现展示了作者模块化形式允许的独特方法,其中 Query 输入 (Q_{dt}) 的离散化因子设置为0。当处理 (dt=0) 的Token时,不会发生隐藏状态更新,如下所示。
这种变化将每个 Query 的操作次数从 (2H(N+1) + \alpha D(3N+H+1)) 减少到 (\alpha D(N+H+1))。另外,XQsSM 的输出仅需用于 Query Token 输入,这使得每张图像的特征复杂度减少到 (2H(N+1) + 2\alpha D N)。总体而言,XQSSM 的计算复杂度为 (2V(H(N+1) + \alpha D N) + M\alpha D(N+H+1)),其中 (M) 是添加到序列中的 Query 数量。此外,在序列形式下的内存复杂度是常数级的,但在并行化时,它会随着序列长度呈线性增长,如 Dao & Gu (2024) 所示。由此得到的矩阵混合器 (\mathcal{M}) 从 ((M+\bar{V} \times M+\bar{V} \times 2H)) 变为 ((M \times V \times 2H)),类似于形状为 ((Q \times K \times H)) 的点积交叉注意力的矩阵混合器。
3.1.4 QUERY AND FEATURE POST-PROCESSING
经过合并输入序列并通过内部SSM模块处理后,需要从输出
中提取queries以获得更新的BEV queries
。在合并操作
过程中生成的 Mask 应用于输出,得到
。然后使用
,每个
中的query在其在
中的原始位置处累积到
。每个
中的元素
的幅度取决于该BEV位置累积的query数量。最终将其投影到
,然后添加到残差BEV Query 网格
中。为了应对由于幅度差异导致的不稳定性,考虑了两种归一化方法:投影前的均值归一化和投影后的RMS归一化。这两种方法及其联合使用的消融实验结果见表7。
BEV 自注意力机制。在行优先顺序下,经过一个单层海德拉层(双向曼巴)对归一化的BEV网格Q进行遍历,然后将其输入到一个变形注意层中以融合多尺度特征。变形注意机制。这是一种稀疏注意方法,在检测任务中表现出色,并且与传统的softmax注意机制相比,在训练时间更短、操作次数更少且具有更好的可扩展性(Zhu et al., 2020)。该方法的基础是在给定图像特征集和作为参考位置的坐标集的情况下,根据预测出的 Query 偏移量,在参考位置周围的区域进行网格采样。
可变形注意力的计算成本分为三个主要部分:每个 Query 计算偏移量的成本
,双线性插值和加权样本求和的成本
,以及每个图像特征层
(形状为
)的值的线性投影成本
。其中,
表示 Query 的数量,
是隐含维度,
是每个 Query 的参考点数量,
是每个参考点的偏移量数量。该操作的整体计算复杂度为
。在假设
的情况下,整体复杂度简化为
。值得注意的是,与图像大小相比,这里的计算复杂度相对较低,因为约束条件中权重较高的部分来源于使用的 Query 数、头数和参考点数。当图像尺寸
相对于 Query 数较小的情况下,可变形注意力在计算复杂度上可能类似于甚至差于像 Mamba 这样的线性注意力方案,但其提供的交叉注意力效果较弱。
4 EXPERIMENTAL SETUP
作者遵循王等人(2022)、李等人(2022a)和杨等人(2023)之前工作的方法。作者在两个 Backbone 网络上进行评估:ResNet101 和 ResNet50,前者在深度预测任务中训练,后者在COCO数据集中训练。在训练过程中, Backbone 网络的第一阶段被冻结,其余阶段以10%的学习率进行训练,从而微调其潜在表示以适应多视角自主驾驶环境。
4.1 DATASET AND METRICS
作者使用 nuScenes 数据集(Caesar等人,2020)进行实验。nuScenes 数据集是一个大型自主驾驶数据集,包含了来自波士顿和新加坡的约1000个驾驶场景。每个场景大约持续20秒。整个数据集中,以每2Hz的频率标注了23种物体类别的3D边界框,其中有10种用于3D检测任务。每个场景由6个视角的摄像头捕捉,视角范围为360度,并包括激光雷达、雷达、GPS、传感器校准和IMU数据。nuScenes devkit 提供了评估指标及其计算框架。用于评估的指标包括:1)平均精度均值(mAP),它使用地面平面上的距离中心误差对预测结果的定位和分类性能在四个不同阈值上进行了评估;以及五种真阳性指标:2)平均平移误差(ATE)、3)平均尺度误差(ASE)、4)平均方向误差(AOE)、5)平均速度误差(AVE)和6)平均属性误差(AAE)。nuScenes 还定义了一种综合评估方法,即nuScenes检测分数(NDS),它是通过结合mAP与五个真阳性指标来计算得出的。在作者的方法和实验中,只有摄像头帧、传感器校准数据和GPS数据被用于生成预测结果。
4.2 IMPLEMENTATION DETAILS
作者使用了学习率为
,并在计划的步骤的前
进行线性 Warm Up,起始值为
。Warm Up 结束后,学习率遵循基于 epoch 的余弦退火调度,最小学习率为
。作者在 30 个 epoch 中使用有效批量大小为 32 的数据进行训练,并在 8 块 A100 GPU 上不累积梯度训练至第 24 个 epoch。从第 100 个步骤开始,应用公式
的指数移动平均方法到所有权重。使用的优化器为具有 0.01 权重衰减的 AdamW,并采用自动混合精度优化器包装器,初始梯度缩放为 512。对 Backbone 网络权重和可变形注意力采样偏移量使用 0.1 的学习率乘数 Zhu et al. (2020)。编码器层使用随机初始化的网络从头开始训练模型。源代码将在发表时提供。
5 RESULTS
作者首先报道了主要结果,并随后进行了消融研究以理解各种模型选择的有效性。除非另有说明,消融研究中均使用了一个时间帧。消融模型训练了12个 epoch,并采用在COCO目标检测数据集上预训练的ResNet50作为主干网络。虽然更多的训练提高了性能,但不同模型参数化并没有显著差异。
主要结果。作者在表1中展示了与兼容参数量和图像输入尺度的最新方法的比较结果。作者仅使用了摄像头特征,并且没有采用杨等(2023)等人工作的辅助损失。作者将tiny和small模型的定义与BEVFormerV1对齐。作者只训练了MamBEV的tiny和small版本。作者的基于CNN的时间成分在small配置中有5700万参数,在tiny配置中有3200万参数。不包括时间部分,MamBEV-Small模型有6500万参数,而MamBEV-Tiny则有3900万参数。相比之下,BEVFormer的小型、小型和 Baseline 模型分别有3400万、6000万和6900万参数。
作者在类似条件下使用ResNet50 Backbone 网络评估以前的最先进技术,报告了与MamBEV-Tiny相当的结果(+0.002 NDS)。作者还报道了MamBEV-Small使用4帧在NDS上比之前的BEVFormerV1-Small模型提高了0.046(+9.6%),比之前的BEVFormerV1-Base模型提高了0.008(+1.5%)。这表明作者的模型可以在降低内存复杂度的同时提高性能。
Spatial Cross Mamba 层的有效性分析。为了验证 Spatial Cross Mamba 层的效果,作者训练了一个不含变形层的较小模型 MamBEV-Small-Pure。在 NDS 和 mAAE 指标上,该配置显示出了良好的性能,并优于 PolarDETR-T 和 BEVFormerV1-Small。然而,使用了变形层的小模型则表现得更好。
效率。在表2中,作者测试了使用作者提出的XQsSM、标准点积注意力或变形注意力的模型配置所占用的内存和计算估计的FLOPs。XQsSM和变形注意力在内存和计算复杂度方面与输入的大小
和
呈线性关系,尽管变形注意力的系数因子较小。
增加时间信息。作者通过提供模型额外的前一帧来研究增加时间信息的效果。使用更多的帧会要求在编码步骤中关注更多信息。由于SSM基注意力机制可以降低这种复杂性,作者能够更好地利用这些信息。作者在表3中报告了实验结果。
如观察所示,增加帧数可以提升整体性能,直到大约5帧为止。使用单个帧进行预测时性能最差。不出所料,mAVE(一种速度预测误差的衡量标准)在模型仅有单个时间步的多视图视频可用时也显著更高。
基于SSM的注意力机制与变形注意力机制的比较。作者对比了作者的空间交叉Mamba形式化表示与变形注意力机制。作者在表4中展示了实验结果。空间交叉Mamba可以作为变形注意力的替代方案,因为两者的性能差异很小。在实验中,作者发现,在训练一个学习更大数量帧的表示且规模更大的网络时,使用混合的空间和变形注意力是有帮助的。
扩展实验。作者进行了实验以了解在 Mamba 模块内部增加通道数量以及状态空间模型内部隐藏状态规模的效果。作者在表5 中报告了隐藏状态规模的实验结果,在表8 中报告了通道数量的实验结果。作者的实验结果显示,在调整这些变量时,观察到的影响非常小。这一发现表明,增加通道特征和隐藏状态的规模未必一定能提高性能。
Query 插入。为了验证作者的BEV位置感知合并(Project)方法的有效性,作者在表6中进行了实验以展示其效果。结果显示,Project是对于作者模型而言最有效的 Query 插入方法,特别是在检测性能和减少关键错误方面。由于Mamba主要通过序列中元素的位置来学习空间关系,因此append和prepend的性能较差是合理的。
6 CONCLUSION
作者的研究提出了一种BEV构建模型。作者在相同的实验条件下评估时,表明该模型在性能上优于先前的工作。
详尽的消融研究证明了该模型对各种变化具有鲁棒性。作者将在未来的研究中开源源代码。
参考
[0]. MAMBEV: ENABLING STATE SPACE MODELS TO LEARN BIRDS-EYE-VIEW REPRESENTATIONS .