MICCAI2022|A Robust Volumetric Transformer for Accurate 3D Tumor Segmentation

MICCAI2022|A Robust Volumetric Transformer for Accurate 3D Tumor Segmentation

用于精确 3D 肿瘤分割的鲁棒体积Transformer

0 摘要

我们提出了一种用于3D分割的Transformer架构，这是一项具有挑战性的任务，需要在编码局部和全局空间线索时保持复杂的平衡，并沿体积的所有轴保存信息。所提出的设计的编码器受益于同时编码局部和全局线索的自注意力机制，而解码器采用并行的自注意力和交叉注意力公式来捕获用于边界细化的精细细节。经验上，我们表明，所提出的设计选择产生了一个计算效率高的模型，在医学分割十项全能（MSD）脑肿瘤分割（BraTS）任务中具有竞争性和有前景的结果。我们进一步表明，通过我们的模型学习的表示对数据损坏是鲁棒的。

1 引言

受计算机视觉中基于Transformer的模型的强大经验结果、其有有前途的泛化能力和鲁棒性特征以及其对长距离交互建模的灵活性的启发，我们提出了一种用于3D医学图像模态（如MRI、CT）分割的体积Transformer架构，称为VT-UNet。早期为3D医学扫描开发基于Transformer的分割模型的努力已显示出优于最先进的CNN的表现。然而，这些方法将3D体积划分为2D切片，并将2D切片作为输入进行处理。因此，对于封装切片间至关重要的大量和潜在的关键体积信息依赖关系丢失。虽然一些混合方法（同时使用卷积块和Transformer）保持3D体积数据的完整，但在文献中尚未探索能够在输入时保持体积数据完整的纯粹基于Transformer的架构的设计。本文在纯Transformer方向设计了一个模型，不仅表现出很好的分割性能，而且对数据伪影有更好鲁棒性。

Transformer模型具有高度动态和灵活的感受野，能够捕捉长距离交互，但为3D分割设计基于Transformer的UNet架构仍然是一项具有挑战性的任务。这是因为：

1. 封装体素信息并捕获体积序列中任意位置之间的连接并不简单。与基于Transformer的2D图像分割方法相比，体积的每个切片中的数据都连接到三个视图，丢弃其中任何一个都是不利的。

2. 在三维中保存空间信息是一项艰巨的任务。即使对于2D图像，当将图像分解为块并将块投影为tokens时，如Vision Transformer（ViT）中所介绍的，局部结构线索也可能丢失，如Tokens-to-token ViT中所示。因此，有效地编码局部线索，同时沿着体积的多个轴捕捉全局交互是一项具有挑战性的任务。

3. 由于自我注意力的平方复杂性和3D体积张量输入的大尺寸，设计基于Transformer的分割模型（计算效率高）需要仔细的设计考虑。

我们提出的VT-UNet模型通过提出多个模块有效地解决了上述设计挑战。在基于UNet的架构中，我们开发了两种类型的Transformer块。

首先，我们的编码器中的块以分层方式直接处理3D体积，以联合捕获局部和全局信息，实质上类似于Swin Transformer块。

其次，对于解码器，我们在扩展路径中引入了并行交叉注意力和自注意力，这在来自解码器的query和来自编码器的key和value之间建立了桥梁。

通过交叉注意力和自注意力的并行化，我们的目标是在解码过程中保持完整的全局上下文，这对于分割任务很重要。由于VT-UNet没有卷积，并且在解码过程中结合了来自两个模块的注意力输出，因此序列的顺序对于获得准确的预测非常重要。受Transformer启发，除了在计算每个Transformer块中的注意力时应用相对位置编码之外，我们增强了解码过程，并注入了从序列中token的融合特征位置提取的补充信息。

贡献：

1. 我们从序列到序列的角度重新制定了3D肿瘤分割，并提出了用于多模态医学图像分割的UNet形体积Transformer。

2. 我们设计了一个具有两个连续的自注意力层的编码器块，以联合捕获局部和全局上下文线索。此外，我们设计了一个解码器块，该解码器块支持基于并行（移位）窗口的自我和交叉注意力。这种并行化使用一个共享的query的projection，并独立地计算交叉和自注意力。为了进一步增强解码中的特征，我们提出了一种凸组合方法以及融合特征位置编码。

3. 结合我们提出的设计选择，我们实质上限制了模型参数，同时与现有方法相比，保持了较低的FLOP（见图1）。

4. 我们进行了广泛的评估，并表明我们的设计实现了SOTA的体积分割结果，同时增强了对数据伪影的鲁棒性。

2 方法

a图 简单解释VT-UNet架构。这里，k表示类的数量。

b图 显示了体积移动窗口的可视化。为了说明，考虑尺寸为D×H×W的MRI体积，D=H=W=8。此外，假设用于分割体积的窗口大小为P×M×M，P＝M＝4。这里，层l在体积Transformer（VT）块的第一步中采用常规窗口分割，这导致2×2×2＝8个窗口。在层l+1内，体积窗口移动（P/2，M/2，M/2）=（2，2，2）tokens。这导致3×3×3=27个窗口。

e图 说明融合模块的结构。

我们分别用粗体小写x和粗体大写X表示向量和矩阵。设X={ x1，x2，··，Xτ}，xi∈ RC是表示感兴趣信号（例如MRI体积）的序列。我们称每个xi为token。我们假设token的原始形式可能不是定义跨度的最佳方式。因此，在自注意力（SA）中，我们通过从输入token学习线性映射来定义跨度。

其中其中 B ∈Rτ×τ是可训练的，作为在体积中的一个相对位置偏差，V = XWV，K = XWK，Q = XWQ。在实践中，多次并行计算多个注意力头的SA被称为多头自注意力（MSA）。公式(1)是我们的体积TRM窗口的多头自我注意（VT-W-MSA）的基本块。
基于体积TRM移位窗口的多头自注意力（VTSW-MSA），接下来讨论。

VT-UNet概述。图2显示了提出的体积Transformer网络（简称VT-UNet）的概念图。

我们模型的输入是一个尺寸为D×H×W×C的3D体积。输出是一个D×H x W×K维体积，表示体素级别类标签的存在/不存在（K是类的数量）。

下面，我们将讨论VTUNet模块的架构形式，并详细解释设计背后的功能和原理。

2.1. The VT Encoder

VT编码器由3D patch分割层、线性嵌入层和3D patch 合并层以及两个连续VT编码器块组成。

3D Patch 切分。基于Transformer的模型使用一系列tokens。VT-UNet的第一个块接受D×H×W×C维医学体积（例如MRI），并通过将3D体积分割成非重叠的3D块来创建一组tokens（见图b）。分区核的大小为P×M×M，导致用
tokens。3Dpadding分割之后是线性嵌入，以将维度为P×M×M的每个tokens映射到C维向量。根据我们的实验，M、P和C的典型值分别为4、4和72。

VT编码器块。在ViT中，tokens由于其构造方式而携带重要的空间信息。最近的几项研究表明，通过在ViT中窗口来执行SA的重要性，最著名的是Swin Transformer。在Swin Transformer的设计中遵循类似的原则，尽管是针对体积数据，我们建议在VT 编码器块（VT Enc Blks）中进行3D窗口操作。特别地，我们提出了两种类型的窗口，即常规窗口和移位窗口，为了简单起见，我们分别通过VT-W-MSA和VT-SW-MSA显示了这两种窗口。图2b提供了VT-W-MSA和VT-SW-MSA的设计细节，而图2（b）解释了窗口操作。

VT-W-MSA和VT-SW-MSA都采用了带有窗口的注意层，然后是两层多层感知器（MLP），其间具有高斯误差线性单元（GELU）添加非线性因素。在每个MSA和MLP之前应用层标准化（LN），在每个模块之后应用残差连接。窗口化使我们能够在建模tokens之间的长距离依赖关系时注入归纳偏置。在VT-W-MSA和VT-SW-MSA中，对窗口内tokens的关注有助于表示学习。在VT-W-MSA中，我们将体积均匀地划分为较小的非重叠窗口，如图2（b）所示。由于VT-W-MSA中相邻窗口中的token无法彼此看到，因此我们使用了VT-SW-MSA中的移位窗口（参见最右侧的面板图2（b）），该窗口桥接了VT-W-MSA中相邻窗口的tokens。窗口化受到Swin Transformer的启发，可以理解为对体积数据的概括。注意，我们工作中的窗口操作类似于Video Swin Transformer，它将窗口的好处从图像扩展到视频。VT Enc Blk集所有功能于一体，实现了以下功能：

3D Patch 合并。 我们使用3D Patch 合并块在VT-UNet的编码器中生成特征层次。具有这样的层次结构对于在密集预测任务的输出中生成更精细的细节至关重要。

在每个VT-Enc-Blk之后，我们以非重叠的方式沿空间轴合并相邻的tokens以产生新的tokens。在这样做时，我们首先连接每组2×2个相邻tokens的特征。生成的向量通过线性映射投影到tokens的通道维度加倍的空间（见图2）。Patch合并的好处不限于功能层次结构。SA的计算复杂性是tokens数量的平方。因此，在每个VT-Enc-Blk之后，Patch合并将VT -UNet的FLOP计数减少16倍。给读者一个更好的想法，正如我们将在章节4中讨论的那样，与完全体积的CNN模型相比，小型VT UNet模型仅使用6.7%的FLOP，同时实现了类似的性能（确实稍好）！请注意，Patch合并块未在瓶颈阶段使用。

2.2.VT Decoder。

在经过VT-Enc-Blk和3D Patch 扩展层组成的瓶颈层之后，VT解码器开始于连续的VT解码器块（VT-Dec-Blk）、3D补丁扩展和分类器最后生成最终预测。VT-Enc-Blk和VT-Dec-Blk之间存在一些基本设计差异，我们将在接下来讨论。

3D Patch Expanding。

此功能用于以某种方式恢复patch合并的效果。换句话说，为了以与输入相同的空间分辨率构造输出，我们需要在解码器中创建新的tokens。

为了便于讨论，考虑瓶颈层之后的Patch Expanding（参见图2的中间部分）。Patch Expanding的输入tokens具有维度8C。在Patch Expanding中，我们首先使用线性映射将输入tokens的维数增加2倍。reshaping之后，我们可以从得到的维度为2×8C的向量中获得维度为4C的2×2标记。这样，我们将沿着空间轴reshaping，因此对于D/4×H/32×W/32×8C，我们创建D/4×H/16×W/16×4C tokens。

VT解码器块。UNet及其变体利用编码器和解码器之间的横向连接来产生精细的详细预测。这是因为当输入通过编码器时，空间信息会丢失，代价是获得更高层次的语义。

UNet中的横向连接使得可以同时获得两个部分的最佳结果，即来自下层的空间信息和来自上层的语义信息（沿着计算图）。考虑到这一点，我们在解码器端提出了SA的混合形式（见图2b）。每个VT Dec Blk从位于VT UNet相同阶段的VT Enc Blk接收其先前VT Dec Blk生成的tokens以及Key和Value 的tokens，见图2a。回想一下，VT Enc Blk有两个SA块，具有常规和移位的窗口操作。VT Dec Blk具有类似的窗口操作，但使用了四个SA块，这些SA块分为SA模块和交叉注意（CA）模块。功能可描述为：

这里，r 和 l 表示解码器模块的左右分支。 SA 的右分支根据等式3作用于先前 VT-Dec-Blk 生成的tokens。 我们通过其中的下标 D强调来自解码器的信息流。然而，CA 的左分支使用解码器生成的query以及从计算图中同一级别的VTEnc-Blk 获得的key和value。这里的想法是使用编码器跨越的思想（由value标识）以及key来从编码器收集的空间信息中获益。这些块还使用常规和移位窗口向模型中注入更多归纳偏差。 请注意，来自具有相同窗口操作的 SA 的values和keys应该组合在一起，因此形成图 2 © 中的交叉连接形式。

备注 1. 有人可能会问为什么要从编码器考虑values和keys。

我们确实研究了其他可能性，例如在解码器生成values的同时使用来自编码器的querys和keys。根据经验，在等式中(3) 描述的形式被观察到提供更好和更稳健的结果，因此我们在 VT-UNet选择该方式

融合模块。 如图 2 (e) 所示，从CA 模块和MSA模块生成的tokens被组合在一起并馈送到下一个 VT-DecBlk，zl使用线性函数计算为：

其中 F(·) 表示傅里叶特征位置编码 (FPE)，α 控制每个 CA 和 MSA 模块的贡献。为了简单起见，在融合由 CA 和 MSA 生成的tokens时，我们使用 α = 0.5 的线性组合。

分类器层。 在解码器中的最终 3D Patch扩展层之后，我们引入了一个分类器层，其中包括一个3D卷积层，用于将深C维特征映射到K分割类。

关于计算复杂性的注释。 式(1)中描述的SA的计算复杂度是由获得Q、K、V，计算QK，以及通过将Softmax的输出（是一个τ×τ矩阵）应用于V而获得所产生的tokens所需要的计算决定的。这加起来为 O（3τC² + 2τ²C)，其中 C 和 τ 分别是维数和token数。根据O(3τC²+ 2τκC)，窗口将减少 SA 的计算负荷。我们假设tokens被分组到 κ 个窗口中，并且在每个窗口中应用 SA。在我们的问题中，tokens是从体积数据 τ 生成，τ远大于 κ ，因此窗口不仅有助于具有更好的辨别能力，而且还有助于减少计算负荷。

3 相关工作

Vision Transformers 在不同的计算机视觉任务 [ViViT,Video Swin Transformer,Training data-efficient image transformers & distillation through attention] 中显示出卓越的实证结果，具有良好的特性。例如，与 CNN 相比，它们更不偏向于纹理 [ Intriguing properties of vision transformers]，并且表现出更好的泛化和鲁棒性。最近还研究了用于图像分割的Transformer[ Swin-Unet,TransUNet,Rethinking semantic segmentation from a sequence-to-sequence perspective with transformers(CVPR)]。 TransUNet是第一个基于 Transformer 的医学图像分割方法。它采用 UNet 结构，并用 ViT替换瓶颈层，其中Patch嵌入应用于从 CNN 编码器生成的特征图（其中输入是 3D 体积的 2D 切片）。与这些混合方法（同时使用卷积和自注意力）不同，Cao 等人 提出了 Swin-UNet，这是一种用于医学图像分割的纯基于transformers的网络。它继承了 swin-transformer 块，并显示出比 TransUNet更好的分割结果。 TransUnet 的 3D 版本称为 TransBTS，具有 CNN 编码器-解码器设计和作为瓶颈层的 Transformer。 Zhou 等人 提出了 nnFormer，其具有基于 3D Swin Transformer 的块作为具有交错卷积干的编码器和解码器。 UNETR中提出了一种使用Transformer作为编码器并通过跳过连接将中间编码器输出直接连接到解码器的模型。 SOTA方法的encoder-decoder结构比较如图2（d）所示。

d图 编码器-解码器结构与其他SOTA方法的比较。所提出的VT-UNet架构没有卷积模块，纯粹基于Transformer块。

上述基于 transformer 的 3D 医学图像分割方法已经显示出它们的前景，与 CNN 对应方法相比，它们具有更好的性能。另一方面，我们提出的模型整体处理体积数据，从而完全编码了切片之间的相互作用。此外，我们提出的模型完全基于 Transformers 构建，并在编码器-解码器设计中引入横向连接来执行 CA 和 SA。这些设计元素有助于实现更好的分割性能以及增强的稳健性。

4 实验

实现细节。 我们使用来自 MSD BraTS 任务的 484 次 MRI 扫描。我们将 484 次扫描分为 80%、15% 和 5%，分别用于训练、验证和测试集。我们使用带有单个 Nvidia A40 GPU 的 PyTorch [Paszke, A., et al.: Automatic differentiation in PyTorch (2017)]。在 ImageNet-22K 上预训练的 Swin-T的权重用于初始化模型。对于训练，我们使用 AdamW 优化器，学习率为 1e−4 1000 个时期，批量大小为 4。我们使用旋转、添加噪声、模糊和添加伽玛作为数据增强技术。

实验结果。 表 1 将 VT-UNet 与最近基于transformer的方法和基于 SOTA CNN 的方法进行了比较。我们通过改变用于模型训练的嵌入式维度的数量来引入 VTUNet 的变体。我们的变体是：

(a) Small VT-UNet-S ：C = 48

(b) Base VT-UNet-B ：C = 72。

我们使用 Dice Srensen系数 (DSC) 和 Hausdorff 距离 (HD) 作为评估指标，并分别为三类计算它们：(1) 增强肿瘤 (ET)，(2) 肿瘤核心 (TC)（添加 ET、NET 和 NCR），以及 (3) 全肿瘤 (WT)（添加 ED 到 TC )，遵循与 [UNETR,nnU-Net,nnFormer] 中类似的评估策略。我们在表 1 中的定量结果表明 VT-UNet 在 DSC 和 HD 中实现了最佳的整体性能。

表 4 显示了未见患者数据的定性分割结果。

我们可以观察到我们的模型可以准确地分割结构并描绘肿瘤的边界。我们相信，捕获相邻切片之间的远程依赖关系对我们模型的性能起着至关重要的作用。我们在表 3 中的实证结果揭示了在 VT-DecBlks 中引入并行 CA 和 SA 以及 FPE 以及凸组合的重要性。我们可以注意到所有这些组件都有助于模型的性能。

鲁棒性分析。 患者的运动和采集条件等因素会给 MRI 带来噪声。在这里，我们调查了VT-UNet 的稳健性，通过推理时将人为噪声引入MR图像。

这些包括 (1) 运动伪影。 (2) 重影。 (3) 尖峰噪声(人字形 噪声)。表 2 比较了我们的方法与 nnFormer的稳健性。结果表明，在存在这些干扰的情况下，VT-UNet 的性能更可靠。我们的发现与 RGB 图像的现有工作一致，其中基于transformer 的模型由于其高度动态和灵活的感受野对遮挡、自然和对抗性扰动已经显示出更好的鲁棒性。

5 结论

本文介绍了一种用于医学图像分割的体积变换网络，它在处理大型 3D 体积方面计算效率高，并学习了对人为噪声具有鲁棒性的表示。我们的结果表明，所提出的模型在体积分割方面比现有的最先进方法取得了一致的改进。我们相信我们的工作可以帮助更好的临床诊断和治疗计划。

Summary

核心：对伪影或者噪声是鲁棒的、架构由纯Transformer构成，编码器能够对全局和局部的信息进行提取，利用3D 架构，能够建模空间信息，具有并行的自注意力以及交叉注意力的解码器对边界细化，同时滑窗能够降低复杂度加快效率