PoinTramba: A Hybrid Transformer-Mamba Framework for Point Cloud Analysis

TL;DR

PoinTramba tackles point cloud analysis by marrying Transformer-based intra-group modeling with a linear-time Mamba inter-group encoder, augmented by a bi-directional importance-aware ordering (BIO) and importance-aware pooling (IAP). The method segments a cloud into groups, uses a Transformer to produce group embeddings, reorders them via BIO to stabilize Mamba processing, and extracts global inter-group features through Mamba before pooling. The training objective combines task performance with specialized losses for importance ordering and embedding alignment, enabling end-to-end optimization. Empirical results on ScanObjectNN, ModelNet40, and ShapeNetPart demonstrate strong performance and efficiency, with BIO and IAP contributing significant gains and achieving competitive or state-of-the-art benchmarks in several settings.

Abstract

Point cloud analysis has seen substantial advancements due to deep learning, although previous Transformer-based methods excel at modeling long-range dependencies on this task, their computational demands are substantial. Conversely, the Mamba offers greater efficiency but shows limited potential compared with Transformer-based methods. In this study, we introduce PoinTramba, a pioneering hybrid framework that synergies the analytical power of Transformer with the remarkable computational efficiency of Mamba for enhanced point cloud analysis. Specifically, our approach first segments point clouds into groups, where the Transformer meticulously captures intricate intra-group dependencies and produces group embeddings, whose inter-group relationships will be simultaneously and adeptly captured by efficient Mamba architecture, ensuring comprehensive analysis. Unlike previous Mamba approaches, we introduce a bi-directional importance-aware ordering (BIO) strategy to tackle the challenges of random ordering effects. This innovative strategy intelligently reorders group embeddings based on their calculated importance scores, significantly enhancing Mamba's performance and optimizing the overall analytical process. Our framework achieves a superior balance between computational efficiency and analytical performance by seamlessly integrating these advanced techniques, marking a substantial leap forward in point cloud analysis. Extensive experiments on datasets such as ScanObjectNN, ModelNet40, and ShapeNetPart demonstrate the effectiveness of our approach, establishing a new state-of-the-art analysis benchmark on point cloud recognition. For the first time, this paradigm leverages the combined strengths of both Transformer and Mamba architectures, facilitating a new standard in the field. The code is available at https://github.com/xiaoyao3302/PoinTramba.

Figures (4)

• Figure 1: The overview of our newly proposed PoinTramba framework (a) and its two main modules, the Intra-group Transformer Encoder (b) and the Inter-group Mamba Encoder (c). Initially, we segment the input point cloud into distinct point groups. Following this, we employ a Transformer encoder to model intra-group relationships and generate group embeddings. An importance-score prediction module is then utilized to predict the importance score for each group embedding. These predicted importance scores are used to reorder the group embeddings. Finally, a Mamba encoder extracts inter-group relationships from the reordered group embeddings, which are subsequently fed into an importance-aware pooling layer. This layer captures the global feature that can be further utilized for various downstream tasks such as classification and segmentation.
• Figure 2: The detailed design of our importance score prediction module (a) and our importance-aware pooling layer (b). The importance score prediction module targets at calculating the similarity between the group embeddings and the global feature, thus predicting the importance scores for group embeddings. The importance-aware pooling layer targets at aggregating the updated group embeddings to obtain the global feature.
• Figure 3: Ablation studies on different ordering strategies and pooling methods. Experiments are conducted on the PB-T50-RS variant of the ScanObjNN dataset. PoinTramba is adopted as the backbone. (a) shows the comparison of different ordering strategies, i.e., random ordering strategy, XYZ ordering strategy, z ordering strategy, Hilbert Ordering strategy, our single-directional importance-aware ordering strategy (SIO) and bi-directional importance-aware ordering strategy (BIO). (b) illustrates the comparison of different pooling methods, i.e., average pooling, max-pooling, weighted sum and our importance-aware pooling (IAP).
• Figure 4: Visualization of the importance scores for the point groups predicted by our PoinTramba model. Samples from various categories in ModelNet40 are used as examples. Red regions indicate higher importance scores, while yellow regions indicate lower importance scores.