Source-linked AI summary
EGE-UNet: an Efficient Group Enhanced UNet for skin lesion segmentation
Jiacheng Ruan, Mingye Xie, Jingsheng Gao, Ting Liu, Yuzhuo Fu
TL;DR
Medical image segmentation models can be too large and computationally costly for mobile health applications. EGE-UNet addresses this with grouped linear-complexity attention and multi-scale feature fusion, achieving state-of-the-art skin lesion segmentation with substantially lower resource requirements. The model is evaluated on ISIC2017 and ISIC2018 and is reported to use about 50KB of parameters.
Problem
Existing segmentation models improve performance with complex modules, but their parameter and computational demands limit suitability for real medical settings.
Method
EGE-UNet combines GHPA, which extracts multi-perspective information with linear-complexity HPA, and GAB, which fuses grouped multi-scale features with decoder masks.
Results
EGE-UNet achieves state-of-the-art segmentation performance on ISIC2017 and ISIC2018 while reducing parameters and computation compared with prior models.
Takeaways & Limitations
The model provides an extremely lightweight approach for skin lesion segmentation with substantially lower resource requirements.
Takeaways & Limitations
EGE-UNet is currently designed only for the skin lesion segmentation task and has not yet been deployed in a real-world environment.
Abstract
from arXiv · showhide
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.
1 Introduction
Medical image segmentation models increasingly use powerful but computationally expensive modules, creating a need for lightweight models suitable for mobile healthcare. EGE-UNet addresses this need with efficient attention and feature-fusion modules while targeting state-of-the-art segmentation performance.
- Automated skin lesion segmentation can help medical professionals identify lesion areas quickly and facilitate subsequent treatment.
- Transformer-based segmentation models improve feature extraction but often introduce substantial parameter and computational costs.Examples include Swin-UNet, TransUNet, TransFuse, UTNetV2, and TransBTS.
- Existing lightweight models reduce resource use, but MALUNet still achieves lower segmentation performance than larger models such as TransFuse.
- The model is designed as an extremely lightweight skin lesion segmentation network with state-of-the-art performance and lower resource requirements.The authors report a parameter count of approximately 50KB.
- EGE-UNet combines GHPA for efficient multi-perspective information integration with GAB for multi-scale feature fusion using auxiliary masks.
2 EGE-UNet
EGE-UNet uses a six-stage U-shaped encoder-decoder with GHPA in deeper encoder stages and GAB between encoder and decoder stages. GHPA provides linear-complexity multi-axis attention, while GAB combines grouped low- and high-level features with decoder masks for multi-scale fusion.
- The overall architecture: EGE-UNet uses a symmetric six-stage encoder-decoder with channel numbers {8, 16, 24, 32, 48, 64}.The first three stages use 3×3 plain convolutions, while the last three use GHPA.
- GAB: GAB replaces simple skip connections and fuses low-level features, resized high-level features, and decoder-generated masks.
- GHPA: GHPA divides features into four channel groups and applies HPA across height-width, channel-height, and channel-width axes.
- GHPA: HPA replaces quadratic self-attention complexity with linear complexity by combining input features with a learned tensor through a Hadamard product.
- GAB: GAB partitions feature maps into four channel groups, applies dilated rates {1, 2, 5, 7}, and combines them with a 1×1 convolution.
- Loss function: Deep supervision generates stage-specific mask information, using a weighted combination of binary cross-entropy and Dice losses.The default stage weights are 1, 0.5, 0.4, 0.3, 0.2, and 0.1.
3 Experiments
Experiments evaluate EGE-UNet on ISIC2017 and ISIC2018 using standard dataset splits and compare its segmentation performance, efficiency, and module ablations. Results report state-of-the-art performance alongside substantial reductions in parameters and computation.
- Datasets: ISIC2017 and ISIC2018 contain 2150 and 2694 dermoscopy images, respectively, and are randomly split into training and testing sets at a 7:3 ratio.
- Comparative results: EGE-UNet achieves comprehensive state-of-the-art performance on ISIC2017 while reducing parameters and computation versus TransFuse by 494x and 160x, respectively.
- Comparative results: EGE-UNet improves mIoU by 1.55% and DSC by 0.97% over UNeXt-S, with parameters and computation reduced to 17% and 72% of UNeXt-S.
- Comparative results: EGE-UNet improves mIoU by 1.03% and DSC by 0.64% over MALUNet, while reducing parameters and computation to 30% and 85% of MALUNet.
- Ablation studies: Replacing baseline convolutions with GHPA improves performance while reducing parameters and computation, whereas replacing skip connections with GAB enables multi-scale feature fusion.
4 Conclusions and Future Works
EGE-UNet combines GHPA and GAB to provide lightweight attention and multi-scale feature fusion for skin lesion segmentation. The authors identify real-world deployment and extension beyond skin lesion segmentation as future directions.
- Conclusions: GHPA simplifies self-attention’s quadratic complexity to linear complexity while grouping features to capture information from different perspectives.
- Conclusions: GAB fuses low-level and high-level features with an auxiliary mask to integrate multi-scale information.
- Future works: The authors plan to deploy EGE-UNet in a real-world environment and extend its lightweight design to other tasks.