[논문 리뷰] Mamba3D: Enhancing Local Features for 3D Point Cloud Analysis via State Space Model

Xu Han, Yuan Tang외 2명|arXiv (Cornell University)|2024. 04. 23.

Remote Sensing and LiDAR ApplicationsEnvironmental Science인용 수 3

한 줄 요약

Mamba3D는 3D 포인트 클라우드 분석을 위한 상태공간 모델을 제안하며, 국소 기하 특징 추출을 위한 명시적 방법으로 국소 노름 풀링(LNP) 블록과 새로운 채널별 역방향 SSM을 통합한 이중방향 SSM(bi-SSM)을 활용하여 글로벌 표현을 향상시킨다. 선형 복잡도를 유지하면서도 최신 기술 수준(SoTA)의 성능을 달성하여, 훈련을 처음부터 시작할 경우 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.

연구 동기 및 목표

• Transformer 기반 3D 포인트 클라우드 모델에서 발생하는 제곱형 복잡도 문제를 해결하여 해상도 및 확장성의 제약을 해소한다.
• 순서가 없는 포인트 클라우드에 직접 Mamba를 적용할 경우 발생하는 가짜 순서 의존성과 명시적 국소 특징 추출의 부재 문제를 해결한다.
• 사전학습에 의존하지 않으면서도 높은 성능을 유지할 수 있고, 소수의 샘플 학습 및 사전학습된 일반화 성능을 지원하는 확장성 있고 효율적인 아키텍처를 개발한다.
• SSM 프레임워크에 국소 기하 특징을 명시적으로 통합하여 포인트 클라우드의 표현 학습을 향상시킨다.
• 특징 채널을 순서가 있는 시퀀스로 간주하여 이중방향 SSM을 설계함으로써 순서가 없는 포인트 클라우드에서 글로벌 특징 학습의 안정성을 높인다.

제안 방법

• 국소 기하 특징 학습을 위해 K-노름(근처 K개의 이웃에 대한 특징 정규화)과 간단하고 미분 가능한 연산자를 사용하는 K-풀링(특징 집계)을 조합한 국소 노름 풀링(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 Transformer의 제곱형 복잡도 문제를 해결하면서도 성능을 유지하거나 향상시킬 수 있는 상태공간 모델이 3D 포인트 클라우드 학습에 효과적으로 적용될 수 있는가?
• RQ2 Mamba처럼 국소 구조에 대한 인덕티브 바이어스가 없는 순서 기반 모델에 국소 기하 특징을 어떻게 명시적으로 통합할 수 있는가?
• RQ3 역방향 SSM에서 특징 채널을 순서가 있는 시퀀스로 간주함으로써 순서가 없는 포인트 클라우드에서 글로벌 특징 학습이 향상되는가?
• RQ4 Mamba3D처럼 선형 복잡도를 가지는 아키텍처가 무거운 사전학습에 의존하지 않고도 포인트 클라우드 분류 작업에서 최신 기술 수준의 성능을 달성할 수 있는가?
• RQ5 제안된 bi-SSM는 포인트 클라우드 벤치마크에서 자기주의(Self-attention)와 단일방향 SSM에 비해 성능과 강건성 측면에서 어떻게 비교되는가?

주요 결과

• Mamba3D는 사전학습 없이 훈련을 처음부터 시작할 경우 ScanObjectNN 데이터셋에서 전체 정확도 92.6%를 달성하여 사전학습이 없는 모델 중 최신 기술 수준(SoTA)을 수립한다.
• 단일 모달 사전학습을 통한 결과, Mamba3D는 ModelNet40에서 전체 정확도 95.1%를 기록하여 단일 모달 사전학습 모델 중 최신 기술 수준(SoTA)을 확립한다.
• 제거 실험 결과, LNP 블록을 제거할 경우 정확도가 1.2% 감소하고, bi-SSM 블록을 제거할 경우 2.3% 감소하여 두 블록이 성능 향상에 핵심적인 역할을 한다는 것을 확인한다.
• bi-SSM를 자기주의 또는 단일방향 SSM으로 대체할 경우 정확도가 1.2% 감소하며, 토큰 뒤집기(variant)는 정확도를 1.1%에서 1.4%까지 떨어뜨려 C-SSM 설계의 효과성을 입증한다.
• K-풀링은 최대풀링, 평균풀링 및 그 조합보다 우수한 성능을 보이며, 이를 대체할 경우 정확도가 0.7%에서 2.8%까지 감소하여 국소 특징 집계에서의 우월성을 입증한다.
• 입력 순서 전략을 전혀 사용하지 않을 경우 성능이 가장 뛰어나며, Mamba3D가 순서가 없는 포인트 클라우드에서 의미론적 특징을 효과적으로 포착함을 시사한다. 특히 k=4와 L=128에서 최고 성능를 기록한다.