UTNet论文笔记

UTNet论文笔记

UTNet: A Hybrid Transformer Architecture for Medical Image Segmentation

Abstract

Transformer架构已经在许多自然语言处理任务中取得了成功。然而，它在医学视觉中的应用仍然没有得到很大的探索。在这项研究中，我们提出了UTNet，一种简单但功能强大的混合Transformer结构，它将自我注意集成到卷积神经网络中，用于增强医学图像分割。UTNet在编码器和解码器中都应用了自我关注模块，以最小的开销捕获不同规模的远程依赖性。为此，我们提出了一种有效的自我注意机制和相对位置编码，将自我注意操作的复杂性从$O（n^2）$显著降低到接近$O（n）$。还提出了一种新的自关注解码器，用于从编码器中跳过的连接中恢复细粒度细节。我们的方法解决了Transformer需要大量数据才能了解视觉感应偏差这一难题。我们的混合层设计允许将Transformer初始化为卷积网络，而无需预先训练。我们在多标签、多供应商心脏磁共振成像队列中评估了UTNet。UTNet显示出优越的分割性能和对最先进方法的鲁棒性，有望在其他医学图像分割中得到很好的推广。代码是可用的https://github.com/yhygao/UTNet。

1 Introduction

卷积网络以其卓越的特征表示能力彻底改变了计算机视觉领域。目前，卷积编码器体系结构在位置敏感任务方面取得了实质性进展，如语义分割[14,11,20,17,6]。所使用的卷积运算通过从邻域像素收集局部信息来捕获纹理特征。为了全局聚合局部滤波器响应，这些模型叠加多个卷积层，并通过下采样扩展感受野。尽管取得了这些进步，但这种模式有两个固有的局限性。首先**，卷积只从邻域像素收集信息，缺乏明确捕获长距离（全局）依赖关系的能力[26,25,5]**。其次，卷积核的大小和形状通常是固定的，因此它们不能适应输入内容[15]。

使用自我注意机制的Transformer体系结构在自然语言处理（NLP）[18]中取得了成功，它能够捕获长期依赖性。自我注意是一种计算原语，通过上下文聚合机制实现成对实体交互，该机制能够捕获远程关联特征。它允许网络根据输入内容动态聚合相关功能。对简单形式的自我注意的初步研究表明，它在分割[4,16]、检测[24]和重建[9]方面很有用。

尽管基于图像的Transformer的应用前景广阔，但Transformer体系结构的培训和部署仍面临一些严峻的挑战。首先，自我注意机制在序列长度方面具有O（n2）时间和空间复杂性，导致大量的训练和推理开销。之前的作品试图降低自我关注的复杂性[10,28]，但仍远未达到完美。由于时间复杂性，标准的自我注意只能应用于面片，例如[3,27]使用16×16平坦的图像面片作为输入序列对图像进行编码，或者在CNN主干的特征图上进行编码，这些特征图已经下采样到低分辨率[4,22]。然而，对于位置敏感的任务，如医学图像分割，高分辨率特征起着至关重要的作用，因为大多数错误分割的区域都位于感兴趣区域的边界附近。其次，Transformer对图像没有感应偏置，在小规模数据集上也不能很好地工作[3]。例如，Transformer可以通过大型数据集（如full JFT300M[3]）进行预培训。但是，即使在ImageNet上进行了预培训，Transformer仍然比ResNet[12,7]更糟糕，更不用说可用医疗数据量少得多的医学图像数据集了。

在本文中，我们提出了一种U-shape hybrid Transformer Network:：UTNet，它将卷积强度和自我注意策略结合起来用于医学图像分割。其主要目标是应用卷积层提取局部强度特征，避免Transformer的大规模预训练，同时利用自我注意捕获远程关联信息。我们遵循UNet的标准设计，但将每个分辨率（最高分辨率除外）的构建块的最后一个卷积替换为提议的Transformer模块。为了提高分割的质量，我们试图运用自我关注在高分辨率特征地图上提取详细的长期关系。为此，我们提出了一种有效的自我注意机制，该机制在时间和空间上将整体复杂度从O（n2）显著降低到接近O（n）。此外，我们在自我注意模块中使用相对位置编码来学习医学图像中的内容-位置关系（bert中的class token）。我们的UTNet在多标签、多供应商心脏磁共振成像队列中展示了优越的分割性能和鲁棒性。鉴于UTNet的设计，我们的框架有望很好地推广到其他医学图像分割。

2 Method

2.1 Revisiting Self-attention Mechanism

该转换器建立在多头自我注意（MHSA）模块[18]的基础上，该模块允许模型从不同的表示子空间联合推断注意。将多个磁头的结果串联起来，然后用前馈网络进行变换。在这项研究中，我们使用了4个头，为了简单起见，在下面的公式和图中没有给出多个头的尺寸$X\in {\mathcal{R}}^{C\times H\times W}$，其中H，W是空间高度、宽度，C是通道数。三个1×1卷积用于投影X以进行查询、键和值嵌入：$\mathbf{Q},\mathbf{K},\mathbf{V}\in {\mathcal{R}}^{d\times H\times W}$，其中d是嵌入每个头部的尺寸。然后将Q，K，V展平并转换成大小为n×d的序列，其中n=HW。自我注意的输出是一个缩放的点积：

注意，$P\in {\mathcal{R}}^{n\times n}$被称为上下文聚合矩阵或相似矩阵。具体来说，第i个查询的上下文聚合矩阵是$P_i=\mathrm{softmax}\left(\frac{{\mathbf{q}}_i{\mathbf{K}}^{\top }}{\sqrt{d}}\right),P_i\in {\mathcal{R}}^{1\times n}$，计算qi和键中每个元素之间的标准化成对点生成。然后使用上下文聚合矩阵作为权重，从值中收集上下文信息。通过这种方式，自我关注在本质上具有全局感受域，并且擅长于捕捉长期依赖。此外，上下文聚合矩阵对输入内容是自适应的，以便更好地进行特征聚合。然而，n×d矩阵的点积导致$O（n^{2}d）$的复杂性。通常，当特征图的分辨率较大时，n远大于d，因此序列长度主导了自我注意计算，并使在高分辨率特征图中应用自我注意变得不可行，例如，对于16×16特征图，n=256，对于128×128特征图，n=16384。

2.2 Efficient Self-attention Mechanism

由于图像是高度结构化的数据，除了边界区域外，本地足迹内的高分辨率特征地图中的大多数像素都具有相似的特征。因此，所有像素之间的成对注意计算是非常低效和冗余的。从理论角度来看，对于长序列而言，自我注意本质上是低等级的[21]，这表明大多数信息集中在最大的奇异值中。受这一发现的启发，我们为我们的任务提出了一种有效的自我注意机制，如图2所示。

其主要思想是使用两个投影来投影关键和价值：$K,V\in {\mathcal{R}}^{n\times d}$ 低维嵌入的$\overline{\mathbf{K}},\overline{\mathbf{V}}\in {\mathcal{R}}^{k\times d}$，其中$k=hw\ll n$, n、h和w是子采样后特征图的缩小尺寸。现在建议的有效自我关注是：

通过这样做，计算复杂度降低到O（nkd）。值得注意的是，对低维嵌入的投影可以是任何下采样操作，例如平均/最大池或跨步卷积。在我们的实现中，我们使用1×1卷积和双线性插值对特征映射进行下采样，下采样率为8。

2.3 Relative Positional Encoding

标准的自我注意模块完全丢弃了位置信息，并且是扰动等变的[1]，这使得它无法对高度结构化的图像内容进行建模。以往文献[13]中的正弦嵌入在卷积层中不具有平移等变的性质。因此，我们通过添加相对高度和宽度信息来使用二维相对位置编码[1]。softmax之前使用像素$i=（i_x，i_y）$和像素之间的相对位置编码的成对注意logit：

其中，$q_i$是像素i的query向量，$k_i$是像素j的key向量，$j,r_{jx}^W-i_x$和$r_{j_y-i_y}^H$是可学习嵌入分别是相对宽度$j_x-i_x$的和相对高度$j_y-i_y$。与有效的自我注意类似，在低维投影后计算相对宽度和高度。包括相对位置嵌入在内的有效自我注意是：

${\mathbf{S}}_H^{rel},{\mathbf{S}}_W^{rel}\in {\mathcal{R}}^{HW\times hw}$沿高度和宽度尺寸的相对位置逻辑矩阵是否满足${\mathbf{S}}_H^{rel}[i,j]=q_i^{\top }{r}_{j_y-i_y}^H,{\mathbf{S}}_W^{rel}[i,j]=q_i^{\top }{r}_{j_x-i_x}^W$

2.4 Network Architecture

图1突出显示了UTNet的体系结构。我们试图结合卷积和自我注意机制的力量。因此，这种混合结构可以利用卷积图像的感应偏差来避免大规模的预训练，以及转换器捕捉远程关系的能力。由于误分割区域通常位于感兴趣区域的边界，因此高分辨率的上下文信息在分割中起着至关重要的作用。因此，我们将重点放在所提出的自我注意模块上，使其能够有效地处理大型特征地图。我们将Transformer模块应用于编码器和解码器的每一级，以从多个尺度收集长期依赖性，而不是天真地将自我注意模块集成到CNN主干的特征映射之上。请注意，我们没有在原始分辨率上应用Transformer，因为在网络的极浅层中添加Transformer模块无助于实验，但会引入额外的计算。一个可能的原因是，网络的浅层更关注细节纹理，而收集全局上下文可能无法提供信息。UTNet的构建块如图1（b）和（c）所示，包括剩余的基本块和Transformer块。对于这两个块，我们在the short cut中使用身份映射的pre-activation setting。这种身份映射在vision[8]和NLP任务[19]中被证明是有效的。

3 Experiment

3.1 Experiment Setup

我们在多标签、多供应商心脏磁共振成像（MRI）挑战队列[2]中系统地评估UTNet，包括左心室（LV）、右心室（RV）和左心室心肌（MYO）的分割。在培训集中，我们有来自两个不同MRI供应商的150张带注释的图像（每个供应商有75张图像），包括：西门子；B:飞利浦。在测试集中，我们有来自4家不同MRI供应商的200张图像（每个供应商50张图像），包括：西门子；B:飞利浦；C:通用电气；D：佳能，供应商C和D在培训集中完全不存在（我们丢弃未标记的数据）。来自不同供应商的MRI扫描在外观上有显著差异，这使我们能够测量模型的稳健性，并在不同设置下与其他模型进行比较。具体来说，我们进行了两个实验，以突出UTNet的性能和健壮性。首先，我们报告了来自同一供应商A的培训和测试数据的初步结果。其次，我们进一步衡量了模型的跨供应商稳健性。这种设置更具挑战性，因为培训和测试数据来自独立供应商。我们报告了每个模型的Dice分数和Hausdorff距离来比较性能。

3.2 Implementation Detail

对于数据预处理，我们将平面间距重新采样为1.2×1.2mm，同时保持沿z轴的间距不变。我们从无到有地训练了150个epochs的所有模特。我们使用指数学习率调度器，基本学习率为0.05。我们在一个GPU上使用批量大小为16的SGD优化器，动量和重量衰减设置为0.9和1e− 分别为4。在模型训练过程中动态应用数据增强，包括随机旋转、缩放、平移、加性噪声和伽马变换。在进入模型之前，所有图像都被随机裁剪为256×256。我们使用骰子损失和交叉熵损失的组合来训练所有网络。

3.3 Segmentation Results

我们将UTNet的性能与多种最先进的分割模型进行了比较。UNet[14]构建在完全卷积网络之上，采用U形架构来捕获上下文信息。ResUNet在架构上类似于UNet，但它使用剩余块作为构建块。CBAM[23]使用两个顺序卷积模块来推断通道和空间注意，以自适应地细化中间特征映射。双注意网络[4]使用两种自我注意来建模语言中的语义相互依赖分别是空间维度和通道维度。为了更好地进行比较，我们在ResUNet主干网中实现了CBAM和双重注意。由于双重注意的二次复杂性，它只适用于4次向下采样后的特征图。

如表1所示，UTNet在所有细分结果（LV、MYO和RV）中表现领先。通过引入剩余连接，ResUNet比最初的UNet稍有改进。CBAM的空间和通道注意是从卷积层推断出来的，它仍然受到有限的感受野的影响。因此，与ResUNet相比，CBAM的改进有限。我们还认识到，双注意方法几乎与ResUNet相同，因为它具有二次复杂性，无法处理更高分辨率的特征图来修复分割边界中的错误。同时，与双注意方法相比，我们的UTNet提供的参数更少，并且它可以从高分辨率特征图中捕获全局上下文信息。

Ablation study

图3（a）显示了不同自我注意位置的表现。x轴上的数字表示自我关注的水平（即经过几层注意力机制），例如，“34”表示进行3次和4次向下采样的水平。随着级别的提高，自我关注可以收集更多细粒度的细节信息，从而提高性能。但是，当添加到原始分辨率时，曲线会饱和。我们之所以这样做，是因为非常浅的层往往更关注局部纹理，而全局上下文信息不再具有信息性。图3（b）显示了有效自我注意的结果缩小尺寸为4、8、16。尺寸减小8倍可获得最佳性能。插值下采样略优于使用最大池。图3（c）显示了使用（a）和（b）中的最佳超参数的变压器编码器、解码器和相对位置编码的效果。变压器编码器和解码器的组合提供了最佳性能。相对位置编码也起着至关重要的作用，因为删除它会导致性能大幅下降。

为了在空间和时间复杂度上与标准自我注意进行头对头比较，我们进一步在四种分辨率（1、2、3、4，与UTNet相同）下应用双重注意，并使用相同的输入图像大小和批量大小（256×256×16）来测试推断时间和内存消耗。UTNet与具有二次复杂度的双重注意相比具有优越的优势，其中GPU内存：3.8 GB vs 36.9 GB，时间：0.146 s vs 0.243 s。

Robustness analysis

表2显示了使用供应商A和B数据的培训模型的结果，然后分别在供应商A、B、C和D上测试模型。在查看C和D供应商的结果时，竞争方法会受到供应商差异的影响，而UTNet则保持竞争性能。这一观察结果可能归因于多层次特征图上的自我注意设计和内容位置注意，使得UTNet能够更好地关注全局上下文信息，而不仅仅是局部纹理。图4进一步显示，UTNet显示了最一致的边界结果，而其他三种方法无法捕捉边界的细微特征，尤其是心脏MRI中的RV和MYO区域。

4 Conclusion

我们提出了一种U-shape hybrid Transformer network （UTNet），以融合卷积层和自我注意机制的优点，用于医学图像分割。我们的混合层设计允许将变压器初始化为卷积网络，而无需预先训练。这种新颖的自我关注使我们能够在编码器和解码器中扩展网络不同级别的操作，以便更好地捕获远程依赖关系。我们相信，这种设计将有助于丰富的参数化变压器模型在医学视觉应用中变得更容易访问。此外，高效处理长序列的能力为UTNet在更下游的医学图像任务中的使用开辟了新的可能性。