​港大开源 PoinTramba | Transformer 与 Mamba 混合让点云模型拥有 Mamba的效率！

点击下方卡片，关注 「AI视界引擎」 公众号

点云分析由于深度学习而取得了实质性进展，尽管先前的基于Transformer的方法在建模这一任务上的长距离依赖性表现出色，但它们的计算需求相当大。相比之下，Mamba提供了更高的效率，但与基于Transformer的方法相比，其潜力有限。

在本文中，作者引入了PoinTramba，这是一个开创性的混合框架，它将Transformer的分析能力与Mamba卓越的计算效率相结合，以增强点云分析。

具体来说，作者的方法首先将点云分割成组，其中Transformer精心捕捉复杂的组内依赖关系并生成组嵌入，而高效的Mamba架构将同时巧妙地捕捉到这些组嵌入之间的组间关系，确保了全面的分析。

与先前的Mamba方法不同，作者引入了一种双向重要性感知排序（BIO）策略来解决随机排序效应的挑战。这种创新策略根据计算的重要性分数智能地重新排序组嵌入，显著提高了Mamba的性能并优化了整体分析过程。作者的框架通过无缝集成这些先进技术，实现了计算效率与分析性能之间的优越平衡，标志着点云分析的一个重大飞跃。在ScanObjectNN、ModelNet40和ShapeNetPart等数据集上的广泛实验证明了作者方法的有效性，为点云识别建立了一个新的最先进的分析基准。

有史以来第一次，这个范式利用了Transformer和Mamba架构的相结合的优势，为该领域树立了一个新标准。

代码可在https://github.com/xiaoyao3302/PoinTramba获取。

1 Introduction

点云作为关键的3D视觉信号，包含了重要的几何信息，已经引起了越来越多的研究兴趣[1; 2]。由于深度学习方法的兴起，点云分析任务，如分类和分割，已经取得了显著的成功。早期方法要么利用辅助数据结构，如 Voxel ，通过3D卷积操作进行处理[3; 4]，要么直接使用多层感知机（MLPs）从原始点集提取视觉表示。然而，这些方法主要依赖于局部依赖关系，往往忽视了全局特征。为了增强点云的全局感知，最近的研究（如PointBERT[8]和PointMAE[9]）引入了带有注意力机制的Transformer[10]来处理长序列。尽管Transformer强大的长距离建模能力直观地提升了分析性能，但全面的机制不可避免地导致了高计算复杂度和内存使用，使得这些方法实用性降低。

另一方面，为了避免在建模长距离关系时产生大的内存负担，引入了一种名为Mamba[11]的新架构，其中包含了状态空间模型[12, 13]模块。这种架构已经在各种自然语言处理[14, 15]和计算机视觉[16, 17]任务中取得了成功。然而，尽管基于Mamba的方法效率高，但在点云分析[18, 19]中，与同等规模的基于Transformer的方法相比，其性能仍然落后。此外，将Mamba应用于点云分析仍需进一步研究。

例如，一个重要的挑战是点云本质上是无序的，而Mamba主要是为处理结构化数据而设计的。因此，Mamba在处理点云数据方面的有效性仍然不确定。

在这项研究中，作者引入了PoinTamba，一个用于点云分析的新型混合框架，它结合了Transformer的强大分析能力和Mamba的效率。具体来说，作者将点云划分为不同的组，并利用Transformer和Mamba分别建模组内和组间关系。最初，利用Transformer在建模长距离依赖方面的卓越能力，作者用它捕捉组内依赖并生成增强的组内特征，即组嵌入。尽管由于组内点数较少，使用Transformer建模组内点云仍然是可行的，但这种方法在建模组间依赖关系时会导致复杂性显著增加，因为组数较多。为了有效地建模组间依赖，作者整合了Mamba，将复杂性降低到线性规模。具体来说，作者将Transformer产生的组嵌入输入到Mamba编码器中。该编码器从每个点云样本中提取组间全局特征，促进分类和分割等分析任务。

特别是，作者没有直接将Mamba应用于无序的点云组，而是引入了一种双向重要性感知排序（BIO）策略。这种方法重新排序组，以减轻随机点云排序对作者Mamba编码器的负面影响。与之前的方法[18, 19]（如使用传统的z排序[20]或Hilbert排序[21]）不同，作者提出了一种算法，为每个组嵌入学习一个“重要性”分数，并据此重新排序组。具体来说，作者将组嵌入和Mamba编码器的全局特征映射到同一特征空间。然后，作者计算每个组嵌入与全局特征之间的余弦相似度。此外，作者利用一个重要性分数预测模块，使每个组嵌入能够预测与计算出的余弦相似度紧密对齐的重要性分数。基于这些重要性分数，作者以双向方式重新排序组嵌入，结合降序和升序。这确保了每个组嵌入从所有其他组嵌入中聚合额外的信息。这种重新排序策略充分利用了Mamba处理结构化数据的潜力，从而产生了更精细的全局组间特征，并显著提升了分析性能。

通过将Transformer和Mamba模块整合到作者的端到端混合分析框架中，作者在保持有效复杂性的同时，在基准数据集上实现了可比较的性能结果。

作者的贡献可以总结如下：

1. 作者提出了PoinTamba，一个结合了Transformer和Mamba的新型混合框架，用于高效有效的点云分析。该框架利用Transformer强大的建模能力产生增强的组内特征，以及Mamba的线性复杂度从大量的组嵌入中生成组间特征。

2. 作者提出了一种新的双向重要性感知排序策略来重新排序组嵌入。这种操作可以充分利用Mamba的递归特性，更好地处理结构化的组顺序，并从所有其他组嵌入中聚合额外的信息。

3. 作者在ScanObjectNN、ModelNet40和ShapeNetPart基准数据集上进行了广泛的点云分类和分割分析实验。这些实验表明，作者的方法取得了可比较的定量结果，并验证了其有效性。

2 相关工作

Deep Learning on Point Cloud Recognition

点云识别的深度学习受到了极大的关注，因为点云包含了丰富的深度信息，可以用于自动驾驶和机器人技术等领域[22]。与2D图像相比，3D点云是稀疏的，且点是无序的，这使得直接处理点云进行识别变得困难。

早期的分析研究直接利用了辅助数据结构，如 Voxel ，这使得传统的深度学习方法，如3D卷积[23; 24]，得以应用。然而，使用这些额外的数据结构导致了显著的计算和内存成本。

面对这个问题，PointNet[5]是一项开创性工作，提出使用MLPs直接处理点集，这种方法简单但取得了有希望的性能，并启发了一系列工作，如PointNet++[6]和DGCNN[7]等。然而，这些方法仅依赖于局部依赖关系，往往忽略了全局特征，因此性能有限。

最近，受到Transformer架构在自然语言处理和2D视觉中对长距离上下文建模成功的启发[10; 25]，各种研究提出了探索Transformer架构在点云[26; 27; 8; 9]中的有效性，并取得了优异的性能。

然而，Transformer中的注意力模块的计算复杂度是，其中表示输入 Token 的数量。因此，尽管Transformer在点云识别方面表现出色，但当输入 Token 的数量增加时，基于Transformer的方法将导致显著的计算开销，限制了这些方法的处理能力。

State Space Models

状态空间模型被广泛应用于通过连接输入和输出并使用潜在状态来模拟系统来处理信息序列[28]。近期，受结构化状态空间（S4）模型[13]的启发，各种研究[29; 30]提出了使用状态空间模型来模拟长距离依赖关系。特别是，新提出的Mamba[11]因其在全球感知方面的巨大潜力而引起了广泛关注。与Transformer的计算复杂度相比，Mamba的计算复杂度为，这使得它在序列长度上具有线性可扩展性。受到Mamba成功的启发，各种研究提出了检查Mamba在2D视觉任务上的有效性，包括图像分类[16; 17]和医学图像分割[31]等。然而，这些基于Mamba的方法仍然无法在与Transformer相当大小的情况下超越基于Transformer的方法[32]。因此，各种研究提出了结合Transformer和Mamba的优点以获得更好的性能和更低计算复杂度，这仍然是一个开放性问题，值得探索[32; 33]。

Mamba on Point Clouds

受到Mamba在2D视觉任务上成功的启发[16; 17; 31]，一些最近的工作提出了检验Mamba在3D点云上的有效性。特别是，PointMamba[18]将Mamba编码器与标准的点云分组划分操作相结合，而PCMamba[19]将Mamba与PointMLP[34]结合，以提取Mamba在点云上的潜力。

与序列语言或图像数据相比，点云高度无序。然而，Mamba是一个循环模型，隐藏状态的确顺序将严重影响Mamba对长距离依赖建模的性能。因此，将Mamba适应于点云的关键问题在于点云的排序策略。为了解决这个问题，PointMamba沿着轴线重新排序点云，而PCMamba根据包括z排序[20]，Hilbert排序[21]等多个顺序重新排序点云。然而，这些排序策略是否适合用Mamba进行点云分析还是启发式的。

3 Method

PoinTramaba

作者的方法流程如图1（a）所示，给定一个输入点云集，作者旨在生成一个全局特征，以用于各种下游任务。作者的方法首先将分割成个点群。对于每个点群，作者采用一个_Intra-group Transformer编码器_，包含个Transformer层，生成个群嵌入，其中表示第个点群，如图1（b）所示。随后，作者引入一种双向重要性感知排序（BIO）算法，以双向方式重新排序群嵌入，将降序和升序结合成。这些重新排序的嵌入随后通过一个_Inter-group Mamba编码器_，如图1（c）所示。经过个Mamba层的更新，作者得到一组更新后的特征，这组特征通过一个重要性感知池化（IAP）操作来生成全局特征。注意，上标表示第个Mamba层，其中表示Mamba层的输入。

Intra-group Transformer and Inter-group Mamba Encoder

组内Transformer编码器。 作者首先将点云分割为组，每组包含个点，正式表示为，_即_，。最初，作者采用最远点采样（FPS）算法选择个关键点。随后，作者使用K最近邻（KNN）算法为每个关键点找到个最近邻。在这种表示法中，下标指的是第组点，而表示第组内的第个点。在将点云分割为组之后，作者利用Transformer在建模长距离依赖方面的卓越能力来捕捉组内依赖并生成增强的组内特征，这些特征被称为组嵌入。具体来说，对于每个点，作者首先将点的坐标，_即_，，通过Transformer的标准嵌入层投影为点嵌入和点位置嵌入[10; 25; 35]。这些点嵌入和点位置嵌入的总和随后被送入包含层的标准Transformer编码器，以建模组内关系并生成组嵌入。

组间Mamba编码器。尽管由于每个组内的点云数量相对较少，使用Transformer对组内点云进行建模是可行的，但是对大量组之间（组间）的依赖关系进行建模将显著增加计算复杂性。为了解决这个问题，作者集成了Mamba编码器，该编码器在建模组间依赖时，能有效地将复杂性降低到线性规模。具体来说，在获得组嵌入之后，作者提出了一种BIO策略来获得重新排序的组嵌入。关于BIO策略的更多细节将在3.3节中提供。同样，作者使用BIO策略重新排序关键点坐标的位置嵌入。通过将重新排序的组和位置嵌入相加，作者得到集合的聚合，该聚合将由层Mamba编码器进行更新。如图1（c）所示，作者采用了来自[18]的Mamba层，这是一个标准的Mamba层，可以表示为：

（

其中 表示第 个Mamba层。这里， 指的是深度可分卷积， 表示层归一化操作， 代表状态空间模型，它作为作者Mamba编码器的基本层，而 表示点积。此外， 表示SiLU [36] 操作。通过 个Mamba层更新组嵌入后，作者得到最终更新的组嵌入 。这些更新的组嵌入随后被送入一个重要性感知池化（IAP）操作以提取全局特征 。通过利用长距离关系建模和在序列长度上的线性可扩展性，作者使Mamba编码器能够进行全局感知。

Importance-aware Ordering and Pooling

双向重要性感知排序. 由于Mamba [11] 的循环特性和点云的无序性，将随机排序的群组嵌入输入到Mamba编码器会显著影响Mamba在长距离依赖建模方面的性能。为了解决这个问题，作者提出了一种新颖的双向重要性感知排序（BIO）策略来重新排序群组嵌入，减轻随机点云排序的负面影响。

具体来说，如图2所示，对于给定的群组嵌入 ，作者采用了两个非线性投影层，即嵌入投影层和全局投影层。这两个层分别将群组嵌入和点云的全局特征 映射到同一个特征空间，即 和 。注意作者稍后会介绍点云的全局特征。因此，作者可以计算 和 之间的余弦相似度 作为重要性得分，可以表示为：

然而，为每个群组嵌入计算 并不可行，因为在排序阶段需要已知全局特征，这是不切实际的。为了克服这个限制，作者使用一个非线性投影层来预测每个群组嵌入的相应重要性得分，即 。作者使用重要性损失 鼓励 接近

重要性感知池化。鉴于更新后的组嵌入，作者执行一个池化操作，从这些嵌入中获取全局特征。回想一下，预测的组嵌入重要性得分表示组嵌入与全局特征之间的余弦相似性，突显了每个组嵌入的重要性。负的重要性得分表明相应的组嵌入对全局特征有负面影响。因此，作者 Proposal 丢弃具有负重要性得分的组嵌入，并使用它们的重要性得分对剩余的嵌入进行重加权。这个过程表示为：

其中是单位阶跃函数。点云的生成全局特征可用于各种下游任务，如分类或分割，使用相应的损失。

Objective Function

作者采用多方面的损失计算方法来进行优化。

具体来说，作者融合了以下损失组件：

1. 任务损失，用于优化特定的下游任务，比如分类或分割，确保性能的稳健。
2. 重要性损失，用于优化不同组嵌入的重要性分数的排序。
3. 对齐损失，用于优化重要性分数的预测，确保模型学习到具有语义意义的嵌入投影层。作者使用三个权衡参数、和来平衡每个损失组件。

因此，作者通过以端到端的方式解决以下优化问题来训练整个网络：

4 实验研究

Experimental Protocols

遵循先前的工作[18]，作者在三个下游任务上评估了作者的方法的有效性，即实物体分类在ScanObjectNN [39]上，合成物体分类在ModelNet40 [4]上，以及部分分割在ShapeNetPart [40]上。数据集的详细信息在附录A.1中提供，实现细节在附录A.2中概述。

Experimental results

在ScanObjectNN上的现实世界目标分类。 作者首先在表1中与其他方法比较了作者的PointTramba在现实世界目标分类基准数据集ScanObjNN上的表现。从表中可以推理，作者的方法明显超过了当前的方法，包括基于MLP的方法如PointNet++和基于Transformer的方法如PointMAE。具体来说，作者的方法分别在三种子ScanObjNN上的表现超过了之前的最佳状态（SOTA）4.0%、3.5%和1.0%。值得注意的是，PointTramba的参数数量甚至少于标准的基于Transformer的方法，如PointBERT和PointMAE。

尽管一些基于MLP的方法，如PointCNN，引入了更少的参数，但它们的性能显著落后于作者的方法。主要原因是作者的方法中的组内Transformer编码器增强了局部特征提取，优于其他依赖MLP完成这一任务的方法。

此外，组内Transformer编码器是轻量级的，确保了PointTramba与PointMamba相比没有引入显著更多的参数。还值得一提的是，PCM使用的参数几乎是作者方法的两倍，但作者的方法仍然优于它，突显了作者方法的有效性。

在ModelNet40上的合成目标分类。 然后，作者在表2上与其他方法比较了作者的PoinTramba在合成目标分类基准数据集ModelNet40上的表现。注意，作者重现了PCM的结果，因为它遵循与作者的不同设置。从表中可以推理，作者的方法实现了SOTA性能，大幅超过了当前的基于Transformer和基于Mamba的方法。

在ShapeNetPart上的部分分割。 最后，作者在部分分割基准数据集ShapeNetPart上与其他方法比较了作者的PoinTramba的表现，见表3。注意，作者重现了PCM的结果，因为它不仅需要坐标，还需要额外的四维信息作为输入。从表中可以推理，尽管ShapeNetPart是一个高度竞争的基准，但作者的方法与之前的SOTA方法取得了相当的性能，验证了作者方法的有效性。

Ablation Study

在本节中，作者强调了模块设计对ScanObjNN数据集的PB-T50-RS变体的贡献，如表4所示。从表中可以看出，仅部署组间Mamba编码器就能使识别性能比PointNet++提高8.2%（变体1与变体0），这主要归功于Mamba的全局感知能力。

将组内Transformer编码器添加到组间Mamba编码器进一步提高了性能0.4%（变体3与变体1），这主要归功于Transformer的长距离建模能力，以捕捉组内依赖。将作者的BIO策略与PoinTramba模型结合使用，带来了2.1%的性能大幅提升（变体6与变体3），突显了作者的排序算法的有效性。此外，重要性感知池化操作帮助模型关注有用的组嵌入，同时忽略不相关的嵌入，与随机排序策略相比，性能提高了0.5%（变体7与变体6）。将BIO和IAP与仅Mamba编码器结合也能达到87.9%的识别准确率，超过Mamba编码器1.8%（变体5与变体1）。此外，考虑到BIO策略包括一个增强局部特征提取的对齐损失，作者进一步验证了的有效性。如表所示，仅采用就能使Mamba模型的性能提高1.4%（变体2与变体1），PoinTramba模型的性能提高1.7%（变体4与变体3）。此外，单独的BIO策略带来了0.4%的性能提升（变体4与变体6）。这些消融研究证明了作者方法中每个组件的有效性。

然后，作者通过采用PoinTramba作为主干网络，验证了排序策略对Mamba模型的重要性，并比较了六种不同的排序策略，即随机排序策略、基于坐标的排序策略（即XYZ排序策略[18]、z排序策略[20]、Hilbert排序策略[21]、单向重要性感知排序策略（SIO）和双向重要性感知排序策略（BIO）。

结果如图3（a）所示。从表中可以看出，随机排序策略比基于坐标的排序策略性能高出0.8%。根据作者的分析，主要原因是随机排序策略将既远离又接近某个嵌入的嵌入放置在该嵌入的邻近区域。这种方式使得嵌入从长距离和短距离聚合信息，从而具有竞争力的全局感知。

值得注意的是，作者的BIO策略比BIO策略高出0.5%，因为BIO策略确保每个嵌入从所有其他嵌入聚合信息，从而实现更好的全局感知性能。此外，作者的BIO策略与其他排序策略相比具有显著优势，表明其有效性。

最后，作者验证了作者的重要性感知池化操作（IAP）的重要性，如图3（b）所示。作者将IAP与几种不同的池化策略进行比较，包括平均池化策略、最大池化策略和加权和策略。从图中可以推理，防止模型受到无用组嵌入的影响，与直接加权和策略相比，性能提高了0.6%，这表明了作者的IAP策略的有效性。

5 Conclusion

在这项工作中，作者介绍了PoinTramba，这是一个开创性的混合框架，它将Transformers强大的建模能力与Mamba的计算效率相结合，用于点云分析。通过整合这两种架构，PoinTramba在计算复杂度与分析性能之间取得了优越的平衡。此外，作者创新性的BIO策略显著地利用了Mamba的性能。

广泛的实验证明了PoinTramba的有效性，在点云分析领域建立了新的最新基准。这种新颖的方法利用了Transformer和Mamba架构的联合优势，是该领域的一个重要进展。对于未来的工作，作者计划探索进一步的优化技术以增强PoinTramba的可扩展性，并研究其在更广泛点云任务中的应用，进一步细化作者的排序策略以提升效率和性能。

局限性。 在本研究中，作者仅关注了重要性感知排序策略。然而，目前尚不清楚这种方法是否是最优的排序算法，能够充分利用Mamba的潜力。此外，需要在更广泛的任务上进行进一步实验，以全面评估作者的PoinTramba的能力。

Datasets

根据先前的工作[18]，作者在三个下游任务上评估了作者的方法的有效性，即实物体分类任务ScanObjectNN [39]，合成物体分类任务ModelNet40 [4]，以及部件分割任务ShapeNetPart [40]。

ScanObjNN数据集[39]是一个具有挑战性的点云物体分类数据集，包含15,000个来自现实世界的点云目标，分为15个类别，其中包括三个变体：OBJ_BG、OBJ_ONLY和PB_T50_RS。

遵循先前的工作[18, 19]，作者将每个样本的点数设置为2048。ModelNet数据集[4]是一个广泛使用的点云物体分类数据集，涵盖40个类别，每个类别包括100个合成样本。遵循先前的工作[18, 19]，作者将每个样本的点数设置为1024。

ShapeNetPart [40]数据集是一个具有挑战性的点云部件分割数据集，包含来自16个不同形状类别的16,880个样本，具有50个部件标签。遵循先前的工作[18, 19]，作者将每个样本的点数设置为2048。

Implementation Details

在作者的所有实验中，作者采用了AdamW[46]优化器以及余弦调度器。作者训练模型300个周期，包括一个10个周期的 Warm up 阶段。所有分类模型的初始学习率设置为3e-4，而部分分割模型的初始学习率设置为2e-4，所有实验的权重衰减为5e-2。

所有任务的批处理大小设置为32。作者的组内Transformer编码器包含4个Transformer层，通道大小为384，不使用分类标记。作者的组间Mamba编码器包含12个Mamba层，通道大小为384，同样不使用分类标记。对于ScanObjectNN和ModelNet数据集上的分类任务，点云样本被划分为256组，每组由16个点组成。对于ShapeNetPart数据集上的部分分割任务，点云样本被划分为128组，每组由32个点组成。

投影层，包括嵌入投影层和全局投影层，是两层卷积层，隐藏维度为128，输出维度为256，其后是批归一化层、ReLU函数和归一化操作。重要性得分预测头是一个两层卷积层，隐藏维度为128，输出维度为1，其后是批归一化层和ReLU函数。超参数 、 和 对于分类任务均设置为1.0，对于部分分割任务则设置为5.0、1.0和1.0。

作者的大部分实验设置来自PointMamba[18]，包括分类头和分割头。

作者的所有分类模型都在配备有十个NVIDIA RTX 3090 GPU的服务器上训练，而所有分割模型则在一个配备有八个NVIDIA A100 GPU（80GB内存）的服务器上训练。

Appendix B Additional Ablation Study

在本节中，作者进行了额外的消融研究，以验证作者的方法的有效性。首先，作者研究了作者的方法对组数和每组点数的敏感性。特别是，作者将点云划分为32到1024个组，每个点组内的点数从128变化到8。结果在表5中报告。如表所示，过小的组大小导致每个点组的语义信息有限。

因此，组内Transformer编码器可能无法编码每个点组的潜在语义特征，导致PoinTramba的识别性能有限。相反，当组大小过大时，轻量级的组内Transformer编码器，即4层Transformer编码器，可能无法为每个点组编码丰富的特征，也会导致识别性能有限。值得注意的是，由于Mamba卓越的长距离建模能力，改变点组数量并不会显著影响作者PoinTramba模型的性能。

最后，作者检查了PoinTramba对Transformer编码器中层数的敏感性。

结果列在表6中。从表中可以推理，当层数较少时，例如1层或2层，组内Transformer编码器过于轻量级，无法捕捉到每个点组的丰富语义特征。在这种配置下，PoinTramba并没有显著优于使用MLPs作为组内编码器的PoinTramba模型，即当时。然而，当使用四层Transformer编码器时，组内编码器有效地建模了每个点组内的长距离依赖关系，从而获得了卓越的识别性能。

点击上方卡片，关注 「AI视界引擎」 公众号