[논문 리뷰] Swin-UMamba: Mamba-based UNet with ImageNet-based pretraining

Jiarun Liu, Hao Yang외 9명|arXiv (Cornell University)|2024. 02. 05.

Vehicle License Plate Recognition인용 수 6

한 줄 요약

Swin-UMamba는 ImageNet-사전학습된 Mamba 블록을 활용하여 여러 데이터셋에서 정확도와 효율성을 향상시키는 2D 의료 영상 분할용 Mamba 기반 UNet이며, 더 가벼운 Mamba 기반 디코더를 사용하는 Swin-UMamba dagger 변형이 있다.

ABSTRACT

Accurate medical image segmentation demands the integration of multi-scale information, spanning from local features to global dependencies. However, it is challenging for existing methods to model long-range global information, where convolutional neural networks (CNNs) are constrained by their local receptive fields, and vision transformers (ViTs) suffer from high quadratic complexity of their attention mechanism. Recently, Mamba-based models have gained great attention for their impressive ability in long sequence modeling. Several studies have demonstrated that these models can outperform popular vision models in various tasks, offering higher accuracy, lower memory consumption, and less computational burden. However, existing Mamba-based models are mostly trained from scratch and do not explore the power of pretraining, which has been proven to be quite effective for data-efficient medical image analysis. This paper introduces a novel Mamba-based model, Swin-UMamba, designed specifically for medical image segmentation tasks, leveraging the advantages of ImageNet-based pretraining. Our experimental results reveal the vital role of ImageNet-based training in enhancing the performance of Mamba-based models. Swin-UMamba demonstrates superior performance with a large margin compared to CNNs, ViTs, and latest Mamba-based models. Notably, on AbdomenMRI, Encoscopy, and Microscopy datasets, Swin-UMamba outperforms its closest counterpart U-Mamba_Enc by an average score of 2.72%.

연구 동기 및 목표

• 의료 영상 분할에서 지역 CNN 수용 영역과 높은 비용의 ViTs를 넘는 장거리 의존성 모델링의 필요성을 동기부여한다.
• 2D 의료 영상에 맞춰 ImageNet 기반 사전학습을 적용한 Mamba 기반 인코더를 도입한다.
• 정확한 분할을 위한 향상된 스킵 연결 및 깊은 감독을 갖춘 U-네트 구조의 디코더를 설계한다.
• Swin-UMamba dagger 변형으로 Patch 확장과 매개변수/ FLOPs 감소를 가진 더 가벼운 디코더를 제안한다.
• Mamba 블록으로의 데이터 효율적 분할을 가능하게 하는 사전학습의 중요성을 입증한다.

제안 방법

• ImageNet에서 사전 학습된 Mamba 기반 인코더를 채택하여 2D 의료 영상의 다중 스케일 특징을 추출한다.
• 2D 시퀀스-샷 VSS (SS2D) 블록을 사용해 장거리 의존성을 가진 2D 시각 데이터를 처리한다.
• 인코더 초기화를 위해 VMamba-Tiny와 가중치를 공유하여 ImageNet 사전학습의 이점을 활용한다.
• 스킵 연결과 깊은 감독을 포함하는 U형 아키텍처의 Swin-UMamba 디코더를 구성한다.
• 패치 확장 및 매개변수/FLOPs 감소를 특징으로 하는 더 가벼운 디코더 변형인 Swin-UMamba dagger를 제공한다.
• Dice + cross-entropy 손실, 깊은 감독, AdamW, 코사인 감쇠, 그리고 사전학습 가중치의 단계적 고정으로 학습한다.

실험 결과

연구 질문

• RQ1 ImageNet 기반의 사전학습이 의료 영상 분할에서 Mamba 기반 모델의 성능을 향상시키는가?
• RQ2 Swin-UMamba가 다양한 데이터셋에서 CNN-, ViT-, 및 다른 Mamba 기반 분할 모델과 어떻게 비교되는가?
• RQ3 매개변수 및 FLOPs가 적은 Swin-UMamba dagger가 경쟁력 있는 결과를 얻을 수 있는가?
• RQ4 사전학습이 Mamba 블록을 이용한 의료 분할에서 수렴 안정성과 데이터 효율성에 미치는 영향은 무엇인가?
• RQ5 Mamba 블록의 장거리 모델링 능력이 2D 의료 영상 작업에 얼마나 잘 적용되는가?

주요 결과

• Swin-UMamba와 Swin-UMamba dagger는 AbdomenMRI, Endoscopy, 및 Microscopy 데이터셋에서 CNN-, ViT-, 및 이전 Mamba 기반 모델보다 우수한 성능을 보인다.
• ImageNet 기반 사전 학습은 Swin-UMamba의 DSC에서 상당한 이점을 제공하며 대략 3.04pp의 상승과 NSD에서 대략 4.19포인트의 이득을 준다.
• 사전학습은 더 빠른 수렴과 학습 안정성을 가능하게 하며, AbdomenMRI에서 기본 모델들보다 훨씬 적은 반복으로 수렴한다.
• Swin-UMamba dagger는 Swin-UMamba(40M, 58.4G) 및 U-Mamba 변형에 비해 매개변수와 FLOPs가 크게 적은 27M 매개변수, 15.0G FLOPs로 경쟁력 있는 결과를 달성한다.
• Endoscopy에서 사전학습된 Swin-UMamba dagger는 비사전학습 변형에 비해 DSC/NSD 개선으로 데이터 효율성의 이점을 강조한다.