[Paper Review] EGE-UNet: an Efficient Group Enhanced UNet for skin lesion segmentation

Jiacheng Ruan, Mingye Xie+3 more|arXiv (Cornell University)|Jul 17, 2023

Cutaneous Melanoma Detection and Management12 citations

TL;DR

EGE-UNet is a lightweight UNet variant using Group multi-axis Hadamard Product Attention (GHPA) and Group Aggregation Bridge (GAB) to achieve state-of-the-art skin lesion segmentation with extremely low parameters (~50KB) and low computation.

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.

Motivation & Objective

• Motivate efficient, mobile-friendly skin lesion segmentation.
• Develop ultra-lightweight architecture with strong performance.
• Introduce GHPA and GAB to extract multi-perspective, multi-scale features while minimizing compute.
• Demonstrate state-of-the-art results on ISIC2017/ISIC2018 with extensive ablations.

Proposed method

• Introduce Group multi-axis Hadamard Product Attention (GHPA) to reduce complexity from quadratic to linear and capture multi-perspective information.
• Introduce Group Aggregation Bridge (GAB) to fuse multi-scale features using low/high-level maps and an auxiliary mask.
• Integrate GHPA and GAB into a U-Net framework with six stages and deep supervision for multi-scale outputs.
• Train end-to-end on ISIC2017 and ISIC2018 with data augmentation, AdamW optimization, and cosine annealing learning rate schedule.
• Use loss combining Binary Cross-Entropy and Dice loss with stage-wise deep supervision weights.
• Evaluate against state-of-the-art methods in terms of parameters, FLOPs, mIoU, and DSC.

Experimental results

Research questions

• RQ1 Can GHPA provide multi-perspective attention with linear complexity suitable for mobile health applications?
• RQ2 Does GAB improve multi-scale feature fusion by incorporating mask guidance in a lightweight framework?
• RQ3 How does EGE-UNet compare to both large Transformers-based methods and other lightweight models on ISIC2017/ISIC2018?
• RQ4 What is the trade-off between extremely low parameter count and segmentation performance?

Key findings

• EGE-UNet achieves state-of-the-art performance on ISIC2017 and ISIC2018 while drastically reducing parameters and computation.
• Compared to large models like TransFuse, EGE-UNet reduces parameters by 494x and FLOPs by 160x.
• EGE-UNet reaches around 50 KB parameter size with high segmentation accuracy.
• GHPA provides linear-complexity attention with multi-axis grouping that enhances features.
• GAB fuses low/high-level features and mask guidance to enable effective multi-scale fusion.
• Ablations confirm the necessity of both GHPA and GAB, and the importance of mask information and dilation in GAB.