3D UX-Net

ICLR2023

摘要

1. 大多数医学图像如 MRI 和 CT 是属于 volumetric data 类型,为了更加充分的利用体素信息，近几年已经提出了不少 3D CNNs 的模型，如 SwinUNETR、UNETR, UNETR++等. 整体来说，这些ViTs模型性能是越来越高，在几个主流的 3D 数据基准测试中也实现了大大小小的 SOTA，特别是 3D 医学图像分割这块。

2. 分层Transformer例如Swin Transformer)重新引入了几个ConvNet的先验，并进一步增强了在3D医疗数据集中适应体积分割的实际可行性。混合方法的有效性主要归功于对非局部自注意的大接受域和大量的模型参数。

3. 提出了一种轻量级的体积ConvNet，称为3D UX-Net，其使用 ConvNet 模块调整分层 Transformer 以实现稳健的体素分割。具体来说，我们重新讨论了具有大内核(LK)大小的体积深度卷积(例如从7 × 7 × 7开始)，以实现更大的全局接受域，灵感来自Swin Transformer。我们进一步用点深度卷积代替Swin Transformer块中的多层感知器(MLP)，并通过更少的归一化和激活层来增强模型性能，从而减少模型参数的数量。

4. 在三个具有挑战性的脑体积和腹部成像公共数据集与当前的 SOTA 模型如 SwinUNETR 对比，同时在以下三大主流数据集性能均达到了最优：

   • MICCAI Challenge 2021 FLARE
   • MICCAI Challenge 2021 FeTA
   • MICCAI Challenge 2022 AMOS

   具体的，相比于 SwinUNETR，3D UX-Net 将 Dice 从 0.929 提高到 0.938 (FLARE2021)，Dice 从 0.867 提高到 0.874 (Feta2021)。此外，为了进一步评估 3D UX-Net 的迁移学习能力 ，作者在 AMOS2022 数据集上训练得到的模型在另一个数据集上取得了 2.27% Dice 的提升(0.880 → 0.900)。

1 引言

1. 首先是先来看看最近提出的一众基于 Transformer 架构的 ViT 模型有什么优劣势。

​ ViT（Vision Transformer）的特征是缺乏图像特定的归纳偏差和可扩展性行为，这些特征会受到大型模型容量和数据集大小的增强。

​ 当涉及到3D数据（如医学图像）时，由于其更高的维度和更大的输入分辨率，传统的2D ConvNet很难应对。最近，基于LK卷积设计的深度卷积技术的进展为2D ConvNet的大感受野提供了可扩展的计算机制。受到这些设计的启发，该研究重新审视了3D体积ConvNet的设计，探索了通过纯ConvNet架构实现SOTA性能的可行性，并比较了与3D ViTs的网络复杂度的差异，提出了一种新的设计3D ConvNet在体积高分辨率任务上的方向。

​ 该研究提出了3D UX-Net，它是一种轻量级的3D ConvNet，旨在通过ConvNet模块来适应Swin Transformer的内在特性，并使用更小的模型容量提高体积分割性能。具体来说，它引入了LK大小的体积深度卷积来模拟Swin Transformer中生成自注意力的大感受野操作。此外，它还引入了点积深度卷积缩放，将每个通道特征独立地分配到更宽的隐藏维度（例如4×输入通道），从而最小化跨通道学习的上下文的冗余，并保持模型性能而不增加模型容量。

​ 3D ViT和3D ConvNet之间的差异，，3D ViT的关键组件是其可扩展性和具有大感受野的全局自注意力，这些特征与卷积有关。由于传统的3D ConvNet无法实现这样的操作，因此提出了层次化Transformer来弥合这些差距。

2. 以 SwinUNETR 为例，其将 ConvNet 的一些先验引入到了 Swin Transformer 分层架构中，进一步增强了在 3D 医学数据集中调整体素分割的实际可行性。

探究通过纯ConvNet架构的（1）实现SOTA性能的可行性，（2）与3D ViTs相比提供更少的网络复杂性（3）提供设计用于体积高分辨率任务的3D ConvNet的新方向。

​ 提出了一种轻量级的体积ConvNet 3D UX-Net，以适应Swin Transformer的内在特性，并使用较小的模型容量增强体积分割性能。具体而言，我们引入具有LK大小的体积深度卷积来模拟生成Swin Transformer中的自注意力操作的大感受野。此外，我们进一步引入点卷积缩放来将每个通道特征独立分布到更宽的隐藏维度中（例如4×输入通道），从而最小化跨通道学习的上下文冗余，保持模型性能而不增加模型容量。

此类 “Conv+Transformer” 的组合拳的有效性在很大程度上归功于以下两个因素：

• 非局部自注意力所带来的大感受野
• 大量的模型参数

3. 本文贡献：

   (1) 提出了3D UX-Net，只使用ConvNet模块来适应Swin Transformer的内在特性.

   (2) 利用具有LK大小的深度卷积作为通用特征提取骨干，并引入点卷积来有效地缩放提取的表示并减少参数。

   (3) 使用三个具有挑战性的公共数据集在直接训练和微调场景中评估3D UX-Net，用于体积多器官/组织分割。 3D UX-Net在两种情况下的所有ConvNets和Transformer SOTA中都持续改善，同时具有较少的模型参数。

2 相关工作

2.1 基于Transformer的分割

​ 许多工作将Swin Transformer和Unet进行结合进行医学图像的3D分割，效果良好，引出否可以使用ConvNets来模拟Transformer行为，以展示高效的特征提取。

2.2 基于深度卷积的分割

（1）深度卷积被证明是标准卷积的一个强大变体，有助于减少参数数量和迁移学习 。相关发展：深度卷积来增强特征，然后在类似UNet的架构中融合解码特征；3D U2-Net利用深度卷积作为领域适配器，提取每个通道的特定于领域的特征 。证明了使用深度卷积增强容积任务的可行性。

（2）LK卷积在医学图像分割中的有效性: LK（7x7）卷积层作为跳跃连接，以解决脾脏增大脾脏分割的解剖的缺陷；将LK和扩张深度卷积用于解码器的容积分割; 然而，LK卷积的使用导致FLOPs显着增加，严重降低了训练和推断效率。为了提高使用LK的模型效率; ConvNeXt作为2D通用骨干网，使用LK深度卷积模拟ViTs的优势，用于自然图像的下游任务, ConvUNeXt则被提出扩展为2D医学图像分割 。通过LK深度卷积带来的大感受野，假设LK深度卷积可以潜在地模拟Transformer的行为，并为分割提供高效的好处。

3 3D UX-Net模型

​ 3D UX-Net 核心理念是设计出一种简单、高效和轻量化的 网络，其适用于 hierarchical transformers 的能力同时保留使用 ConvNet 模块的优势，如归纳偏置。具体地，其编码器模块设计的基本思想可分为: 1) block-wise(微观层面) 及 2) layer-wise(宏观层面)【可以类比下 ConvNeXt】。

block-wise部分

首先从下列三个不同的视角进行讨论:

(1) Patch-wise Features Projection(基于patch的特征投影): 对比 ConvNets 和 ViTs 之间的相似性，两个网络都使用一个共同的基础模块将patch特征表示缩小到特定的尺寸。以往的方法大都没有将图像块展平为具有线性层的顺序输入，因此作者采用具备大卷积核的投影层来提取 patch-wise 特征作为编码器的输入。

(2) Volumetric Depth-wise Convolution with LKs(体积深度卷积与LK): Swin transformer 的内在特性之一是用于计算非局部 MSA 的滑动窗口策略。总的来说，有两种分层方法来计算 MSA：基于窗口的 MSA (W-MSA) 和 移动窗口 MSA (SW-MSA)。这两种方式都生成了跨层的全局感受野，并进一步细化了非重叠窗口之间的特征对应关系。受深度卷积思想的启发，作者发现自注意力中的加权和方法与每通道卷积基础之间的相似之处，其认为使用 LK 大小的深度卷积可以在提取类似于 MSA 块的特征时提供大的感受野。因此，本文建议采用 LK 大小（例如，从 7 × 7 × 7 开始）通过体素深度卷积压缩 Swin transformer 的窗口移动特性。如此一来便可以保证每个卷积核通道与相应的输入通道进行卷积运算，使得输出特征与输入具备相同的通道维度。

(3) Inverted Bottleneck with Depthwise Convolutional Scaling(带深度卷积扩展比例的反向瓶颈): Swin transformer 的另一个固有结构是，它们被设计为 MLP 块的隐藏层维度比输入维度宽四倍，如下图所示。有趣的是，这种设计与 ResNet 块中的扩张率相关。因此，我们利用 ResNet 块中的类似设计并向上移动深度卷积来计算特征。此外，通过引入了具有 1 × 1 × 1 卷积核大小的深度卷积缩放(DCS)，以独立地线性缩放每个通道特征。通过独立扩展和压缩每个通道来丰富特征表示，可以最小化跨通道上下文产生的冗余信息，同时在每个阶段增强了与下采样块的跨通道特征对应。最后，通过使用 DCS，可以进一步将模型复杂度降低 5%，并展示了与使用 MLP 模型的架构相当的结果。

​ 图1：我们提出的设计卷积块以模拟Swin Transformer的行为概述。我们利用深度卷积和点乘缩放来适应大感受野，并通过扩展独立通道来丰富特征。我们进一步比较了不同的体积卷积神经网络和Swin Transformer块架构。黄色虚线展示了网络瓶颈中扩展特征通道的空间位置差异。

layer-wise

宏观设计中的卷积块展示了与swin transformer相比提取特征的大感受野和类似操作的可能性。进一步研究ConvNets和swin transformer在层次设置上的差异，并优化模型架构以更好地模拟ViTs在宏观层面上的特征。

还是从三个视角来看：

（1）Applying Residual Connections(应用残差连接)：从下图左上角可以看到，标准的 3D U-Net 模型内嵌的模块为 2 个 3 x 3 x 3 卷积的堆叠，其展示了使用小卷积核提取具有增加通道的局部表示的朴素方法；而其右手边的 SegResNet 则应用了类似 3D 版本的瓶颈层，先降维再升维最后再接残差表示；紧接着左下角的便是 Swin Transformer，其基于窗口注意力+MLP层的组合；最后右下角便是本文所提出的模块，其在最后一个缩放层之后应用输入和提取特征之间的残差连接。此外，在残差求和前后并没有应用到任何的归一化层和激活层。

（2）Adapting Layer Normalization(调整层归一化)：我们知道，在卷积神经网络中，BN 是一种常用策略，它对卷积表示进行归一化以增强收敛性并减少过拟合。然而，之前的工作已经证明 BN 会对模型的泛化能力产生不利影响。因此作者这里跟 ConvNeXt 一致，将 BN 替换为 LN。

（3）Using GELU as the Activation Layer(使用GELU作为激活层)：ReLU 是个好东西，几乎是现代 CNNs 模型的首选激活函数。作者在这里提倡使用 GELU，这是一种基于高斯误差的线性变换单元，相对 ReLU 更加平滑，也是其中一种变体，解决 ReLU 因负梯度被硬截断而导致的神经元失活问题。

4 3D UX-Net方法

​ 可以看出，3D UX-Net整体的设计思路非常简洁，框架也是一目了然，就一个标准的 3D U-Net 架构，由编码器-解码器组成，同时结合长跳跃连接操作进一步利用来将多分辨率特征连接到基于卷积的解码器网络，帮助网络更好的恢复空间细节的定位。图2展示了3D UX-Net的完整架构。我们在本节中进一步描述编码器和解码器的细节。

4.1 Depth-wise convolution encoder(深度卷积编码器)

​ 首先，输入部分应用一个大卷积核的 Projection 层将原始图像映射到一个低维空间分辨率的潜在空间特征表示，一来降低显存参数量计算量等，二来可以增大网络感受野，顺便再处理成编码器输入所需的格式，总之好处多多，不过都是基操。

处理完之后就输入到编码器中进行主要的特征提取，该编码器共4个stage，也是标准的16倍下采样。每个stage由多个不同的 3D UX-Net Block 构成，具体长看下左图。主要的特点就是主打轻量化和大感受野。

细节：

采用LK卷积层来计算分割特征映射，其大小为H/2 × W/2 × D/2，并将其投影到C = 48维空间中，而不是将补丁压平,并使用线性层进行投影。为了适应计算局部自注意力的特点，我们使用深度卷积（DWC）作为卷积核大小，从7×7×7开始，填充为3，作为“移位窗口”，均匀划分特征图。由于在Swin Transformer中提取大量补丁的全局自注意力通常不可行，我们假设使用LK大小的深度卷积可以有效地提取具有全局感受野的特征。因此，我们定义编码器块在层l和l + 1中的输出如下：

DCS代表深度卷积缩放,与Swin Transformer相比，用两个DWC层替换了常规的窗口分区多头自注意力模块W-MSA和SW-MSA。

编码器的完整架构包括4个阶段，每个阶段包括2个LK卷积块（即L=8个总层），在每个块内，DCS层在DWC层之后。DCS层有助于在不增加模型参数的情况下扩展特征图的维度（扩展到输入通道大小的4倍），并最小化跨通道学习到的体积上下文的冗余。为了在通道之间交换信息，我们使用了一个标准的卷积块，内核大小为2×2×2，步幅为2，以缩小特征分辨率2倍。相同的过程在第2阶段、第3阶段和第4阶段继续进行，这样的分层表示在多尺度设置中在每个阶段中提取，并进一步用于学习密集的体积分割。

4.2 DECODER

​ 编码器中每个阶段的多尺度输出通过长跳跃连接链接到基于 ConvNet 的解码器，并形成一个类似U形的网络以用于下游的分割任务。这一块就跟标准的 3D U-Net 几乎没啥两样了。具体而言，我们提取编码器每个阶段i（i∈0、1、2、3、4）的输出特征映射，并进一步利用一个残差块包括两个后归一化的3×3×3卷积层和实例归一化来稳定提取的特征。然后使用转置卷积层上采样处理过的每个阶段的特征，并与前一阶段的特征连接起来。为了进行下游体积分割，我们还将输入补丁的残差特征与上采样的特征进行拼接，并将特征输入带有softmax激活的1×1×1卷积层的残差块中，以预测分割概率。

5 实验

5.1 数据集

• MICCAI Challenge 2021 FLARE
• MICCAI Challenge 2021 FeTA
• MICCAI Challenge 2022 AMOS

5.2 实现细节

在两种情况下进行评估：1）直接监督训练和2）使用预训练权重的迁移学习。FeTA2021和FLARE2021数据集用于直接训练场景的评估，而AMOS数据集用于迁移学习场景的评估。我们对FeTA2021和FLARE2021数据集进行了五倍交叉验证，使用80％（训练）/ 10％（验证）/ 10％（测试）的划分。对于迁移学习场景，我们利用FLARE2021中训练的最佳折叠模型的预训练权重，并对AMOS2022上的模型权重进行微调，以评估3D UX-Net的微调能力，使用相同的训练/验证/测试划分。总体而言，我们通过与目前的体积 Transformer和ConvNet SOTA方法进行比较，评估3D UX-Net在完全监督的情况下用于体积分割的性能。我们使用Dice相似系数作为评估指标，以比较预测和真实标签之间的重叠区域。此外，我们进行了消融研究，以研究不同卷积核大小和使用深度卷积替换线性层进行特征提取的可变性对性能的影响。

5.3 实验结果

从上述表格可以看出，3D UX-Net 在所有分割任务中均展示出最佳性能，并且 Dice 分数有了显着提高、（FeTA2021：0.870 到 0.874，FLARE2021：0.929 到 0.934）。

（1）FeTA 和 FLARE 的评估

表1显示了当前在体积分割医学图像中使用的Transformer和ConvNets SOTA的结果比较。3D UX-Net在所有分割任务中展现出最佳性能，并且Dice得分有显着提高（FeTA2021：从0.870提高到0.874，FLARE2021：从0.929提高到0.934）

观察到3D UX-Net在使用FeTA2021数据集进行训练时展现出最快的收敛速度。有趣的是，随着训练样本数量的增加，SwinUNETR和3D UX-Net的收敛效率开始变得相容。除了定量的表示，图3进一步提供了额外的信心，展示了3D UX-Net在分割质量方面的改进。器官和组织的形态与真实标签相比得到很好的保留。

（2）AMOS的迁移学习

使用AMOS 2022数据集进一步研究了3D UX-Net与Transformer SOTA的迁移学习能力的比较。我们观察到3D UX-Net的微调性能显著优于其他Transformer网络，平均Dice为0.900（提高2.27%），大多数器官分割展现出了一致的质量提高。

虽然每个变压器网络的收敛曲线与FLARE2021训练模型的相当，但3D UX-Net进一步展示了其快速收敛和通过微调提高模型鲁棒性的能力。此外，图3中的定性表示展示了在保持相邻器官之间的边界和最小化过分割其他器官区域的可能性方面有显著改进

（3）消融实验

利用FeTA2021和FLARE2021数据集对不同模块进行消融实验。所有的消融实验都是在7×7×7的内核大小的情况下进行的，除了评估内核大小的可变性研究。

1）与标准卷积比较：调查了标准卷积和深度卷积在初始特征提取方面的效果。使用标准卷积，它表现出略微的提高。然而，使用标准卷积的模型参数约为使用深度卷积的3.5倍，而使用深度卷积的分割性能在两个数据集中仍然表现出可比性。

2）核大小的变化：从表3中可以看到，7×7×7的卷积核对于FeTA2021数据集的工作最佳，而FLARE2021的分割性能则是在13×13×13的核大小下表现最佳的。使用13×13×13的核大小对FLARE2021的显著改进可能是因为其提供了更大的感受野，增强了腹部区域内多个相邻器官之间的特征对应关系。对于FeTA2021数据集，只有小婴儿的大脑可以被很好地定位为前景，而7×7×7的核大小是提取组织对应关系的最佳接受野。

3）调整DCS：我们发现，在没有使用MLP进行特征缩放的情况下，性能有显著下降。使用线性缩放，在FLARE2021中性能显著提高，而在FeTA2021中略微提高。有趣的是，在使用1×1×1的深度卷积进行缩放时，对于FeTA2021和FLARE2021数据集都表现出略微的性能提高。此外，模型参数从56.3M降至53.0M，而没有牺牲模型性能。

6 讨论

​ 提出了一种基于纯卷积模块模拟 Swin Transformer 行为的块状设计。通过跳跃连接将设计进一步适应为通用的编码器骨干，用于体积分割的“U-Net”类似架构。我们发现，提高性能的关键组成部分可以分为两个主要方面：1）计算 MSA 的滑动窗口策略和2）扩大计算特征通道的反向瓶颈结构。W-MSA 增强了在每个窗口内学习特征对应关系的能力，而 SW-MSA 则加强了不同不重叠窗口之间的特征级别交叉连接。这种策略将 ConvNet 先验知识集成到 Transformer 网络中，并扩大了特征提取的感受野。然而，我们发现深度卷积可以展示类似于 Swin Transformer 块中计算 MSA 的操作。在深度卷积中，我们将每个输入通道与单个卷积滤波器卷积，并将卷积输出堆叠在一起，这与 Swin Transformer 中的特征输出的补丁合并层是可比的。此外，采用带有 LK 滤波器的深度卷积与 W-MSA 和 SW-MSA 具有相似之处，可以学习一个大感受野内的特征连接。 我们的设计提供了类似于 Swin Transformer 的功能，并且通过使用 ConvNet 模块减少了模型参数数量的优势。另一个有趣的差异是反向瓶颈结构。图1 显示了 Swin Transformer 和一些标准 ConvNet 都具有特定的瓶颈结构（黄色虚线）。 Swin Transformer 瓶颈中的独特组成部分是将通道大小保持为输入维度的四倍，并将 MSA 层的空间位置保持不变。我们遵循 Swin Transformer 块中的反向瓶颈结构，并将深度卷积移动到类似于 MSA 层的顶部。我们引入了在点卷积设置中使用深度卷积的思想来扩展具有更宽通道的密集特征。有趣的是，我们发现在数据集之间显示了轻微的性能改善（FeTA2021：0.872 至 0.874，FLARE2021：0.933 至 0.934），但具有较少的模型参数。由于每个编码器块仅包含两个缩放层，缩放块的数量有限可能会对性。

7 总结

​ 本文为大家介绍了 3D UX-Net，这是第一个将分层 Transformer 的特征与用于医学图像分割的纯 ConvNet 模块相适应的网络结构哦。具体地，本文重新设计了具有深度卷积的编码器块，以更低的代价实现与 Transformer 想媲美的能力。最后，通过在三个具有挑战性的公共数据集上进行的广泛实验表明所提方法的高效性。