nnWNet: 重新思考 Transformer 在生物医学图像分割中的应用及统一评估基准

基于 Transformer 的生物医学图像分割

Transformer 模型将图像视为小图像块的序列。它采用基于全连接网络的自注意力机制来捕获整个图像上的长程依赖关系，并使用位置编码来表示图像块之间的空间关系。一些研究探索了将 Transformer 整合到 U 型架构中。常见策略包括：无卷积层的纯 Transformer UNet，如 SwinUNet、D-Former、nnFormer；用 Transformer 替换卷积编码器或解码器，如 TC-CoNet、ConvTransSeg、UNETR、Swin UNETR、3D UX-Net；在编码器和解码器之间引入基于 Transformer 的桥接或瓶颈，如 CoTr、MissFormer、DCA、TransAttUNet；串联或并联堆叠卷积层和 Transformer 层，如 ConvFormer、H2Former、DuAT、MaxViT-UNet。此外，一些研究专注于优化自注意力机制以降低计算成本和内存消耗，如 BRAU-Net++、UNETR++。然而，这些架构设计存在全局和局部特征交替传递和混淆的缺陷。我们在第 3.1 节中更详细地讨论了这个问题。

nnUNet

nnUNet 是一个自动配置的分割框架，它设定了图像预处理、数据增强和训练超参数。它不仅作为一个开箱即用的分割解决方案，还为评估模型性能提供了一个统一的基准。在本研究中，我们基于 nnUNet 框架进行了全面且公平的比较。

方法

在第 3.1 节中，我们回顾并总结了各种编码器 - 解码器架构，分析了当前结合 Transformer 和卷积的设计中的矛盾。在第 3.2 节中，我们对全卷积 UNet 进行轻微调整，提出了 WNet，它整合了 Transformer 并解决了这些矛盾。在第 3.3 节中，我们将 WNet 整合到 nnUNet 框架中，用于统一的基准测试。

3.1 编码器 - 解码器架构探讨

如图 1 所示，当前的编码器 - 解码器模型可以总结并抽象为三种主要架构：编码器 - 解码器（encoder-decoder）、编码器 - 瓶颈 - 解码器（encoder-bottleneck-decoder）以及编码器 - 桥接 - 解码器（encoder-bridge-decoder）。这三种架构由四个模块组成：编码器、解码器、瓶颈和桥接。这四个模块不仅可以由纯卷积层或纯 Transformer 层顺序连接而成，也可以是卷积层和 Transformer 层的串联或并联。

表 1 列出了这些架构及其对应的一些模型。卷积专注于捕获局部细节，而 Transformer 则强调全局关系。它们的结合有潜力增强特征表示能力。

然而，当前的组合策略是矛盾的。无论是用 Transformer 作为编码器、卷积作为解码器（或反之），还是交替堆叠卷积层和 Transformer 层，它们都迫使 Transformer 层利用局部特征作为输入来捕获长程依赖，并迫使卷积层使用全局特征作为输入来提取局部细节。这破坏了全局和局部特征在整个模型中连续稳定的流动，导致训练不稳定和显著的性能下降。我们在第 4.3 节的对比实验中观察到了这一现象。

3.2 无矛盾地整合 Transformer

为了将 Transformer 整合到编码器 - 解码器架构中，并确保全局和局部特征能够不间断、不混淆地连续传递，我们轻微调整了 UNet，提出了 WNet。图 2 展示了 WNet 的概览。WNet 由两个级联的编码器 - 解码器架构组成，对输入图像进行重叠补丁嵌入，使用局部范围块构建多尺度编码和解码层，并在每个尺度的编码层和解码层之间引入全局范围桥。我们采用卷积层作为 LSB 来提取局部细节，并使用 Transformer 层作为 GSB 来捕获长程依赖关系。LSB 和 GSB 在每个编码 - 解码尺度上进行特征融合以交换信息。我们在补充材料中详细解释了架构设计的动机。

图 2：WNet 概览。OPE 表示重叠补丁嵌入。LSB 表示局部范围块。GSB 表示全局范围桥。

（符号）表示局部和全局特征融合。

H 和 W 表示特征图尺寸。

C1、C2、C3 和 C4 表示不同尺度特征图的通道数。

具体来说，重叠特征嵌入通过步长为 1 的 3×3 卷积和下采样实现。LSB 通过步长为 1 及下采样/上采样的 3×3 残差块实现。GSB 通过具有全局注意力机制的 Transformer 实现。特征融合通过通道拼接实现。我们在第 4.6 节的消融研究中比较了各种注意力机制和融合策略。请注意，基于 LSB 的局部特征包含隐式的位置关系，并且由于 LSB 和 GSB 可以交换信息，我们没有在 GSB 中引入额外的位置编码。

如何解决矛盾？

从图 3 可以看出，WNet 中的全局和局部特征可以在整个模型中连续流动，并在每个尺度上相互交换信息，这不仅解决了矛盾，而且有效地将全局和局部特征融合为一个统一的表示。

图 3：全局和局部特征的传递。(a) 局部特征的传递。(b) 全局特征的传递。两种特征连续流动和融合，没有间断。

我们在图 4 中比较了不同尺度下 LSB 和 GSB 的有效感受野。LSB 的 ERF 较小，而 GSB 的 ERF 较大，它们是互补的。这表明 LSB 主要捕获局部特征，而 GSB 专注于非局部特征。关于 ERF 可视化的更多细节可在补充材料中找到。此外，WNet 只是 U 型架构的扩展，没有任何复杂操作。其参数量少于 nnUNet。我们在第 4.3 节比较了不同模型的参数数量和性能。

图 4：LSB 和 GSB 在 ISIC-2017 [9]、Kvasir-SEG [27]、CREMI [12] 上的有效感受野（100 张图像平均）。顶行：LSB 中的残差块。底行：GSB 中的 11×11 深度可分离卷积自注意力。(a) ISIC-2017 尺度 1。(b) ISIC-2017 尺度 2。(c) ISIC-2017 尺度 3。(d) ISIC-2017 尺度 4。(e) Kvasir-SEG 尺度 4。(f) CREMI 尺度 4。

3.3 呼吁建立统一的评估基准

由于生物医学图像在图像尺寸、维度、模态、成像区域、像素（体素）大小、分割对象、数据规模和病理变化方面具有多样性，即使是针对相同临床应用的数据集之间也存在显著差异。这导致了缺乏标准化的训练和验证集划分以及一致的性能评估指标。一些研究尝试细化图像预处理（如图像重采样和感兴趣区域裁剪）并优化关键超参数（如图像块大小和训练轮数），以确保其模型表现更好。

总的来说，由于缺乏统一的评估基准，对比实验已不仅仅是模型架构的竞争。nnUNet [25, 26] 是一个自动配置的分割框架。更重要的是，它为纯粹、公平地比较各种架构提供了一个统一的基准。我们将 WNet 整合到 nnUNet 框架中以标准化实验设置。

实验

4.1 数据集

我们在四个 2D 数据集（DRIVE [44]、ISIC-2017 [9]、Kvasir-SEG [27] 和 CREMI [12]）和四个 3D 数据集（Parse2022 [35]、AMOS22 [28]、BTCV [31]、ImageCAS [53]）上评估我们的模型。表 2 展示了它们实验设置的详细信息。这八个数据集的补充信息在补充材料中提供。从表 2 可以看出，这些数据集涵盖了不同的模态、维度、成像区域、数据规模和分割对象，评估模型性能更具代表性和说服力。

4.2 统一评估基准

我们通过 nnUNet 框架（版本 2.3.1）实现我们的模型。所有模型的训练和推理均在单张 NVIDIA RTX A6000 GPU 上进行。所有模型均从头开始训练。为了平衡训练时间和效果，我们分别对 2D 和 3D 数据集使用 2d 和 3d_lowres 训练器。2D 数据集的训练轮数设为 200，3D 数据集设为 500。我们使用五折交叉验证来评估分割结果，并使用 Jaccard 指数、Dice 系数、95% 豪斯多夫距离和平均表面距离作为性能指标。Jaccard 和 Dice 强调像素级精度，而 95HD 和 ASD 强调边界精度。此外，我们仅使用模型的输出进行评估，不做任何额外的后处理。其他配置（如学习率、损失函数、图像块大小和优化器）均采用 nnUNet 框架中的默认设置。

4.3 与最先进模型的比较

我们将 nnWNet 与最先进的模型进行了广泛比较，包括纯卷积、纯 Transformer 以及混合模型。表 3 展示了在 2D 和 3D 数据集上的比较结果。nnWNet 大幅超越了先前的最先进模型；这可能是因为我们的模型能够同时捕获长程依赖关系和局部细节，并确保它们在多个特征尺度上连续传递和融合。我们还注意到，一些纯 Transformer 模型未能收敛；可能是因为它们过度强调长程依赖而忽略了关键的局部特征。一些混合模型表现出显著的性能下降；这可能归因于全局和局部特征传递过程中的混淆和矛盾。此外，在统一基准下，一些模型的表现不如一个精心设计的 UNet。

还在表 3 中比较了模型大小。我们的模型在所有编码器 - 解码器架构中保持了相对较少的参数数量。

表 3：在 2D 和 3D 数据集上与最先进模型的比较。

注：此处仅保留表格说明，因表格内容复杂且为实验数据摘要，建议用户查阅原文 PDF 或图片以获取准确表格。表中绿色表示性能优于 nnUNet，橙色表示性能差于 nnUNet，红色表示训练失败或性能显著下降。粗体表示所有模型中的最佳性能。

4.4 统计分析

表 4 显示了 nnWNet 与其他模型在 Jaccard 指标上的统计显著性。我们的模型取得了 87.50% 的显著率，在 32 次 t 检验中成功 28 次（p<0.05）。这表明 nnWNet 具有统计上显著的改进。

4.5 定性结果

图 5 和图 6 展示了不同模型的各种定性结果。基于全局和局部特征的连续流动与多尺度融合，nnWNet 实现了更好的像素级精度和轮廓精度。更多定性结果可在补充材料中找到。

4.6 消融研究

为了验证每个组件的有效性，我们在四个 2D 数据集上进行了以下消融研究。

不同自注意力机制的比较。

表 5 比较了在 GSB 中添加不同自注意力机制的性能，包括大核深度卷积、池化和高效多头自注意力。池化取得了更好的结果。我们在表 3 的所有相关实验中使用池化作为自注意力机制。

融合策略的比较。

表 6 比较了全局和局部特征的不同融合策略。我们发现原始 UNet 使用的通道拼接取得了更好的性能。我们将其应用于表 3 的所有相关实验。

结合卷积与 Transformer 的有效性。

我们将 LSB 替换为 PoolFormer，将 GSB 替换为残差块，分别生成纯 Transformer WNet 和纯卷积 WNet。表 7 比较了混合、纯卷积和纯 Transformer 架构的性能。总体而言，混合架构表现出优越的结果。请注意，纯卷积架构可能获得更好的轮廓精度，因为卷积专注于局部细节并对边界敏感。

全局和局部特征连续传递的必要性。

我们替换原始 UNet 中的编码器和解码器，以证明连续传递全局和局部特征的重要性。具体来说，我们实现了纯卷积、纯 Transformer、卷积与 Transformer 串联以及卷积与 Transformer 并联来构建编码器和解码器。表 8 比较了这些架构的性能。我们的发现表明串联连接有负面影响。这可能是因为串联连接迫使全局和局部特征交替生成，阻碍了连续传递。相比之下，并联连接以及在每个下采样前的融合产生了积极结果，因为它避免了从局部特征捕获长程依赖的同时又从全局特征提取局部细节的冲突。使用卷积编码器与 Transformer 解码器的性能与使用 Transformer 编码器与卷积解码器相似，这两种架构的性能介于串联和并联之间。此外，我们发现 Transformer 与卷积的组合优于纯卷积架构，进一步证明了同时关注全局和局部特征的必要性。

结论

我们系统性地总结了基于 Transformer 的生物医学图像分割编码器 - 解码器架构，并分析了其设计中的矛盾。为了整合 Transformer 并解决这些矛盾，我们轻微调整了 UNet 并提出了 WNet。WNet 能够同时捕获长程依赖关系和局部细节，并确保它们的连续传递和多尺度融合。我们将 WNet 整合到 nnUNet 框架中以进行统一的基准测试。我们的模型实现了最先进的生物医学图像分割性能，在四个 2D 数据集和四个 3D 数据集上的大量实验证明了其有效性。