速度快4倍 | MIT&交大&清华联合提出FlatFormer,一个非常高效的Transformer方法

简介: 速度快4倍 | MIT&交大&清华联合提出FlatFormer,一个非常高效的Transformer方法

Transformer,作为CNN的替代品,已在许多模态(例如,文本和图像)中证明其有效性。对于3D点云Transformer,现有研究主要集中在将它们的准确性推至最先进水平。然而,它们的延迟落后于稀疏卷积模型(3倍 slower),这阻碍了它们在资源受限、延迟敏感的应用中的使用(例如,自动驾驶)。这种低效性源于点云的稀疏和不规则性质,而Transformer是为了密集、规则的工作负载而设计的。

本文提出了FlatFormer,通过以空间邻近度为代价换取更好的计算规律来弥补这一延迟差距。作者首先对点云进行窗口排序,并将其分成大小相等的组,而不是形状相等的窗口。这样,可以有效地避免昂贵的结构化和填充开销。然后,在组内应用自注意力来提取局部特征,交替排序轴来收集不同方向的特征,并移动窗口以跨组交换特征。

FlatFormer在Waymo Open Dataset上实现了最先进的准确性,同时比(基于Transformer的)SST快4.6倍,比(稀疏卷积的)CenterPoint快1.4倍。这是第一个在边缘GPU上实现实时性能的点云Transformer,同时比稀疏卷积方法快,且在大规模基准测试上达到与大规模基准测试相当或甚至更好的准确性。

1 Introduction

Transformer已成为自然语言处理(NLP)中的首选模型,作为许多成功的大语言模型(LLMs)的 Backbone 。最近,其影响力进一步扩展到视觉领域,视觉Transformer(ViTs)已在许多视觉模态(例如,图像和视频)中展现出与卷积神经网络(CNNs)相当或甚至更好的性能。然而,3D点云尚未成为其中的一部分。

与图像和视频不同,3D点云固有地稀疏且不规则。大多数现有的点云模型基于3D稀疏卷积,它只在非零特征上计算卷积。它们需要专用的系统支持才能在并行硬件(例如,GPUs)上实现高利用率。

许多研究致力于点云 Transformer (PCTs),以探索它们作为稀疏卷积的替代品的潜力。全局点云 Transformer得益于自注意力的一致计算模式,但受到与点数平方成正比的计算成本(相对于点数)的严重影响。局部点云 Transformer将自注意力应用于定义在类似基于点模型的局部邻域,因此受到昂贵邻居收集的限制。这些方法仅适用于单个物体或部分室内扫描(少于4k个点)且无法有效地扩展到室外场景(超过30k个点)。

受到Swin Transformer的启发,窗口点云 Transformer在窗口 Level 计算自注意力,在大规模3D检测基准测试上实现了惊人的准确性。尽管这些窗口在空间上具有规律性,但由于稀疏性,它们的点数可能会有极大的不同(相差超过80倍)。这种严重的失衡导致冗余计算和低效的填充和分区开销。因此,窗口点云 Transformer 比稀疏卷积模型慢3倍(如图1所示),限制了它们在资源受限、延迟敏感场景(例如,自动驾驶,增强现实)中的应用。

本文引入了FlatFormer来缩小这一巨大的延迟差距。在构建基于窗口的点云 Transformer (window PCTs)的基础上,FlatFormer通过将3D点云分成“相等大小的组”而不是“相等形状的窗口”来换取更好的计算规律。在组内应用自注意力来提取局部特征,交替排序轴来收集不同方向的特征,并移动窗口以跨组交换特征。

得益于规则的计算模式,FlatFormer在Waymo Open Dataset上实现了最先进的准确性,同时比(基于Transformer的)SST快4.6倍,比(稀疏卷积的)CenterPoint快1.4倍。这是第一个在边缘GPU上实现实时性能的点云 Transformer ,同时比稀疏卷积方法快,且在大规模基准测试上达到与大规模基准测试相当或甚至更好的准确性。

到作者为止,FlatFormer是第一个在较低延迟下达到稀疏卷积方法相等或更好的准确性,同时比卷积神经网络(CNN)快得多的点云 Transformer 。它也是第一个在边缘GPU上实现实时性能的点云 Transformer 。在更好地支持 Transformer (例如,NVIDIA Hopper)的硬件支持下,点云 Transformer 将具有巨大的潜力成为3D深度学习中的首选模型。

2 Why are Point Cloud Transformers Slow?

尽管点云 Transformer (PCTs)正在逐步赶上稀疏卷积检测器的准确性,但它们与最快的PCT(SST)和稀疏卷积CenterPoint之间的延迟仍存在3倍差距(图1)。在本节中剖析了PCTs效率的瓶颈,这为FlatFormer的设计奠定了坚实的基础。

Global PCTs

受到ViT的启发,点云上 Transformer 的最简单和直接的设计是全局PCTs,其中每个点都是一个标记。它们在整个点云上利用多头自注意力(MHSA)进行全局处理。虽然全局PCTs在小型3D物体上有效,但由于其的复杂度,在扩展到大场景时表现不佳,其中是标记数,是通道数。

如图2所示,全局PCTs的运行时间随着输入点数增长而呈平方增长。例如,在32k输入点*的情况下,该模型在NVIDIA A6000 GPU上执行需要近一秒钟,比CenterPoint慢66倍。

Local PCTs

研究行人提出了局部PCTs来解决全局PCTs的可扩展性问题。它们将多头自注意力(MHSA)应用于每个点的邻居,而不是整个点云。因此,它们的计算复杂度为,其中是点数,是每个点的邻居数,是通道数。由于对于实际工作负载,的范围在20k到100k之间,小于64,因此它们的理论成本比全局PCTs低得多。

然而,局部PCTs受到邻居准备开销的影响很大。由于点云稀疏且不规则,需要首先找到每个点的邻居,然后将数据从格式重构为格式,以便应用MHSA。这两步很慢,对于VoTr在Waymo上的单场景执行,需要22毫秒(即36%的总运行时间),这已经比整个CenterPoint模型慢。对于Point Transformer(PT)[98],准备邻居的开销可能占运行时间的70%。这样的开销在单层中就已经大于CenterPoint的总运行时间!

Window PCTs

Swin Transformers在各种视觉识别任务上的巨大成功催生了窗口PCTs的设计,其中,SST是代表性工作。它首先将点云投影到鸟瞰空间(BEV空间),然后将BEV空间划分为大小相等、非重叠的窗口,并在每个窗口内应用多头自注意力(MHSA)。与Swin Transformer类似,SST使用窗口平移来实现窗口间的信息交换。

与图像不同,点云是稀疏且非均匀地分布在空间中的。因此,每个窗口内的点数不同,可以相差两个数量级(图4)。由于基本的MHSA Kernel 无法有效地支持可变序列长度,SST将大小相似的窗口分组,并将每个分组内的所有窗口都填充到该组内的最大值(图5)。

然后,在单独的每个分组内应用MHSA。在实际应用中,这种填充引入了1.7倍的计算开销。更糟糕的是,将点分成大小相等的窗口也引入了显著的延迟开销:在Waymo上,每个场景需要18毫秒,甚至比CenterPoint的总运行时间还要慢。总之,填充和分组的额外开销使得SST在硬件友好性方面比稀疏卷积方法差。

4 FlatFormer

在第3节中学到的所有经验教训,作者将设计作者的点云 Transformer 以实现可扩展性和高效性。

Overview

FlatFormer的基本构建模块是Flattened Window Attention (FWA)。如图5所示,FWA采用基于窗口的排序来平铺点云并将其划分为大小相等的组,而不是形状相等的窗口。这自然地解决了组大小不平衡问题并避免了填充和分组开销。

FWA然后在每个组内应用自注意力来提取局部特征,交替排序轴来聚合不同方向的特征,并将窗口移动来跨组交换特征。最后,作者提供了一种改进的FWA实现,进一步提高了其效率并最小化了开销。

Flattened Window Attention (FWA)

4.2.1 Sorting & Grouping

基于窗口的排序。对于一个点云,作者首先对每个点的坐标进行量化。

脚注†:作者假设点云在2D中,以便于表示,但作者的方法可以应用于3D或更高维的点云。

其中表示窗口形状。接下来,作者首先按窗口坐标对所有点进行排序,然后按窗口内局部坐标进行排序。这一步将无序的点云变为有序的点云,其中同一个窗口内的点将相邻排列。

按相等大小的组划分。传统的窗口PCTs会将同一个窗口内的点分组在一起。然而,如第3节所述,每个组中的点数可以因固有的稀疏性而有很大的不同。

为了克服填充开销,作者根据排序序列将点云分成大小相等的组。这一步使得后续的组注意力具有完全规律的工作负载。从实现角度来看,作者的分组只涉及简单的张量重排(由于它不改变布局,所以是免费的),比SST中的窗口分组的效率更高。

交替排序轴。在两个轴之间,具有更高的优先级。因此,具有相同的点将相邻排列,而具有相同的点在排序序列中可能非常远离彼此,破坏了几何局部性。

为了解决这种不平等,作者在不同的FWA块中交替使用和作为排序轴。这与将3x3核分解为3x1和1x3核的空间可分卷积非常相似。堆叠具有不同排序轴的FWA块使得模型能够从不同方向聚合局部特征。

相等大小 vs 相等窗口。作者做出的关键设计选择是将点云分成相等大小的组而不是相等形状的窗口。存在权衡:相等窗口组保持完美的空间邻近性(即每个组具有相同的半径),但破坏了计算规律性,而相等大小组确保了计算工作负载的平衡(即每个组具有相同数量的点数),但无法保证几何局部性。

作者在第5.3节中证明,计算规律性更重要,因为空间不规则性可以通过作者的算法设计部分解决:即窗口排序提供了一个相当好的空间排序,而自注意力对异常值(即远离点对)具有鲁棒性。

4.2.2 Group Attention

在每个组内,作者应用自注意力来提取局部特征。对于具有坐标和特征的每个点组,作者遵循标准的 Transformer 块设计:

其中MHSA,FFN和LN分别对应多头自注意力,前馈层和层归一化,而PE则给出全局绝对位置编码。在这里,作者使用标准的 Transformer 块设计MHSA。作者的方法将受益于更高效的注意力变体,例如线性注意力[30],作者将留待未来的工作中使用。

窗口平移。得益于非重叠设计,基于窗口的注意力通常具有比卷积(如大小为3的稀疏卷积中的27个邻居)更大的感受野(例如,作者的FWA中的69个邻居,而窗口平移的27个邻居)。然而,其建模能力有限,因为不同组之间没有特征交换。

类似于Swin Transformer,作者采用平移窗口的方法,在连续的FWA块中交替更改排序配置。具体而言,作者在平移窗口的FWA块中,将所有点的坐标进行排序。这种机制在保持组间特征通信的同时,有效地维持了组间工作负载的独立性。值得注意的是,交替排序轴也可以实现特征交换。

Efficient Implementation

除了算法设计外,作者还提供了一种实现方式,可以提高MHSA和FFN的效率,并最小化排序和掩码开销。所有的优化都是为了作者的点云 Transformer 设计而专门设计的,并且不适用于稀疏卷积模型。

高效的MHSA。在MHSA中, Query ,Key 和Value 首先会被单独的线性层进行变换。作者将这三个线性投影打包成一个批处理的矩阵乘法(由于在作者的FWA中,,和具有相同形状),以提高并行性。此外,标准的注意力实现将和实现。作者利用最近一种高效的保持函数不变实现(FlashAttention),它使用分块来减少内存读写次数,从而实现更好的效率。

高效的FFN。FFN由两个带有GELU激活的中间层的线性层组成。作者在Triton中实现了一个融合的线性 Kernel (将激活函数集成到层中,以避免将中间结果写入DRAM),并观察到作者的线性 Kernel (由Triton优化)比cuBLAS更有效,这可能是由于不寻常的高宽比工作负载。

重复排序。对所有点的坐标进行排序是一项不可忽视的开销。由于坐标保持不变(无下采样),作者使用相同的轴和窗口重复使用排序结果(即每个点的排名)。在实践中,这将作者的模型中的排序开销降低了50%。

丢弃残差。输入点云的大小可能不能被组大小整除,在分区后可能会产生一个组内点数较少的组。这种微小的不规则性仍然会导致自注意力中的一些开销,因为作者需要引入掩码来正确地将其置为0。相反,作者直接丢弃非满组的最终结果。这只对应所有点的不到0.1%,对模型的性能几乎没有影响(0.1%)。

5 Experiments

Setup

数据集。作者在Waymo Open Dataset(WOD)上进行了作者的实验,该数据集包含1150个LiDAR点云序列。每个序列包含200帧,由360度视野的LiDAR传感器以每秒10帧的速度收集。有四个前景类别,其中三个(车辆、行人和骑行者)用于检测指标评估。

指标。作者遵循Waymo的官方指标来计算所有方法的标准3D mAP和 Head 加权3D mAP(mAPH)。车辆、行人和骑行者三个类别中的匹配IoU阈值设置为默认值(0.7,0.5和0.5)。物体分为两个难度 Level ,小于五个激光点或标记为困难的物体被归类为 Level 2(L2),其他物体定义为 Level 1(L1)。作者在主论文中主要报告L2指标,并在附录中提供详细指标。

模型。基于FWA,作者提供了一个实现FlatFormer的3D检测版本。遵循PointPillars的设计,首先将点云进行0.32m0.32m的voxel化,并在每个voxel上使用MLP。然后,应用8个连续的FWA块,每个块使用交替的排序轴(例如,或)和窗口平移配置(例如,开或关)。所有FWA块的窗口形状为99,组大小为69。遵循SST,作者没有应用任何空间降采样,这对小物体有益。最后,应用常规BEV编码器和基于中心检测Head,如CenterPoint。

为了执行所有稀疏卷积检测器的3D编码器。对于基于 Transformer 的检测器,作者使用它们的官方实现来测量运行时间。所有在3D编码器之后的模块(例如,BEV编码器和检测Head)都使用TensorRT 8.4执行。作者将在前1000个样本上执行所有方法50次。作者报告平均延迟(排除异常值)。作者不包括数据加载和后处理时间。

结果。与Table 1中的结果类似,作者的FlatFormer在稀疏卷积和基于 Transformer 的检测器之间实现了稳定的性能改进,同时具有更好的效率。对于单帧模型,FlatFormer比SST,SST-Center和最近的VoxSet快4.2倍3.3倍3.7倍。它还与强大的稀疏卷积 Baseline 相比具有优势:比CenterPoint快1.4倍,具有1.7更高的L2 mAPH,并达到与PillarNet相当的性能。在两帧设置中,FlatFormer的性能优势更加明显。

具体而言,FlatFormer比CenterPoint快1.4倍,具有2.8 L2 mAPH更高的准确性,并以类似的延迟超过PillarNet1.3 L2 mAPH。对于三输入帧,FlatFormer比SST和SST-Center分别快4.6倍3.2倍。它还实现了更好的延迟-准确性权衡(比CenterPoint++快1.1倍,准确率提高0.4%)。

值得注意的是,尽管FlatFormer需要比CenterPoint++多1.7倍的MACs,但它仍然更快。这表明作者的设计比稀疏卷积 Baseline 更符合硬件需求。

部署。作者在NVIDIA Jetson AGX Orin上部署了作者的FlatFormer。这是资源受限的边缘GPU平台,广泛用于实际自动驾驶汽车。如图6所示,FlatFormer在16 FPS下运行,即1.2倍

5.2.2 Two-Stage Detectors

模型。为了验证FlatFormer的泛化性,在FSD中的3D Backbone 上进行了替换,FSD是一种最先进的两阶段检测器,并将其结果与先前的两阶段模型进行了比较。作者在单阶段实验中保持相同的网格分辨率、窗口形状和组大小。

基准与延迟。作者在Table 2中与最先进的两阶段检测器进行了比较。遵循相同的延迟测量协议。对于CenterFormer,作者调整了官方实现以支持SpConv v2.2.3后端在FP16精度下进行公平比较。

结果。所有现有高性能的两阶段检测器都是稀疏卷积的,而FlatFormer是唯一一种基于 Transformer 的实现,实现了最先进的准确性。与CenterFormer相比,FlatFormer在输入帧数方面的可扩展性更好。请注意,作者的论文主要关注优化3D Backbone 的延迟。然而,两阶段检测器通常在运行时由第二阶段瓶颈,这超出了作者的范围。作者期望FlatFormer的延迟可以受益于更高效的第二阶段设计。

Analysis

在本节中,作者提出了一些分析来验证作者设计选择的有效性。所有实验基于作者使用20%数据训练的单帧模型。

5.3.1 Flattened Window Attention

在图7中,作者可视化了作者在FWA中学习的注意力权重。颜色表示注意力权重的规模,较暖的颜色表示较大的注意力权重。黑色点对应 Query 点,灰色点对应权重小于阈值的点。

对于向前移动的车辆,转向和停车的车辆,车辆上的 Query 点始终高度关注附近同一车辆上的点,而远离车辆的点具有非常小的学习注意力权重。这样的观察可以部分解释FWA的有效性:即,即使等大分组不会创建空间规律的窗口,模型也可以学习抑制背景中的异常点的重要性,并专注于每个组内的更重要前景点。

5.3.2 Ablation Studies on Model Design

排序策略。作者在Table 3中分析了作者基于窗口的轴交替排序策略的有效性。将点随机分组而不进行任何空间排序,与FlatFormer相比,性能将降低约4%。此外,由于边界区域上的空间不连续性,直接按坐标排序或按单个轴排序窗口,都会提供次优的结果。作者还发现,窗口平移可以带来0.5%的性能提升。

值得注意的是,尽管作者的排序策略不能保证窗口在几何上与SST中的窗口一样规则,但FlatFormer仍然在所有三个类别中始终优于SST。

窗口形状。FlatFormer在不同的窗口形状下都能实现鲁棒性能。根据Table 4中的结果,作者在所有实验中选择窗口形状为99(2.88m2.88m,与车辆的大小相同),并将组大小固定为窗口大小的85%。

组大小。作者进一步研究了组大小的选择,如Table 5所示。作者将窗口形状固定为99(2.88m2.88m,与车辆的大小相同),并将组大小在50%,85%和125%的窗口形状范围内进行扫描。

结果表明,将组大小设置为窗口大小的85%时,性能最佳。直观上,如果组大小太小,FlatFormer将无法拥有足够大的感受野(例如,组大小=1,FWA将退化为MLP)。当组大小过大(例如,组大小=整个点云)时,每个组内将存在大量异常值,FlatFormer将像全局PCT一样行为,这是不希望的。

输入分辨率。从Table 6中,作者可以看到0.32m0.32m的输入分辨率是FlatFormer在延迟-准确性权衡中的最佳选择,而进一步增加输入大小只会降低效率,而没有性能改进。

5.3.3 Breakdown for System Optimizations

在图8中,作者分析了作者在第4.3节中提出的系统优化方法的有效性。高效的MHSA实现(FlashAttention)将作者的推理延迟提高了1.7倍。将的计算放入单个线性 Kernel 中,将带来1.3倍的加速。

将线性层和激活层(在FFN中)融合,将带来1.2倍的加速。最后,丢弃非满窗口将提高作者的推理延迟,带来1.1倍的改进。总之,作者的系统优化将FlatFormer的延迟提高了2.9倍,使其背

讨论。CenterPoint的背后是SpConv,它是一个高度优化的稀疏卷积推理库,基于CUTLASS实现。然而,FlatFormer在NVIDIA GPU上仍然实现了最佳的效率。作者部分地将作者的效率优势归因于FlatFormer中相等的组,这不仅给作者提供了最佳计算规律,还消除了计算开销。

另一方面,SpConv实现了一种3D稀疏卷积,使用带掩码的隐式GEMM算法,当一个线程块内的点没有完全相同的邻居模式时,必然会引入计算开销。因此,尽管FlatFormer在GPU上能够击败稀疏卷积模型,尽管它们进行了系统优化。

6 Conclusion

本文介绍了FlatFormer,用于弥合点云 Transformer 和稀疏卷积模型之间的巨大效率差距。它采用相等大小的组划分而不是相等窗口的组划分,以换取计算规律性。

FlatFormer在Waymo Open Dataset上实现了最先进的准确性,比之前的点云 Transformer 快4.6倍。作者希望FlatFormer可以激发未来研究设计高效且准确的点云 Transformer 。

参考

[1].FlatFormer: Flattened Window Attention for Efficient Point Cloud Transformer.

相关实践学习
部署Stable Diffusion玩转AI绘画(GPU云服务器)
本实验通过在ECS上从零开始部署Stable Diffusion来进行AI绘画创作,开启AIGC盲盒。
相关文章
|
22天前
|
机器学习/深度学习 人工智能 算法
Optima:清华联合北邮推出优化通信效率和任务有效性的训练框架
Optima是由清华大学和北京邮电大学联合推出的一个优化通信效率和任务有效性的训练框架。该框架通过迭代生成、排名、选择和训练范式,显著提高了基于大型语言模型(LLM)的多智能体系统(MAS)的通信效率和任务效果。Optima不仅减少了令牌使用,还为改进推理时间扩展法则提供了新的可能性。
41 6
Optima:清华联合北邮推出优化通信效率和任务有效性的训练框架
|
2月前
|
人工智能 语音技术 UED
仅用4块GPU、不到3天训练出开源版GPT-4o,这是国内团队最新研究
【10月更文挑战第19天】中国科学院计算技术研究所提出了一种名为LLaMA-Omni的新型模型架构,实现与大型语言模型(LLMs)的低延迟、高质量语音交互。该模型集成了预训练的语音编码器、语音适配器、LLM和流式语音解码器,能够在不进行语音转录的情况下直接生成文本和语音响应,显著提升了用户体验。实验结果显示,LLaMA-Omni的响应延迟低至226ms,具有创新性和实用性。
69 1
|
5月前
|
人工智能 自然语言处理 网络架构
单一作者论文,谷歌提出百万专家Mixture,超越密集前馈、稀疏MoE
【7月更文挑战第27天】谷歌提出了一种名为“百万专家Mixture”的神经网络架构,旨在解决Transformer模型处理大规模数据时面临的计算和内存效率问题。该架构通过利用“产品键”技术实现从大规模专家池中的高效检索,相较于传统密集前馈网络和稀疏MoE模型,在性能-计算权衡方面展现出明显优势。尽管如此,模型训练的复杂性和大规模模型的有效管理仍然是挑战。[链接](https://arxiv.org/abs/2407.04153)
68 2
|
7月前
|
数据采集 人工智能 自然语言处理
谷歌DeepMind发布Gecko:专攻检索,与大7倍模型相抗衡
【4月更文挑战第12天】谷歌DeepMind的Gecko模型以小巧身形(256维)展现出媲美大型语言模型的检索性能,且在MTEB基准测试中超越768维模型。采用两步蒸馏法训练,适用于多任务及硬件环境,尤其在多语言处理上表现出色。尽管训练成本高、泛化能力待优化,但其创新为文本嵌入技术带来新可能。
124 7
谷歌DeepMind发布Gecko:专攻检索,与大7倍模型相抗衡
|
机器学习/深度学习 缓存 分布式计算
Transformer取代者登场!微软、清华刚推出RetNet:成本低、速度快、性能强
Transformer取代者登场!微软、清华刚推出RetNet:成本低、速度快、性能强
129 0
Transformer取代者登场!微软、清华刚推出RetNet:成本低、速度快、性能强
|
机器学习/深度学习 自然语言处理 并行计算
大幅优化推理过程,字节高性能Transformer推理库获IPDPS 2023最佳论文奖
大幅优化推理过程,字节高性能Transformer推理库获IPDPS 2023最佳论文奖
133 0
|
机器学习/深度学习 存储 缓存
LLM推理提速2.8倍,CMU清华姚班校友提出「投机式推理」引擎SpecInfer,小模型撬动大模型高效推理
LLM推理提速2.8倍,CMU清华姚班校友提出「投机式推理」引擎SpecInfer,小模型撬动大模型高效推理
290 0
|
存储 并行计算 算法
斯坦福大学CS博士新作:新型Attention提速2-4倍,BERT单节点训练最快
斯坦福大学CS博士新作:新型Attention提速2-4倍,BERT单节点训练最快
276 0
|
自然语言处理 算法 数据可视化
加速100倍,性能媲美SOTA,浙大提出无数据知识蒸馏新方法FastDFKD
加速100倍,性能媲美SOTA,浙大提出无数据知识蒸馏新方法FastDFKD
140 0
|
编解码 自然语言处理 算法
谷歌、MIT「迭代共同认证」视频问答模型:SOTA性能,算力少用80%
谷歌、MIT「迭代共同认证」视频问答模型:SOTA性能,算力少用80%
118 0
下一篇
DataWorks