极市导读
本文提出了第一个点云分析的状态空间模型,保留了线性复杂度和全局建模能力,同时避免复杂的结构设计。>>加入极市CV技术交流群,走在计算机视觉的最前沿
首个基于状态空间模型的点云分析方法PointMamba再次迎来更新,带来了架构和性能的全面升级。本文着眼于讨论Mamba相比Transformer的架构优势,PointMamba在保持线性复杂度和全局信息建模的条件下,使用Vanilla Mamba即在点云分析数据集上达到了优越的性能,在长序列输入下相比Transformer方法在计算开销等有显著的提升,为后续研究提供了一个简单且高效的新基线。
论文标题:PointMamba: A Simple State Space Model for Point Cloud Analysis
论文地址:https://arxiv.org/abs/2402.10739
代码地址:https://github.com/LMD0311/PointMamba
项目主页:https://lmd0311.github.io/projects/PointMamba/
PointMamba做了什么?
目前大量的工作使用Transformer结合自监督学习范式,不断将点云分析任务的性能提升到新高度。作为Transformer中的核心部分,attention机制能有效地建模全局信息,但其二次复杂度也带来了额外的计算开销,极大地限制了模型拓展到长序列的能力。关注到最近NLP领域Mamba架构带来的线性复杂度优势,我们期望为点云分析设计一种线性复杂度且保留全局建模能力的点云分析模型。在本文中,我们提出了第一个点云分析的状态空间模型,保留了线性复杂度和全局建模能力,同时避免复杂的结构设计;具备许多优点:结构简洁(使用Vanilla Mamba)、低计算开销、知识可迁移(支持自监督学习)等。
PointMamba怎么做的?
PointMamba 的设计遵循奥卡姆剃刀原则, 确保了结构的尽可能简洁。首先, 输入点云 , 通过 FPS 选取 个关键点, 得到 。因为在 FPS 中 个关键点为随机选取, 可能为单向建模的 Mamba 模型带来挑战, 因此选择使用空间填充曲线来扫描并序列化点云。具体而言, 使用三维希尔伯特曲线 (Hilbert) 和其转置变体 (Trans-Hilbert)在三维空间中按顺序遍历全部关键点, 得到 和 。接着使用通用的 Point tokenizer, 依据前述两组关键点得到两组序列化的 Point tokens, 记为 和 。此外, 引入了一个极其简单的顺序指示器(order indicator)以区分不同扫描策略的 token 特征, , 其中 分别为缩放因子。将变换后的两组特征进行拼接得到 Mamba encoder 的输入 。
Mamba encoder为N个Vanilla Mamba block组成。为了更好地理解所提出的PointMamba具有的全局建模能力,我们提供了一个直观的示意图。如图所示,在对第一组token(即基于Hilbert的标记)建模后,积累的全局信息可以改进下一组标记(即基于Trans-Hilbert的标记)的序列化过程。这种机制可确保Trans-Hilbert序列中的每个序列化点都能从之前处理过的Hilbert序列的整个历史中获得信息,从而使建模过程具有更丰富的上下文和全局建模处理。
模型的pipeline如下:
此外,PointMamba提出了序列化的掩码建模预训练方式,通过随机选取一种空间填充曲线(Hilbert或Trans-Hilbert)序列化点云,并使用对应的order indicator得到encoder输入。经过Vanilla Mamba encoder的特征编码,在被mask掉token的地方插入mask token,经过Vanilla Mamba decoder 和一个投影层将mask token恢复为原始点云。自监督预训练pipeline如下:
PointMamba的效果如何?
在真实世界点云分类数据集ScanObjectNN、模拟数据集ModelNet40、Fewshot learning、部件分割中,PointMamba取得了优异的性能,具体情况如下:
在效率方面,得益于Mamba的架构优势,除去tokenizer只测试encoder的情况下,当token序列增长到较长序列(如32,768)时,相较Point-MAE(Vanilla Transformer)在推理速度、推理显存占用、FLOPs上提升30.2×,24.9×,5.2×。
此外,以下是掩码建模和部件分割的可视化:
总结
总结来说,PointMamba作为首个状态空间的点云分析模型,通过设计输入序列化为单向建模的Mamba具备高效的全局建模能力,用最简洁的Vanilla架构验证了新架构相比Transformer的优势,在点云分析任务中取得优越的性能并展现了巨大的潜力。

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

点击阅读原文进入CV社区
收获更多技术干货

