Mamba3D: Enhancing Local Features for 3D Point Cloud Analysis via State Space Model

TL;DR

Mamba3D introduces a linear-complexity, SSM-based backbone for 3D point clouds that explicitly integrates local geometry via Local Norm Pooling and global context via bidirectional SSM (bi-SSM). Patch embeddings create patch-level tokens, fed into an encoder that alternates a token-level LNP with a channel-level bi-SSM, aided by a [CLS] token and positional encoding. Local propagation and aggregation are handled by K-norm and K-pooling within LNP, while a forward L+SSM and backward C-SSM handle global features without relying on token order. Pre-training with Point-BERT and Point-MAE strategies enhances transfer to downstream tasks, yielding SoTA results on ScanObjectNN and ModelNet40 and strong few-shot performance, all with near-linear FLOPs. Overall, Mamba3D delivers a robust, scalable alternative to Transformers for 3D point cloud understanding, enabling large-scale pre-training and efficient inference at scale.

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.

Figures (3)

• Figure 1: Illustration of Mamba3D. (a) Overview. We first segment input point cloud into $L$ patches using FPS & KNN, and then obtain the initial patch embeddings via a light-PointNet. After adding a $\mathtt{[CLS]}$ token, we apply standard positional encoding to all patch embeddings, which are then fed into Mamba3D Encoder. Finally we use a task header to fit downstream tasks. (b-c) Mamba3D Encoder details. The illustration of the Mamba3D Encoder is inspired by ViT.
• Figure 2: Illustration of C-SSM. Instead of horizontal token flip, we propose a vertical feature flip to alleviate the pseudo-order reliance.
• Figure 3: Ablation on (a) the parameter $k$ in the LNP block, (b) input patch sequence length $L$, and (c) ordering strategy. The overall accuracy (%), training time/epoch (s) and FLOPs (G) are reported.