VT-UNet——基于transformer的医学3D分割网络

文章提出了一种用于体积分割的Transformer架构，这是一项具有挑战性的任务，需要在局部和全局空间编码上保持复杂平衡，以及沿体积的所有轴保留信息。所提出设计的编码器受益于自我注意机制以同时编码局部和全局特征，而解码器采用并行自我和交叉注意公式来捕获精细细节以进行边界细化。提出的设计是一个计算效率高的模型，在医学分割十项全能 (MSD) 脑肿瘤分割 (BraTS) 任务上具有竞争力和有希望的结果。进一步表明，对数据损坏的情况我们的模型学习到的表示是具有鲁棒性的。

一、模型结构

VT-UNet结构如上图所示，模型的输入是大小为D×H×W×C的3D体积，输出是一个D×H×W×K的3D体积，表示体素级类标签。

VT编码器。VT编码器由 3D补丁分区层和线性嵌入层组成，3D补丁合并层后跟着两个连续的VT编码器块。

a、3D补丁分区。基于Transformer的模型使用一系列标记。VT-UNet的第一个块接受D × H × W × C维医学体数据（例如 MRI），并通过将3D体积拆分为不重叠的3D块来创建一组标记（见图 b)。分区内核的大小为 P × M × M，因此通过τ = [D/P]×[H/M]×[W/M]个标记来描述体积。3D补丁分区之后是线性嵌入，以将维度为P × M × M的每个标记映射到C维向量。根据我们的实验，M、P 和 C 的典型值分别为 4、4和72。

b、VT 编码器块。在 ViT 中，标记由于其构造方式而携带重要的空间信息。遵循 Swin Transformers设计中的类似原则，尽管是针对体积数据，建议在VT编码器块 (VT-Enc-Blks)中进行3D窗操作。特别是，提出了两种窗类型，即常规窗和移位窗，为简单起见，分别用 VTW-MSA 和 VT-SW-MSA 表示。图b提供了VT-W-MSA和VT-SW-MSA的设计细节，而图b 说明了窗操作。VT-W-MSA和VT-SW-MSA的基本构建模块是由自注意力（SA）来构建的，SA计算公式如下所示。

VT-W-MSA和VT-SW-MSA都采用带窗的注意层，然后跟着2层多层感知器 (MLP)，其中间使用非线性的高斯误差线性单元 (GELU)。在每个MSA和MLP之前应用层规范化 (LN)，并在每个模块之后应用残差连接。在VT-W-MSA中，将体积均匀地分成更小的非重叠窗口，如图 b所示。由于相邻窗口中的标记无法通过VT-W-MSA 相互看到，使用VT-SW-MSA中的移位窗口（参见图b最右侧），其可以桥接 VT-W-MSA 相邻窗口中的标记。将所有模块结合起来，VT-Enc-Blk实现了以下功能，其中z^l和zl分别表示块l的VT-W-MSA模块和MLP模块的输出特征。

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

VT解码器。网络瓶颈层由VT-Enc-Blk和3D补丁扩展层组成，VT解码器从连续的VT解码器块(VT-Dec-Blks)、3D补丁扩展层和以产生最终预测的分类器组成。

a、3D补丁扩展。此功能用于以某种方式恢复补丁合并的效果。为了构建与输入具有相同空间分辨率的输出，需要在解码器中创建新的标记。瓶颈层补丁扩展的输入标记的维度为8C。在补丁扩展中，首先使用线性映射将输入标记的维数增加两倍。在重塑之后，可以从维数为2×8C的结果向量中获得维数为4C的2×2个标记。再沿空间轴重塑，对于D/4×H/32×W/32×8C，创建D/4×H/16×W/16×4C标记。

b、VT 解码器块。为了融合来自较低层的空间信息和来自较高层的语义信息，解码器端采用SA混合形式（参见图c的说明）。每个VT-Dec-Blk 接收其前面VT-Dec-Blk生成的标记和来自位于VT-UNet同一级的VT-Enc-Blk 的键（KE）和值（VE）标记，见图a。VT-Enc-Blk有两个SA块，具有常规和移位窗操作。VT-Dec-Blk享有类似的窗操作，但使用分为SA模块和交叉注意力（CA）模块的四个SA块，可以描述为：

这里，r和l表示解码器模块的左右分支。SA的右分支根据等式3作用于先前VT-Dec-Blk 生成的标记。通过其中的下标D来强调来自解码器的信息流。然而，CA的左分支使用解码器生成的查询以及从计算图中同一级别的VT-Enc-Blk获得的键和值。这里的想法是使用编码器跨越的基础（由值来标识）以及键来获取编码器中收集的空间信息。这些块也使用常规和移位窗向模型中注入更多的归纳偏差。来自具有相同窗操作的SA的值和键应该合并，例如图c中的交叉连接形式。文中作者表明公式3中的SA模块和CA模块可以观察到更好、更稳健的结果，推测这是由于在反向传播期间从解码器到编码器有额外的连接，这可能促进梯度传输。

c、融合模块。如图e所示，从CA模块和MSA模块生成的标记组合在一起并馈送到下一个VT-Dec-Blk，zl使用线性函数计算为：

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

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

二、实验细节

使用MSD BraTS数据集。将484例MRI分为 80%、15% 和 5%，分别用于训练、验证和测试集。使用单卡Nvidia A40（48G显存）GPU和PyTorch框架。在ImageNet-22K上预训练的Swin-T的权重用于初始化模型。对于训练，使用AdamW优化器，学习率为1e−4 ，1000个epoch，批量大小为 4。使用旋转、添加噪声、模糊和添加伽玛作为数据增强技术。

三、结论

(1)、从序列到序列的角度重新制定了体分割，并提出了一种用于多模态医学图像分割的UNet形状体Transformer。

(2)、设计了一个编码器块，其中包含两个连续的自注意力层来共同捕获局部和全局上下文信息。此外，设计了一个解码器块，它支持基于并行（移位）窗口的自我和交叉注意力。这种并行化使用查询的一个共享投影并独立计算交叉和自我注意力。为了进一步增强在解码中的功能，提出了一种傅立叶位置编码的凸组合方法。

(3)、结合提出的设计选择，大大限制了模型参数，同时与现有方法相比保持较低的FLOP。

(4)、进行了广泛的评估，表明模型实现了最先进的体分割结果，同时也增强了对数据伪影的鲁棒性。