[论文解读] 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 Management被引用 12

一句话总结

EGE-UNet 是一种轻量级的 UNet 变体，使用 Group multi-axis Hadamard Product Attention (GHPA) 和 Group Aggregation Bridge (GAB) 来实现皮肤病变分割的最先进性能，参数极少（约 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.

研究动机与目标

• 推动高效、适合移动端的皮肤病变分割。
• 开发 ultra-lightweight 架构以实现强性能。
• 引入 GHPA 和 GAB，在尽量减少计算的同时提取多视角、多尺度特征。
• 在 ISIC2017/ISIC2018 上通过广泛消融实验展示最前沿结果。

提出的方法

• 引入 Group multi-axis Hadamard Product Attention (GHPA) 将复杂度从二次方降低到线性并捕获多视角信息。
• 引入 Group Aggregation Bridge (GAB) 以使用低层/高层特征和辅助掩码对多尺度特征进行融合。
• 将 GHPA 和 GAB 集成到六阶段的 U-Net 框架中，采用深度监督实现多尺度输出。
• 在 ISIC2017 和 ISIC2018 上进行端到端训练，采用数据增强、AdamW 优化与余弦退火学习率调度。
• 使用将二元交叉熵损失与 Dice 损失结合的损失函数，并配合阶段性深度监督权重。
• 在参数、FLOPs、mIoU 与 DSC 方面与最先进方法进行比较评估。

实验结果

研究问题

• RQ1 GHPA 是否在移动健康应用中提供线性复杂度的多视角注意力？
• RQ2 GAB 是否通过掩码引导在轻量框架中改进多尺度特征融合？
• RQ3 EGE-UNet 相较于大型 Transformer 基模型和其他轻量模型在 ISIC2017/ISIC2018 的表现如何？
• RQ4 极低参数数量与分割性能之间的权衡是什么？

主要发现

• EGE-UNet 在 ISIC2017 和 ISIC2018 上达到最先进的性能，同时显著降低参数量和计算量。
• 相较于像 TransFuse 这样的大模型，EGE-UNet 将参数量降低 494 倍、FLOPs 降低 160 倍。
• EGE-UNet 的参数规模约为 50 KB，分割精度很高。
• GHPA 提供线性复杂度的多轴分组注意力，提升特征表达。
• GAB 将低层/高层特征与掩码引导进行融合，从而实现有效的多尺度融合。
• 消融分析证实 GHPA 与 GAB 的必要性，以及 GAB 中掩码信息和扩张在性能中的重要性。