UNETR:用于3D图像分割的Transformer

论文：UNETR: Transformers for 3D Medical Image
Segmentation
论文地址：https://arxiv.org/pdf/2103.10504.pdf
摘要
近年来，具有收缩路径和扩展路径（例如，编码器和解码器）的全卷积神经网络（FCNN）在各种医学图像分割应用中表现出突出的地位。在这些体系结构中，编码器通过学习全局上下文信息成为一个不可或缺的角色，而此过程中获取的全局上下文表示形式将被解码器进一步用于语义输出预测。尽管取得了成功，但作为FCNN的主要构建模块的卷积层的局限性，限制了在此类网络中学习远程空间相关性的能力。受自然语言处理（NLP）转换器在远程序列学习中的最新成功的启发，我们将体积（3D）医学图像分割的任务重新设计为序列到序列的预测问题。特别是，我们介绍了一种称为UNEt变压器（UNETR）的新颖架构，该架构利用纯变压器作为编码器来学习输入量的序列表示并有效地捕获全局多尺度信息。转换器编码器通过不同分辨率的跳跃连接直接连接到解码器，以计算最终的语义分段输出。我们已经使用医学分割十项全能（MSD）数据集广泛验证了我们提出的模型在不同成像方式（即MR和CT）上对体积脑肿瘤和脾脏分割任务的性能，并且我们的结果始终证明了良好的基准。

1 引言
医学图像分割在许多临床诊断方法中起着不可或缺的作用，并且通常是对解剖结构进行定量分析的第一步。 自从深度学习问世以来，FCNN，尤其是编码器-解码器体系结构[19,13,14,12]在各种医学语义分割任务[1,22,11]中已取得了最新的成果。 在典型的U-Net [21]体系结构中，

• 编码器负责通过逐渐降低采样的特征来学习全局上下文表示；
• 而解码器则将采样的表示上采样至输入分辨率，以进行像素/体素语义预测；
• 另外，跳跃连接以不同的分辨率合并了编码器和解码器的输出，因此可以恢复在下采样期间丢失的空间信息。
  尽管此类基于FCNN的方法具有强大的表示学习功能，但它们在学习远程依存关系方面的性能仅限于其局部接收域。结果，在捕获多尺度上下文信息中的这种缺陷导致具有可变形状和尺度（例如，具有不同大小的脑损伤）的结构的分割。通过采用无意识的卷积层[4,15,10]，已经做出了一些努力来缓解这个问题。但是，由于CNN的局限性，它们的接收场仍然局限于一个小区域。
  在NLP领域，基于变压器的模型[24,6]在各种任务中都达到了最新的基准。变压器中的自我注意机制使他们能够动态地突出显示单词序列的重要特征并了解其长期依赖性。最近，通过引入Visual Transformer（ViT）[7]，该概念已扩展到计算机视觉。在ViT中，图像表示为一系列补丁嵌入，这些补丁嵌入将用于直接预测类别标签以进行图像分类。
  在这项工作中，我们建议利用变压器进行体积医学图像分割，并为此目的引入一种被称为UNETR的新颖架构。特别是，我们将3D分割的任务重新设计为1D序列到序列的预测问题，并使用纯转换器作为编码器从嵌入的输入色块中学习上下文信息。从变压器编码器提取的表示通过多个分辨率的跳过连接与解码器合并，以预测分段输出。
  我们已经在MSD数据集中广泛验证了我们的UNETR对脑肿瘤和脾脏分割任务的有效性[22]，并且与我们的验证集中的其他模型相比，我们的实验证明了良好的性能。据我们所知，我们是第一个提出用于体积医学图像分割的完全基于变压器的编码器的公司。考虑到体积数据在医学成像中的盛行及其在分割中的广泛应用，我们认为我们的UNETR为可用于各种应用的新型基于变压器的分割模型铺平了道路。

2 相关工作
基于CNN的分割网络：自从开创性的U-Net [21]以来，基于CNN的网络已在各种2D和3D各种医学图像分割任务上取得了最新的成果[8,29,25,9 , 16,28]。 尽管取得了成功，但这些网络的局限性在于它们在学习全局上下文和长期空间依赖方面的表现不佳，这可能严重影响具有挑战性的任务的分割性能。
Visual Transformers：Visual Transformers最近在各种计算机视觉任务中获得了关注。 Dosovitskiy等。 [7]通过对纯变压器的大规模预训练和微调，在图像分类数据集上展示了最新的性能。在目标检测中，基于端到端变压器的模型在多个基准测试中显示出突出的地位[2,30]。最近，一些努力[27,3,23,26]已经探索了使用基于变压器的模型进行2D图像分割的可能性。 Chen等。 [3]通过在U-Net的瓶颈中采用变压器作为层，提出了一种用于多器官分割的2D方法。另外，Zhang等。 [26]建议在分开的流中使用CNN和变压器，并对它们的输出进行融合。瓦拉纳拉苏（Valanarasu）等。 [23]提出了一种基于变压器的轴向注意机制，用于2D医学图像分割。
我们的模型与这些工作之间存在三个主要区别：
（1）UNETR专为3D分割量身定制，并直接利用体积数据；
（2）UNETR使用变压器作为分段网络的主要编码器，并通过跳过连接将其直接连接到解码器，而不是将其用作分段网络中的关注层；
（3）UNETR不依赖于主干CNN来生成输入序列，而是直接利用标记化补丁。

3 方法论
3.1 架构
我们在图1中介绍了所提出模型的概述。UNETR利用收缩-扩展模式，该模式由一堆变压器组成，作为编码器，该编码器通过跳过连接与解码器连接。我们首先描述变压器编码器的工作机制。正如NLP中常用的那样，这些变压器以一维输入嵌入序列工作。在我们的UNETR中，通过将3D输入体积$x\in R^{H\times W\times D\times }$