[논문 리뷰] EGE-UNet: an Efficient Group Enhanced UNet for skin lesion segmentation

Jiacheng Ruan, Mingye Xie외 3명|arXiv (Cornell University)|2023. 07. 17.

Cutaneous Melanoma Detection and Management인용 수 12

한 줄 요약

EGE-UNet은 GHPA 및 GAB를 사용한 경량 UNet 변형으로 매개변수 매우 적은 (~50KB) 및 계산이 낮은 상태에서 최첨단 피부 병변 분할을 달성합니다.

ABSTRACT

Transformer and its variants have been widely used for medical image segmentation. However, the large number of parameter and computational load of these models make them unsuitable for mobile health applications. To address this issue, we propose a more efficient approach, the Efficient Group Enhanced UNet (EGE-UNet). We incorporate a Group multi-axis Hadamard Product Attention module (GHPA) and a Group Aggregation Bridge module (GAB) in a lightweight manner. The GHPA groups input features and performs Hadamard Product Attention mechanism (HPA) on different axes to extract pathological information from diverse perspectives. The GAB effectively fuses multi-scale information by grouping low-level features, high-level features, and a mask generated by the decoder at each stage. Comprehensive experiments on the ISIC2017 and ISIC2018 datasets demonstrate that EGE-UNet outperforms existing state-of-the-art methods. In short, compared to the TransFuse, our model achieves superior segmentation performance while reducing parameter and computation costs by 494x and 160x, respectively. Moreover, to our best knowledge, this is the first model with a parameter count limited to just 50KB. Our code is available at https://github.com/JCruan519/EGE-UNet.

연구 동기 및 목표

• 효율적이고 모바일 친화적인 피부 병변 분할을 고무한다.
• 강력한 성능을 가진 초경량 아키텍처를 개발한다.
• 계산을 최소화하면서 다각적·다중 스케일 특징을 추출하기 위해 GHPA와 GAB를 도입한다.
• ISIC2017/ISIC2018에서 광범위한 ablations로 최첨단 결과를 입증한다.

제안 방법

• Group multi-axis Hadamard Product Attention (GHPA)을 도입하여 복잡도를 제곱에서 선형으로 줄이고 다-perspective 정보를 포착한다.
• Group Aggregation Bridge (GAB)을 도입하여 저/고레벨 맵과 보조 마스크를 사용해 다중 스케일 특징을 융합한다.
• GHPA와 GAB를 6단계의 U-Net 프레임워크에 통합하고 다중 스케일 출력을 위한 깊은 감독을 적용한다.
• ISIC2017 및 ISIC2018에서 데이터 증강, AdamW 최적화 및 코사인 감소 학습률 스케줄로 엔드-투-엔드로 학습한다.
• Binary Cross-Entropy와 Dice 손실을 스테이지별 깊은 감독 가중치와 결합해 사용한다.
• 매개변수, FLOPs, mIoU 및 DSC 측면에서 최첨단 방법과의 비교 평가를 수행한다.

실험 결과

연구 질문

• RQ1 GHPA가 모바일 건강 애플리케이션에 적합한 선형 복잡도의 다-perspective 주의력을 제공할 수 있는가?
• RQ2 GAB가 마스크 가이던스를 도입하여 경량 프레임워크에서 다중 스케일 특징 융합을 개선하는가?
• RQ3 ISIC2017/ISIC2018에서 EGE-UNet은 대형 Transformer 기반 방법 및 다른 경량 모델과 어떻게 비교되는가?
• RQ4 매우 낮은 매개변수 수와 분할 성능 간의 트레이드오프는 어느 정도인가?

주요 결과

• EGE-UNet은 ISIC2017 및 ISIC2018에서 최첨단 성능을 달성하면서 매개변수 및 계산을 대폭 감소시킨다.
• TransFuse 같은 대형 모델과 비교하면 EGE-UNet은 매개변수를 494배, FLOPs를 160배 감소시킨다.
• EGE-UNet은 대략 50 KB의 매개변수 용량에서 높은 분할 정확도를 달성한다.
• GHPA는 다축 그룹화로 선형 복잡도의 주의력을 제공하고 특징을 강화한다.
• GAB는 저/고수준 특징과 마스크 가이던스를 융합하여 효과적인 다중 스케일 융합을 가능하게 한다.
• Ablation 결과는 GHPA와 GAB의 필요성과 GAB에서 마스크 정보와 확장의 중요성을 보여준다.