[论文解读] Mamba3D: Enhancing Local Features for 3D Point Cloud Analysis via State Space Model

Xu Han, Yuan Tang+2 more|arXiv (Cornell University)|Apr 23, 2024

Remote Sensing and LiDAR ApplicationsEnvironmental Science被引用 3

一句话总结

Mamba3D 提出了一种专为3D点云分析设计的状态空间模型，通过引入局部归一化池化（LNP）模块显式提取局部几何特征，并采用新型通道级反向状态空间模型（C-SSM）的双向状态空间模型（bi-SSM），以增强全局表征能力。该方法在保持线性复杂度的同时实现了最先进性能，在从零开始训练的 ScanObjectNN 上达到 92.6% 的准确率，在单模态预训练的 ModelNet40 上达到 95.1% 的准确率。

ABSTRACT

Existing Transformer-based models for point cloud analysis suffer from quadratic complexity, leading to compromised point cloud resolution and information loss. In contrast, the newly proposed Mamba model, based on state space models (SSM), outperforms Transformer in multiple areas with only linear complexity. However, the straightforward adoption of Mamba does not achieve satisfactory performance on point cloud tasks. In this work, we present Mamba3D, a state space model tailored for point cloud learning to enhance local feature extraction, achieving superior performance, high efficiency, and scalability potential. Specifically, we propose a simple yet effective Local Norm Pooling (LNP) block to extract local geometric features. Additionally, to obtain better global features, we introduce a bidirectional SSM (bi-SSM) with both a token forward SSM and a novel backward SSM that operates on the feature channel. Extensive experimental results show that Mamba3D surpasses Transformer-based counterparts and concurrent works in multiple tasks, with or without pre-training. Notably, Mamba3D achieves multiple SoTA, including an overall accuracy of 92.6% (train from scratch) on the ScanObjectNN and 95.1% (with single-modal pre-training) on the ModelNet40 classification task, with only linear complexity. Our code and weights are available at https://github.com/xhanxu/Mamba3D.

研究动机与目标

• 解决基于Transformer的3D点云模型中二次方复杂度的瓶颈问题，以提升分辨率与可扩展性。
• 克服直接将Mamba应用于无序点云时存在的伪序依赖性与缺乏显式局部特征提取的问题。
• 设计一种可扩展且高效的架构，在不依赖预训练的前提下保持高性能，同时支持强大的 few-shot 与预训练泛化能力。
• 显式地将局部几何特征集成到SSM框架中，以提升点云表示学习能力。
• 设计一种双向SSM，将特征通道视为有序序列，以在无序点云中稳定全局特征学习。

提出的方法

• 提出局部归一化池化（LNP）模块，结合K-归一化（相对于K个最近邻归一化特征）与K-池化（通过简单且可微的操作聚合特征），实现局部几何特征学习。
• 引入双向状态空间模型（bi-SSM），由一个标记前向SSM（L+）与一种新型通道级反向SSM（C-SSM）组成，后者沿通道维度处理特征，以稳定全局上下文学习。
• 设计C-SSM使其在特征通道上作为伪有序序列运行，降低对任意点序的依赖，提升对无序输入的鲁棒性。
• 将LNP模块实现为仅含0.3M参数，以保持模型效率与可扩展性。
• 采用简单且参数高效的K-池化操作，避免使用昂贵的MLP或池化层，从而减少FLOPs与参数量，同时保持性能。
• 将LNP与bi-SSM模块整合为统一架构，支持标准协议下的端到端训练与预训练（如Point-BERT与Point-MAE）。

实验结果

研究问题

• RQ1 状态空间模型能否被有效适配于3D点云学习，以克服Transformer模型的二次方复杂度，同时保持或提升性能？
• RQ2 如何在Mamba这类缺乏局部结构归纳偏差的序列模型中，显式注入局部几何特征？
• RQ3 在反向SSM中将特征通道视为有序序列，是否能改善无序点云中的全局特征学习？
• RQ4 像Mamba3D这样线性复杂度的架构，能否在不依赖重型预训练的情况下实现点云分类任务的最先进性能？
• RQ5 与自注意力机制和单向SSM相比，所提出的bi-SSM在点云基准测试中的鲁棒性与性能表现如何？

主要发现

• 在从零开始训练的情况下，Mamba3D在ScanObjectNN数据集上实现了92.6%的整体准确率，成为无预训练模型中的新SOTA。
• 在单模态预训练设置下，Mamba3D在ModelNet40数据集上达到95.1%的整体准确率，成为单模态预训练模型中的新SOTA。
• 消融实验表明，移除LNP模块导致准确率下降1.2%，移除bi-SSM模块导致准确率下降2.3%，证实二者在性能中的关键作用。
• 用自注意力或单向SSM替代bi-SSM导致准确率下降1.2%；而使用点顺序翻转（token-flip）变体则导致性能下降1.1%至1.4%，证明C-SSM设计的有效性。
• K-池化优于最大池化、平均池化及其组合，当被替换时准确率下降0.7%至2.8%，表明其在局部特征聚合中的优越性。
• 模型在不使用任何输入排序策略时表现最佳，表明Mamba3D能有效从无序点云中捕捉语义特征，且在k=4与L=128时达到峰值性能。