Attention U-Net: Learning Where to Look for the Pancreas

Abstract

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

1 Introduction

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

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

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

1.1 RelatedWork

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

1.2 Contributions

在本文中，我们提出了一种新颖的自我关注门控模块，可用于基于CNN的标准图像分析模型，用于密集标签预测。 此外，我们探索AGs对医学图像分析的益处，特别是在图像分割的背景下。 这项工作的贡献可归纳如下：

•  我们进一步采用[11]中提出的注意力方法，提出基于网格的门控，使注意系数更加特定于局部区域。 与基于全局特征向量的门控相比，这提高了性能。 此外，我们的方法可用于密集预测，因为我们不执行自适应池池化。
• 我们提出了应用于医学成像任务的前馈CNN模型中的软注意技术的第一个用例之一。 提出的注意力门可以取代图像分类[36]中使用的硬注意方法和图像分割框架中的外部器官定位模型。
• 提出了标准U-Net模型的扩展，以提高模型对前景像素的灵敏度，而无需复杂的启发式算法。 通过实验观察到U-Net的准确度改进在不同的图像数据集中是一致的。

2 Methodology

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

卷积层通过逐层处理局部信息逐步提取更高维度的图像表示（x1）。 最终，这根据它们的语义分离高维空间中的像素。 通过这个顺序过程，模型预测以从大的感受野收集的信息为条件。 因此，通过顺序地应用线性变换，然后是非线性激活函数，在层1的输出处获得特征映射x1。 通常选择整流线性单元：其中i和c分别表示空间和通道尺寸。 特征激活可以表示为：，其中*表示卷积运算，并且为了符号清晰度，在公式中省略了空间下标（i）。 在卷积层1中应用的函数由可训练的核参数l表征。通过最小化训练目标来学习参数，例如， 交叉熵损失，使用随机梯度下降（SGD）。 在本文中，我们在标准U-Net架构之上构建我们的注意力模型。 U-Nets通常用于图像分割任务，因为它们具有良好的性能和GPU内存的有效使用。 后一个优点主要与在多个图像尺度上提取图像特征有关。 粗糙特征图捕获上下文信息并突出显示前景对象的类别和位置。 多尺度提取的特征图之后通过跳跃连合并，以组合粗略和精细级别的预测，如图1所示。

用于图像分析的注意门：为了捕获足够大的感知野并因此捕获语义上下文信息，特征图网格在标准CNN架构中逐渐下采样。通过这种方式，粗糙空间网格水平模型的位置和全局尺度内组织之间的关系。然而，仍然难以减少对于显示大的形状可变性的小物体的假阳性预测。为了提高准确性，当前的分割框架[14,26,27]依赖于附加的目标定位模型来将任务简化为单独的定位和后续分割步骤。在这里，我们证明通过在标准CNN模型中集成注意力门（AG）可以实现相同的目标。这不需要训练多个模型和大量额外的模型参数。与多级CNN中的定位模型相反，AG在不相关的背景区域中逐渐抑制特征响应，而不需要在网络之间裁剪ROI。

注意系数，，识别显着图像区域和修剪特征响应，以仅保留与特定任务相关的激活，如图3a所示。 AG的输出是输入特征图和注意力系数的元素乘法：。在默认设置中，针对每个像素矢量计算单个标量注意值，其中对应于层l中的特征图的数量。在多个语义类的情况下，我们建议学习多维注意系数。这受[29]的启发，其中多维注意系数用于学习句子嵌入。因此，每个AG都学会专注于目标结构的子集。如图2所示，门控矢量用于每个像素i以确定聚焦区域。门控向量包含上下文信息，以修剪[32]中建议的较低级别的特征响应，其使用AG进行自然图像分类。我们使用加性注意[2]来获得门控系数。虽然这在计算上更昂贵，但实验证明它可以获得比乘法注意更高的准确度[19]。加性注意力公式表示如下：

其中对应于sigmod激活函数。AG的特征在于一组参数包含：线性变换，和偏置项。 对于输入张量，使用通道1x1x1卷积计算线性变换。 在其他上下文[33]中，这被称为基于向量级联的注意力，其中级联特征和被线性地映射到维度中间空间。在图像字幕[1]和分类[11]任务中，softmax激活函数用于归一化注意系数（2）; 然而，顺序使用softmax会在输出端产生较稀疏的激活。 因此，我们选择了sigmoid激活函数。 这实验性地导致AG参数的更好的训练收敛。 与[11]相反，我们提出了一种网格注意力技术。 在这种情况下，门控信号不是用于所有图像像素的全局单个矢量，而是用于图像空间信息的网格信号。 更重要的是，每个跳跃连接的门控信号聚合来自多个成像尺度的信息，如图1所示，这增加了查询信号的网格分辨率并实现了更好的性能。 最后，我们要指出，AG参数可以使用标准的反向传播更新进行训练，而无需在硬注意力中使用基于采样的更新方法[21]。

U-Net模型中的注意门：提出的AG被合并到标准U-Net架构中，以突出通过跳过连接传递的显着特征，参见图1.从粗略提取的信息用于门控以消除不相关和嘈杂的歧义在跳过连接中的响应。 这在concatenation操作之前执行以仅合并相关激活。 另外，AG在正向传播期间以及反向传播期间过滤神经元激活。 源自背景区域的梯度在反向传播期间向下加权。 这允许较浅层中的模型参数主要基于与给定任务相关的空间区域来更新。 层中卷积参数的更新规则可以表述如下：

右侧的第一个梯度项用缩放。在多维AG的情况下，对应于每个网格尺度的矢量。在每个子AG中，提取并融合补充信息以定义跳过连接的输出。为了减少可训练参数的数量和AG的计算复杂度，在没有任何空间支持（1x1x1卷积）的情况下执行线性变换，并且输入特征映射被下采样到门控信号的分辨率，类似于non-local blocks[33]。相应的线性变换将特征图解耦并将它们映射到用于门控操作的较低维空间。如[11]中所建议的，在门控功能中不使用低级特征图，即第一个跳跃连接，因为它们不表示高维空间中的输入数据。我们使用深度监督[16]来强制中间特征图在每个图像尺度上在语义上有区别。这有助于确保不同尺度的注意力单元能够影响对大范围图像前景内容的响应。因此，我们防止从跳跃连接的小子集重建密集预测。

3 Experiments and Results

提出的AG模型是模块化的，与应用类型无关; 因此，它可以很容易地适应分类和回归任务。 为了证明其对图像分割的适用性，我们在具有挑战性的腹部CT多标签分割问题上评估Attention U-Net模型。 特别是，由于形状变化和组织对比度差，胰腺边界描绘是一项艰巨的任务。 我们的模型在分割性能，模型容量，计算时间和内存要求方面与标准3D U-Net进行了比较。

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

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

注意力图分析：从测试图像获得的注意力系数相对于训练时期可视化（参见图4）。 我们通常观察到AG最初具有均匀分布并且在所有位置都通过特征。 这是逐步更新和定位到目标器官边界。 此外，在较粗糙的尺度上，AG提供了粗略的器官轮廓，这些器官在更精细的分辨率下逐渐细化。 此外，通过在每个图像尺度上训练多个AG，我们观察到每个AG学习专注于器官的特定子集。

分割实验：所提出的Attention U-Net模型是针对多级腹部CT分割的标准U-Net进行基准测试的。我们使用CT-150数据集进行训练（120）和测试（30）。相应的dice score（DSC）和表面距离（S2S）在表1中给出。胰腺预测的结果表明，注意力门（AG）通过改善模型的表达能力来增加召回率，因为它依赖于AG用于定位前景像素。使用这两个模型获得的预测之间的差异在图3b中进行了定性比较。在第二个实验中，使用较少的训练图像训练相同的模型（30），以表明对于不同大小的训练数据，性能改进是一致的和显着的。对于这两种方法，我们观察到随着训练规模减小，脾脏DSC的性能下降。提出的框架下降幅度不大。对于肾脏分割，由于组织对比度较高，因此模型具有相似的准确度。

在表1中，我们还报告了两种模型的可训练参数的数量。 我们观察到，通过在标准U-Net中增加8％的额外容量，就DSC而言，性能可提高2 3％。 为了公平比较，我们还培训更高容量的U-Net模型，并与具有更小网络规模的建议模型进行比较。 表2中显示的结果表明AG的添加不仅仅是简单地增加（均匀地）跨网络的所有层的模型容量。 因此，在使用AG来减少训练多个单独模型的冗余的情况下，AG应该使用额外的容量来定位组织。