[论文解读] SAUNet: Shape Attentive U-Net for Interpretable Medical Image Segmentation

Jesse Sun, Fatemeh Darbehani+2 more|arXiv (Cornell University)|Jan 21, 2020

Advanced Neural Network Applications参考文献 39被引用 22

一句话总结

SAUNet 提出了一种形状感知的 U-Net 架构，通过并行学习结构特征与传统的基于纹理的特征，提升了医学图像分割的可解释性与鲁棒性。该模型在 SUN09 和 AC17 心脏 MRI 数据集上实现了最先进性能，且在推理过程中生成多分辨率显著性图，无需依赖事后基于梯度的可解释性方法。

ABSTRACT

Medical image segmentation is a difficult but important task for many clinical operations such as cardiac bi-ventricular volume estimation. More recently, there has been a shift to utilizing deep learning and fully convolutional neural networks (CNNs) to perform image segmentation that has yielded state-of-the-art results in many public benchmark datasets. Despite the progress of deep learning in medical image segmentation, standard CNNs are still not fully adopted in clinical settings as they lack robustness and interpretability. Shapes are generally more meaningful features than solely textures of images, which are features regular CNNs learn, causing a lack of robustness. Likewise, previous works surrounding model interpretability have been focused on post hoc gradient-based saliency methods. However, gradient-based saliency methods typically require additional computations post hoc and have been shown to be unreliable for interpretability. Thus, we present a new architecture called Shape Attentive U-Net (SAUNet) which focuses on model interpretability and robustness. The proposed architecture attempts to address these limitations by the use of a secondary shape stream that captures rich shape-dependent information in parallel with the regular texture stream. Furthermore, we suggest multi-resolution saliency maps can be learned using our dual-attention decoder module which allows for multi-level interpretability and mitigates the need for additional computations post hoc. Our method also achieves state-of-the-art results on the two large public cardiac MRI image segmentation datasets of SUN09 and AC17.

研究动机与目标

• 解决基于深度学习的医学图像分割模型在可解释性和鲁棒性方面的不足。
• 通过学习更具鲁棒性的形状特征，而非基于纹理的特征，提升模型泛化能力。
• 通过将多分辨率可解释性直接嵌入模型，消除对计算成本较高的事后显著性方法（如 SmoothGrad）的依赖。
• 在大型公开心脏 MRI 分割基准（SUN09 和 AC17）上实现最先进性能。
• 开发一种临床可采纳的分割框架，支持诊断应用中的透明性与可靠性。

提出的方法

• 提出一种双流编码器-解码器架构，包含主纹理流和次门控形状流，以并行方式学习形状相关特征。
• 采用双注意力解码器模块，在推理过程中生成多分辨率的空间注意力图与形状注意力图，实现内置可解释性。
• 使用双任务损失目标，联合优化分割准确率与形状特征学习，提升模型鲁棒性。
• 在编码器与解码器路径之间应用跳跃连接，以保留空间上下文信息，同时实现跨尺度的特征融合。
• 通过不同解码器层级的注意力图生成显著性图，避免了 SmoothGrad 等方法所需的额外前向-反向传播计算。
• 在 SUN09 和 AC17 数据集上训练时，采用 RAdam 优化器，初始学习率为 5e-4，批量大小为 4。

实验结果

研究问题

• RQ1 与标准 CNN 相比，显式学习形状特征的双流 U-Net 架构是否能提升分割的鲁棒性？
• RQ2 集成形状流是否能减少从训练数据（SUN09）迁移到不同测试集（AC17）时的性能下降？
• RQ3 内置注意力图是否能提供无需依赖事后基于梯度方法（如 SmoothGrad）的多分辨率可解释性？
• RQ4 形状感知特征在双心室心脏 MRI 分割任务中，能在多大程度上提升分割准确率？
• RQ5 所提出方法是否能在保持内在可解释性的同时实现最先进性能？

主要发现

• 包含形状流的模型在 SUN09 训练集上达到 90.84% 的 mIoU，优于无形状流的基线模型（90.20%）。
• 在 AC17 验证集上，包含形状流的模型达到 81.73% 的 mIoU，相比基线模型（79.33%）的性能下降更小，差距为 1.76%。
• 形状流提升了泛化能力，表现为从 SUN09 到 AC17 的 mIoU 下降幅度减小了 1.76%，表明模型更具鲁棒性。
• 所提方法在 20 秒内即可完成整个 AC17 验证集的多分辨率显著性图生成，而 SmoothGrad 方法需 24 分钟，显著提升了计算效率。
• 形状流生成的注意力图清晰地聚焦于左心室和右心室的解剖形状，证实模型学习到了有意义的结构特征。
• 双注意力解码器实现了分层可解释性，高分辨率模块关注细节，低分辨率模块捕捉全局形状先验。