[Paper Review] 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 Applications21 citations

TL;DR

PointMamba introduces a simple state space model with a reordering strategy and Mamba blocks to achieve global, linear-complexity point cloud analysis, outperforming transformer baselines with fewer parameters and FLOPs.

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}.

Motivation & Objective

• Motivate a linear-complexity global modeling approach for 3D point clouds to overcome quadratic attention costs.
• Adapt state space models (SSMs) to irregular, unordered point data via a simple reordering strategy.
• Demonstrate that PointMamba can outperform transformer-based methods on classification and part segmentation tasks.
• Show that PointMamba achieves competitive or superior accuracy while reducing parameters and FLOPs.

Proposed method

• Embed point patches into tokens using a lightweight PointNet-based tokenizer with FPS and KNN to form patches.
• Apply a reordering strategy that sorts tokens along x, y, z axes and triples the sequence to enable causal modeling in Mamba.
• Process the reordered tokens through a series of Mamba blocks with layer norm, SSM, depth-wise conv, and residual connections.
• Train with Point-MAE style pre-training using masked patch reconstruction and Chamfer distance as the reconstruction loss.
• Optionally ablate design choices such as CLS token usage and feature dimension to study their impact on performance.

Experimental results

Research questions

• RQ1 Can a state space model with linear complexity capture global context in point clouds as effectively as attention-based transformers?
• RQ2 Does a simple geometric reordering of point tokens enable causal modeling in Mamba for non-sequential point data?
• RQ3 What is the impact of design choices (CLS token, token dimensionality) on PointMamba’s performance in classification and segmentation?
• RQ4 How does PointMamba perform in pre-training+downstream settings compared to transformer baselines on standard 3D benchmarks?

Key findings

• PointMamba outperforms Transformer-based counterparts on multiple point cloud tasks while using fewer parameters and FLOPs (≈44.3% fewer parameters and ≈25% fewer FLOPs).
• In ScanObjectNN, PointMamba achieves higher OBJ-BG and OBJ-ONLY scores than the Transformer-based baseline when trained from scratch, and remains competitive with pre-training variants.
• A simple reordering strategy improves Mamba’s ability to model non-causal, unordered point clouds, yielding measurable accuracy gains over the 1x (no reordering) setting.
• Omitting the CLS token and using average pooling of the final block’s output yields the best classification performance for PointMamba.
• PointMamba demonstrates competitive Part Segmentation results on ShapeNetPart with substantially fewer parameters than many baselines, and exhibits favorable memory efficiency for long sequences.