【论文翻译-3】Attention U-Net: Learning Where to Look for the Pancreas
Attention U-Net: Learning Where to Look for the Pancreas
- 阅读日期：2020年11月25日
Abstract
- 我们提出新型注意力门控（attention gate, AG）模型用于医学成像，它可以自动学习并关注不同形状和大小的目标结构。使用AGs训练的模型能隐式学习抑制输入图像中不相关的区域，同时强调对特定任务有用的显著特征。这使得我们无需使用级联级联卷积神经网络来定位外部组织/器官。AGs能容易集成到标准CNN架构，例如U-Net，并且计算开销小，同时提高了模型的灵敏度和预测精度。在两个大型腹部CT数据集中用多类图像分割任务评估Attention U-Net。实验结果表明AGs在保持计算效率的同时，在不同数据集和训练规模的情况下均提高了U-Net的性能。代码开源：https://github.com/ozan-oktay/Attention-Gated-Networks
1 Introduction
-
由于对大量医学图像进行手工密集标注是一项繁琐且容易出错的任务，自动医学图像分割在图像分析领域得到了广泛的研究。精确可靠的解决方案是希望提高临床工作流程效率和通过快速和自动提取定量测量方法来支持决策。
-
随着卷积神经网络的出现，自动化医学图像分析任务（包括心脏MR分割、癌性肺节点检测）的性能可以达到接近放射科医师的水平。表达能力高，快速推理和滤波器共享特征使得CNNs成为图像分割的标准。全卷积网络FCNs和U-Net是两种常用的架构。尽管它们的表示能力很高，但是当患者之间目标器官的形状和大小差异较大时，这些架构依赖于多级级联CNNs。级联框架提取感兴趣区域（ROI），并对特定的ROI进行密集预测。应用领域包括心脏MRI[14]、心脏CT[23]、腹部CT[26,27]分段、肺CT结节[17]的检测。但是，这种方法会导致计算资源和模型参数的过度和冗余使用；例如，类似的低级特征会被级联中的所有模型重复提取。为了解决这一普遍问题，我们提出了一个简单而有效的解决方案，即“注意力门”(AGs)。带有AGs的CNN模型可以像训练FCN模型一样以标准的方式从头开始训练，并且AGs可以自动学习集中在目标结构上而不需要额外的监督。在测试时，这些门会产生动态的软区域建议，并突出对特定任务有用的显著特征。此外，它们不会引入显著的计算开销，也不需要像多模型框架那样需要大量的模型参数。反过来，所提出的AGs通过抑制不相关区域的特征激活，提高了稠密标签预测模型的敏感性和准确性。这样，在保持较高预测精度的同时，消除了使用外部器官定位模型的必要性。类似的注意机制也被用于自然图像分类和视频字幕，自适应特征池化，其中模型在选定图像区域的一个子集进行。在本文中，我们概括了这种设计，并提出了基于图像网格的门控，使注意系数能够特定于局部区域。此外，我们的方法可以用于基于注意力的密集预测。
-
我们在标准U-Net架构(Attention U-Net)中演示了AG的实现，并将其应用于医学图像。我们选择具有挑战性的CT胰腺分割问题来作为评估。该问题困难之处在于组织对比度低，器官形状和大小的变异性大。我们通过两个常用的基准来评估我们的执行情况：TCIA Pancreas CT-82和multi-class abdominal CT-150。结果表明，在不需要多个CNN模型的情况下，AGs一致性地提高了不同数据集和训练规模的预测精度。
1.1 Related Work
-
CT Pancreas Segmentation：早期对腹部CT中胰腺分割的研究使用了统计形状模型或多图谱技术，特别是，atlas方法受益于人工注释传播控制的隐式形状约束。然而，在公共基准测试中，比如TCIA数据集，基于atla的框架的骰子相似系数(DSC)在69.6%到73.9%之间。在[39]中，提出了一种基于分类的图像配准框架来消除图像集对配准的依赖。最近，人们提出了级联多阶段CNN模型[26,27,38]来解决这个问题。这里，使用一个初始的粗糙模型(例如U-Net或回归森林)来获得一个ROI，然后使用一个被裁剪的ROI经过第二个模型分割细化。类似地，**在【4】中使用2D-FCN和递归神经网络(RNN)模型的组合来利用相邻轴向切片之间的依赖关系。这些方法在TCIA基准(81.2% - 82.4% DSC)中实现了最先进的性能。如果不使用级联框架，性能会下降2.0%到4.4%。【37】最近的工作提出了一个迭代的两阶段模型，递归地更新局部和全局的预测，两个模型都是端到端训练的。**除了标准FCNs外，在CT胰腺分割问题中还应用了稠密连接和稀疏卷积。密集连接和稀疏核通过减少非零参数的数量来降低计算复杂度
-
Attention Gates：AGs通常用于自然图像分析、知识图谱和自然语言处理(NLP)（用于图像字幕、机器翻译和分类任务）。最初的工作是通过解释相对于输入图像的输出类分数的梯度来探索注意力图。另一方面，可训练的注意力是设计出来的，分为硬注意力和软注意力。Hard attention【21】，迭代建议的区域和裁剪，通常是不可微的，需要依靠强化学习来更新参数，这使得模型训练更加困难。[36]中使用递归集中注意力来检测胸部x射线扫描的异常。相反，Soft attention是概率性的，它使用标准的反向传播而不需要蒙特卡罗抽样。例如，加法软注意被用于句子到句子的翻译，最近又被应用于图像分类。在[10]中，channel-wise attention被用来突出重要的特征维度，这是ILSVRC 2017年图像分类挑战中表现最好的。Self-attention技术【11,33】提出用来消除对外部门控信息的依赖。例如，[33]中使用了非局部自我注意来捕获长期依赖关系。在[11,32]中，利用self-attention对特定的类进行池化，提高了图像分类精度和鲁棒性。
1.2 Contributions
- 在本文中，我们提出了一种新的自注意门控模块，可用于基于CNN的标准图像分析模型的密集标签预测。此外，我们探讨了AGs对医学图像分析的好处，特别是在图像分割的背景下。这项工作的贡献可以总结如下:
- 我们进一步采用了【11】（self-attention）中提出的注意方法，提出了基于网格的门控，使注意系数更聚焦于局部区域。与基于全局特征向量的门控相比，这提高了性能。此外，我们的方法可以用于密集预测，因为我们不执行自适应池。
- 我们提出的soft-attention技术是第一个用于医学成像任务的前馈CNN模型。提出的注意力门可以替代图像分类[36]和图像分割框架中的外部器官定位模型中使用的hard-attention方法。
- 提出了一种对标准U-Net模型的扩展，在不需要复杂启发式的情况下提高模型对前景像素的灵敏度。通过实验观察，对比U-Net，在不同的成像数据集中精度都提高了。
2 Methodology
- Fully Convolutional Network（FCN）：在公共基准数据集上，卷积神经网络(CNNs)在医学图像分析方面优于传统方法，同时也比图切和多图谱分割技术快一个数量级。这主要是因为(I)使用随机梯度下降(SGD)优化来学习区域特定的图像特征，(II)学习到的滤波器参数在所有像素之间共享，(III)图像卷积操作很好地利用了医学图像中的结构信息。特别是，全卷积网络(FCN)[18]，如U-Net[24]、DeepMedic[13]和整体嵌套网络[16,35]已被证明在各种任务中，包括心脏MR[3]、脑瘤[12]和腹部CT[26, 27]的图像分割任务中，能够实现稳健和准确的性能。
- 卷积层通过逐层处理局部信息来逐渐提取高维图像表示 （ x l ） （x^{l}） （xl）。最终，这将根据它们的语义在高维空间中分离像素。通过这个连续的过程，模型的预测是基于从一个大的感受野收集的信息。因此，特征图 ( x l ) (x^{l}) (xl)的获得方式如下，第 l l l层的输出依次使用一个线性变换，后接一个非线性激活函数（通常选择ReLU： σ 1 ( x i , c l ) = m a x ( 0 , x i , c l ) \sigma_{1}(x_{i,c}^{l}) = max(0, x_{i,c}^{l}) σ1(xi,cl)=max(0,xi,cl)），其中 i i i和 c c c分别表示空间和通道维度。特征激活公式为： x c l = σ 1 ( ∑ c ′ ∈ F l x c ′ l − 1 ∗ k c ′ , c ) x_{c}^{l} = \sigma_{1}(\sum _{c^{'} \in F_{l}} x_{c^{'}}^{l-1} * k_{c^{'}, c}) xcl=σ1(∑c′∈Flxc′l−1∗kc′,c)，其中 ∗ * ∗表示卷积操作，这里空间下标 ( i ) (i) (i)省略了。函数 f ( x l ; Φ l ) = x ( l + 1 ) f(x^{l}; \Phi^{l}) = x^{(l+1)} f(xl;Φl)=x(l+1)，其中 Φ l 是 网 络 学 习 的 核 参 数 \Phi^{l}是网络学习的核参数 Φl是网络学习的核参数。参数是通过最小化训练目标，例如交叉熵损失，使用随机梯度下降（SGD）。本文在标准的U-Net架构上建立了注意力模型。U-Net网络性能良好，能高效利用GPU内存，因此被广泛用于图像分割任务。优势主要与多尺度图像特征的提取有关。粗特征图捕捉上下文信息，突出显示前景对象的类别和位置。在多个尺度上提取的特征图稍后将通过跳跃连接合并，结合粗级和细级的稠密预测，如图1所示。
- Attention Gates for Image Analysis：为了获得足够大的感受野，从而获得语义上下文信息，在标准的CNN架构种，特征图会逐渐下采样。这样，在全局尺度上，粗空间网络层次的特征可以定位和理清组织之间的关系。然而，对于具有较大形状差异的小物体，仍然很难减少假阳性的预测。为了提升精度，当前的分割框架依赖额外的目标物体先定位模型，为的是将任务简化为定位和分割。这里，我们证明将AGs加入到标准CNN模型中能够达到相同的结果。这并不需要对多个模型进行训练，也不需要大量额外的模型参数。与多阶段CNNs中的定位模型相比，AGs逐渐抑制不相关背景区域的特征响应，而不需要在网络之间裁剪ROI。
-
**注意力系数 α ∈ [ 0 , 1 ] \alpha \in [0, 1] α∈[0,1]**识别显著的图像区域，修剪特征响应，仅仅保留与特定任务相关的响应，如图3a所示。AGs的输出计算方法是输入特征图与注意力系数的元素乘法： x ^ i , c l = x i , c l ⋅ α i l \hat{x}_{i, c}^{l} = x_{i, c}^{l} \cdot \alpha_{i}^{l} x^i,cl=xi,cl⋅αil。默认地，单标量的注意值是由每个像素向量 x i l ∈ R F l x_{i}^{l} \in \mathbb{R}^{F_{l}} xil∈RFl计算得到的，其中 F l F_{l} Fl表示在第 l l l层的特征图数量。在存在多个语义类别的情况下，我们提出学习多维注意力系数。这是受到了[29]的启发，多维注意力系数用于学习句子嵌入。因此，每个AG学会关注目标结构的一个子集。如图2所示，每个像素 i i i使用一个门控向量 g i ∈ R F g g_{i} \in \mathbb{R}^{F_{g}} gi∈RFg来确定聚焦的区域。门控向量包含上下文信息，如[32]（使用AGs用于自然图像分类）说明的那样修剪低水平的特征响应。我们使用加性注意力【2】来获得门控系数。虽然这样做计算量更大了，但是与乘法注意力【19】相比，前者能获得更高的精度。加性注意力的公式如下：
q a t t l = Ψ T ( σ 1 ( W x T x i l + W g T g i + b g ) ) + b Ψ (1) q_{att}^{l} = \Psi^{T}(\sigma_{1}(W_{x}^{T}x_{i}^{l} + W_{g}^{T}g_{i} + b_{g})) + b_{\Psi} \tag{1} qattl=ΨT(σ1(WxTxil+WgTgi+bg))+bΨ(1)
α i l = σ 2 ( q a t t l ( x i l , g i ; Θ a t t ) ) (2) \alpha_{i}^{l} = \sigma_{2}(q_{att}^{l}(x_{i}^{l}, g_{i}; \Theta_{att})) \tag{2} αil=σ2(qattl(xil,gi;Θatt))(2)
其中， σ 2 ( x i , c ) = 1 1 + e x p ( − x i , c ) \sigma_{2}(x_{i, c}) = \frac{1}{1+exp(-x_{i, c})} σ2(xi,c)=1+exp(−xi,c)1是sigmoid激活函数。AG的一组参数 Θ a t t \Theta_{att} Θatt包括：线性变换 W x ∈ R F t × F i n t W_{x} \in \mathbb{R}^{F_{t} \times F_{int}} Wx∈RFt×Fint， W g ∈ R F g × F i n t W_{g} \in \mathbb{R}^{F_{g} \times F_{int}} Wg∈RFg×Fint， Ψ ∈ R F i n t × 1 \Psi \in \mathbb{R}^{F_{int} \times 1} Ψ∈RFint×1和偏置项 b Ψ ∈ R b_{\Psi} \in \mathbb{R} bΨ∈R， b g ∈ R F i n t b_{g} \in \mathbb{R}^{F_{int}} bg∈RFint。线性变换是通过对输入张量在通道方向进行 1 × 1 × 1 1\times 1\times 1 1×1×1卷积来计算的。在其他方法【33】中，这被称为基于向量连接的注意力（vector concatenation-based attention），其中连接的特征 x l x^{l} xl和 g g g被线性映射到 R F i n t \mathbb{R}^{F_{int}} RFint维度空间中。在图像字幕和分类任务中，softmax函数用于将注意力系数 ( σ 2 ) (\sigma_{2}) (σ2)归一化；但是，连续使用softmax会导致输出的激活值很稀疏。因此，我们选择sigmoid激活函数。这使得训练AG参数时能更好收敛。与【11】相比，我们提出了网格注意力技术。在这种情况下，门控信号不是所有图像像素的全局单一向量，而是适应于图像空间信息的网格信号。更重要的是，每个跳跃连接的门控信号是从多个图像尺度聚合的信息，如图1所示，这增加了查询信号的网格分辨率，并实现了更好的性能。最后，我们想指出的是，AG参数可以通过标准的反向传播更新进行训练，而不需要使用在hard-attention[21]中使用的基于采样的更新方法。 -
Attention Gates in U-Net Model：提出的AGs被合并到标准U-Net架构中，以突出通过skip连接的显著特征，见图1。将从粗尺度提取出的信息应用到门控中，可以消除跳跃连接产生的不相关和嘈杂的响应。这（由粗尺度提取的信息应用到门控中）是在连接操作之前执行的，目的是合并相关的激活。此外，AGs在前向传播和反向传播时对神经元的激活进行过滤。从背景区域产生的梯度在反向传播时减少权重。这使得在浅层的模型参数更新是基于给定任务相关的空间区域。在 l − 1 l-1 l−1层卷积参数的更新方式如下：
∂ ( x ^ i l ) ∂ ( Φ l − 1 ) = ∂ ( α i l f ( x i l − 1 ; Φ l − 1 ) ) ∂ ( Φ l − 1 ) = α i l ∂ ( f ( x i l − 1 ; Φ l − 1 ) ) ∂ ( Φ l − 1 ) + ∂ ( α i l ) ∂ ( Φ l − 1 ) x i l (3) \frac{\partial(\hat{x}_{i}^{l})}{\partial(\Phi^{l-1})} = \frac{\partial(\alpha_{i}^{l}f(x_{i}^{l-1}; \Phi^{l-1}))}{\partial(\Phi^{l-1})} = \alpha_{i}^{l} \frac{\partial(f(x_{i}^{l-1}; \Phi^{l-1}))}{\partial(\Phi^{l-1})} + \frac{\partial(\alpha_{i}^{l})}{\partial(\Phi^{l-1})}x_{i}^{l} \tag{3} ∂(Φl−1)∂(x^il)=∂(Φl−1)∂(αilf(xil−1;Φl−1))=αil∂(Φl−1)∂(f(xil−1;Φl−1))+∂(Φl−1)∂(αil)xil(3) -
右边的第一个梯度项乘以 α i l \alpha_{i}^{l} αil，对于多维度的AGs， α i l \alpha_{i}^{l} αil对应每个网格尺度上的一个向量。在每个子AG中，提取并融合补充的信息得到跳跃连接的输出。为了减少AGs训练参数数量和计算复杂度，进行线性变换（ 1 × 1 × 1 1 \times 1 \times 1 1×1×1卷积），将输入的特征图向下采样到门控信号的分辨率，类似于非局部块【33】。相应的线性变换将特征映射解耦，并将其映射到更低维的空间用于门控操作。如[11]中所建议的，低级特征映射(即第一个跳跃连接)在门控函数中不使用，因为它们不表示高维空间中的输入数据。我们使用深度监督[16]强制中间特征图在每个图像尺度上具有语义区分。这有助于确保不同尺度的注意单元能够影响对大范围图像前景内容的响应。这样能够防止从跳跃连接的小子集重建密集预测
3 Experiments and Results
-
提出的AG模型是模块化的，与应用类型无关;因此，它可以很容易地适应分类和回归任务。为了证明它在图像分割中的适用性，我们评估了注意力U-Net模型在一个具有挑战性的腹部CT多标签分割问题。特别是胰腺边界的划定是一项困难的任务，由于形状变异和组织造影剂差。我们的模型与标准的3D U-Net在分割性能、模型容量、计算时间和内存需求方面进行了比较。
-
Evaluation Datasets：实验中使用了两种不同的CT腹部数据集：(I) 150例胃癌患者腹部3D CT扫描（CT-150），在所有的图像中，胰腺、肝脏和脾脏的边界由三名训练有素的研究人员半自动划定，并由临床医生手动验证。[27]中使用了相同的数据集来作为U-Net模型在胰腺分割中的基准。(II)第二个dataset (CT-82)包括82张增强3D CT扫描，逐片进行胰腺手工标注。此数据集(NIH-TCIA)【25】是公开可用的，通常用于CT胰腺分割框架的基准。由于大图像大小和硬件内存的限制，将两个数据集的图像向下采样到各向同性2.00毫米分辨率。
-
Implementation Details：与最先进的CNN分割框架[4,26]相比，我们提出了一个3d模型来捕捉足够的语义上下文。梯度更新计算使用小批量大小为2到4个样本。对于较大的网络，梯度平均在多个前向和后向通道上使用。所有的模型都使用Adam optimiser[15]、批量归一化、深度监控【16】和标准的数据增强技术(仿射变换、轴向翻转、随机裁剪)进行训练。强度值线性缩放得到正态分布N(0,1)。使用定义在所有语义类上的Sorensen-Dice loss【20】对模型进行训练，实验表明该方法对类别失衡不那么敏感。初始化门控参数，使注意力门控经过所有空间位置的特征向量。此外，我们不需要基于hard-attention方法中所使用的多个训练阶段，因此简化了训练过程。代码基于PyTorch并且开源。
-
Attention Map Analysis：从测试图像中得到的注意系数与训练时间有关(见图4)。观察到AGs最初具有均匀分布，并在所有位置通过特征。逐渐更新和定位目标器官的边界。此外，在较粗的尺度上，AGs提供器官的粗略轮廓，这些轮廓在较细的分辨率下逐渐细化。此外，通过在每个图像尺度上训练多个AGs，我们观察到每个AG都学会了关注特定的器官子集。
- Segmentation Experiments：本文提出的Attention U-Net模型是以腹部CT多类分割的标准U-Net为基准的。我们使用CT-150数据集进行训练(120)和测试(30)。对应的骰子分数(DSC)和表面距离(S2S)如表1所示。胰腺预测的结果表明，注意力门(AGs)通过提高模型的表达能力来提高召回值(p = .005)，因为它依赖AGs来定位前景像素。图3b定性地比较了这两种模型的预测结果之间的差异。在第二个实验中，同样的模型使用较少的训练图像进行训练(30)，表明在不同大小的训练数据下，性能的提高是一致的和显著的(p = .01)。对于这两种方法，我们观察到脾脏DSC的性能下降，训练规模减少。所提出的框架得到的下降幅度较小。对于肾脏分割，由于组织对比度较高，模型达到了相似的准确性。
- 在表1中，我们还报告了两个模型的可训练参数的数量。我们观察到，通过向标准U-Net增加8%的额外容量，就DSC而言，性能可以提高2-3%。为了进行公平的比较，我们还训练了大容量的U-Net模型，并将其与网络规模较小的所提出的模型进行比较。表2所示的结果表明，AGs的添加不仅仅是简单地(均匀地)增加网络各层的模型容量(p = .007)。因此，当AGs用于减少训练多个、单模型的冗余度时，应该使用额外的容量来对组织进行定位。
- Comparision to State-of-the-Art：提出的架构在公共TCIA CT胰腺基准上进行评估，以比较其性能与最先进的方法。首先，我们将在CT-150数据集上训练的模型直接应用到CT-82数据集上，观察两种模型在不同数据集上的适用性。相应结果(BFT)见表3。U-Net模型优于传统的atlas技术[34]，尽管它是在一个不相交的数据集上训练的。此外，在不同数据集上，attention模型在胰腺分割方面的表现更加一致。这些模型随后在TCIA数据集(61次训练，21次测试)的子集上进行微调(AFT)。在输出softmax计算中，排除脾脏和肾脏对应的输出节点，仅对背景和胰腺标签进行梯度更新。表3和表4的结果显示，与连接的多模型CNN方法相比[4,26,38]，由于增加了训练数据和更丰富的语义信息(如脾脏标签)，性能有所提高。此外，我们使用从CT-82数据集中随机选取的61张训练图像从零开始训练这两个模型(SCR)。与CT-150数据集的结果相似，AGs提高了分割精度，降低了表面距离(p = .03)，提高了胰腺像素的召回率(p = .09)。
- 表4总结了最先进的CT胰腺分割模型的结果，以便进行比较。由于这些模型是在同一个训练数据集上训练的，因此这种比较可以让我们了解到注意力模型与相关文献的比较。值得注意的是，在我们的框架中并没有使用后处理(例如条件随机场)，因为实验主要关注在一个孤立环境下AGs带来的性能改进的量化。同样的，残差和密集连接可以像在[6]中一样，结合AGs来改善分割结果。在这方面，我们的3D Attention U-Net模型的表现与最先进的技术差不多，尽管输入的图像被降低采样到较低的分辨率。更重要的是，与基于单模型的分割框架相比，我们的方法显著提高了结果(见表4)。我们不需要多个CNN模型来定位和分割对象边界。最后，我们使用Attention U-Net对CT-82数据集进行5倍交叉验证以进行更好的比较，结果胰腺标签的DSC为81.48±6.23。
4 Discussion and Conclusion
-
提出了一种新的用于医学图像分割的注意力门控模型。我们的方法消除了使用额外目标定位模型的必要性。所提出的方法是通用的和模块化的，因此它可以很容易地应用于图像分类和回归问题，如在自然图像分析和机器翻译的例子。实验结果表明，所提出的AGs对组织/器官的识别和定位非常有利。对于可变的小尺寸器官，如胰腺，这一点尤其正确，而对于全局的分类任务，预期也会有类似的行为。
-
迁移学习和多阶段训练方案均可改善AGs的训练效果。例如，预先训练好的U-Net权值可以用来初始化注意力网络，在微调阶段可以对gates进行相应的训练。同样，在机器学习方面也有大量的文献在探索不同的门控架构。例如，高速公路网络[7]利用gate block周围的残差连接来实现更好的梯度反向传播和更柔和的注意机制。虽然我们的残差连接实验没有提供任何显著的性能改善，未来的研究将集中在这方面，以获得更好的训练行为。最后，我们注意到随着GPU计算能力和内存的提高，更大容量的3D模型可以在不需要图像降采样的情况下以更大的批处理规模进行训练。这样，我们就不需要使用特别的后处理技术来进一步改善最先进的结果。类似地，注意力U-Net的性能可以通过使用高分辨率输入批次而不需要额外的启发式来进一步增强。