Attention U-Net笔记

原文地址：Learning Where to Look for the Pancreas

Abstract

我们提出了一种新的用于医学成像的attention gate（AG）模型，该模型能够自动学习聚焦不同形状和大小的目标结构。使用AGs训练的模型隐式学习抑制输入图像中的无关区域，同时突出对特定任务有用的显著特征。这使我们能够消除使用级联卷积神经网络（CNN）的显式外部组织/器官定位模块的必要性。AGs可以很容易地集成到标准的CNN架构中，比如U-Net模型，只需最小的计算开销，同时提高模型的灵敏度和预测精度。在两个用于多类图像分割的大型CT腹部数据集上对所提出的注意力U-Net结构进行了评估。实验结果表明，AGs在保持计算效率的同时，在不同的数据集和训练规模下持续提高了U-Net的预测性能。提议的架构的源代码是公开的。

1、Introduction

自动医学图像分割在图像分析领域得到了广泛的研究，这是因为对大量医学图像进行手动密集标记是一项繁琐且容易出错的任务。通过快速自动提取定量测量值，需要准确可靠的解决方案来提高临床工作流程效率，并支持决策。

随着卷积神经网络（CNN）的出现，自动医学图像分析任务（包括心脏MR分割[3]和肺癌结节检测[17]）可以实现接近放射科医生水平的性能。高表示能力、快速推理和滤波器共享特性使CNN成为事实上的图像分割标准。全卷积网络（FCN）[18]和U-Net[24]是两种常用的体系结构。尽管这些结构具有良好的代表性，但当靶器官在形状和大小方面表现出较大的患者间差异时，它们依赖于多级级联CNN。级联框架提取region of interest（感兴趣区域）（ROI），并对该特定ROI进行密集预测。应用领域包括心脏MRI[14]、心脏CT[23]、腹部CT[26,27]分割和肺部CT结节检测[17]。然而，这种方法会导致计算资源和模型参数的过度和冗余使用；例如，级联中的所有模型都会重复提取类似的低级特征。为了解决这个普遍的问题，我们提出了一个简单而有效的解决方案，即注意门（AGs）。带有AGs的CNN模型可以用类似于FCN模型训练的标准方式从头开始训练，AGs可以自动学习关注目标没有额外监督的结构。在测试时，这些门会动态地隐式生成软区域建议，并突出对特定任务有用的显著特征。此外，它们不会引入显著的计算开销，也不需要像多模型框架那样需要大量的模型参数。作为回报，所提出的AGs通过抑制无关区域的特征激活，提高了密集标签预测的模型敏感性和准确性。通过这种方式，可以消除使用外部器官定位模型的必要性，同时保持较高的预测精度。natural image classification[11]和captioning[1]也提出了类似的注意机制，以执行自适应特征池，其中模型预测仅以选定图像区域的子集为条件。在本文中，我们推广了这种设计，并提出了基于图像网格的选通，允许注意系数特定于局部区域。此外，我们的方法可以用于基于注意力的密集预测。

我们演示了AG在标准U-Net体系结构（Attention U-Net）中的实现，并将其应用于医学图像。我们选择具有挑战性的CT胰腺分割问题，为我们提出的贡献提供实验证据。由于组织对比度低，器官形状和大小变化大，这个问题构成了一个困难的任务。我们在两个常用的基准上评估我们的实现：TCIA胰腺CT-82[25]和多类腹部CT-150。结果表明，AGs在不需要多个CNN模型的情况下，能够在不同数据集和训练规模下持续提高预测精度，同时实现最先进的性能。

1.1 、Related Work

CT Pancreas Segmentation:

腹部CT胰腺分割的早期工作使用了统计形状模型[5,28]或多图谱技术[22,34]。特别是，atlas方法受益于人工注释传播所施加的隐式形状约束。然而，在TCIA数据集[25]等公共基准中，基于atlas的框架的Dice相似系数（DSC）的范围为69.6%到73.9%[22,34]。[39]提出了一种基于分类的框架，以消除atlas对图像配准的依赖性。最近，人们提出了级联多阶段CNN模型[26、27、38]来解决这个问题。在这里，使用初始粗层次模型（例如U-Net或回归森林）来获得ROI，然后使用裁剪的ROI通过第二个模型进行细分细化。类似地，在[4]中，2D-FCN和递归神经网络（RNN）模型的组合被用于利用相邻轴向切片之间的相关性。这些方法在TCIA基准中实现了最先进的性能（81.2%− 82.4%（DSC）。如果不使用级联框架，性能将下降2.0%到4.4%。最近的工作[37]提出了一个迭代的两阶段模型，该模型递归地更新局部和全局预测，并且两个模型都是端到端训练的。除了标准的FCN，密集连接[6]和稀疏卷积[8,9]也被应用于CT胰腺分割问题。密集连接和稀疏核通过减少非零参数的数量来降低计算复杂度。

Attention Gates:

AGs通常用于自然图像分析、知识图和语言处理（NLP），用于图像字幕[1]、机器翻译[2,30]和分类[11,31,32]任务。最初的工作是通过解释输出类分数相对于输入图像的梯度来探索注意图。另一方面，可训练注意力是通过设计来实现的，分为硬注意和软注意。硬注意[21]，例如迭代区域建议和裁剪，通常是不可微的，并且依赖强化学习来更新参数，这使得模型训练更加困难。在[36]中，递归硬注意用于检测胸部X光扫描中的异常。相反，软注意是概率的，并且使用标准的反向传播，不需要蒙特卡罗抽样。例如，（additive soft attention）加法软注意用于句子间的翻译[2,29]，最近用于图像分类[11,32]。在[10]中，渠道关注被用来突出重要的特征维度，这是ILSVRC 2017图像分类挑战赛中表现最好的。Self-attention技术[11,33]已经被提出，以消除对外部门控信息的依赖。例如，在[33]中，非局部自我注意被用来捕捉长期依赖。在[11,32]中，自我关注被用于执行特定于类的池，这会导致更准确和更稳健的图像分类性能。

1.2 、Contributions

在本文中，我们提出了一种新的自我注意门控模块，可用于基于CNN的密集标签预测标准图像分析模型。此外，我们还探讨了AGs对医学图像分析的好处，尤其是在图像分割方面。这项工作的贡献总结如下：

1）我们进一步采用了[11]中提出的注意方法，提出了grid-based gating，允许注意系数更具体地针对局部区域。与基于全局特征向量的选通相比，这提高了性能。此外，我们的方法可以用于密集预测，因为我们不执行自适应池。

2）我们提出了一个应用于医学成像任务的前馈CNN模型中的第一个 soft-attention技术用例。提出的注意门可以取代图像分类[36]中使用的硬注意方法，以及图像分割框架中的外部器官定位模型[14、22、26、27]。

3）提出了对标准U-Net模型的扩展，以提高模型对前景像素的敏感性，而不需要复杂的启发式算法。通过实验观察到，U-Net的精度改进在不同的成像数据集中是一致的。（泛化性较好）

2、Methodology

Fully Convolutional Network (FCN):

卷积神经网络（CNN）在公共基准数据集上的医学图像分析中优于传统方法[14,17]，同时比图切割和多图谱分割技术[34]快一个数量级。这主要归因于以下事实：（I）使用随机梯度下降（SGD）优化学习特定于域的图像特征；（II）学习的核在所有像素上共享；（III）图像卷积操作充分利用了医学图像中的结构信息。特别是，U-Net[24]、DeepMedic[13]和整体嵌套网络[16,35]等完全卷积网络（FCN）[18]已被证明在各种任务中，包括心脏MR[3]、脑肿瘤[12]和腹部CT[26,27]图像分割任务中，能够实现稳健而准确的性能。

这一段公式多，放原文，主要内容还是再讲卷积神经网络的线性回归和激活函数，以及整个卷积神经网络的基本流程，之后有介绍了他们怎样将AGS和U-Net网络结合，粗略的特征地图捕捉上下文信息，并突出显示前景对象的类别和位置。在多个尺度上提取的特征图随后通过skip connections进行合并，以结合粗、细密度预测，结构图如图一：

Attention Gates for Image Analysis:

为了获得足够大的感受野，从而获得语义上下文信息，在标准的CNN架构中，特征图网格逐渐减少采样。通过这种方式，特征在粗略的空间网格级别模型上定位，并在全局范围内确定组织之间的关系。然而，对于形状变化较大的小物体，仍然很难减少假阳性预测。为了提高精度，当前的分割框架[14,26,27]依赖于其他先前的对象定位模型，将任务简化为单独的定位和后续分割步骤。在这里，我们证明了通过将注意力门（AGs）整合到标准CNN模型中也可以实现同样的目标。这不需要训练多个模型和大量额外的模型参数。与多阶段CNN中的定位模型不同，AGs逐步抑制无关背景区域中的特征响应，而无需在网络之间裁剪ROI。

注意系数αi∈ [0,1]，识别显著图像区域并修剪特征响应，以仅保留与特定任务相关的激活，如图3a所示。AGs的输出是输入特征映射和注意系数的元素相乘：ˆxli，c=xli，c·αli。在默认设置中，为每个像素向量xli计算单个标量注意值∈ RFl，其中Fl对应于l层中特征映射的数量。在多个语义类的情况下，我们建议学习多维注意系数。这是受到[29]的启发，其中多维注意系数用于学习句子嵌入。因此，每个AG学习关注目标结构的子集。如图2所示，一个选通向量gi∈ RFg用于每个像素i以确定聚焦区域。如[32]所述，选通向量包含上下文信息，用于修剪较低级别的特征响应，它使用AGs进行自然图像分类。我们使用附加注意[2]来获得选通系数。虽然这在计算上更昂贵，但实验表明，它比乘法注意具有更高的精确度[19]。附加注意的公式如下：（这里放原文，有一部分公式）

softmax激活函数用于归一化注意系数（σ2）；然而，顺序使用softmax会在输出端产生更稀疏的激活。因此，我们选择了一个S形激活函数。实验结果表明，AG参数的训练收敛性更好。与[11]相比，我们提出了一种网格注意技术。在这种情况下，选通信号不是所有图像像素的全局单一向量，而是一个适应于图像空间信息的栅格信号。更重要的是，每个skip connections的选通信号聚合了来自多个成像比例的信息，如图1所示，这提高了查询信号的网格分辨率并实现了更好的性能。最后，我们想指出的是，AG参数可以通过标准的反向传播更新进行训练，而无需在硬注意中使用基于采样的更新方法[21]。

Attention Gates in U-Net Model:

建议的AGs被纳入标准U-Net架构中，以突出通过skip connections的显著特征，见图1。从粗尺度提取的信息用于选通，以消除skip connections中不相关和噪声响应的歧义。这是在连接操作之前执行的，以仅合并相关激活。此外，AGs过滤前向传递和后向传递期间的神经元激活。来自背景区域的渐变在后退过程中被向下加权。这使得较浅层中的模型参数主要基于与给定任务相关的空间区域进行更新。l层卷积参数的更新规则l − 1可表述如下：

右手边的第一个梯度项用αli缩放。对于多维AGs，αLi对应于每个网格尺度上的向量。在每个子AG中，提取互补信息并进行融合，以定义skip connections的输出。为了减少AGs的可训练参数数量和计算复杂度，在没有任何空间支持（1x1x1卷积）的情况下执行线性变换，并将输入特征映射降采样到选通信号的分辨率，类似于非局部块[33]。相应的线性变换将特征映射解耦，并将其映射到低维空间以进行选通操作。如[11]所述，低级特征映射，即第一个skip connections，不用于选通功能，因为它们不代表高维空间中的输入数据。我们使用深度监督[16]来强制中间特征映射在每个图像尺度上具有语义区分性。这有助于确保不同尺度的注意力单元能够影响对大范围图像前景内容的反应。因此，我们防止从skip connections的小子集重建密集预测。

3、Experiments and Results

提出的AG模型是模块化的（可以应用在任何图像领域），与应用类型无关；因此，它可以很容易地适用于分类和回归任务。证明其在图像分割中的适用性，我们在一个具有挑战性的腹部CT多标签分割问题上评估了注意力U-Net模型。特别是，由于形状变异性和组织对比度差，胰腺边界划定是一项困难的任务。我们的模型在分割性能、模型容量、计算时间和内存需求方面与标准3D U-Net进行了比较。

Evaluation Datasets:

在实验中，使用了两种不同的腹部CT数据集：（I）150次腹部3D CT扫描，这些扫描来自被诊断为胃癌的患者（CT-150）。在所有图像中，胰腺、肝脏和脾脏的边界由三名训练有素的研究人员半自动划定，并由临床医生手动验证。[27]中使用了相同的数据集，对胰腺分割中的U-Net模型进行基准测试。（二）第二个数据集1（CT-82）由82个对比增强3D CT扫描组成，其中胰腺手动注释逐层执行。该数据集（NIH-TCIA）[25]是公开的，通常用于基准CT胰腺分割框架。由于较大的图像尺寸和硬件内存限制，来自两个数据集的图像被降采样到各向同性2.00 mm分辨率。

Implementation Details:

与最先进的CNN分割框架[4,26]相比，我们提出了一个3D模型来捕捉足够的语义上下文。梯度更新使用2到4个样本的小批量计算。对于较大的网络，在多个正向和反向过程中使用梯度平均。所有模型都使用Adam Optimizer[15]、批量归一化、深度监督[16]和标准数据增强技术（仿射变换、轴向翻转、随机裁剪）进行训练。强度值线性缩放以获得正态分布N（0，1）。这些模型是使用Sorensen Dice loss[20]对所有语义类进行训练的，实验表明，Sorensen Dice loss[20]对类不平衡不太敏感。门控参数被初始化，以便注意门通过所有空间位置的特征向量。此外，我们不需要像基于注意力的方法那样需要多个训练阶段，因此简化了训练过程。我们使用PyTorch的实现是公开的。

Attention Map Analysis:

从测试图像中获得的注意力系数根据训练时间进行可视化（见图4）。我们通常观察到AGs最初在所有位置都具有均匀分布和通行特征。这是逐步更新和定位的目标器官边界。此外，在较粗的尺度下，AGs提供了器官的大致轮廓，这些器官在较细的分辨率下逐渐细化。此外，通过在每个图像尺度上训练多个AGs，我们观察到每个AGs学习聚焦于特定器官子集。

Segmentation Experiments:

本文提出的注意力U-Net模型在腹部CT多类分割中与标准U-Net模型进行了对比。我们使用CT-150数据集进行训练（120）和测试（30）。表1给出了相应的骰子分数（DSC）和表面距离（S2S）。胰腺预测的结果表明，注意门（AGs）通过提高模型的表达能力来增加回忆值（p=0.005），因为它依赖AGs来定位前景像素。图3b定性比较了使用这两个模型得出的预测之间的差异。在第二个实验中，使用较少的训练图像（30）训练相同的模型，以表明对于不同大小的训练数据，性能改善是一致且显著的（p=.01）。对于这两种方法，我们都观察到性能下降随着训练规模的减小，训练也会增加。在拟议的框架下，下降的幅度较小。对于肾脏分割，由于组织对比度更高，模型达到了类似的精度。

在表1中，我们还报告了两种模型的可训练参数数量。我们观察到，通过向标准U-Net添加8%的额外容量，性能可以在DSC方面提高2-3%。为了进行公平比较，我们还训练了更大容量的U-Net模型，并与较小网络规模的拟议模型进行了比较。表2所示的结果表明，AGs的添加不仅仅是简单地（均匀地）增加网络所有层的模型容量（p=.007）。因此，当AGs用于减少训练多个单独模型的冗余时，AGs应使用额外的能力来定位组织。

（注：这里表二第一列应该是写错了，应该U-Net (120/30) 和U-Net (30/120)）

Comparison to State-of-the-Art:

在公共TCIA CT胰腺基准上对所提出的架构进行评估，以将其性能与最先进的方法进行比较。最初，在CT-150数据集上训练的模型直接应用于CT-82数据集，以观察两个模型在不同数据集上的适用性。表3给出了相应的结果**（BFT）。尽管U-Net模型是在不相交的数据集上训练的，但它的性能优于传统的atlas技术[34]。此外，注意力模型在不同数据集的胰腺分割中表现一直更好。这些模型随后在TCIA数据集（61 train，21 test）的子集上进行微调（AFT**）。与脾和肾对应的输出节点被排除在输出softmax计算之外，梯度更新仅针对背景和胰腺标签进行计算。结果见表3和表4与串联多模型CNN方法[4,26,38]相比，由于额外的训练数据和更丰富的语义信息（例如脾脏标签），显示出更好的性能。此外，我们从CT-82数据集中随机选取61张训练图像，对这两个模型进行从头开始（SCR）训练。与CT-150数据集上的结果类似，AGs提高了分割精度并降低了 the surface distances（表面距离）（p=0.03），因为胰腺像素的召回率（p=0.09）。

为了便于比较，表4总结了最先进的CT胰腺分割模型的结果。由于这些模型是在同一个训练数据集上训练的，因此这种比较可以深入了解注意力模型与相关文献的比较情况。需要注意的是，我们的框架中没有使用后处理（如条件随机场），因为实验主要集中在量化AGs在孤立环境中带来的性能改善。类似地，如[6]中所述，residual and dense connections剩余和密集连接可与AGs结合使用，以改善分割结果。在这方面，我们的3D Attention U-Net模型的性能与最先进的类似，尽管输入图像被降采样到较低的分辨率。更重要的是，与基于单一模型的细分框架相比，我们的方法显著改善了结果（见表4）。我们不需要多个CNN模型来定位和分割对象边界。最后，我们使用Attention U-Net对CT-82数据集进行了5倍交叉验证，以进行更好的比较，胰腺标签的DSC值为81.48±6.23。

4、Discussion and Conclusion

在本文中，我们提出了一种新的用于医学图像分割的注意门模型。我们的方法消除了应用外部对象定位模型的必要性。该方法具有通用性和模块化的特点，可以很容易地应用于自然图像分析和机器翻译中的图像分类和回归问题。实验结果表明，所提出的AGs对组织/器官识别和定位非常有益。对于胰腺等大小不等的器官尤其如此，在全球分类任务中也会出现类似的情况。

AGs的train行为可以受益于转移学习和多阶段train计划。例如，预先训练的U-Net权重可以用来初始化注意网络，而gates可以在微调阶段进行相应的训练。类似地，机器学习领域也有大量文献探索不同的门控体系结构。例如，highway networks[7]利用门块周围的残余连接，允许更好的梯度反向传播和稍微柔和的注意机制。尽管我们对residual connections的实验没有提供任何显著的性能改进，但未来的研究将集中在这方面，以获得更好的训练行为。最后，我们注意到，随着GPU计算能力和内存的提高，更大容量的3D模型可以用更大的批量进行训练，而无需图像下采样。这样，我们就不需要使用特殊的后处理技术来进一步改善最先进的结果。类似地，注意力U-Net的性能可以通过使用精细分辨率的输入批次而进一步增强，而无需额外的heuristics。最后，我们要感谢Salim Arslan和Dan Busbridge对这项工作的有益评论。