[논문 리뷰] MALUNet: A Multi-Attention and Light-weight UNet for Skin Lesion Segmentation

Jiacheng Ruan, Suncheng Xiang외 3명|arXiv (Cornell University)|2022. 11. 03.

Cutaneous Melanoma Detection and Management인용 수 6

한 줄 요약

이 논문은 네 가지 새로운 주의 모듈(DGA, IEA, CAB, SAB)을 갖춘 경량 U-Net 아키텍처인 MALUNet을 제안한다. 이는 특징 추출과 샘플 간 관계를 향상시키면서도 파라미터 수와 계산 비용을 최소화하도록 설계되었다. 표준 U-Net 대비 44배 적은 파라미터와 166배 적은 FLOPs를 기록하며 ISIC2017 및 ISIC2018에서 최신 기준 성능을 달성했으며, mIoU는 2.39% 향상되고 DSC는 1.49% 향상되었다.

ABSTRACT

Recently, some pioneering works have preferred applying more complex modules to improve segmentation performances. However, it is not friendly for actual clinical environments due to limited computing resources. To address this challenge, we propose a light-weight model to achieve competitive performances for skin lesion segmentation at the lowest cost of parameters and computational complexity so far. Briefly, we propose four modules: (1) DGA consists of dilated convolution and gated attention mechanisms to extract global and local feature information; (2) IEA, which is based on external attention to characterize the overall datasets and enhance the connection between samples; (3) CAB is composed of 1D convolution and fully connected layers to perform a global and local fusion of multi-stage features to generate attention maps at channel axis; (4) SAB, which operates on multi-stage features by a shared 2D convolution to generate attention maps at spatial axis. We combine four modules with our U-shape architecture and obtain a light-weight medical image segmentation model dubbed as MALUNet. Compared with UNet, our model improves the mIoU and DSC metrics by 2.39% and 1.49%, respectively, with a 44x and 166x reduction in the number of parameters and computational complexity. In addition, we conduct comparison experiments on two skin lesion segmentation datasets (ISIC2017 and ISIC2018). Experimental results show that our model achieves state-of-the-art in balancing the number of parameters, computational complexity and segmentation performances. Code is available at https://github.com/JCruan519/MALUNet.

연구 동기 및 목표

• 높은 계산 및 메모리 요구량으로 인해 자원이 제한된 임상 환경에 복잡한 딥러닝 모델을 구현하는 데 도전하는 문제를 해결하기 위해.
• 모델 복잡도를 증가시키지 않고도 국소 및 전역 특징 표현을 향상시켜 피부 병변 분할 성능을 향상시키기 위해.
• 자기주의 기반 모델에서 자주 간과되는 의료 영상 분할에서의 샘플 간 관계 수립을 위해.
• 채널 및 공간 축을 모두 고려해 다단계 및 다스케일 특징을 효과적으로 융합하여 파라미터 추가 비용을 최소화하면서도 분할 정확도를 향상시키기 위해.

제안 방법

• 확장된 컨볼루션과 게이트드 주의를 조합한 확장된 게이트드 주의(DGA) 블록을 도입하여, 파라미터 수를 줄이며 전역 및 국소적 맥락을 포착한다.
• 전체 데이터셋 수준의 표현을 학습함으로써 샘플 간 관계를 모델링하는 역전형 외부 주의(IEA) 블록을 제안한다.
• 채널 축과 공간 축에서 각각 주의 맵을 생성하고 다단계 특징을 융합하기 위해 채널 주의 브리지(CAB) 및 공간 주의 브리지(SAB) 모듈을 설계한다.
• 모든 네 가지 모듈을 유효한 여섯 단계의 U-Net 아키텍처에 통합하여 MALUNet을 구성함으로써 모델 크기를 최소화하면서 성능을 극대화한다.
• DGA에서 깊이 분리형 컨볼루션을 사용하여 계산 복잡도를 줄이고 효율성을 유지한다.
• 각 단계의 채널 수를 8, 16, 24, 32, 48, 64로 줄여 compact하면서도 강력한 백본을 달성하기 위해 아키텍처를 최적화한다.

실험 결과

연구 질문

• RQ1 경량 U-Net 아키텍처가 파라미터 수와 FLOPs를 크게 줄이며 최신 기준 분할 성능을 달성할 수 있는가?
• RQ2 의료 영상 분할에서 내부 샘플 및 외부 샘플 특징 관계를 향상시키기 위해 주의 메커니즘을 설계할 수 있는가?
• RQ3 채널 및 공간 수준에서의 다단계 특징 융합 기법은 분할 정확도 향상에 얼마나 효과적인가?
• RQ4 전역(IEA를 통해) 및 국소(DGA를 통해) 주의를 통합함으로써 모델 복잡도를 증가시키지 않고 성능 향상을 이룰 수 있는가?
• RQ5 제안된 아키텍처는 ISIC2017 및 ISIC2018와 같은 다양한 피부 병변 데이터셋에서 높은 성능을 유지하는가?

주요 결과

• MALUNet은 표준 U-Net 대비 44배 적은 파라미터와 166배 적은 계산 복잡도(GFLOPs)를 기록하면서 ISIC2018에서 mIoU는 2.39% 향상되고 DSC는 1.49% 향상되었다.
• DGA 블록은 FLOPs와 파라미터 수를 줄이며 기준 U-Net(BasicUNet) 대비 mIoU를 0.90% 향상시키고 DSC는 0.66% 향상시켰다.
• IEA 블록은 파라미터 수가 약간 증가했음에도 불구하고 표준 EA 블록을 초월하여 mIoU는 0.46% 향상되고 DSC는 0.29% 향상되었다.
• DGA와 IEA를 CAB 및 SAB와 융합한 조합이 가장 뛰어난 성능을 보였으며, ISIC2018에서 mIoU는 80.23%, DSC는 89.14%를 달성했다.
• 절단 실험 결과 각 모듈이 성능 향상에 독립적으로 기여하며, CAB 및 SAB는 단계 간 특징 융합을 향상시킨다.
• MALUNet은 모델 효율성과 분할 정확도 사이의 최신 기준 균형을 달성하여, ISIC2017 및 ISIC2018 데이터셋에서 기존의 경량 모델을 모두 능가하는 성능을 보였다.