[论文解读] MALUNet: A Multi-Attention and Light-weight UNet for Skin Lesion Segmentation

Jiacheng Ruan, Suncheng Xiang+3 more|arXiv (Cornell University)|Nov 3, 2022

Cutaneous Melanoma Detection and Management被引用 6

一句话总结

本文提出 MALUNet，一种轻量级 U-Net 架构，包含四个新型注意力模块——DGA、IEA、CAB 和 SAB——旨在增强特征提取和样本间关系，同时最小化参数量和计算成本。与标准 U-Net 相比，该模型在 ISIC2017 和 ISIC2018 数据集上实现了最先进性能，参数量减少 44 倍，FLOPs 减少 166 倍，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，对架构进行优化，实现紧凑而强大的主干网络。

实验结果

研究问题

• RQ1 轻量级 U-Net 架构能否在参数量和 FLOPs 显著减少的情况下实现最先进分割性能？
• RQ2 注意力机制能否被设计为同时增强医学图像分割中的样本内与样本间特征关系？
• RQ3 在通道与空间层面的多阶段特征融合机制在提升分割精度方面的有效性如何？
• RQ4 通过 IEA 实现的全局注意力与通过 DGA 实现的局部注意力相结合，能否在不增加模型复杂度的前提下提升性能？
• RQ5 所提出的架构能否在 ISIC2017 和 ISIC2018 等多样化的皮肤病变数据集上保持高性能？

主要发现

• MALUNet 相较于标准 U-Net，参数量减少 44 倍，计算复杂度（GFLOPs）减少 166 倍，同时在 ISIC2018 上 mIoU 提升 2.39%，DSC 提升 1.49%。
• 与基线 U-Net（BasicUNet）相比，DGA 模块使 mIoU 提升 0.90%，DSC 提升 0.66%，且 FLOPs 和参数量更低。
• IEA 模块优于标准 EA 模块，mIoU 提升 0.46%，DSC 提升 0.29%，且参数量仅略有增加。
• DGA 与 IEA 结合 CAB 和 SAB 可实现最佳性能，在 ISIC2018 上达到 80.23% 的 mIoU 和 89.14% 的 DSC。
• 消融实验证实，各模块对性能提升均有独立贡献，其中 CAB 和 SAB 显著增强了阶段间特征融合。
• MALUNet 在模型效率与分割精度之间实现了最先进平衡，在 ISIC2017 和 ISIC2018 数据集上均优于现有轻量级模型。