[논문 리뷰] PointMamba: A Simple State Space Model for Point Cloud Analysis

Dingkang Liang, Xin Zhou외 6명|arXiv (Cornell University)|2024. 02. 16.

Remote Sensing and LiDAR Applications인용 수 21

한 줄 요약

PointMamba는 재배치 전략과 Mamba 블록을 갖춘 간단한 상태공간 모델을 도입하여 전역적이며 선형 복잡도의 포인트 클라우드 분석을 달성하고, 더 적은 매개변수와 FLOPs로 트랜스포머 기반 비교대상보다 성능이 우수하다.

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 포인트 클라우드에서 제곱적 주의 비용을 극복하기 위해 선형 복잡도의 글로벌 모델링 접근법을 제시한다.
• 간단한 재배치 전략을 통해 불규칙하고 비순차적인 포인트 데이터에 상태공간 모델(SSMs)을 적용한다.
• 분류와 파트 세그먼테이션 작업에서 PointMamba가 트랜스폼 기반 방법보다 더 나은 성능을 낼 수 있음을 보여준다.
• 매개변수와 FLOPs를 줄이면서 PointMamba가 경쟁력 있거나 더 나은 정확도를 달성함을 보여준다.

제안 방법

• FPS와 KNN을 사용한 경량 PointNet 기반 토크나이저로 포인트 패치를 토큰으로 임베딩하여 패치를 형성한다.
• x, y, z 축을 따라 토큰을 정렬하고 시퀀스를 3배로 늘리는 재배치 전략을 적용하여 Mamba에서 인과 모델링을 가능하게 한다.
• 정렬된 토큰들을 계층 정규화, SSM, 깊이별 컨벌루션 및 잔차 연결을 갖춘 Mamba 블록 시퀀스를 통해 처리한다.
• 재구성 손실로 Chamfer 거리와 마스킹된 패치 재구성을 사용하는 Point-MAE 스타일 프리트레이닝으로 학습한다.
• 성능에 미치는 영향을 연구하기 위해 CLS 토큰 사용 및 특징 차원과 같은 설계 선택을 선택적으로 제거(ablate)한다.

실험 결과

연구 질문

• RQ1 선형 복잡도의 상태공간 모델이 포인트 클라우드의 글로벌 컨텍스트를 주의 기반 트랜스포머만큼 효과적으로 포착할 수 있는가?
• RQ2 포인트 토큰의 간단한 기하학적 재배치가 비순차 포인트 데이터에서 Mamba의 인과 모델링을 가능하게 하는가?
• RQ3 설계 선택(CLS 토큰, 토큰 차원)이 분류 및 세그먼테이션에서 PointMamba의 성능에 미치는 영향은 무엇인가?
• RQ4 표준 3D 벤치마크에서 프리트레이닝+다운스트림 설정에서 PointMamba의 성능은 트랜스포머 대조군과 비교해 어떤가?

주요 결과

• PointMamba는 다수의 포인트 클라우드 작업에서 트랜스포머 기반 동료들보다 더 나은 성능을 보이면서 매개변수와 FLOPs를 더 적게 사용한다(매개변수 약 44.3% 감소, FLOPs 약 25% 감소).
• ScanObjectNN에서 PointMamba는 처음부터 학습할 때 트랜스포머 기반 기준선보다 OBJ-BG 및 OBJ-ONLY 점수가 높으며, 프리트레이닝 변형에서도 경쟁력을 유지한다.
• 간단한 재배치 전략이 비인과적이고 비순차적인 포인트 클라우드를 모델링하는 Mamba의 능력을 향상시켜, 1x(재배치 없음) 설정 대비 명확한 정확도 향상을 가져온다.
• CLS 토큰을 생략하고 마지막 블록 출력의 평균 풀링을 사용하는 것이 PointMamba의 분류 성능을 최상으로 만든다.
• PointMamba는 ShapeNetPart에서 많은 베이스라인에 비해 현저히 적은 매개변수로 경쟁력 있는 파트 세그먼테이션 결과를 보여주며, 긴 시퀀스에 대한 메모리 효율성도 우수하다.