[논문 리뷰] SegMamba: Long-range Sequential Modeling Mamba For 3D Medical Image Segmentation

Zhaohu Xing, Ye Tian외 3명|arXiv (Cornell University)|2024. 01. 24.

Medical Image Segmentation Techniques인용 수 9

한 줄 요약

SegMamba는 3D 의학 영상의 장거리 의존성을 모델링하기 위해 U자 형태의 아키텍처에서 Mamba 상태 공간 블록을 사용하여 CNN 및 트랜스포머 기반 기준보다 효율적인 추론으로 BraTS2023 분할 성능을 달성합니다.

ABSTRACT

The Transformer architecture has shown a remarkable ability in modeling global relationships. However, it poses a significant computational challenge when processing high-dimensional medical images. This hinders its development and widespread adoption in this task. Mamba, as a State Space Model (SSM), recently emerged as a notable manner for long-range dependencies in sequential modeling, excelling in natural language processing filed with its remarkable memory efficiency and computational speed. Inspired by its success, we introduce SegMamba, a novel 3D medical image extbf{Seg}mentation extbf{Mamba} model, designed to effectively capture long-range dependencies within whole volume features at every scale. Our SegMamba, in contrast to Transformer-based methods, excels in whole volume feature modeling from a state space model standpoint, maintaining superior processing speed, even with volume features at a resolution of {$64 imes 64 imes 64$}. Comprehensive experiments on the BraTS2023 dataset demonstrate the effectiveness and efficiency of our SegMamba. The code for SegMamba is available at: https://github.com/ge-xing/SegMamba

연구 동기 및 목표

• 고해상도 3D 의학 영상에서 전역적이고 장거리 의존성을 모델링해야 할 동기를 제시한다.
• 다중 스케일의 전체 볼륨 특징을 포착하는 Mamba 기반 인코더를 제안한다.
• 정확한 분할을 위한 스킵 연결이 있는 CNN 기반 디코더를 통합한다.
• 대용량 3D 부피에서 학습과 추론의 효율성을 유지한다.

제안 방법

• 다음의 세 부분으로 구성된 SegMamba 아키텍처를 도입한다: 다수의 블록이 있는 Mamba 기반 인코더, CNN 기반 디코더, 그리고 U-자형 스킵 연결.
• Mamba 블록을 효율적으로 순차 모델링할 수 있도록 3D 특징을 1D 긴 시퀀스로 평탄화한 다음, 시그마 연산으로 3D로 복원한다.
• Section 2.1에 정의된 식들에 따라 잔차와 같은 블록 구조에서 레이어 정규화(LN), Mamba 블록, MLP를 사용한다.
• 줄기(stem)를 7x7x7 깊이 방향 합성곱으로 다운샘플링하여 다중 스케일 특징 z^0를 생성하고 Mamba 블록을 통해 점진적으로 처리한다.
• 교차 엔트로피 손실, SGD 옵티마이저, 다항 스케일의 학습률 스케줄, 표준 데이터 증강으로 학습하고, 추론 시 테스트 타임 증강을 적용한다.

실험 결과

연구 질문

• RQ1 SegMamba가 Transformer 기반 접근법과 비교하여 Mamba를 사용해 3D 의학 부피의 장거리 의존성을 효율적으로 모델링할 수 있는가?
• RQ2 SegMamba 인코더–디코더 아키텍처가 BraTS2023에서 분할 정확도를 향상시키면서 고해상도(예: 64x64x64 특징 맵)에서 추론 속도를 유지하는가?

주요 결과

• SegMamba는 BraTS2023에서 WT, TC, ET 모두에서 Dice 점수 최상위에 도달했고 HD95는 WT: 3.37, TC: 3.85, ET: 3.48로 나타난다(또한 WT: 93.61, TC: 92.65, ET: 87.71).
• BraTS2023의 평균 Dice는 91.32%이고 Avg HD95는 4.01로, 보고된 결과에서 CNN 기반 및 Transformer 기반 기준보다 우수하다.
• 표 1에 따르면 평균 Dice에서 UX-Net 및 SwinUNETR-V2보다 각각 1.63% 및 1.93% 높게 앞선다.
• 3D 의학 영상 분할에 대한 Mamba 기반 장거리 모델링의 효율성을 보여주며, 고해상도(64x64x64)에서도 경쟁력 있는 속도를 유지한다.