[논문 리뷰] MambaU-Lite: A Lightweight Model based on Mamba and Integrated Channel-Spatial Attention for Skin Lesion Segmentation

Thi‐Nhu‐Quynh Nguyen, Quang Tri Ho외 4명|arXiv (Cornell University)|2024. 12. 02.

Cutaneous Melanoma Detection and Management인용 수 4

한 줄 요약

MambaU-Lite는 Mamba의 효율적인 시퀀스 모델링을 컨볼루션 네트워크(CNN)와 함께 통합하고, 통합된 채널-공간 주의 메커니즘을 가진 새로운 P-Mamba 블록을 포함한 경량 의료 영상 분할 모델이다. 이 모델은 단지 420만 개의 파라미터와 1.25G FLOPs로 ISIC2018 및 PH2 데이터셋에서 최신 기술 수준(SOTA) 성능을 달성하며, 계산 및 메모리 비용을 최소화하면서도 기존의 경량 모델보다 DSC 및 IoU 점수에서 뛰어난 성능을 보였다.

ABSTRACT

Early detection of skin abnormalities plays a crucial role in diagnosing and treating skin cancer. Segmentation of affected skin regions using AI-powered devices is relatively common and supports the diagnostic process. However, achieving high performance remains a significant challenge due to the need for high-resolution images and the often unclear boundaries of individual lesions. At the same time, medical devices require segmentation models to have a small memory foot-print and low computational cost. Based on these requirements, we introduce a novel lightweight model called MambaU-Lite, which combines the strengths of Mamba and CNN architectures, featuring just over 400K parameters and a computational cost of more than 1G flops. To enhance both global context and local feature extraction, we propose the P-Mamba block, a novel component that incorporates VSS blocks along-side multiple pooling layers, enabling the model to effectively learn multiscale features and enhance segmentation performance. We evaluate the model's performance on two skin datasets, ISIC2018 and PH2, yielding promising results. Our source code will be made publicly available at: https://github.com/nqnguyen812/MambaU-Lite.

연구 동기 및 목표

• 계산 및 메모리 비용을 최소화하면서도 높은 정확도를 유지하는 경량 딥 러닝 모델을 개발하여 피부 병변 분할을 수행하는 것.
• 모호한 경계를 가진 병변과 고해상도 입력 영상을 처리하기 위한 효율적인 아키텍처로 도전 과제를 해결하는 것.
• Mamba의 장거리 모델링 능력을 CNN과 통합하여 전반적인 맥락 정보와 국소 특징 추출 능력을 향상시키는 것.
• 자원 제약이 있는 의료 기기에서의 구현을 위해 모델 크기와 추론 비용을 줄이는 것.
• 새로운 하이브리드 아키텍처를 통해 ISIC2018 및 PH2와 같은 벤치마크 데이터셋에서의 분할 성능을 향상시키는 것.

제안 방법

• Mamba와 CNN 성분을 융합한 하이브리드 U-Net 유사 아키텍처를 제안하며, 핵심 모듈로 P-Mamba 블록을 사용한다.
• P-Mamba 블록을 도입하여 VSS(Vector-Selective Scan) 블록과 다중 풀링 레이어를 통합함으로써 다중 해상도 특징을 캡처한다.
• 스킵 커넥션에 CBAM(Convolutional Block Attention Module)을 활용하여 공간적 및 채널 기반 특징 표현을 향상시킨다.
• 분할 성능을 최적화하기 위해 Dice 손실과 Tversky 손실을 조합한 이중 손실 함수를 사용하며, 초모수 γ₁=0.7과 γ₂=0.3을 설정한다.
• U-Lite를 영감으로 삼은 경량 설계를 통해 파라미터 수를 0.42M으로 줄이고 메모리 사용량을 1.67MB로 감소시킨다.
• 학습에 Adam 옵timizer를 사용하고, 코즈인 애너일링 학습률 스케줄링과 데이터 증강 기법을 적용한다.

Figure 1 : The architecture of the proposed MambaU-Lite model

실험 결과

연구 질문

• RQ1 하이브리드 Mamba-CNN 아키텍처는 피부 병변 분석에 있어 최소한의 파라미터와 FLOPs로 높은 분할 정확도를 달성할 수 있는가?
• RQ2 제안된 P-Mamba 블록은 표준 Mamba 또는 CNN 블록에 비해 특징 학습 능력을 어떻게 향상시키는가?
• RQ3 스킵 커넥션에 채널-공간 주의 메커니즘을 통합하면 도전적인 피부 병변 데이터셋에서 분할 성능이 향상되는가?
• RQ4 기존의 경량 모델들과 비교했을 때 MambaU-Lite는 정확도, 파라미터 수, 추론 비용 측면에서 어떻게 성능을 내는가?
• RQ5 최소한의 아키텍처 복잡성으로도 MambaU-Lite는 ISIC2018 및 PH2와 같은 다양한 피부 병변 데이터셋에서 잘 일반화되는가?

주요 결과

• ISIC2018 데이터셋에서 MambaU-Lite는 DSC(Dice Similarity Coefficient) 0.9057과 IoU 0.8361을 기록하여 U-Net, Attention U-Net, UNeXt, DCSAU-Net, U-Lite를 모두 능가했다.
• PH2 데이터셋에서 MambaU-Lite는 DSC 0.9572와 IoU 0.9189를 기록하여 경량 모델 중 최신 기술 수준(SOTA) 성능을 달성했다.
• 모델은 단지 0.42만 개의 파라미터와 1.67MB의 메모리 사용량을 가지며, U-Lite(0.88M 파라미터, 3.51MB) 및 기타 모델보다 훨씬 작다.
• U-Lite(0.69G FLOPs)보다 높은 FLOP 수(1.25G)를 가짐에도 불구하고 MambaU-Lite는 더 뛰어난 정확도를 달성하여 더 나은 효율-정확도 트레이드오프를 보였다.
• P-Mamba 블록은 다중 해상도 특징을 효과적으로 캡처하여 경계가 모호한 병변에서의 분할 성능 향상에 기여했다.
• 스킵 커넥션에 CBAM을 통합함으로써 특징 정제가 향상되어 더 정밀한 마스크 예측이 가능해졌다.

Figure 2 : The main components’ architectures of the proposed MambaU-Lite model