[논문 리뷰] U-Mamba: Enhancing Long-range Dependency for Biomedical Image Segmentation

Jun Ma, Li Fei-Fei외 1명|arXiv (Cornell University)|2024. 01. 09.

AI in cancer detection인용 수 213

한 줄 요약

U-Mamba는 선형 확장으로 긴 범위 의존성을 포착하고 데이터셋 간 자동 구성하는 하이브리드 CNN-SSM 네트워크로, CNN- 및 Transformer 기반 경쟁자들을 여러 과제에서 능가하는 생의학 영상 분할 모델입니다.

ABSTRACT

Convolutional Neural Networks (CNNs) and Transformers have been the most popular architectures for biomedical image segmentation, but both of them have limited ability to handle long-range dependencies because of inherent locality or computational complexity. To address this challenge, we introduce U-Mamba, a general-purpose network for biomedical image segmentation. Inspired by the State Space Sequence Models (SSMs), a new family of deep sequence models known for their strong capability in handling long sequences, we design a hybrid CNN-SSM block that integrates the local feature extraction power of convolutional layers with the abilities of SSMs for capturing the long-range dependency. Moreover, U-Mamba enjoys a self-configuring mechanism, allowing it to automatically adapt to various datasets without manual intervention. We conduct extensive experiments on four diverse tasks, including the 3D abdominal organ segmentation in CT and MR images, instrument segmentation in endoscopy images, and cell segmentation in microscopy images. The results reveal that U-Mamba outperforms state-of-the-art CNN-based and Transformer-based segmentation networks across all tasks. This opens new avenues for efficient long-range dependency modeling in biomedical image analysis. The code, models, and data are publicly available at https://wanglab.ai/u-mamba.html.

연구 동기 및 목표

• 생의학 영상 분할에서 긴 범위 의존성 모델링의 도전 과제 해결.
• 로컬 특징 추출과 긴 범위 컨텍스트를 결합한 일반-purpose CNN-SSM 백본(U-Mamba) 제안.
• 다양한 데이터셋에 자동으로 적응하는 자체 구성 메커니즘 도입.
• 4개의 다양한 생의학 분할 작업(3D 복부 CT/MRI, 내시경 도구, 현미경 세포)에서 U-Mamba를 평가하고 CNN 및 Transformer 기반 기준선과 비교.
• 재현성 확보를 위한 오픈 소스 코드, 모델, 데이터 제공

제안 방법

• 긴 범위 의존성을 선형 확장으로 포착하는 Scan(S6)를 가진 선택적 구조화 상태 공간 시퀀스 모델(Mamba) 도입.
• 로컬 정보와 글로벌 정보를 융합하기 위해 두 개의 잔차 블록이 먼저 Mamba 블록을 앞에 두는 U-Net 유사 인코더-디코더(U-Mamba) 구성.
• Mamba 블록에 맞게 3D/2D 특징 맵을 시퀀스 친화적인 형태로 평탄화한 뒤 두 병렬 가지의 Hadamard 곱 융합.
• 두 가지 네트워크 변형 제공: U-Mamba_Bot(병목만 Mamba) 및 U-Mamba_Enc(인코더 블록 전체에 Mamba)
• 아키텍처 깊이와 패치 크기를 데이터셋별로 자동으로 조정하는 nnU-Net의 자체 구성 프레임워크와 통합.
• 기존 기준선 간 공정 비교를 보장하기 위한 Dice와 cross-entropy 손실의 고정 학습 프로토콜 사용

실험 결과

연구 질문

• RQ1 U-Mamba가 CNN 및 Transformer를 넘어 생의학 영상 분할에서 장기 의존성 모델링을 개선할 수 있는가?
• RQ2 선형 확장 Mamba 블록이 다양한 2D/3D 생의학 분할 작업에서 실용적 이점을 제공하는가?
• RQ3 복부 장기, 기구, 세포 분할에서 강력한 CNN- 및 Transformer 기반 기준선과 U-Mamba의 비교는 어떠한가?
• RQ4 U-Mamba의 자체 구성 메커니즘( nnU-Net 이용)이 매뉴얼 조정 없이 데이터셋 간 견고한 적응을 가능하게 하는가?

주요 결과

방법	복부 CT의 장기 DSC	복부 CT의 장기 NSD	복부 MRI의 장기 DSC	복부 MRI의 장기 NSD
nnU-Net	0.8615 ± 0.0790	0.8972 ± 0.0824	0.8309 ± 0.0769	0.8996 ± 0.0729
SegResNet	0.7927 ± 0.1162	0.8257 ± 0.1194	0.8146 ± 0.0959	0.8841 ± 0.0917
UNETR	0.6824 ± 0.1506	0.7004 ± 0.1577	0.6867 ± 0.1488	0.7440 ± 0.1627
SwinUNETR	0.7594 ± 0.1095	0.7663 ± 0.1190	0.7565 ± 0.1394	0.8218 ± 0.1409
U-Mamba_Bot	0.8683 ± 0.0808	0.9049 ± 0.0821	0.8453 ± 0.0673	0.9121 ± 0.0634
U-Mamba_Enc	0.8638 ± 0.0908	0.8980 ± 0.0921	0.8501 ± 0.0732	0.9171 ± 0.0689

• U-Mamba_Bot가 복부 CT에서 DSC 0.8683 및 NSD 0.9049를 달성하며 CNN 및 Transformer 기반 기준선을 능가합니다.
• U-Mamba_Enc가 CT에서 DSC 0.8638 및 NSD 0.8980, 복부 MRI에서 DSC 0.8501 및 NSD 0.9171로 교차 모달 성능이 강하게 나타납니다.
• 2D 과제에서 MRI의 기관, 내시경의 기구, 현미경의 세포에 대해 기저 모델보다 더 우수한 평균 성능을 보입니다.
• 데이터셋 전반에서 세그먼트 아웃라이어가 줄고 간 liver, 담낭 등 도전적 기관과 도구의 경계 정의가 향상됩니다.