PointMamba: 点云分析的简单状态空间模型

在这里向大家介绍我们的最新工作：PointMamba: A Simple State Space Model for Point Cloud Analysis。状态空间模型作为序列建模的新方法最近在许多领域有了积极的进展，在本文中我们构建了一个具有全局建模和线性复杂性的3D点云分析新框架，性能优于基于 vanilla Transformer 的同类工作，同时显著节省了约 44.3% 的参数和 25% 的 FLOPs。据我们所知，PointMamba是第一个把状态空间模型运用到点云分析任务上的文章。论文、模型、代码均已开源，欢迎大家引用和star✨。

论文：

https//arxiv.org/abs/2402.10739

代码：

https//github.com/LMD0311/PointMamba

引言

点云数据的不规则性和稀疏性一直是3D视觉领域中的难题。虽然Transformer基于其强大的全局信息建模能力，在点云分析任务中展现了潜力，但其计算复杂度随着输入长度的增加而显著增长，限制了其在长序列模型上的应用。在此背景下，我们提出了PointMamba，一个结合了状态空间模型（SSM）优势的框架，旨在通过具有线性复杂度的全局建模方法，克服现有技术的限制。

模型架构

Point Tokenizer: 与Point-BERT/Point-MAE类似，我们首先采用一个简单的Point Tokenizer来生成Point Patches，通过最远点采样（FPS）和K最近邻（KNN）算法将输入点云划分为一系列局部区域，然后将这些区域映射到特征空间中，形成Point tokens。

Reordering Strategy: 为了让SSM能够更有效地捕获点云的全局结构，我们提出了一种重排序策略。这一策略通过基于Point Patches的几何中心坐标，沿 x,y,zx,y,z 轴顺序的几何顺序重新组织Point tokens，以逻辑上的扫描顺序来增强模型的全局建模能力。

Mamba Block: 重排序后的Point Tokens将被送入一系列Mamba block进行处理。Mamba block结合了深度可分离卷积和selective SSM，通过有效捕获点云的全局和局部特征，对点云结构进行建模。

预训练策略

为了进一步提升PointMamba的性能，我们还探索了基于Point-MAE的预训练策略，通过mask部分Point Patches并训练模型去重建它们，模型学习到了丰富的点云表示。这一步不仅提高了模型对点云结构的理解能力，也为后续的下游任务提供了强大的特征表示。具体公式为：

值得注意的是，Decoder的输入将mask token 拼接在visible token 之后，满足了SSM的单向建模特点。

实验验证

我们在多个点云分析任务上验证了PointMamba的性能，包括物体分类和部分分割。实验结果证明，PointMamba不仅在效率上具有显著优势，同时在准确度上也超过了许多现有的基于Transformer的模型。

结论

总结来说，PointMamba作为一个新颖的点云分析工具，通过其高效的全局建模能力和简洁的模型结构，在3D视觉任务中展现了巨大的潜力。我们希望这项工作能够激励更多的研究，共同推动点云分析领域的进步。

公众号后台回复“数据集”获取100+深度学习各方向资源整理