[논문 리뷰] MedMamba: Vision Mamba for Medical Image Classification

Yubiao Yue, Zhenzhang Li|arXiv (Cornell University)|2024. 03. 06.

COVID-19 diagnosis using AI인용 수 65

한 줄 요약

MedMamba는 SS-Conv-SSM 블록으로 로컬 및 장거리 특징을 포착하는 Vision Mamba 기반 아키텍처를 의료 영상 분류에 도입했으며, 여러 모달리티에서 경쟁력 있는 결과로 평가되었습니다.

ABSTRACT

Since the era of deep learning, convolutional neural networks (CNNs) and vision transformers (ViTs) have been extensively studied and widely used in medical image classification tasks. Unfortunately, CNN's limitations in modeling long-range dependencies result in poor classification performances. In contrast, ViTs are hampered by the quadratic computational complexity of their self-attention mechanism, making them difficult to deploy in real-world settings with limited computational resources. Recent studies have shown that state space models (SSMs) represented by Mamba can effectively model long-range dependencies while maintaining linear computational complexity. Inspired by it, we proposed MedMamba, the first Vision Mamba for generalized medical image classification. Concretely, we introduced a novel hybrid basic block named SS-Conv-SSM, which purely integrates the convolutional layers for extracting local features with the abilities of SSM to capture long-range dependencies, aiming to model medical images from different image modalities efficiently. By employing the grouped convolution strategy and channel-shuffle operation, MedMamba successfully provides fewer model parameters and a lower computational burden for efficient applications without sacrificing accuracy. We thoroughly evaluated MedMamba using 16 datasets containing ten imaging modalities and 411,007 images. Experimental results show that MedMamba demonstrates competitive performance on most tasks compared with the state-of-the-art methods. This work aims to explore the potential of Vision Mamba and establish a new baseline for medical image classification, thereby providing valuable insights for developing more powerful Mamba-based artificial intelligence algorithms and applications in medicine. The source codes and all pre-trained weights of MedMamba are available at https://github.com/YubiaoYue/MedMamba.

연구 동기 및 목표

• 의료 영상에서 장거리 의존성 포착의 한계를 CNN으로 해결하려 한다.
• 로컬 및 장거리 모델링을 결합한 경량의 효율적 아키텍처를 제안한다.
• 다양한 의료 영상 모달리티에서 MedMamba를 평가하여 SSM 기반 의료 영상 분류의 새로운 기준선을 확립한다.

제안 방법

• SS-Conv-SSM 모듈을 도입하여 컨볼루션 로컬 특징과 상태-공간 모델의 장거리 의존성을 융합한다.
• SS2D(2D-선택적 스캔)를 개발하여 다방향 특징 시퀀스를 스캔하고 병합한다.
• 패치 임베딩 계층(4x4 패치) 후 채널 확장과 함께 네 단계의 SS-Conv-SSM 블록을 사용한다.
• 레이어 정규화, 이중 분기 아키텍처(Conv-분기와 SSM-분기) 및 상호 작용을 위한 채널 셔플을 적용한다.
• 데이터 증강이나 사전 학습 없이 훈련하여 아키텍처 주도 이점을 보여준다.
• 이미지 데이터의 엔드 투 엔드 모델링을 위한 이산화된 상태 공간 모델링(Abar, Bbar) 및 컨볼루셔널 커널 어셈블리(Kbar) 이론적 근거를 제공한다.

실험 결과

연구 질문

• RQ1 상태 공간 모델과 컨볼루션 특징을 통합하는 것이 모달리티 전반에 걸친 의료 영상 분류에 어떤 영향을 미치는가?
• RQ2 MedMamba가 사전 학습이나 데이터 증강 없이도 다양한 의료 데이터세트에서 경쟁력 있는 성능을 달성할 수 있는가?
• RQ3 SS2D 선택적 스캐닝 메커니즘이 특징 추출 및 장거리 의존성 모델링에 미치는 영향은 무엇인가?

주요 결과

• MedMamba는 CNN- 및 Transformer 기반 기준선과 비교하여 경쟁적이며 일부 데이터셋에서 우수한 성능을 보인다.
• 해당 아키텍처는 피부 병변, 초음파, 흉부 X-레이, 내시경, MedMNIST 유도 작업 등 여러 모달리티에서 잘 작동한다.
• SS-Conv-SSM은 국소(컨볼루션) 및 전역(SSM) 특징을 효과적으로 결합하여 의료 영상에서 미세한 정보와 거친 정보를 모두 포착한다.
• 공개/비공개 다섯 개의 의료 데이터셋에 대한 광범위한 실험은 MedMamba가 SSM 기반 의료 영상 분류의 새로운 기준선을 확립할 수 있음을 시사한다.
• 복제 및 추가 연구를 위한 소스 코드는 공개되어 있다.