UTNet：A Hybrid Transformer Architecture for Medical Image Segmentation

UTNet：A Hybrid Transformer Architecture for Medical Image Segmentation

摘要

1. 提出UTNet，这是一种简单而强大的混合Transformer架构，它将自我注意力集成到卷积神经网络中，用于增强医学图像分割。
2. UTNet在编码器和解码器中都应用了自我注意模块，以最小的开销捕获不同规模的远程依赖。
3. 提出了一种有效的自注意机制和相对位置编码，使自注意操作的复杂度从O($n^2$)显著降低到近似O(n)
4. 提出了一种新的自注意解码器，从编码器中跳过的连接中恢复细粒度的细节
5. 混合层设计允许将transformer初始化为卷积网络，而不需要预先训练。我们已经在多标签、多供应商的心脏磁共振成像队列中对UTNet进行了评估。UTNet表现出卓越的分割性能和相对于最先进的方法的健壮性。

Introduction

背景：

第一段：卷积神经网络的编码器解码器架构在位置敏感任务（语义分割）上取得实质性的进展，其中的卷积运算通过从邻域像素收集局部信息来捕获纹理特征。为了全局聚合局部滤波器响应，这些模型叠加多个卷积层，并通过下采样扩展感受野。但是仍然存在其固有的局限性：卷积只从邻域像素收集信息，缺乏明确捕获长距离（全局）依赖关系的能力；卷积核的大小和形状通常是固定的，因此它们不能适应输入内容。

第二段：使用自注意力机制的Transformer的结构由于具有捕获远程依赖性的能力在NLP上取得了成功。自注意力是通过上下文聚合机制实现成对实体交互，该机制能够捕获远程关联特征。

提出问题：

第三段：Transformer架构仍存在一些问题：

1. 自我注意机制在序列长度方面具有O（$n^2$）时间和空间复杂性，导致大量的训练和推理开销。对于位置敏感的任务，如医学图像分割，高分辨率特征起着至关重要的作用，因为大多数错误分割的区域都位于感兴趣区域的边界附近
2. Transformer对图像没有感应偏置，在医学图像这种小规模数据集上也不能很好地工作。

解决方法：

第四段：因此本文作者提出UTNet网络架构，它将卷积强度和自我注意策略结合起来用于医学图像分割。其主要目标是应用卷积层提取局部强度特征，避免Transformer的大规模预训练，同时利用自我注意捕获远程关联信息。

创新点：

遵循UNet的标准设计，但将每个分辨率（最高分辨率除外）的构建块的最后一个卷积替换为Transformer模块。为了提高分割的质量，我们试图运用自我关注在高分辨率特征地图上提取详细的长期关系。为此，我们提出了一种有效的自我注意机制，该机制在时间和空间上将整体复杂度从O（$n^2$）显著降低到接近O（n）。此外，我们在自我注意模块中使用相对位置编码来学习医学图像中的内容-位置关系（bert中的class token）。

方法

重新审视自注意力机制

Transformer是建立在多头自我注意力（MHSA）模块的基础上，它允许模型从不同的表示子空间联合推断注意力，并且将多个头部的结果串联起来，然后使用前馈网络进行变换。多头自注意力的计算公式如下：

P∈R^n×n被称为上下文聚合矩阵或相似矩阵,将这个相似矩阵用作权重，以从中收集上下文信息。通过这种方式自注意力机制具有全局感受野，并且善于捕获长距离依赖。

高效的自注意力机制

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

其主要思想是用两个投影来投影关键和价值：K , V ∈ R ^nxd低维嵌入的k，v ∈R^kxd

在这里插入图片描述

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

相对位置编码

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

其中，q_i是像素i的query向量，k_i是像素j的key向量，j，r_jx^W-i_x和r_jy-iy^H是可学习嵌入分别是相对宽度,j_x-i_x和相对高度j_y-i_y

与有效的自我注意类似，在低维投影后计算相对宽度和高度。包括相对位置嵌入在内的有效自我注意是：

网络架构

图1是UTNet的体系结构,结合了卷积和自我注意机制，这种混合结构可以利用卷积图像的感应偏差来避免大规模的预训练，以及转换器捕捉远程关系的能力。由于误分割区域通常位于感兴趣区域的边界，因此高分辨率的上下文信息在分割中起着至关重要的作用。因此，将重点放在所提出的自我注意模块上，使其能够有效地处理大型特征地图。UTNet的构建块如图1（b）和（c）所示，包括Residual Basic Block和Transformer Encoder块。

U-Net中蓝色箭头是卷积操作（conv3X3,ReLU）在UTNet中替换成Residual Basic Block操作和绿色箭头的Transformer Encoder的操作，crop and copy的操作换成第一层的跳跃连接，其他层换成了 Transformer Decoder操作。

编码器部分：

以单通道输入图片为例，原始分辨率上通过两个Residual Basic Block操作后相应地进行变化，下采样（maxpool的池化操作2X2 ）高和宽就会减半，编码器下面的操作都是一个Residual Basic Block操作，一个Transformer Encoder的操作，channel会相应的增加。没有在原始分辨率上应用Transformer，因为在网络的极浅层中添加Transformer模块无助于实验，但会引入额外的计算。一个可能的原因是，网络的浅层更关注细节纹理，而收集全局上下文可能无法提供信息。

解码器部分：

在通过和编码器部分得到的特征图进行一个 Transformer Decoder操作融合信息，特征层的高宽都会放大为原来的两倍，在进行一个Residual Basic Block操作得到的特征图继续以上操作，最后的黄色箭头部分是上采样（转置卷积），原始分辨率层在通过Residual Basic Block操作输出最后的分割结果。

实验结果

在多标签、多供应商心脏磁共振成像（MRI）挑战队列中系统地评估UTNet，包括左心室（LV）、右心室（RV）和左心室心肌（MYO）的分割。

在培训集中，有来自两个不同MRI供应商的150张带注释的图像（每个供应商有75张图像），包括：西门子；B:飞利浦。在测试集中，有来自4家不同MRI供应商的200张图像（每个供应商50张图像），包括：西门子；B:飞利浦；C:通用电气；D：佳能，供应商C和D在培训集中完全不存在（丢弃未标记的数据）。来自不同供应商的MRI扫描在外观上有显著差异，这使得能够测量模型的稳健性，并在不同设置下与其他模型进行比较。

具体来说，进行了两个实验，以突出UTNet的性能和健壮性。首先，报告了来自同一供应商A的培训和测试数据的初步结果。其次，进一步衡量了模型的跨供应商稳健性。这种设置更具挑战性，因为培训和测试数据来自独立供应商。报告了每个模型的Dice分数和Hausdorff距离来比较性能。

数据预处理

将面内间距重新采样为1.2×1.2 mm，同时保持沿z轴的间距不变。从头开始训练所有型号150个epoch。使用基本学习率为0.05的指数学习率调度器。在一个GPU上使用了批次大小为16的SGD优化器，动量衰减和权重衰减分别设置为0.9%和1e−4。

数据增强是在模型训练过程中动态应用的，包括随机旋转、缩放、平移、加性噪声和Gamma变换。在进入模型之前，所有图像都被随机裁剪到256×256。使用Dice损失和交叉熵损失相结合的方法来训练所有的网络。

分割结果

比较了UTNet和多种最先进的分割模型的性能，UNET建立在完全卷积网络之上，采用U型架构来捕获上下文信息。ResUNet在架构上类似于UNET，但它使用残差块作为构建块。CBAM使用两个连续的卷积模块来推断通道和空间注意力，以自适应地细化中间特征地图。双重注意网络使用两种自我注意分别在空间和通道维度上对语义相互依赖进行建模。为了更好地进行比较，在ResUNet主干中实现了CBAM和Double Attendence。由于其二次复杂性，双重注意仅适用于4次下采样后的特征映射。

如表1所示，UTNet在所有细分结果(LV、MYO和RV)中都表现出领先的表现。通过引入残差连接，ResUNet比原来的UNET略有改进。CBAM的空间和通道注意都是从卷积层推断出来的，但它仍然受制于有限的接受野。因此，与ResUNet相比，CBAM的改进有限。双注意方法与ResUNet方法几乎相同，因为它存在二次复杂度，无法处理更高分辨率的特征图来修复分割边界中的错误。同时，UTNet比双注意力方法提供的参数更少，并且可以从高分辨率的特征地图中捕获全局上下文信息。

消融实验：

图3.消融研究。(a)不同自我注意位置的影响。(b)缩小规模和有效自我注意投射的效果。© transformer编码器、transformer解码器的作用及相对位置编码。

图3(a)显示了不同自我注意位置的表现。X轴上的数字表示自我关注的级别，例如，‘34’表示执行3次和4次下采样的级别。随着水平的提高，自我注意可以收集到更细粒度的细节信息，表现也会提高。但是，当添加到原始分辨率时，曲线会饱和。这是因为非常浅的层倾向于更多地关注局部纹理，其中全局上下文信息不再是信息性的。图3（c)显示了有效的自我注意减少4、8、16的结果。减少的8会产生最好的表现。插值下采样略好于使用最大合并。图3（c）显示了使用(a)和(b)中的最佳超参数的transformer编码器、解码器和相对位置编码的效果。transformer编码器和解码器的结合提供了最佳的性能。相对位置编码也起着至关重要的作用，因为删除它会导致较大的性能下降。

稳健性分析：

表2显示了带有来自供应商A和B的数据的培训模型的结果，然后分别在供应商A、B、C和D上测试这些模型。在查看C和D供应商的结果时，竞争方法会受到供应商差异的影响，而UTNet则保持有竞争力的业绩。这种观察可能归因于自我注意在多个层次的特征地图上的设计和内容-位置注意，使得UTNet能够更好地专注于全局语境信息，而不仅仅是局部纹理。

表2.鲁棒性比较，用Dice分数衡量。所有模型都使用供应商A、B的数据进行训练，并使用供应商A、B、C、D的数据进行测试。C和D括号中的数字表示与A和B的平均值相比性能下降。

图4进一步显示UTNet显示最一致的边界结果，而其他三种方法无法捕捉边界的细微特征，特别是心脏MRI中的RV和MYO区域。

结论

提出了一种U型混合变换网络(UTNet)来融合卷积层和自注意力机制的优点，用于医学图像分割。混合层设计允许将transformer初始化为卷积网络，而不需要预先训练。新的自注意结构允许在编码器和解码器中扩展不同级别的网络操作，以便更好地捕获远程依赖关系。

这种设计将有助于丰富参数化的transformer模型在医学视觉应用中变得更容易使用。此外，高效处理长序列的能力为UTNet在更多下游医学图像任务中的使用开辟了新的可能性。