UTNet: A Hybrid Transformer Architecture for Medical Image Segmentation

UTNet: A Hybrid Transformer Architecture for Medical Image Segmentation
UTNet：一种用于医学图像分割的混合transformer结构
论文链接：UTNet: A Hybrid Transformer Architecture for Medical Image Segmentation

Abstract

transformer架构已经在许多自然语言处理任务中取得了成功。然而，它在医学视觉上的应用在很大程度上还没有被探索。在这项研究中，我们提出了UTNet，一个简单但功能强大的混合transformer架构，它将自我关注集成到卷积神经网络中，以增强医学图像分割。UTNet在编码器和解码器中都应用了自我注意模块，以便以最小的开销捕获不同规模的远程依赖。为此，我们提出了一种有效的自我注意机制，并结合相对位置编码，将自我注意操作的复杂度从O(N2)大幅降低到近似O(N)。为了从编码器中跳过的连接中恢复细粒度的细节，还提出了一种新的自注意解码器。我们的方法解决了transformer需要大量数据来学习视觉感应偏差的两难境地。我们的混合层设计允许将transformer初始化为卷积网络，而不需要预先训练。我们已经在多标签、多供应商的心脏磁共振成像队列中对UTNet进行了评估。UTNet表现出卓越的分割性能和相对于最先进的方法的健壮性，有望在其他医学图像分割中得到很好的推广。

Introduction

卷积网络以其卓越的特征表示能力给计算机视觉领域带来了革命性的变化。目前，卷积编解码结构在位置敏感任务中取得了实质性进展，如语义分割。所使用的卷积运算通过从邻域像素收集局部信息来捕获纹理特征。为了全局聚合局部滤波响应，这些模型堆叠了多个卷积层，并通过下采样扩展了接收范围。尽管取得了这些进步，但这一范式有两个固有的局限性。首先，卷积仅从邻域像素收集信息，缺乏显式捕获远程(全局)依赖关系的能力。其次，卷积核的大小和形状通常是固定的，因此它们不能适应输入内容。
使用自我注意机制的transformer架构已经在自然语言处理(NLP)中获得成功，因为它具有捕获远程依赖的能力。自我注意是一种计算原语，它通过上下文聚合机制实现成对的实体交互，该机制具有捕获远程关联特征的能力。它允许网络根据输入内容动态聚合相关特征。对简单形式的自我注意的初步研究表明，它在分割、检测和重建方面是有用的。
虽然基于图像的transformer应用前景看好，但培训和部署transformer架构面临着一些艰巨的挑战。首先，自我注意机制相对于序列长度具有O(N2)的时间和空间复杂度，导致大量的训练和推理开销。以前的研究试图降低自我注意的复杂性，但仍远未达到完美。由于时间复杂性，标准的自我注意只能应用于patch方式，例如[3，27]使用16×16平坦的图像patch作为输入序列来编码图像，或者在来自CNN主干的特征地图上进行编码，这些特征地图已经被下采样到低分辨率[4，22]。然而，对于像医学图像分割这样的位置敏感任务，高分辨率特征起着至关重要的作用，因为大多数误分割区域都位于感兴趣区域的边界附近。其次，transformer对图像没有感应偏向，在小规模的数据集上表现不佳。例如，Transformer可以通过像Full JFT300M[3]这样的大型数据集从预培训中受益。但即使有了ImageNet的预培训，Transformer仍然比ResNet差，更不用说医学图像数据集，可用的医疗数据量要少得多。
在本文中，我们提出了一种U型混合变换网络：UTNet，它综合了卷积和自关注策略的优点，用于医学图像分割。该方法的主要目的是利用卷积分层提取局部强度特征，避免变压器的大规模预训练，同时利用自我注意捕捉远距离联想信息。我们遵循UNET的标准设计，但将每个分辨率(最高分辨率除外)中构建块的最后一次卷积替换为建议的变压器模块。为了提高分割质量，我们试图应用自我注意在高分辨率特征地图上提取详细的长期关系。为此，我们提出了一种有效的自我注意机制，在时间和空间上将整体复杂度从O(N2)大幅降低到接近O(N)。此外，我们在自我注意模块中使用相对位置编码来学习医学图像中的内容-位置关系。我们的UTNet在多标签、多供应商的心脏磁共振成像队列中显示出卓越的分割性能和稳健性。考虑到UTNet的设计，我们的框架有望很好地推广到其他医学图像分割上。

Method

Revisiting Self-attention Mechanism

transformer建立在多头自我注意(MHSA)模块的基础上，该模块允许模型从不同的表示子空间联合推断注意力。来自多个头部的结果被连接起来，然后用前馈网络进行转换。在这项研究中，我们使用4个头部，为了简单起见，在下面的公式和图中没有给出多头的尺寸。考虑一个输入特征映射X∈R^C×H×W，其中H、W是空间高度、宽度，C是通道数。使用3个1×1卷积来投影X到查询、键、值的嵌入：Q，K，V∈R^d×H×W，其中d是每个头中嵌入的维度。然后将Q、K、V展平并转置成大小为n×d的序列，其中n=HW。自我注意的输出是按比例的点积:

请注意，P∈R^n×n称为上下文聚集矩阵，或相似度矩阵。具体地说，第i个查询的上下文聚合矩阵为Pi=Softmax(qiK^T /√d)，Pi∈R^1×n，它计算qi和键中的每个元素之间的归一化成对点生成。然后，将上下文聚集矩阵用作权重，以从这些值收集上下文信息。通过这种方式，自我关注本质上具有全局感受野，并善于捕捉长期依赖。此外，上下文聚合矩阵自适应于输入内容，以实现更好的特征聚合。然而，n×d矩阵的点积运算复杂度为O(N2d)。通常，当特征映射的分辨率较大时，n远大于d，因此序列长度决定了自我注意的计算，使得在高分辨率特征映射中应用自我注意是不可行的，例如，16×16特征映射n=256时，128×128特征映射n=16384。

Efficient Self-attention Mechanism

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

图2.提出的高效多头自我注意(MHSA)。(a)transformer编码器中使用的MHSA。(B)transformer解码器中使用的MHSA。它们具有相似的概念，但(b)接受两个输入，包括来自编码器跳过连接的高分辨率特征和来自解码器的低分辨率特征。
其==主要思想是使用两个投影将键和值K，V∈R^n×d投影到低维嵌入
中：，其中k=hw<<n，h和w是经过二次采样后变小了的特征地图的大小。提出的高效自我关注现在是：

通过这样做，计算复杂度降低到O(NKD)。值得注意的是，到低维嵌入的投影可以是任何下采样操作，例如平均/最大池化或strided卷积。在我们的实现中，我们使用1×1卷积，然后是双线性插值来对特征地图进行下采样，缩减的尺寸是8。

Relative Positional Encoding

标准的自我注意模块完全丢弃了位置信息，并且是扰动等变的[1]，对于高度结构化的图像内容的建模是无效的。前人工作[13]中的正弦嵌入在卷积层中不具有平移等变性质。因此，我们通过添加相对高度和宽度信息来使用二维相对位置编码。在Softmax之前使用像素i=(ix，iy)和像素j=(jx，jy)之间的相对位置编码的成对注意logit：

其中qi是像素i的查询向量，kj是像素j的键向量，rjx−ix^W和rjy−iy^H分别是相对宽度jx−ix和相对高度jy−iy的可学习嵌入。与有效的自我注意类似，在低维投影后计算相对宽度和高度。包括相对位置嵌入在内的有效自我注意是：

其中，S^relH，S^relW∈R^HW×hw是沿高度和宽度维度的相对位置对数矩阵满足S^relH[i，j]=q^Ti r^Hjy−iy，S^relW[i，j]=q^ti r^Wjx−ix。

Network Architecture

图1突出显示了UTNet的架构。我们寻求将卷积机制和自我注意机制的优势结合起来。因此，该混合结构可以利用卷积产生的图像的感应偏差来避免大规模的预训练，以及Transformer捕获远程关系的能力。由于误分割区域通常位于感兴趣区域的边界，因此高分辨率的上下文信息在分割中起着至关重要的作用。因此，我们将重点放在所提出的自我注意模块上，使其能够有效地处理大尺寸的特征地图。我们没有天真地将自我注意模块集成在CNN主干上的特征地图上，而是将Transformer模块应用到编码器和解码器的每一级，以从多个尺度收集远程依赖关系。请注意，我们不会在原始分辨率上应用Transformer，因为在网络的非常浅的层中添加Transformer模块对实验没有帮助，反而会引入额外的计算。一个可能的原因是，网络的浅层更多地关注细节纹理，而收集全局上下文可能不会提供信息。UTNet的模块如图1(b)和©所示，包括剩余基本模块和transformer模块。对于这两个块，我们在快捷方式中使用身份映射的预激活设置。这种身份映射在vision[8]和NLP任务[19]中被证明是有效的。

图1 (a)提出的UTNet的混合体系结构。所提出的高效的自注意机制和相对位置编码允许我们在编码器和解码器中使用Transformer来聚合来自多个尺度的全局上下文信息。(B)激活前残差基础块。©transformer编码块的结构。

Experiment

Experiment Setup

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

Implementation Detail

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

Segmentation Results

我们比较了UTNet和多种最先进的分割模型的性能。UNET[14]建立在完全卷积网络之上，采用U型架构来捕获上下文信息。ResUNet在架构上类似于UNET，但它使用残差块作为构建块。CBAM[23]使用两个连续的卷积模块来推断通道和空间注意力，以自适应地细化中间特征地图。双重注意网络[4]使用两种自我注意分别在空间和通道维度上对语义相互依赖进行建模。为了更好地进行比较，我们在ResUNet主干中实现了CBAM和Double Attendence。由于其二次复杂性，双重注意仅适用于4次下采样后的特征映射。
如表1所示，UTNet在所有细分结果(LV、MYO和RV)中都表现出领先的表现。通过引入残差连接，ResUNet比原来的UNET略有改进。CBAM的空间和通道注意都是从卷积层推断出来的，但它仍然受制于有限的接受野。因此，与ResUNet相比，CBAM的改进有限。我们还认识到，双注意方法与ResUNet方法几乎相同，因为它存在二次复杂度，无法处理更高分辨率的特征图来修复分割边界中的错误。同时，我们的UTNet比双注意力方法提供的参数更少，并且可以从高分辨率的特征地图中捕获全局上下文信息。
表1.根据Dice分数和效率比较的分割性能。所有模型都使用供应商A的数据进行训练和测试。Hausdorff的远程结果在补充资料中报告。
Ablation study
图3(a)显示了不同自我注意位置的表现。X轴上的数字表示自我关注的级别，例如，‘34’表示执行3次和4次下采样的级别。随着水平的提高，自我注意可以收集到更细粒度的细节信息，表现也会提高。但是，当添加到原始分辨率时，曲线会饱和。我们认为这是因为非常浅的层倾向于更多地关注局部纹理，其中全局上下文信息不再是信息性的。图3©显示了有效的自我注意减少4、8、16的结果。减少的8会产生最好的表现。插值下采样略好于使用最大合并。图3©显示了使用(a)和(b)中的最佳超参数的transformer编码器、解码器和相对位置编码的效果。transformer编码器和解码器的结合提供了最佳的性能。相对位置编码也起着至关重要的作用，因为删除它会导致较大的性能下降。

图3.消融研究。(a)不同自我注意位置的影响。(b)缩小规模和有效自我注意投射的效果。©transformer编码器、transformer解码器的作用及相对位置编码。
为了与标准自我注意在空间和时间复杂度上进行正面比较，我们进一步在四个分辨率(1，2，3，4，与UTNet相同)上应用双重注意，并使用相同的输入图像大小和批次大小(256×256×16)来测试推理时间和内存消耗。UTNet通过二次复杂度获得了优于双重关注的优势，其中图形处理器内存：3.8G比36.9G，时间：0.146秒比0.243秒。
Robustness analysis
表2显示了带有来自供应商A和B的数据的培训模型的结果，然后分别在供应商A、B、C和D上测试这些模型。在查看C和D供应商的结果时，竞争方法会受到供应商差异的影响，而UTNet则保持有竞争力的业绩。这种观察可能归因于自我注意在多个层次的特征地图上的设计和内容-位置注意，使得UTNet能够更好地专注于全局语境信息，而不仅仅是局部纹理。图4进一步显示UTNet显示最一致的边界结果，而其他三种方法无法捕捉边界的细微特征，特别是心脏MRI中的RV和MYO区域。
表2.鲁棒性比较，用Dice分数衡量。所有模型都使用供应商A、B的数据进行训练，并使用供应商A、B、C、D的数据进行测试。C和D括号中的数字表示与A和B的平均值相比性能下降。

图4.供应商C和D看不见的测试数据上的硬案例可视化。前两行和后两行分别表示结果和供应商C和D的放大视图。轮廓指出了ground truth的注解。最好用L V(绿色)、MYO(黄色)和R V(红色)颜色观看。供应商C的测试用例由于运动伪影而模糊，而供应商D的测试用例噪声大，边界对比度低。只有UTNet提供一致的分段，这证明了它的健壮性。补充部分提供了更多的分割结果可视化和注意力图。

Conclusion

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