轻量级SegFormer3D | 重塑三维图像分割的内存高效Transformer架构，助力医学成像 ！

基于视觉Transformer（ViTs）架构的采用，在三维医学图像（MI）分割领域代表了重大进步，它通过增强全局上下文理解，超越了传统的卷积神经网络（CNN）模型。尽管这种范式转变显著提高了三维分割性能，但现有最先进的架构需要极其大型且复杂的结构，以及大规模的计算资源进行训练和部署。 此外，在经常遇到的医学成像有限数据集的背景下，更大的模型可能在模型泛化和收敛上都存在障碍。为了应对这些挑战，并且为了证明轻量级模型在三维医学成像研究中具有价值，作者提出了SegFormer3D，这是一个分层Transformer，能够在多尺度体积特征上计算注意力。 此外，SegFormer3D避免了复杂的解码器，并使用全多层感知器（MLP）解码器来聚合局部和全局注意力特征，以产生高度精确的分割 Mask 。所提出的内存高效Transformer在紧凑的设计中保持了较大模型的性能特点。 SegFormer3D通过提供一个参数减少

，GFLOPS减少

的模型，使深度学习在三维医学图像分割领域民主化，与当前最先进（SOTA）模型相比。作者在三个广泛使用的Synapse、BRaTs和ACDC数据集上对SegFormer3D进行了基准测试，取得了具有竞争力的结果。 代码：https://github.com/OSUPCVLab/SegFormer3D.git。

1 Introduction

深度学习在医疗领域的出现具有变革性，提供了前所未有的学习和分析复杂医学数据模式的能力。在医学图像分析中，一个基本的任务是三维体积图像分割，这对于诊断和治疗中的肿瘤和多器官定位等应用至关重要。传统的方法包括采用编码器-解码器架构，首先将图像转换成低维表示，然后解码器将这个表示映射到 Voxel 级的分割 Mask 。然而，这些架构由于有限的感受野，难以生成准确的分割 Mask 。

近来，基于Transformer的技术由于ViT利用注意力层捕捉全局关系的能力，展示了卓越的分割性能。这与展现出局部诱导偏置特性的CNN形成了鲜明对比。

图1：参数数量与BraTs上的性能对比 作者将SegFormer3D与现有的3D体积图像分割架构进行比较，评估模型性能与参数数量的关系。绿色条形代表模型参数，而紫色曲线显示了每种架构的平均面部性能。作者证明，在450万参数时，SegFormer3D是一个极具竞争力的轻量级架构，适用于3D医学图像分割。

在面对小规模数据集时，由于缺乏归纳偏置，通常需要依赖大规模数据集进行预训练[7]，这在医学图像领域并不常见。此外，ViTs的计算效率受到浮点运算数量和多头自注意力块中的逐元素函数的限制[17]。在3D医学成像任务中，这个问题更为突出，因为转换后的3D体积输入序列的长度明显较长。此外，医学成像数据经常表现出重复结构[6]，这表明它可以被压缩，这一点通常被医学领域的3D SOTA ViT架构忽略。

遵循TransUnet[5]和UNFTR[11]的开创性工作之后，医学界大量研究致力于设计基于Transformer的架构，这些架构利用ViTs强大的编码能力以及CNN在解码阶段的特征细化能力。例如，结合了卷积的局部感受野和全局注意力。尽管具有优势，但从零开始训练时，ViTs无法匹配CNN的泛化能力。本文提出SegFormer3D，一个体积分层ViT，它将[26]扩展到3D医学图像分割任务。

与在固定尺度上渲染特征图的普通ViT[7]不同，Segformer3D根据Pyramid Vision Transformer[25]对输入体积的不同尺度的特征图进行编码。作者的设计使Transformer能够捕捉到输入的各种从粗到细粒度的特征。SegFormer3D还使用了高效的自我注意力模块[25]，该模块将嵌入序列压缩到固定比例，从而显著降低模型复杂度而不牺牲性能。

此外，SegFormer3D还使用了[26]中的重叠 Patch 嵌入模块，该模块保持了输入 Voxel 的空间连续性。这种嵌入使用了位置无关编码[14]，防止在训练和推理期间因分辨率不匹配而导致的精度损失，这在医学图像分割中是常见的情况。为了高效生成高质量的分割 Mask ，SegFormer3D使用了[26]中引入的全MLP解码器。在三个基准数据集-Synapse[15]，ACDC[1]，BRaTs[20]上的综合实验验证了SegFormer3D的质量和定量效果。

作者的贡献可以总结为：

1. 作者引入了一种轻量级、内存高效的分割模型，该模型保留了大型模型在3D医学成像中的性能特点。
2. 拥有450万的参数和17 GFLOPS，Segformer3D相较于现有技术水平（SOTA）实现了参数数量和模型复杂度分别减少 34 倍和 13 倍。
3. 作者展示了在没有预训练的情况下具有高度竞争力的结果，强调了轻量级ViTs的泛化能力，并探索如Segformer3D这样的架构在医学成像领域中是一块值得研究的有价值的领域。

2 相关工作

在Unet被引入之后，已经提出了许多用于医学图像分析的方法，例如Dense-unet 和深度监督CNN。Unet也被扩展到3D医学图像分析中，例如3D-Unet，V-net，nn-Unet和。研究行人还设计了分层架构以捕获上下文信息。在中，Milletari等人使用V-net将体积分辨率降低以保留有益的图像特征。Cicek等人[6]替换了3D-unet中的2D到3D卷积。Isensee等人[13]提出了nn-Unet通用分割架构，该架构可以在多个尺度上提取特征。在中，PGD-UNet使用可变形卷积处理不规则器官形状和肿瘤的医学图像分割。

几篇最近的论文研究了Transformer-卷积架构，例如TransUnet[5]，Unetr[11]，SwinUnet[10]，TransFuse[29]，nnFormer[31]。TransUnet[5]将Transformer与U-Net结合，用于编码图像块并通过高分辨率上采样的CNN特征进行解码以实现定位。Hatamizadeh等人[11]提出了UNETR，这是一个3D模型，它将Transformer的长距离空间依赖特性与“U形”编码器-解码器结构中的固有CNN归纳偏置相结合。在UNETR中，Transformer块编码能捕获一致的全局表示的特征，并随后在基于CNN的解码器内跨各种分辨率进行整合。Zhou等人[31]提出了nnFormer，这是一种源自Swin-UNet[3]架构的方法。Wang等人[24]提出了TransBTS，该方法使用常规的卷积编码器-解码器架构和一个Transformer层作为瓶颈。

3 Method

Transformer （Transformers）的采用大大提高了立体医学图像分割的性能。然而，目前高性能的架构为了模型性能而优先考虑过参数化，牺牲了效率。为了展示在不影响性能的情况下轻量级且高效的 Transformer 的优势，作者引入了Segformer3D。拥有450万个参数和17 GFLOPS，作者展示了参数数量和复杂度分别减少了34倍和13倍，凸显了所提出架构在3D医学图像分割中的重要性。

编码器：在Transformer框架内使用3D医学图像会导致序列长度过长，从而增加模型的计算复杂性。例如，一个标准的3D MRI体积，其维度为128

，将产生一个32,768的序列长度，而一个典型的2D RGB图像，其维度为256

，则产生一个256的序列长度。作者的分层Transformer融入了三个关键元素来提高计算效率并减少总参数数量，同时保持了SOTA（当前最佳技术水平）的性能。首先，作者采用了重叠的 Patch 合并技术，以克服在 Voxel 生成过程中邻近信息丢失的问题。这种技术与在ViT[7]中看到的 Patch 机制相比，使模型能够更好地理解 Voxel 之间的过渡点，并且已经被证明可以提高整体分割精度[26]。接下来，为了在不牺牲性能的情况下解决序列长度瓶颈，作者集成了一种有效的自注意力机制[25]。这种方法使模型能够更有效地捕获长距离依赖性，从而提高了可扩展性和性能。传统的自注意力机制将一组形状为[Batch, Sequence, Features]的向量序列作为输入，并生成3个独特的投影： Query （Query）、键（Key）和值（Value）向量。一旦生成，注意力分数的计算公式为

。由于操作

，原始分割过程的计算复杂度为

。尽管在2D图像中这种复杂性可能被忽略，但对于长的3D序列来说，它证明了对于高效架构设计是一个挑战。在[25, 26]中引入了高效注意力。

该方法显著降低了由3D体积张量产生的计算复杂度，从

降低到

。作者在编码器的四个阶段中将降低参数

设定为

，

，

，

。

表1：Segformer3D与SOTA在大小（M）和复杂性方面的对比。Segformer3D在不过度牺牲性能的情况下，显著降低了参数数量和计算复杂性。

最后，作者的方法通过采用混合ffn模块[26]来解决在ViTs中调整体积成像大小与固定位置编码之间的挑战。这个模块能够自动学习位置线索，消除对固定编码的需求，确保了卓越的可扩展性和性能。解码器：在基于编解码器设计的医学图像分割中，解码阶段发挥着关键作用，这种设计在UNET为基础的架构中广泛采用[10, 11]。这个框架既用于基于CNN的编码器，也用于基于Transformer的编码器。在3D医学图像的背景下，有效的解码通常需要连续的3D卷积，但作者反而证明了线性层的集成是一种针对医学图像分割高度有效的解码策略。作者的方法简化了解码过程，确保在多样化数据集中高效且一致地解码体积特征，而无需过度参数化。简单的解码过程是：

第一步：

(1) 第二步：

(2) 第三步：

(3) 第四步：

(4)

第一步：

(1) 第二步：

(2) 第三步：

(3) 第四步：

(4)

类似于在UNET中引入的跳跃连接，每个阶段的特征被整理起来，并生成一个固定维度的投影。一旦所有维度都标准化后，作者对每个特征进行上采样并进行拼接，然后进行融合操作。融合后的特征被输入到一个线性投影 Head （3D 1x1卷积）以生成最终的分割 Mask 。

4 Experimental Results

图2：在BRaTs上的定性结果。每一行是MRI序列中的一个独立帧，而每一列是三维体积图像分割解决方案。作者定性展示了与现有技术水平（SOTA）方法相比具有高度准确的分割性能，同时保持了轻量级和高效的架构。

遵循最新的三维体积分割SOTA架构，作者使用了相同的数据库和评估方法，以确保在所有架构之间进行公平和一致的比较。作者在三个广泛使用的数据库上训练和评估所提出的模型，并未使用外部数据进行预训练。这些数据库分别是脑肿瘤分割（BraTS）[20]、Synapse多器官分割（Synapse）[15]，以及自动心脏诊断（ACDC）[1]数据库。

所有实验，包括训练、实时增强和推理，都是在单个Nvidia RTX 3090 GPU上使用PyTorch完成的。每个模型都采用相同的的学习率策略，包括学习率 Warm up 阶段，在这个阶段作者线性地将学习率从

增加到

，之后采用PolyLR衰减策略。广泛采用的AdamW优化器[19]被用于学习率为

。对于损失函数，采用等权重Dice-交叉熵损失组合，以结合每个损失函数在优化过程中的优点，从而改善收敛性。作者将批量大小设置为4，并且像SOTA架构一样，每个模型训练1000个周期。此外，所有实验都是在不使用复杂预训练努力的情况下进行的，以展示所提出架构在实际医疗数据集上的性能，而不需要额外的数据。

Results on Brain Tumor Segmentation (BraTs)

图3：在Synapse上的定性结果。每一行是CT序列中的一个独立帧，而每一列是不同的3D体积图像分割解决方案。每个器官 Mask 都使用独特的颜色代码进行高亮。作者定性展示了与已确立的SOTA方法相比具有高度准确的分割性能，同时保持了轻量级设计。

BraTs [20] 是一个来自MRI扫描的医疗图像分割数据集。该数据集包含484个MRI图像，涵盖四种模式，分别是FLAIR、T1w、T1gd和T2w。数据来自19个机构，带有三种肿瘤子区域（水肿（ED）、增强型肿瘤（ET）和非增强型肿瘤（NET））的 GT 标签。遵循包括nnFormer [31]在内的主要出版物中的相同数据准备、增强和报告策略，作者在整个肿瘤（WT）、增强型肿瘤（ET）和肿瘤核心（TC）上报告了作者的结果。在表2(a)中，作者展示了Segformer3D作为一个强大的轻量级分割架构，在参数保持450万、计算复杂度为17.5 GFLOP的情况下，明显优于更大规模且广泛采用的卷积神经网络（CNN）和Transformer架构。这证明了高效的自注意力模块在表示学习能力上的优势，它分析整个图像块序列而不进行压缩。

最后，在图2(a)中，作者展示了所提出架构的极具竞争力的定性结果，表明其性能优于已确立的架构。

Results on Multi-Organ CT Segmentation (Synapse)

Synapse数据集[15]提供了30张标注的CT图像。作者在数据处理、训练和数据划分（定义于[31]中）之后生成作者的结果。由于标注内容的多样性，覆盖了包括脾脏、胰腺、胆囊等多个器官，Synapse数据集是一个复杂的多类别分割挑战。表3(a)中的定量结果显示，Segformer3D在参数量为1.5亿的nnFormer[31]架构之后，排名第二。此外，在图3(a)中，作者展示了与当前SOTA架构相比，高精度的器官分割 Mask 定性性能结果，验证了所提方法的视觉一致性。最后，与广泛使用的架构[10, 11, 12]相比，Segformer3D仅用450万的参数生成了具有竞争力的结果，并且表明过度参数化并不会导致性能的大幅提升，特别是在数据受限的情况下。

Results on Automated cardiac diagnosis (ACDC)

图4：在ACDC数据集上的定性结果。每一行是cine-MRI序列中的一个独立帧，而每一列是不同的3D体积图像分割解决方案。作者的方法展示了与现有技术水平（SOTA）方法相比的高度准确的分割结果，同时保持了轻量级和高效的结构。

ACDC[1]是一个包含100名患者的数据集，用于对左心室（LV）和右心室（RV）以及心肌（Myo）[1]进行三维体积分割。为了与已发表的研究进行一对一的比较，作者遵循[31]中指定的相同的训练和推理流程，并使用Dice指标来衡量分割的准确性。表4（a）定性地展示了所提出架构与大且高度复杂的解决方案相比具有很高的竞争力。所提出的模型在平均参数数量多出

，计算复杂性高出

的情况下，与SOTA性能的差距在1%以内。图4（a）中可视化了定性结果的比较，展示了在不进行大规模预训练小数据集的情况下，性能也非常具有竞争力。

Conclusion

结构如UNETR、TransUNet和nnFormer等，已经使用ViT框架革新了三维体积医疗图像分割。这种范式转变显著增强了模型在上下文理解能力方面，相比于传统的纯卷积神经网络（CNN）对应物。然而，这一改进是以参数数量和模型复杂性的大幅增加为代价的，这归因于自注意力模块的复杂性质。

除了模型大小和复杂性之外，大型模型阻止了那些计算资源有限的医学研究者有效地训练并将这些模型整合到他们的工作流程中。此外，大型模型可能会在模型泛化和收敛方面带来挑战，特别是在医学成像中常见的数据集有限的情况下。

为了在不牺牲性能的情况下克服这些限制，作者引入了Segformer3D，这是一个轻量级结构，其参数和计算复杂度分别只有最先进（SOTA）架构的

和

。作者将作者的解决方案与当前的SOTA解决方案以及其他高被引作品进行了基准测试，并展示了轻量级高效架构可以帮助显著提高性能，超过更大的模型，而无需额外的预训练，并且只需最小的计算资源。

最后，作者主张将研究努力导向开发高性能轻量级架构，特别是在像医学成像这样具有实质性现实世界优势的领域，这不仅拓宽了可访问性，也促进了在现实世界场景中这些架构的实际应用。

参考

[1].SegFormer3D: an Efficient Transformer for 3D Medical Image Segmentation.