[论文解读] PointMamba: A Simple State Space Model for Point Cloud Analysis

Dingkang Liang, Xin Zhou+6 more|arXiv (Cornell University)|Feb 16, 2024

Remote Sensing and LiDAR Applications被引用 21

一句话总结

PointMamba 引入一个简单的状态空间模型，带有重新排序策略和 Mamba 模块，实现全局、线性复杂度的点云分析，凭借更少的参数和 FLOPs 超越 Transformer 基线。

ABSTRACT

Transformers have become one of the foundational architectures in point cloud analysis tasks due to their excellent global modeling ability. However, the attention mechanism has quadratic complexity, making the design of a linear complexity method with global modeling appealing. In this paper, we propose PointMamba, transferring the success of Mamba, a recent representative state space model (SSM), from NLP to point cloud analysis tasks. Unlike traditional Transformers, PointMamba employs a linear complexity algorithm, presenting global modeling capacity while significantly reducing computational costs. Specifically, our method leverages space-filling curves for effective point tokenization and adopts an extremely simple, non-hierarchical Mamba encoder as the backbone. Comprehensive evaluations demonstrate that PointMamba achieves superior performance across multiple datasets while significantly reducing GPU memory usage and FLOPs. This work underscores the potential of SSMs in 3D vision-related tasks and presents a simple yet effective Mamba-based baseline for future research. The code will be made available at \url{https://github.com/LMD0311/PointMamba}.

研究动机与目标

• 为 3D 点云引入线性复杂度的全局建模方法，以克服二次注意力成本。
• 通过简单的重新排序策略，将状态空间模型 (SSMs) 适应于不规则、无序的点数据。
• 证明 PointMamba 在分类和部分分割任务上可以超过基于 Transformer 的方法。
• 展示 PointMamba 在降低参数和 FLOPs 的同时达到具有竞争力甚至更高的准确性。

提出的方法

• 使用基于轻量级 PointNet 的分词器，通过 FPS 和 KNN 将点补丁嵌入为标记单元以形成补丁。
• 应用一种重新排序策略，对标记按 x、y、z 轴排序并将序列三倍扩展，以在 Mamba 中实现因果建模。
• 通过一系列带有层归一化、SSM、深度卷积和残差连接的 Mamba 块处理重新排序后的标记。
• 采用 Point-MAE 风格的预训练进行训练，使用掩码补丁重构和 Chamfer 距离作为重构损失。
• 可选地消融设计选择，如 CLS token 的使用和特征维度，以研究它们对性能的影响。

实验结果

研究问题

• RQ1 具有线性复杂度的状态空间模型是否能像基于注意力的 Transformer 那样有效地捕捉点云的全局上下文？
• RQ2 对点标记进行简单的几何重新排序是否能够在非序列化的点数据中实现 Mamba 的因果建模？
• RQ3 设计选择（CLS token、标记维度）对 PointMamba 在分类与分割任务中的性能有何影响？
• RQ4 在标准三维基准上，与 Transformer 基线相比，PointMamba 在预训练+下游任务设置中的表现如何？

主要发现

• PointMamba 在多个点云任务上优于基于 Transformer 的对手，同时使用更少的参数和 FLOPs（约少 44.3% 的参数和约 25% 的 FLOPs）。
• 在 ScanObjectNN 中，PointMamba 在从零训练时获得比基线 Transformer 更高的 OBJ-BG 和 OBJ-ONLY 分数，并且在预训练变体中保持竞争力。
• 一种简单的重新排序策略提升了 Mamba 对非因果、无序点云的建模能力，带来相对于 1x（无重新排序）设置的可测量准确性提升。
• 省略 CLS token，并对最终块输出进行平均池化，获得 PointMamba 的最佳分类性能。
• PointMamba 在 ShapeNetPart 上的部分分割结果具有竞争力，参数显著少于许多基线，并且在长序列上表现出较好的内存效率。