[论文解读] Attention Swin U-Net: Cross-Contextual Attention Mechanism for Skin Lesion Segmentation

Ehsan Khodapanah Aghdam, Reza Azad+2 more|arXiv (Cornell University)|Jan 1, 2022

Cutaneous Melanoma Detection and Management被引用 6

一句话总结

本文提出 Att-SwinU-Net，一种增强版的 Swin U-Net 架构，在跳跃连接中引入一种新颖的跨上下文注意力机制，用于皮肤病变分割。通过用可学习的注意力模块替代标准的拼接操作，整合编码器特征并实现跨上下文的特征优化，该模型在多个基准数据集上实现了最先进（SOTA）的性能，优于基于 CNN 和基于 Transformer 的模型。

ABSTRACT

Melanoma is caused by the abnormal growth of melanocytes in human skin. Like other cancers, this life-threatening skin cancer can be treated with early diagnosis. To support a diagnosis by automatic skin lesion segmentation, several Fully Convolutional Network (FCN) approaches, specifically the U-Net architecture, have been proposed. The U-Net model with a symmetrical architecture has exhibited superior performance in the segmentation task. However, the locality restriction of the convolutional operation incorporated in the U-Net architecture limits its performance in capturing long-range dependency, which is crucial for the segmentation task in medical images. To address this limitation, recently a Transformer based U-Net architecture that replaces the CNN blocks with the Swin Transformer module has been proposed to capture both local and global representation. In this paper, we propose Att-SwinU-Net, an attention-based Swin U-Net extension, for medical image segmentation. In our design, we seek to enhance the feature re-usability of the network by carefully designing the skip connection path. We argue that the classical concatenation operation utilized in the skip connection path can be further improved by incorporating an attention mechanism. By performing a comprehensive ablation study on several skin lesion segmentation datasets, we demonstrate the effectiveness of our proposed attention mechanism.

研究动机与目标

• 解决 U-Net 架构在医学图像分割中局部感受野受限的问题。
• 通过用基于注意力的机制替代标准拼接操作，提升跳跃连接中的特征重用性。
• 通过增强 Swin U-Net 中的长距离依赖建模能力，提升对复杂皮肤病变的分割性能。
• 通过混合注意力机制，在皮肤病变分割基准上实现最先进性能。

提出的方法

• 在跳跃连接中提出两级注意力机制：首先，将注意力权重从编码器模块传递，以突出重要标记；其次，通过跨上下文注意力机制，对不同尺度的特征进行优化。
• 将所提出的注意力模块集成到 Swin U-Net 的跳跃连接路径中，用可学习的基于注意力的特征融合替代标准拼接操作。
• 采用双分支注意力机制，分别建模空间和通道依赖关系，以提升特征表示能力。
• 利用 Swin Transformer 的移位窗口机制，在编码器和解码器路径中实现局部与全局特征建模。
• 通过消融实验评估注意力模块位置、输入分辨率、模型规模以及组件移除的影响。
• 在多个皮肤病变数据集上使用标准分割损失函数端到端训练模型。

实验结果

研究问题

• RQ1 跳跃连接中的注意力机制是否能提升 U-Net 架构中特征的重用性和分割精度？
• RQ2 所提出的跨上下文注意力机制与跳跃连接中的标准拼接操作相比表现如何？
• RQ3 在跳跃连接中，注意力模块的最佳数量和位置是什么，以实现皮肤病变分割的最优性能？
• RQ4 所提出方法是否在皮肤病变分割基准上优于现有的基于 CNN 和基于 Transformer 的模型？
• RQ5 空间注意力和通道注意力组件对最终性能的独立贡献是什么？

主要发现

• 在 ISIC2017 数据集上，Att-SwinU-Net 的 Dice 相似度系数（DSC）达到 0.9240，优于先前的 SOTA 方法（TMU-Net）0.0076。
• 在 ISIC2018 数据集上，模型的 DSC 达到 0.9105，在 PH2 数据集上达到 0.9504，表明在所有数据集上均表现出一致的优越性。
• 消融实验表明，在跳跃连接中使用三个注意力模块可获得最佳性能，DSC 从 0.8987 提升至 0.9240。
• 移除空间注意力模块或跨上下文注意力模块，DSC 分别下降 0.0039 和 0.0044，证实了二者各自的贡献。
• 在 ISIC2017 数据集上，模型的准确率达到 0.9656，表明分割具有一致性且误报率较低。
• 定性结果表明，与 TransUNet 和 Swin U-Net 相比，该模型生成的分割掩码更平滑、更精确，尤其在减少欠分割伪影方面表现更优。