TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation

Abstract

• 医学图像分割是开发医疗保健系统的重要先决条件，尤其是对于疾病诊断和治疗计划。在各种医学图像分割任务中，u 形架构（也称为 U-Net）已成为事实上的标准并取得了巨大的成功。然而，由于卷积运算的固有局部性，U-Net 通常在显式建模长程依赖性方面表现出局限性。专为序列到序列预测而设计的 Transformers 已成为具有先天全局自注意机制的替代架构，但由于低级细节不足，可能导致定位能力有限。 在本文中，我们提出了 TransUNet，它兼具 Transformers 和 U-Net 的优点，是医学图像分割的强大替代方案。一方面，Transformer 对来自卷积神经网络 (CNN) 特征图的标记化图像块进行编码作为提取全局上下文的输入序列。另一方面，解码器对编码的特征进行上采样，然后将其与高分辨率 CNN 特征图相结合以实现精确定位。
• 我们认为 Transformers 可以作为医学图像分割任务的强大编码器，与 U-Net 结合使用，通过恢复局部空间信息来增强更精细的细节。TransUNet 在包括多器官分割和心脏分割在内的不同医学应用上的表现优于各种竞争方法。代码和模型可在 https://github.com/Beckschen/TransUNet 上找到。
• 论文地址：[2102.04306] TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation
• 医学图像分割（Medical image segmentation）是图像分割在医学领域的重要应用，具有非常高的落地价值，基本上都以 U-Net 为主，考虑到 CNN 的卷积核对长距离依赖(long-range dependency)信息建模上的不足，很适合用 Transformer 来弥补，但是如果只用 Transformer 处理图像，其对局部信息的处理特别是细节上又有缺陷。本文作者提出了 TransUNet，结合了 Transformer 和 U-Net，来进行医学图像分割。
  • 首先用 CNN 提取 feature map，然后将patch size设置为 1∗1 ，说白了就是把 feature map 给展开成像素级序列，作为 Transformer 的输入。
  • 在 Decoder 阶段，为了弥补 Transformer 在局部信息处理上的不足，又在不同阶段（层）将 CNN encoder 阶段提取的 feature map 和 Transformer hidden feature 结合，这样解码时即有全局信息又有局部细节信息。Encoder 结合了 ResNet-50 和 ViT，这两个模型都是在 ImageNet 上预训练过的。

Introduction

• 卷积神经网络 (CNN)，尤其是全卷积网络 (FCN) ，已成为医学图像分割领域的主流。在众多变体中，U-Net 已成为事实上的选择，它由对称编码器-解码器网络组成，并带有跳跃连接以增强细节保留。基于这种方法，已在广泛的医学应用中取得了巨大成功，例如从磁共振 (MR) 中进行心脏分割、从计算机断层扫描 (CT) 中进行器官分割以及从结肠镜检查视频中进行息肉分割 。

• 尽管基于 CNN 的方法具有出色的表征能力，但由于卷积运算的固有局部性，它们在建模显式长程关系方面通常存在局限性。因此，这些架构通常性能较弱，尤其是对于在纹理、形状和大小方面表现出较大患者间差异的目标结构。为了克服这一限制，现有研究提出建立基于 CNN 特征的自注意机制。另一方面，为序列到序列预测而设计的 Transformers 已成为替代架构，它们完全采用分配卷积运算符，而仅依赖注意机制 。与之前基于 CNN 的方法不同，Transformers 不仅在建模全局上下文方面功能强大，而且在大规模预训练下还表现出对下游任务的卓越可迁移性。在机器翻译和自然语言处理 (NLP) 领域 ，Transformers 的成功得到了广泛见证。最近，各种图像识别任务的尝试也达到了甚至超过了最先进的性能 。

• 在本文中，我们首次介绍了 Transformer 在医学图像分割领域的潜力。然而，有趣的是，我们发现简单的使用方式（即使用 Transformer 对标记化的图像块进行编码，然后直接将隐藏的特征表示上采样为全分辨率的密集输出）无法产生令人满意的结果。

• 这是因为 Transformers 将输入视为 1D 序列，并在所有阶段专注于对全局上下文进行建模，因此导致低分辨率特征缺乏详细的定位信息。而且这些信息无法通过直接上采样到全分辨率来有效恢复，因此导致分割结果粗糙。另一方面，CNN 架构（例如 U-Net ）提供了一种提取低级视觉线索的途径，可以很好地弥补这种精细的空间细节。

• 为此，我们提出了第一个医学图像分割框架TransUNet，该框架从序列到序列预测的角度建立了自注意力机制。为了弥补Transformers带来的特征分辨率损失，TransUNet采用了混合CNN-Transformer架构，既能利用来自CNN特征的详细高分辨率空间信息，也能利用Transformers编码的全局上下文。受U形架构设计的启发，Transformers编码的自注意力特征随后被上采样，与编码路径中跳过的不同高分辨率CNN特征相结合，以实现精确定位。我们表明，这样的设计使我们的框架能够保留Transformers的优势，同时也有利于医学图像分割。实证结果表明，与以前基于CNN的自注意力方法相比，我们基于Transformer的架构提供了一种更好的利用自注意力的方法。此外，我们观察到，更深入地结合低级特征通常会带来更好的分割精度。大量实验证明了我们的方法在各种医学图像分割任务上优于其他竞争方法。

Related Works

• 将 CNN 与自注意力机制相结合。各种研究都试图通过基于特征图对所有像素的全局交互进行建模，将自注意力机制集成到 CNN 中。例如，Wang 等人设计了一个非局部运算符，可以将其插入多个中间卷积层 。基于编码器-解码器 U 形架构，Schlemper 等人 提出了集成到跳过连接中的附加注意力门模块。与这些方法不同，我们使用 Transformers 在我们的方法中嵌入全局自注意力。

• Transformers。Transformers 最早 提出用于机器翻译，并在许多 NLP 任务中建立了最先进的技术。为了使 Transformers 也适用于计算机视觉任务，人们对其进行了一些修改。例如，Parmar 等人仅在每个查询像素的局部邻域中应用自注意力，而不是全局应用。Child 等人提出了稀疏 Transformer，它采用可扩展的近似值来实现全局自注意力。最近，Vision Transformer (ViT) 通过将具有全局自注意力的 Transformers 直接应用于全尺寸图像，在 ImageNet 分类中取得了最先进的水平。 据我们所知，所提出的 TransUNet 是第一个基于 Transformer 的医学图像分割框架，它建立在非常成功的 ViT 之上。

• u-net

  • 在该图中，我们看到了专为生物医学图像分割设计的 U-Net 架构。该模型因其独特的 “U ”形而闻名，这种形状源于编码器-解码器设置，其跳过连接有助于保留重要细节。

  • 

  • U-Net 的左侧是编码器，它使用的是卷积神经网络 (CNN)，类似于图像分类中使用的卷积神经网络。它的工作原理是逐步缩小输入图像的****空间维度*（使其尺寸更小），同时增加*深度*（通道数）。这一过程输出的*特征图****可以捕捉到从图像中提取的信息。编码器中的每一步都包括

    • 从图像中提取特征的卷积操作 （紫色箭头）。
    • ReLU激活，有助于模型学习非线性特征。
    • 最大池化操作（红色箭头），通过选择前一个特征图中最突出的特征来降低特征图的维度。

  • 右侧是解码器，执行与编码器相反的功能。它逐步恢复空间尺寸，使其与原始图像的尺寸相匹配，从而精确定位物体的位置。解码器的每个步骤包括

    • 上卷积操作（绿色箭头），可放大特征图。
    • 与相应编码器特征图的连接 （灰色箭头）。这一步骤至关重要，而直接将特征从编码器带入解码器的跳转连接则为这一步骤提供了便利。这些连接有助于模型记住图像结构的精确细节，而这些细节可能会在编码过程中丢失。

  • 最后，卷积将特征图转换成最终的分割图，根据模型的训练将输入图像的每个像素分为不同的类别。

• 虽然 U-Net 在许多任务中都很强大，但它也有局限性。由于 U-Net 主要关注局部细节，可能会丢失全局信息，因此在****长距离依赖关系*建模方面比较吃力。此外，对于*病人之间差异****较大的目标结构，U-Net 的性能也可能较弱，例如在医学图像中，不同病人的肿瘤在纹理、形状和大小上可能存在显著差异。

• 变换器最初是为自然语言处理（NLP）而设计的，现在已被调整用于图像处理，成为[视觉变换器（ViT）*](https://arxiv.org/pdf/2010.1192。ViT 通过将图像视为像素序列或小图像片段，将Transformer架构调整为视觉数据。

  • 

• ViT 的主要组成部分：

  • 自注意力机制：ViT 的核心优势在于其自我注意机制，它动态权衡图像各个部分的重要性。这使 ViT 能够捕获“全局上下文”，突出特定任务所需的关键特征。
  • MSA（多头自我注意）：ViT 使用 MSA 同时关注多个图像部分。通过将输入划分为多个“头”，它可以并行处理不同的方面，从而捕获更广泛的信息，例如特征类型和空间关系。
  • MLP（多层感知器）：自我注意后，ViT 通过 MLP 处理特征，MLP 由能够学习非线性特征组合的全连接层组成。MLP 整合了注意力头之间的信息，增强了模型辨别复杂模式和关系的能力。
  • Transferability : ViT 的一个主要优势是它的**Transferability **，由于其丰富的广义特征表示，允许在大型数据集上训练的模型以最少的调整有效地适应新的相关任务。

• 尽管 ViT 具有优势，但它在定位方面遇到了困难，由于其 1D 序列处理方法，通常无法精确定位图像中的确切对象位置。当单独与 CNN 解码器一起使用时，这会导致分辨率输出较低。TransUNet 架构通过将 ViT 的全局关注与 CNN 的局部精度相结合来解决这个问题。TransUNet 是一种混合架构，它结合了 CNN 和 Transformer 技术，以利用来自 CNN 的高分辨率空间信息和来自 Vision Transformer （ViT） 的全局上下文洞察。这种集成允许精确的图像分割，这在医学成像中特别有用。

  • 编码器：该编码器采用双重方法，从 CNN 编码器开始，将输入图像处理成特征图。然后，此特征图被标记化为补丁，由 Transformer 编码器处理。此设置可以有效地捕获局部详细信息和全局模式。
  • 译码器：解码器使用 Transformer 编码的特征并对其进行上采样，将它们与 CNN 中的高分辨率特征图集成。此步骤对于分割区域的精确定位至关重要。

Method

• 给定一个图像 $x\in {\mathbb{R}}^{H\times W\times C}$，其空间分辨率为 H×W，通道数为 C。我们的目标是预测相应的像素级标签图，大小为 H×W。最常见的方法是直接训练 CNN（例如 UNet），首先将图像编码为高级特征表示，然后将其解码回完整空间分辨率。与现有方法不同，我们的方法通过使用 Transformer 将自注意力机制引入编码器设计中。我们将首先在第 3.1 节中介绍如何直接应用 Transformer 对分解图像块中的特征表示进行编码。然后，将在第 3.2 节中详细阐述 TransUNet 的整体框架。
  • 初始 CNN 处理： 输入图像首先由 CNN 处理，CNN 生成维度为 H×W 的特征图。此特征图从图像中捕获本地化特征，为进一步处理提供了丰富的基础。
  • **标记化：**特征图被划分为大小为 P×P 的块。然后将每个补丁展平为一维序列，从而产生一个包含 N 个标记的数组，标记表示原始图像的各个部分，但采用适合 transformer 处理的格式。
  • 补丁嵌入：每个标记 $x_{i，p}$ 都使用可训练的线性投影（通常称为嵌入矩阵）转换为 D 维嵌入空间
    • 本质上是将输入标记重塑并投影到更高维的空间中。位置嵌入 Epos 被添加到这些嵌入中，以合并不同补丁之间的空间关系，形成 transformer 的初始序列。此序列 z0 现在以保留内容和位置上下文的形式表示图像，从而促进Transformer的全局上下文处理。
  • **变压器加工：**嵌入式补丁 z0 序列被馈送到 transformer 的层中，每个层都包含多头自我注意 （MSA） 和多层感知器 （MLP） 块：
    • **MSA 层：**在每一层 L ，序列都会经历自我注意，其中模型计算所有补丁之间的注意力分数，使其能够根据图像与分割任务的相关性专注于图像的不同部分。每个标记的 MSA 层的输出是其自身特征和其他标记的特征的组合，按其计算的注意力分数加权
    • **MLP 层：**在自我注意之后，每个 token 的特征都会通过 MLP 传递，MLP 可以通过应用非线性变换来进一步细化特征
• 本文讨论了通过利用双线性上采样对解码器使用“朴素”方法的可能性，并表明这会导致低级细节的丢失。因此，在 TransUNet 中使用了级联上采样器 （cascade upsampler CUP）。
  • 由多个上采样步骤组成，用于解码隐藏特征，输出最终的分割掩码。在重塑隐藏特征序列后，我们通过级联多个上采样块来实例化 CUP，以达到从 $\frac{H\ast P}{W\ast P}$ 到 H*W 的全分辨率，其中每个块包括：
    • 两个 upsampling 运算符。
    • 一个 3×3 的卷积层。
    • ReLU 激活函数。
  • 该研究战略性地选择特定参数设置来优化 TransUNet 的性能（这些参数是根据它们在详细图像处理与计算效率之间取得平衡的能力来选择的。）：
    • 默认图像大小：224x224 像素
    • 补丁大小： 16
    • 批量大小： 24
• 增加跳过连接的数量可以提高模型性能。之所以进行这种增强，是因为更多的连接允许更好地集成局部详细信息和全局上下文。通过在级联上采样器 （CUP） 的所有三个中间上采样阶段（不包括输出层）插入 skip 连接来获得最佳结果。这可以通过使用加法 Transformer 加以改进，但由于其计算成本较高，TransUNet 架构并未采用。
• 输入分辨率：使用更高的分辨率可以更好地显示可用于分割的细节，从而获得更准确的结果。但是，这一优势伴随着计算需求的增加和资源消耗的增加，因此需要仔细平衡。
• 补丁大小：减小补丁大小会增加补丁的数量，从而增加模型中的连接数量，从而增强其捕获复杂细节的能力。这可以提高分割准确性。
• 模型缩放：通过添加更多层或增加参数来扩大模型通常会获得更好的性能。但是，较大的配置需要更多的计算资源，这使得它们对于某些应用程序的可行性较低。

Transformer as Encoder

• 图像序列化。按照[An image is worth 16x16 words: Transformers for image recognition at scale]的方法，我们首先通过将输入x重塑为扁平的二维块序列 { x p i ∈ R P 2 ⋅ C ∣ i = 1 , . . . , N } \{x^i_p\in\R^{P^2·C}|i=1,...,N\} {xpi​∈RP2⋅C∣i=1,...,N} 来执行标记化，其中每个块的大小为P×P，$N=\frac{HW}{P^2}$是图像块的数量（即输入序列长度）。

• 补丁嵌入。我们使用可训练的线性投影将矢量化补丁 xp 映射到潜在的 D 维嵌入空间中。为了对补丁空间信息进行编码，我们学习特定位置嵌入，并将其添加到补丁嵌入中以保留位置信息，如下所示：

  • $$z_0=[x_p^{1}E;x_p^{2}E;...;x_p^{N}E]+E_{pos}$$

  • 其中 $E\in \mathbb{R}(P^2\cdot C)\times D$ 是块嵌入投影，$E_{pos}\in R^{N\times D}$ 表示位置嵌入。

• Transformer 编码器由 L 层多头自注意力 (MSA) 和多层感知器 (MLP) 模块组成（等式 (2)(3)）。因此，第 ’ 层的输出可以写成如下形式：

  • $$\begin{gathered} z_l^{\prime }=MSA(LN(Z_{l-1}))+z_{l-1} \\ {z}_l=MLP(LN(z_l^{\prime }))+z_l^{\prime } \end{gathered}$$

  • 其中 LN(·) 表示层归一化运算符，zL 是编码图像表示。Transformer 层的结构如图 1(a) 所示。

  • 

  • 图 1：框架概述。 (a) Transformer 层示意图；(b) 所提出的 TransUNet 的架构。

TransUNet

• 出于分割目的，一个直观的解决方案是简单地将编码的特征表示 $z_L\in {\mathbb{R}}^{\frac{HW}{P^2}\times D}$ 上采样到全分辨率，以预测密集输出。在这里，为了恢复空间顺序，编码特征的大小应该首先从 $\frac{HW}{P^2}$ 重塑为 $\frac{H}{P}\times \frac{W}{P}$ 。我们使用 1 × 1 卷积将重塑特征的通道大小减小到类数，然后将特征图直接双线性上采样到全分辨率 H × W 以预测最终的分割结果。在第 4.3 节的后续比较中，我们在解码器设计中将这种简单的上采样基线表示为“None”。

• 尽管将 Transformer 与简单的上采样相结合已经产生了不错的性能，但如上所述，这种策略并不是 Transformer 在分割中的最佳用法，因为 $\frac{H}{P}\times \frac{W}{P}$ 通常比原始图像分辨率 H × W 小得多，因此不可避免地会导致低级细节（例如器官的形状和边界）的丢失。因此，为了弥补这种信息损失，TransUNet 采用混合 CNN-Transformer 架构作为编码器以及级联上采样器来实现精确定位。图 1 显示了所提出的 TransUNet 的概览。

• CNN-Transformer 混合模型作为编码器。TransUNet 不使用纯 Transformer 作为编码器（第 3.1 节），而是采用 CNN-Transformer 混合模型，其中首先使用 CNN 作为特征提取器来为输入生成特征图。块嵌入应用于从 CNN 特征图中提取的 1×1 块，而不是从原始图像中提取的块。我们选择这种设计是因为

  • 1）它允许我们在解码路径中利用中间高分辨率 CNN 特征图；
  • 2）我们发现混合 CNN-Transformer 编码器的性能比简单地使用纯 Transformer 作为编码器更好。

• 级联上采样器。我们引入了一个级联上采样器 (CUP)，它由多个上采样步骤组成，用于解码隐藏特征并输出最终的分割掩码。在将隐藏特征 $z_L\in {\mathbb{R}}^{\frac{HW}{P^2}\times D}$ 序列重塑为 $\frac H P ×\frac W P× D $ 形状后，我们通过级联多个上采样块来实例化 CUP，以达到从 H P × W P 到 H×W 的全分辨率，其中每个块依次由 2× 上采样运算符、3×3 卷积层和 ReLU 层组成。我们可以看到，CUP 与混合编码器一起形成了一个 U 形架构，可通过跳跃连接实现不同分辨率级别的特征聚合。CUP 的详细架构以及中间的跳跃连接可以在图 1(b) 中找到。

Experiments and Discussion

Dataset and Evaluation

• Synapse 多器官分割数据集Synapse | Sage Bionetworks.我们使用 MICCAI 2015 多图谱腹部标记挑战赛中的 30 次腹部 CT 扫描，共计 3779 张轴向增强腹部临床 CT 图像。每个 CT 体积由 85 ∼ 198 个 512 × 512 像素的切片组成，体素空间分辨率为 ([0:54 ∼ 0:54] × [0:98 ∼ 0:98] × [2:5 ∼ 5:0])mm3。 我们报告了平均 DSC 和平均豪斯多夫距离 (HD)对8个腹部器官（主动脉、胆囊、脾脏、左肾、右肾、肝脏、胰腺、脾脏、胃）进行实验，随机分成18个训练案例（2212个轴向切片）和12个验证案例。

• 自动心脏诊断挑战赛。ACDC 挑战赛收集了来自 MRI 扫描仪的不同患者的检查结果。电影 MRI 图像是在屏住呼吸时获取的，一系列短轴切片覆盖了从左心室底部到心尖的心脏，切片厚度为 5 至 8 毫米。短轴平面内空间分辨率从 0.83 到 1.75 平方毫米/像素。

• 每位患者的扫描结果均手动标注左心室 (LV)、右心室 (RV) 和心肌 (MYO) 的真实值。我们报告了平均 DSC，其中随机分为 70 个训练案例（1930 个轴向切片）、10 个用于验证的案例和 20 个用于测试的案例。

Implementation Details

• 对于所有实验，我们都应用简单的数据增强，例如随机旋转和翻转。对于纯 Transformer 编码器，我们简单采用具有 12 个 Transformer 层的 ViT。对于混合编码器设计，我们将 ResNet-50 和 ViT 结合起来，在本文中表示为 “R50-ViT”。所有 Transformer 主干（即 ViT）和 ResNet-50（表示为 “R-50”）均在 ImageNet 上进行预训练。除非另有说明，否则输入分辨率和块大小 P 设置为 224×224 和 16。因此，我们需要在 CUP 中连续级联四个 2× 上采样块以达到全分辨率。并且对于模型使用 SGD 优化器进行训练，学习率为 0.01，动量为 0.9，权重衰减为 1e-4。默认批量大小为 24，默认训练迭代次数分别为 ACDC 数据集的 20k 和 Synapse 数据集的 14k。所有实验均使用单个 Nvidia RTX2080Ti GPU 进行。所有3D体积都以逐片的方式推断，并将预测的2D切片堆叠在一起，以重建3D预测以供评估。

Comparison with State-of-the-arts

• 我们通过将我们的 TransUNet 与四个先前的最先进技术进行比较，在 Synapse 多器官分割数据集上进行主要实验：1）V-Net ；2）DARR ；3）U-Net 和 4）AttnUNet 。

• 为了证明我们的 CUP 解码器的有效性，我们使用 ViT [4] 作为编码器，并分别比较使用简单上采样（“\None”）和 CUP 作为解码器的结果；为了证明我们的混合编码器设计的有效性，我们使用 CUP 作为解码器，并分别比较使用 ViT 和 R50ViT 作为编码器的结果。为了与 ViThybrid 基线（R50-ViT-CUP）和我们的 TransUNet 进行公平的比较，我们还用 ImageNet 预训练的 ResNet-50 替换了 U-Net 和 AttnUNet 的原始编码器。表 1 报告了 DSC 和平均 hausdorff 距离（以毫米为单位）的结果。

  • 

  • 表 1：Synapse 多器官 CT 数据集上的比较（平均骰子得分百分比和平均 hausdorff 距离（毫米）以及每个器官的骰子得分百分比）。

• 首先，我们可以看到，与 ViT-None 相比，ViT-CUP 在平均 DSC 和豪斯多夫距离方面分别提高了 6.36% 和 3.50 毫米。这一改进表明我们的 CUP 设计提供了比直接上采样更好的解码策略。同样，与 ViT-CUP 相比，R50-ViT-CUP 还在 DSC 中额外提高了 3.43%，在豪斯多夫距离中额外提高了 3.24 毫米，这证明了我们的混合编码器的有效性。我们的 TransUNet 以 R50-ViT-CUP 为基础，并配备了跳跃连接，在基于 Transformer 的不同模型变体中取得了最佳效果。

• 其次，表 1 还显示，与现有技术相比，我们所提出的 TransUNet 有显著的改进，例如，考虑到平均 DSC，性能提升范围为 1.91% 到 8.67%。特别是，直接将 Transformers 应用于多器官分割可以得到合理的结果（ViT-CUP 的 DSC 为 67.86%），但无法与 U-Net 或 attnUNet 的性能相媲美。这是因为 Transformers 可以很好地捕获有利于分类任务的高级语义，但缺乏用于分割医学图像精细形状的低级线索。另一方面，将 Transformers 与 CNN 相结合，即 R50-ViTCUP，表现优于 V-Net 和 DARR，但仍然比纯基于 CNN 的 R50-U-Net 和 R50-AttnUNet 效果差。最后，当通过跳跃连接与 U-Net 结构相结合时，所提出的 TransUNet 创下了新的最佳成绩，比 R50-ViT-CUP 和之前最好的 R50-AttnUNet 分别高出 6.19% 和 1.91%，表明 TransUNet 具有强大的学习高级语义特征和低级细节的能力，这在医学图像分割中至关重要。平均豪斯多夫距离也呈现出类似的趋势，这进一步证明了我们的 TransUNet 相对于这些基于 CNN 的方法的优势。

Analytical Study

• 为了彻底评估所提出的 TransUNet 框架并验证不同设置下的性能，进行了各种消融研究，包括：1）跳过连接的数量；2）输入分辨率；3）序列长度和补丁大小以及 4）模型缩放。

• 跳跃连接的数量。如上所述，集成 U-Netlike 跳跃连接有助于通过恢复低级空间信息来增强更精细的分割细节。此消融的目标是测试在 TransUNet 中添加不同数量的跳跃连接的影响。通过将跳跃连接的数量改为 0 (R50-ViT-CUP)/1/3，图 2 总结了所有 8 个测试器官的平均 DSC 分割性能。请注意，在“1-skip”设置中，我们仅在 1/4 分辨率比例下添加跳跃连接。我们可以看到，添加更多跳跃连接通常会带来更好的分割性能。通过在 CUP 的所有三个中间上采样步骤（输出层除外）中插入跳跃连接，即在 1/2、1/4 和 1/8 分辨率比例下（如图 1 所示），可实现最佳平均 DSC 和 HD。因此，我们为 TransUNet 采用了这种配置。还值得一提的是，较小器官（即主动脉、胆囊、肾脏、胰腺）的性能提升比较大器官（即肝脏、脾脏、胃）的性能提升更明显。这些结果强化了我们最初的将类似 U-Net 的跳过连接集成到 Transformer 设计中的直觉，能够学习精确的低级细节。

  • 

  • 图 2：TransUNet 中跳过连接数量的消融研究。

• 作为一项有趣的研究，我们在跳跃连接中应用了附加 Transformers，类似于 [Attention gated networks]，并发现这种新型跳跃连接甚至可以进一步提高分割性能。由于 GPU 内存限制，我们在 1/8 分辨率比例跳跃连接中使用了一个轻量级 Transformer，同时保持其他两个跳跃连接不变。结果，这一简单的改变使 DSC 的性能提升了 1.4%。

• 关于输入分辨率的影响。TransUNet 的默认输入分辨率为 224×224。这里，我们还提供了在高分辨率 512×512 上训练 TransUNet 的结果，如表 2 所示。当使用 512×512 作为输入时，我们保持相同的块大小（即 16），这会导致 Transformer 的序列长度增加约 5 倍。如 [ViT] 所示，增加有效序列长度可带来显著的改进。对于 TransUNet，将分辨率比例从 224×224 更改为 512×512 可使平均 DSC 提高 6.88%，但计算成本会大得多。因此，考虑到计算成本，本文中的所有实验比较均以默认分辨率 224×224 进行，以证明 TransUNet 的有效性。

  • 

  • 表 2：输入分辨率影响的消融研究。

• On the Influence of Patch Size/Sequence Length.我们还研究了补丁大小对 TransUNet 的影响。结果总结在表 3 中。我们观察到，补丁大小越小，分割性能通常越高。请注意，Transformer 的序列长度与补丁大小的平方成反比（例如，补丁大小 16 对应的序列长度为 196，而补丁大小 32 对应的序列长度较短，为 49），因此减小补丁大小（或增加有效序列长度）可带来稳健的改进，因为 Transformer 会针对较长的输入序列编码每个元素之间更复杂的依赖关系。按照 ViT 中的设置，我们在本文中使用 16×16 作为默认补丁大小。

  • 

  • 表 3：patch块大小和序列长度的消融研究。

• 模型缩放。最后但并非最不重要的一点是，我们对 TransUNet 的不同模型大小进行了消融研究。具体来说，我们研究了两种不同的 TransUNet 配置，即“Base”和“Large”模型。对于“base”模型，隐藏大小 D、层数、MLP 大小和头数分别设置为 12、768、3072 和 12，而“large”模型的这些超参数分别为 24、1024、4096 和 16。从表 4 中我们得出结论，更大的模型可带来更好的性能。考虑到计算成本，我们对所有实验都采用“Base”模型。

  • 

  • 表4：模型规模的消融研究。

Visualizations

• 我们在 Synapse 数据集上提供了定性比较结果，如图 3 所示。可以看出：1）纯基于 CNN 的方法 U-Net 和 AttnUNet 更容易对器官进行过度分割或分割不足（例如，在第二行中，脾脏被 AttnUNet 过度分割，而被 UNet 分割不足），这表明基于 Transformer 的模型，例如我们的 TransUNet 或 R50-ViT-CUP 具有更强的能力来编码全局上下文并区分语义。 2）第一行的结果表明，我们的 TransUNet 预测的假阳性比其他方法少，这表明 TransUNet 在抑制那些噪声预测方面比其他方法更有利。3）在基于 Transformer 的模型中进行比较，我们可以观察到 R50-ViT-CUP 在边界和形状方面的预测往往比 TransUNet 更粗糙（例如，第二行对胰腺的预测）。

• 此外，在第三行中，TransUNet 正确预测了左肾和右肾，而 R50-ViT-CUP 错误地填充了左肾的内孔。 这些观察结果表明，TransUNet 能够进行更精细的分割并保留详细的形状信息。原因是 TransUNet 既享有高级全局上下文信息，又享有低级细节的好处，而 R50-ViT-CUP 仅依赖于高级语义特征。这再次验证了我们最初的直觉，即将类似 U-Net 的跳跃连接集成到 Transformer 设计中以实现精确定位。

  • 

  • 图 3：通过可视化对不同方法进行定性比较。从左到右：(a) Ground Truth，(b) TransUNet，© R50-ViT-CUP，(d) R50AttnUNet，(e) R50-U-Net。我们的方法预测的假阳性更少，并保留了更精细的信息。

Generalization to Other Datasets

• 为了展示我们的 TransUNet 的泛化能力，我们进一步评估了其他成像模式，即旨在实现自动心脏分割的 MR 数据集 ACDC。我们观察到 TransUNet 相对于纯 CNN 方法（R50-UNet 和 R50-AttnUnet）和其他基于 Transformer 的基线（ViT-CUP 和 R50-ViT-CUP）有持续改进，这与 Synapse CT 数据集上的先前结果相似。

  • 

  • 表 5：DSC 中 ACDC 数据集的比较（%）。

Conclusion

• Transformer 被称为具有强大先天自注意力机制的架构。在本文中，我们首次研究了 Transformer 在一般医学图像分割中的应用。为了充分利用 Transformer 的强大功能，我们提出了 TransUNet，它不仅通过将图像特征视为序列来编码强大的全局上下文，而且还通过 U 形混合架构设计充分利用了低级 CNN 特征。作为主流基于 FCN 的医学图像分割方法的替代框架，TransUNet 的性能优于各种竞争方法，包括基于 CNN 的自注意力方法。
• TransUNet 架构主要处理 2D 图像，在合成分割输出之前独立管理每个图像切片。这种 2D 方法虽然有效，但可能会导致信息丢失——在肿瘤分割等应用中尤其明显，其中切片之间的连续性可以表明肿瘤的存在。为了解决这一限制，还可以考虑了一种 3D 处理方法，该方法处理子卷或补丁中的数据。这种方法不仅可以保持空间连续性，还可以有效地管理计算成本。3D TransUNet 。