[论文解读] Attention U-Net: Learning Where to Look for the Pancreas

Ozan Oktay, Jo Schlemper+10 more|arXiv (Cornell University)|Apr 11, 2018

Advanced Neural Network Applications参考文献 31被引用 4,524

一句话总结

该论文将注意门（Attention Gates）整合到 U-Net（Attention U-Net）中，以学习聚焦目标结构，在 CT 中改进胰腺分割，不需要外部定位模块，开销极小。

ABSTRACT

We propose a novel attention gate (AG) model for medical imaging that automatically learns to focus on target structures of varying shapes and sizes. Models trained with AGs implicitly learn to suppress irrelevant regions in an input image while highlighting salient features useful for a specific task. This enables us to eliminate the necessity of using explicit external tissue/organ localisation modules of cascaded convolutional neural networks (CNNs). AGs can be easily integrated into standard CNN architectures such as the U-Net model with minimal computational overhead while increasing the model sensitivity and prediction accuracy. The proposed Attention U-Net architecture is evaluated on two large CT abdominal datasets for multi-class image segmentation. Experimental results show that AGs consistently improve the prediction performance of U-Net across different datasets and training sizes while preserving computational efficiency. The code for the proposed architecture is publicly available.

研究动机与目标

• 推动在医学图像中对形状/大小变化较大的器官进行准确分割。
• 通过引入注意门（AGs）消除对多阶段级联定位与分割的需求。
• 证明 AGs 在具有挑战性的腹部 CT 数据集上提高分割性能，同时保持计算效率。

提出的方法

• 提出基于网格的软性叠加注意门（soft additive attention gates），产生逐像素的注意系数，并以较粗尺度的门控信号为条件。
• 将 AGs 集成到标准三维 U-Net 的跳跃连接中，在特征融合之前抑制无关激活。
• 使用 1x1x1 卷积实现高效、低参数的门控；为多类分割采用多维注意。
• 通过标准端到端反向传播进行训练，使用 Dice 损失和跨尺度的深度监督，以促进语义层面有意义的门控。

实验结果

研究问题

• RQ1 在 U-Net 中的注意门是否能提高对高变异性器官（如胰腺）的分割精度？
• RQ2 AGs 是否降低对外部定位模块的依赖，同时仍能维持或提高准确性和效率？
• RQ3 Attention U-Net 在不同数据集和训练规模下的表现如何？
• RQ4 在多类腹部分割中，AGs 对召回率和边界准确性的影响是什么？

主要发现

• Attention U-Net 在两个腹部 CT 数据集上持续优于标准 U-Net 的分割性能（例如，胰腺 Dice 分数在 120/30 的训练/测试拆分中从 0.814 提高到 0.840）。
• AGs 提高胰腺召回率和整体分割准确性，同时没有显著的计算开销或模型参数的大幅增加。
• 即使在训练数据减少时加入 AGs 也能带来改进（例如从 120/30 分割减少到 30/120），显示出对数据集规模的鲁棒性。
• 注意力图在训练过程中对胰腺、肾和脾逐步定位，表明跨尺度具有可解释的门控行为。
• 在 TCIA Pancreas-CT 基准测试中，Attention U-Net 在无需多模型级联、参数比高容量的 Vanilla U-Nets 更少的情况下，取得了具有竞争力且通常更优的结果。