Exploring contextual modeling with linear complexity for point cloud segmentation

TL;DR

It is shown that: 1) Spatial locality and robust contextual understanding are critical for strong performance, and 2) Mamba features linear computational complexity, offering superior data and inference efficiency compared to Transformers, while still being capable of delivering strong contextual understanding.

Abstract

Point cloud segmentation is an important topic in 3D understanding that has traditionally has been tackled using either the CNN or Transformer. Recently, Mamba has emerged as a promising alternative, offering efficient long-range contextual modeling capabilities without the quadratic complexity associated with Transformer's attention mechanisms. However, despite Mamba's potential, early efforts have all failed to achieve better performance than the best CNN-based and Transformer-based methods. In this work, we address this challenge by identifying the key components of an effective and efficient point cloud segmentation architecture. Specifically, we show that: 1) Spatial locality and robust contextual understanding are critical for strong performance, and 2) Mamba features linear computational complexity, offering superior data and inference efficiency compared to Transformers, while still being capable of delivering strong contextual understanding. Additionally, we further enhance the standard Mamba specifically for point cloud segmentation by identifying its two key shortcomings. First, the enforced causality in the original Mamba is unsuitable for processing point clouds that have no such dependencies. Second, its unidirectional scanning strategy imposes a directional bias, hampering its ability to capture the full context of unordered point clouds in a single pass. To address these issues, we carefully remove the causal convolutions and introduce a novel Strided Bidirectional SSM to enhance the model's capability to capture spatial relationships. Our efforts culminate in the development of a novel architecture named MEEPO, which effectively integrates the strengths of CNN and Mamba. MEEPO surpasses the previous state-of-the-art method, PTv3, by up to +0.8 mIoU on multiple key benchmark datasets, while being 42.1% faster and 5.53x more memory efficient.

Figures (23)

• Figure 2: Proposed meta-architecture and various block options used for analysis. The model that exclusively uses choice A is called Pure CNN, the model that exclusively uses choice B is called Pure Mamba, and the model that exclusively uses choice C is called Pure Transformer.
• Figure 4: Analysis of a representative Transformer-based model, PTv3, demonstrates that additional context beyond a certain amount is unnecessary.
• Figure 6: Our proposed architecture, Meepo, integrates CNN-Mamba blocks throughout the proposed meta-architecture to harness their strengths in local and contextual modeling. To optimize for point cloud segmentation, Meepo modifies the standard Mamba by replacing causal convolutions with regular convolutions, preserving critical spatial information. Additionally, it introduces a novel Bidirectional Strided SSM, which enhances contextual modeling by minimizing directional bias.
• Figure 7: Comparison between Meepo's and PTv3's ptv3 predictions. Black color are unlabelled points. Red boxes with dash-dotted lines are wrong predictions by PTv3.
• Figure : (a) mIoU v.s. Latency on V100