[논문 리뷰] Mamba-UNet: UNet-Like Pure Visual Mamba for Medical Image Segmentation

Ziyang Wang, Jian-Qing Zheng외 3명|arXiv (Cornell University)|2024. 02. 07.

Medical Image Segmentation Techniques인용 수 79

한 줄 요약

Mamba-UNet은 UNet 유사 인코더-디코더 내에서 순수한 Visual Mamba 블록을 사용하여 장거리 특징 모델링을 개선하고 MRI 심장 데이터에서 UNet 및 Swin-UNet보다 더 높은 분할 정확도를 달성합니다.

ABSTRACT

In recent advancements in medical image analysis, Convolutional Neural Networks (CNN) and Vision Transformers (ViT) have set significant benchmarks. While the former excels in capturing local features through its convolution operations, the latter achieves remarkable global context understanding by leveraging self-attention mechanisms. However, both architectures exhibit limitations in efficiently modeling long-range dependencies within medical images, which is a critical aspect for precise segmentation. Inspired by the Mamba architecture, known for its proficiency in handling long sequences and global contextual information with enhanced computational efficiency as a State Space Model (SSM), we propose Mamba-UNet, a novel architecture that synergizes the U-Net in medical image segmentation with Mamba's capability. Mamba-UNet adopts a pure Visual Mamba (VMamba)-based encoder-decoder structure, infused with skip connections to preserve spatial information across different scales of the network. This design facilitates a comprehensive feature learning process, capturing intricate details and broader semantic contexts within medical images. We introduce a novel integration mechanism within the VMamba blocks to ensure seamless connectivity and information flow between the encoder and decoder paths, enhancing the segmentation performance. We conducted experiments on publicly available ACDC MRI Cardiac segmentation dataset, and Synapse CT Abdomen segmentation dataset. The results show that Mamba-UNet outperforms several types of UNet in medical image segmentation under the same hyper-parameter setting. The source code and baseline implementations are available.

연구 동기 및 목표

• 의료 영상 분할에서 장거리 의존성 모델링의 향상을 동기화한다.
• Encoder, bottleneck, decoder에 대해 Visual Mamba 블록(VSS)을 사용하는 UNet에서 영감을 받은 아키텍처를 제안한다.
• 스킵 연결 및 VMamba 기반 프레임워크 내 패치 병합/확장을 통해 공간 세부 정보를 보존한다.
• 공개 MRI 심장 데이터셋에서 분할 성능을 평가하고 베이스라인과 비교한다.

제안 방법

• Pure Visual Mamba 블록을 핵심 구성 블록으로 사용하는 UNet 유사 인코더–디코더를 채택한다.
• 입력 이미지를 패치 토큰으로 표현한 다음 계층적 VSS 블록을 패치 병합/확장을 통해 처리한다.
• 에너더와 디코더 간의 스킵 연결을 사용하여 다중 스케일 특징을 융합한다.
• 초기화 향상을 위해 인코더에 pretrained VMamba-Tiny를 로드한다.
• 고정 하이퍼파라미터하에 SGD로 학습하고 표준 분할 지표로 평가한다.

Figure 1: A brief introduction of the evolution of recent developments of UNet with incorporation of Transformer and State Space Models (SSM) for medical image segmentation.

실험 결과

연구 질문

• RQ1 VMamba 기반 블록이 일반적인 UNet 및 ViT 기반 접근법과 비교하여 의료 영상 분할에서 장거리 의존성 모델링을 개선할 수 있는가?
• RQ2 Pure VMamba UNet이 MRI 데이터에서 계산 효율성을 유지하면서 우수한 분할 정확도를 달성하는가?
• RQ3 동일한 학습 설정에서 Mamba-UNet이 UNet 및 Swin-UNet과 비교한 성능은 어떠한가?

주요 결과

• Mamba-UNet은 MRI 심장 테스트 세트에서 Dice 0.9281, IoU 0.8698, 정확도 0.9972를 달성한다.
• 동일 하이퍼파라미터에서 Mamba-UNet은 UNet보다 Dice 및 IoU에서 우수하다.
• Mamba-UNet은 Baseline과 비교했을 때 HD(95%) 및 ASD 점수에서 경쟁력을 보이며 HD 2.4645 및 ASD 0.7677를 기록한다.
• Swin-UNet은 Dice 0.9188 및 IoU 0.8545를 생성하여 Mamba-UNet보다 낮고, UNet은 Dice 0.9248을 나타낸다.
• Mamba-UNet은 보고된 지표에서 경계 정밀도 우수성을 보여준다.

Figure 2: The architecture of Mamba-UNet, which is composed of encoder, bottleneck, decoder and skip connections. The encoder, bottleneck and decoder are all constructed based on Visual Mamba block.