[论文解读] Swin Deformable Attention Hybrid U-Net for Medical Image Segmentation

Lichao Wang, Jiahao Huang+2 more|arXiv (Cornell University)|Feb 28, 2023

Radiomics and Machine Learning in Medical ImagingMedicine被引用 3

一句话总结

本文提出 SDAH-UNet，一种混合 U-Net 架构，整合了 Swin 可变形多头自注意力（SDMSA）与并行卷积分支，以提升医学图像分割的准确性和可解释性。通过实现动态、可变形的注意力机制，精准聚焦于心肌或肿瘤等显著区域，该模型减少了计算冗余，并通过形变点提供视觉解释，进而在 ACDC 和 BraTS2020 数据集上实现最先进性能，左心室分割的 DSC 达 92.23%。

ABSTRACT

Medical image segmentation is a crucial task in the field of medical image analysis. Harmonizing the convolution and multi-head self-attention mechanism is a recent research focus in this field, with various combination methods proposed. However, the lack of interpretability of these hybrid models remains a common pitfall, limiting their practical application in clinical scenarios. To address this issue, we propose to incorporate the Shifted Window (Swin) Deformable Attention into a hybrid architecture to improve segmentation performance while ensuring explainability. Our proposed Swin Deformable Attention Hybrid UNet (SDAH-UNet) demonstrates state-of-the-art performance on both anatomical and lesion segmentation tasks. Moreover, we provide a direct and visual explanation of the model focalization and how the model forms it, enabling clinicians to better understand and trust the decision of the model. Our approach could be a promising solution to the challenge of developing accurate and interpretable medical image segmentation models.

研究动机与目标

• 解决混合 CNN-Transformer 模型在医学图像分割中可解释性不足的问题。
• 通过在 Swin Transformer 框架中引入可变形采样，减少标准多头自注意力机制中的计算冗余。
• 通过 SDMSA 实现全局结构建模与并行卷积分支实现局部纹理学习，提升分割准确性。
• 通过形变点提供模型注意力的内在、可视化解释，使临床医生能够理解模型决策过程。
• 通过可解释的聚焦机制，开发在高性能与临床可信度之间取得平衡的模型。

提出的方法

• 所提出的 SDAH-UNet 采用 U-Net 架构，包含三个编码器与解码器模块，每个模块均集成一种新型 SDAPC 模块，结合可变形注意力与卷积特征。
• SDAPC 模块整合了 Swith Deformable Multi-Head Self-Attention (SDMSA) 模块，可学习动态、非均匀的采样点，从而聚焦于左心室壁等不规则形状的解剖结构。
• 并行卷积分支保留局部纹理细节，并提供归纳偏置，与全局注意力机制形成互补。
• 模型采用小卷积核与 GELU 激活函数的卷积嵌入模块，以保持空间分辨率与感受野。
• 反卷积扩展模块对特征进行上采样，以生成掩码预测，同时保持高分辨率输出。
• 通过形变点、注意力分数热力图与基于梯度的可视化方法（SEG-GRAD-CAM）增强模型可解释性，其中形变点在目标区域实现更精确的聚焦。

实验结果

研究问题

• RQ1 在基于 Swin Transformer 的架构中引入可变形注意力，是否能提升在具有复杂不规则结构的医学图像上的分割准确性？
• RQ2 将可变形注意力与卷积分支结合，是否能在提升特征学习的同时减少计算冗余？
• RQ3 注意力机制中的形变点是否能提供比传统基于梯度的方法更准确、更鲁棒的模型聚焦解释？
• RQ4 所提出的混合架构是否在解剖结构与病灶分割任务上均优于现有最先进模型？
• RQ5 该模型的注意力机制是否能够以支持临床信任与理解的方式实现可视化解释？

主要发现

• SDAH-UNet 在 ACDC 数据集的左心室分割任务中取得了 92.23% 的 Dice 相似系数（DSC），优于最先进模型。
• 当使用四个 SDAPC 模块时，该模型在 ACDC 数据集中左心室的 DSC 达到 96.91%，HD95 为 1.22 mm。
• 消融实验表明，将标准模块替换为 SDAPC 模块后，DSC 平均提升约 1.2%。
• 在 SDAPC 模块中移除 SDMSA 或卷积分支，DSC 均下降约 2%，证实了两个组件的贡献。
• 与注意力分数热力图或基于梯度的方法相比，形变点能提供更清晰、更鲁棒的模型聚焦可视化解释，尤其在捕捉肿瘤边界与增强区域方面表现更优。
• 该模型在解剖结构（ACDC）与病灶（BraTS2020）分割任务中均表现出优越性能，证实了其良好的泛化能力。