[論文レビュー] 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 Applications被引用数 21

ひとこと要約

PointMamba は再配置戦略と Mamba ブロックを備えた簡易な状態空間モデルを導入し、全局的で線形計算量の点群解析を実現。パラメータ数と FLOPs が少ない Transformer のベースラインを上回る。

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

研究の動機と目的

• 3D 点群の二次元注意コストを克服するための線形計算量によるグローバルモデリング手法を動機づける。
• 単純な再配置戦略を介して状態空間モデル（SSM）を不規則で順序付けられていない点データに適用する。
• PointMamba が分類とパートセグメンテーションのタスクで Transformer ベースの方法を上回ることを示す。
• PointMamba がパラメータと FLOPs を削減しつつ競争力のある、または優れた精度を達成することを示す。

提案手法

• 軽量な PointNet ベースのトークナイザを使用して FPS と KNN でパッチを形成し、点パッチをトークンに埋め込む。
• x、y、z 軸に沿ってトークンを並べ替え、連続を三分する再配置戦略を適用して Mamba に因果モデリングを可能にする。
• 再配置されたトークンを、レイヤーノルム、SSM、深さ方向畳み込み、残差接続を備えた一連の Mamba ブロックを通して処理する。
• マスク付きパッチ再構成と再構成損失として Chamfer 距離を用いた Point-MAE スタイルの事前学習で訓練する。
• CLS トークンの使用や特徴次元など設計上の選択を排除・有無することで、性能への影響を研究する。

実験結果

リサーチクエスチョン

• RQ1 線形計算量を持つ状態空間モデルは、注意機構を用いたトランスフォーマーと同等に点群のグローバルな文脈を捉えられるのか。
• RQ2 点トークンの単純な幾何再配置は、非連続の順序付けられていない点データに対して Mamba に因果モデリングを可能にするのか。
• RQ3 設計上の選択（CLS トークン、トークン次元） on PointMamba の分類とセグメンテーションの性能にどのような影響を与えるのか。
• RQ4 標準的な3D ベンチマークで transformer ベースラインと比較した場合、事前学習+下流設定で PointMamba の性能はどうなるのか。

主な発見

• PointMamba は、パラメータ数と FLOPs が少ない状態で、複数の点群タスクにおいて Transformer ベースの対向手法を上回る（パラメータ約44.3%、FLOPs約25%の削減）。
• ScanObjectNN では、スクラッチから訓練した場合、Transformer ベースのベースラインより OBJ-BG および OBJ-ONLY スコアが高く、事前学習変種でも競争力を維持する。
• 単純な再配置戦略は、非因果・順序付けられていない点群をモデル化する Mamba の能力を高め、1x の再配置なし設定よりも精度の改善をもたらす。
• CLS トークンを省略し、最終ブロック出力の平均プーリングを用いると、PointMamba の分類性能が最も良くなる。
• PointMamba は ShapeNetPart でのパートセグメンテーション結果が多くのベースラインよりもパラメータが大幅に少なく、長いシーケンスでのメモリ効率も良好である。