[論文レビュー] Mamba3D: Enhancing Local Features for 3D Point Cloud Analysis via State Space Model

Xu Han, Yuan Tang+2 more|arXiv (Cornell University)|Apr 23, 2024

Remote Sensing and LiDAR ApplicationsEnvironmental Science被引用数 3

ひとこと要約

Mamba3Dは、局所的な幾何的特徴抽出のためのLocal Norm Pooling（LNP）ブロックと、新しいチャネル単位の逆方向SSMを備えた双方向SSM（bi-SSM）を用いた、3次元点群解析に特化した状態空間モデルを提案する。線形計算量で最先端の性能を達成し、学習から再訓練した場合にScanObjectNNで92.6%の精度を達成し、単一モodalの事前学習を用いた場合にModelNet40で95.1%の精度を達成した。

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.

研究の動機と目的

• 変換器ベースの3次元点群モデルにおける2次関数的計算量のボトル neck を解消し、解像度とスケーラビリティを向上させること。
• 順序なし点群への直接的なMambaの適用における限界、特に擬似順序依存性と明示的な局所的特徴抽出の欠如を克服すること。
• 事前学習に依存せず、高い性能を維持しつつ、スケーラブルで効率的なアーキテクチャを構築すること。同時に、少サンプル学習および事前学習による汎化性能を強化すること。
• 点群の表現学習を向上させるために、局所的な幾何的特徴を状態空間モデルフレームワークに明示的に統合すること。
• 特徴チャネルを順序付きシーケンスとして扱うことで、順序なし点群におけるグローバル特徴学習の安定化を図る双方向SSMを設計すること。

提案手法

• K近傍の特徴を正規化するK-ノルムと、単純で微分可能演算子を用いたKプーリングを組み合わせたLocal Norm Pooling（LNP）ブロックを提案し、局所的幾何的特徴学習を実現する。
• トークン前方SSM（L+）と、チャネル次元に沿って処理する新しいチャネル単位の逆方向SSM（C-SSM）から構成される双方向SSM（bi-SSM）を導入し、グローバルコンテキスト学習の安定化を図る。
• C-SSMを特徴チャネルを仮想的な順序付きシーケンスとして処理することで、任意の点の順序に依存せず、順序なし入力に対してより高い耐性を発揮するように設計した。
• LNPブロックをたった0.3Mパラメータで実装し、モデルの効率性とスケーラビリティを維持した。
• 高価なMLPやプーリング層を回避する単純でパラメータ効率の良いKプーリング演算子を採用し、FLOPsとパラメータ数を削減しながら性能を保持した。
• LNPブロックとbi-SSMブロックを統合したアーキテクチャを設計し、標準的なプロトコル（例：Point-BERTやPoint-MAE）に準拠した学習から再訓練および事前学習の両方をサポートした。

実験結果

リサーチクエスチョン

• RQ1 変換器の2次関数的計算量の問題を克服しつつ、性能を維持または向上させるために、状態空間モデルを3次元点群学習に効果的に適応できるか？
• RQ2 Mambaのような系列ベースのモデルは、局所構造に対するインダクティブバイアスを欠いているが、局所的幾何的特徴をどのように明示的に統合できるか？
• RQ3 逆方向SSMにおいて特徴チャネルを順序付きシーケンスとして扱うことで、順序なし点群におけるグローバル特徴学習が向上するか？
• RQ4 Mamba3Dのような線形計算量アーキテクチャが、重い事前学習に依存せず、点群分類タスクで最先端の性能を達成できるか？
• RQ5 点群ベンチマークにおいて、提案されたbi-SSMは自己自己注意機構や一方向SSMと比較して、どのように耐性と性能を発揮するか？

主な発見

• 学習から再訓練した場合、Mamba3DはScanObjectNNデータセットで92.6%の全体精度を達成し、事前学習を行わないモデルの中で新たな最先端を記録した。
• 単一モダリティの事前学習を用いた場合、Mamba3DはModelNet40で95.1%の全体精度を達成し、単一モダリティ事前学習モデルの中で新たな最先端を樹立した。
• アブレーションスタディの結果、LNPブロックを削除すると精度が1.2%低下し、bi-SSMブロックを削除すると2.3%低下した。これにより、両ブロックが性能向上に不可欠であることが確認された。
• bi-SSMを自己自己注意機構や一方向SSMに置き換えると1.2%の精度低下が生じた一方、トークンフリップ変種では1.1%～1.4%の低下が生じた。これにより、C-SSMの設計の有効性が裏付けられた。
• KプーリングはMaxpooling、Avgpooling、およびそれらの組み合わせを上回り、置き換えにより精度が0.7%～2.8%低下した。これにより、局所的特徴集約においてKプーリングの優位性が示された。
• 入力の順序戦略を一切使用しない場合に最も高い性能を発揮し、Mamba3Dが順序なし点群からも意味的特徴を効果的に捉えられることを示した。k=4およびL=128でピーク性能を記録した。