Swin-UMamba：结合基于ImageNet的预训练和基于Mamba的UNet模型
摘要
https://arxiv.org/pdf/2402.03302v1.pdf
准确的医学图像分割需要整合从局部特征到全局依赖的多尺度信息。然而，现有方法在建模长距离全局信息方面面临挑战，其中卷积神经网络（CNNs）受限于其局部感受野，而视觉转换器（ViTs）则受到其注意力机制高二次复杂度的困扰。最近，基于Mamba的模型因其在长序列建模中的出色能力而备受关注。多项研究表明，这些模型在各种任务中可以超越流行的视觉模型，提供更高的准确性、更低的内存消耗和更少的计算负担。然而，现有的基于Mamba的模型大多是从头开始训练的，并没有探索预训练的力量，这已被证明对于数据高效的医学图像分析是非常有效的。本文介绍了一种新型的基于Mamba的模型Swin-UMamba，专门设计用于医学图像分割任务，利用基于ImageNet的预训练优势。我们的实验结果表明，基于ImageNet的训练在提升基于Mamba的模型性能方面发挥着至关重要的作用。Swin-UMamba与CNNs、ViTs和最新的基于Mamba的模型相比，表现出卓越的性能。值得注意的是，在AbdomenMRI、Encoscopy和Microscopy数据集上，Swin-UMamba比其最接近的对应模型U-Mamba平均高出3.58%。Swin-UMamba的代码和模型可在https://github.com/JiarunLiu/Swin-UMamba上公开获取。
关键词：医学图像分割；基于ImageNet的预训练；长距离依赖建模。
1、引言
医学图像分割在现代临床实践中扮演着重要角色，如辅助诊断、制定治疗方案和实施治疗等 [1,2,3]。典型的分割过程依赖于经验丰富的医生，这既劳动密集又耗时。此外，由于主观解读和观察者间变异性，专家之间的分割一致性可能会有所不同 [4,5]。这突显了需要自动化分割方法来提高医学图像分析的效率、准确性和一致性，以便做出准确和快速的诊断 [6,7]。
近年来，深度学习在医学图像分割领域取得了显著的进展 [8,9,10,11]。然而，准确的医学图像分割需要将局部特征与其相应的全局依赖关系相结合 [12]。从图像数据中高效捕获复杂和长距离的全局依赖关系仍然是一个挑战。两种流行的方法，卷积神经网络（CNNs）和视觉转换器（ViTs），在长距离依赖关系建模方面都存在自己的局限性。如SegResNet [13]、U-Net [8]和nnU-Net [11]等CNN在医学图像分割中得到了广泛应用。它们在提取局部特征方面很有效，但可能难以捕获全局上下文和长距离依赖关系。这是因为CNNs本质上受其局部感受野的限制 [14]，这限制了它们从图像中的远距离区域捕获信息的能力。另一方面，ViTs已经显示出处理全局上下文和长距离依赖关系的能力 [15,16]。然而，由于ViTs的注意力机制，它们在处理长序列建模时受到高二次复杂度的限制 [17]，而在医学领域，高分辨率图像并不罕见（例如全幻灯片病理图像 [18]、高分辨率MRI/CT扫描 [19]）。尽管存在复杂性，但转换器在处理有限数据集时容易过拟合 [20]，这表明它们对数据的需求很高。
然而，现有的基于Mamba的模型大多是从头开始训练的。对于基于Mamba的模型在医学图像分割任务中的预训练影响尚不清楚，而对于数据高效的医学图像分析，CNN [10]和ViT [28]的预训练已被证明是相当有效的。这在医学领域尤为重要，因为医学图像数据集通常在规模和多样性上受到限制 [29,30]。了解预训练基于Mamba的模型在医学图像分割中的有效性，可以为提高深度学习模型在医学成像应用中的性能提供有价值的见解。
有几个挑战需要解决。首先，由于Mamba块在视觉领域的应用相对较新，需要进一步的实验验证来评估其性能，并确定最适合医学图像分割任务的Mamba块结构。另一个挑战在于，现有的用于医学图像分割的基于Mamba的模型没有考虑到从ImageNet预训练模型的迁移性。因此，需要重新设计网络结构以集成预训练模型。第三，对于现实世界中的部署，特别是资源受限的环境中（这在医学实践中很常见），需要基于Mamba的模型的可扩展性和效率 [31]。
在本文中，我们提出了一种基于Mamba的网络Swin-UMamba，用于二维医学图像分割。Swin-UMamba使用通用编码器将预训练视觉模型的强大功能与专为医学图像分割任务设计的解码器相结合。此外，我们提出了一种基于Mamba的解码器的变体结构SwinUMamba†，该结构具有较少的参数和较低的FLOPs，适用于高效应用。我们的贡献可以总结如下：
- 据我们所知，我们是首次尝试探索预训练基于Mamba的网络在医学图像分割任务中的影响。我们的实验验证了基于Mamba的网络在医学图像分割中，使用ImageNet进行预训练的重要作用，有时这是至关重要的。
- 我们提出了一种新的基于Mamba的网络Swin-UMamba，专门设计用于整合预训练模型的强大功能。此外，我们提出了一种变体结构SwinUMamba†，该结构在保持竞争力的同时，具有较少的网络参数和较低的FLOPs。
- 我们的结果表明，Swin-UMamba和Swin-UMamba†都能以前所未有的优势超越以前的分割模型，包括CNN、ViT和最新的基于Mamba的模型，这突出了在医学图像分割任务中基于ImageNet的预训练和所提出架构的有效性。
2、方法
我们在图1中展示了Swin-UMamba的整体架构。它主要由三部分组成：1) 基于Mamba的编码器，该编码器在大规模数据集（即ImageNet）上进行预训练，以提取不同尺度的特征；2) 具有多个上采样块的解码器，用于预测分割结果；3) 跳跃连接，用于弥合低级细节和高级语义之间的差距。我们将在以下部分中介绍Swin-UMamba的详细结构。
2.1、基于Mamba的VSS块
Mamba在自然语言处理中的最新进展[17]使用空间状态序列模型（SSMs）[21]将长序列建模的注意力复杂度从二次降低到线性。在视觉任务中使用Mamba的优点是显而易见的[26]：更高的准确性、更低的计算负担和更少的内存消耗。然而，2D视觉数据和1D语言序列之间的区别需要仔细考虑。例如，在视觉任务中，2D空间信息至关重要[24]，而在1D序列建模中则不是主要关注点。直接将Mamba应用于扁平化的图像不可避免地会导致受限的接收域，其中无法估计未扫描的块之间的关系。
在[24]的见解基础上，我们将视觉状态空间（VSS）块作为Swin-UMamba的基本单元。VSS块通过使用二维选择性扫描（SS2D）来解决与二维图像数据相关的挑战。这种方法沿着四个方向展开图像块，创建四个不同的序列。然后，每个特征序列将通过SSM进行处理。最后，将输出特征合并以形成完整的二维特征图。给定输入特征z，SS2D的输出特征 z ˉ \bar{z} zˉ可以表示为：
z v = expand ( z , v ) z ˉ v = S6 ( z v ) z ˉ = merge ( z ˉ 1 , z ˉ 2 , z ˉ 3 , z ˉ 4 ) \begin{aligned} z_{v} & =\operatorname{expand}(z, v) \\ \bar{z}_{v} & =\operatorname{S6}\left(z_{v}\right) \\ \bar{z} & =\operatorname{merge}\left(\bar{z}_{1}, \bar{z}_{2}, \bar{z}_{3}, \bar{z}_{4}\right) \end{aligned} zvzˉvzˉ=expand(z,v)=S6(zv)=merge(zˉ1,zˉ2,zˉ3,zˉ4)
在公式中， v ∈ V = { 1 , 2 , 3 , 4 } v \in V=\{1,2,3,4\} v∈V={1,2,3,4} 表示四个不同的扫描方向。扩展（expand）和合并（merge）操作对应于[24]中的扫描扩展和扫描合并操作。选择性扫描空间状态序列模型（S6）是VSS块的核心SSM操作符。它使一维数组（例如文本序列）中的每个元素能够通过压缩的隐藏状态与任何先前扫描的样本进行交互。有关S6的更多详细信息，请参见[24]。VSS块的整体结构如图1所示。
2.2、结合基于ImageNet的预训练
主要的挑战在于如何将通用的预训练模型有效地集成到医学图像分割任务中。先前的研究[23]通常使用具有Mamba块的特定架构，而没有考虑从通用视觉模型的迁移性。为了解决这个问题，我们构建了一个编码器，它与最新的基于Mamba的视觉方法VMamba-Tiny[24]具有相似的结构。该模型在广泛的ImageNet数据集上进行预训练，具有多尺度特征，使我们能够整合通用视觉模型的强大功能来提取具有长程建模能力的信息，降低过拟合的风险，并为Swin-UMamba建立一个稳健的初始化。
Swin-UMamba的编码器可分为五个阶段。第一阶段是起始阶段，包含一个用于2倍下采样的卷积层，该卷积层使用7x7大小的卷积核，填充大小为3，步长为2。卷积层之后进行了二维实例归一化。Swin-UMamba的第一阶段与VMamba不同，因为我们更喜欢逐步下采样的过程，每个阶段都进行2倍下采样。这种策略旨在保留低级别细节，这对于医学图像分割非常重要[8,32]。第二阶段使用了一个2x2大小的块嵌入层，保持了特征分辨率为原始图像的1/4，这与VMamba中的嵌入特征相同。后续阶段遵循VMamba-Tiny的设计，每个阶段都由一个用于2倍下采样的块合并层和若干个用于高级特征提取的VSS块组成。与ViTs不同，由于VSS块的因果性质[24]，我们在Swin-UMamba中没有采用位置嵌入。
在Swin-UMamba的第二阶段到第五阶段，VSS块的数量分别为{2,2,9,2}。每个阶段后的特征维度相对于之前的阶段都呈二次增长，具体为D={48,96,192,384,768}。我们使用了VMamba-Tiny在ImageNet上预训练的权重来初始化VSS块和块合并层，如图1所示。值得注意的是，由于块大小和输入通道的差异，块嵌入层并没有使用预训练的权重进行初始化。
2.3、Swin-UMamba 解码器
我们遵循常用的U型架构，并采用了密集的跳跃连接来构建Swin-UMamba。U-Net及其变体在医学图像分割任务中已显示出显著的效率。这种架构利用跳跃连接来恢复低级别细节，并采用编码器-解码器结构进行高级信息提取。为了增强U-Net中的原生上采样块，我们引入了两项改进：1）一个额外的带有残差连接的卷积块，用于处理跳跃连接特征；2）在每个尺度上增加一个额外的分割头，用于深度监督[33]。
上采样块的结构如图1所示。给定来自阶段l的跳跃连接特征 z l ′ z_{l}^{\prime} zl′和来自上一个上采样块的特征 z l + 1 z_{l+1} zl+1，第l个上采样块的输出特征 z l z_{l} zl和阶段l的分割图 y l ∈ R h l × w l × K y_{l} \in R^{h_{l} \times w_{l} \times K} yl∈Rhl×wl×K可以表述如下：
z ^ l = Res l ( 2 ) ( Cat ( z l + 1 , Res l ( 1 ) ( z l ′ ) ) ) z l = DeConv l ( z ^ l ) , y l = Conv l ( z ^ l ) \begin{array}{l} \hat{z}_{l}=\operatorname{Res}_{l}^{(2)}\left(\operatorname{Cat}\left(z_{l+1}, \operatorname{Res}_{l}^{(1)}\left(z_{l}^{\prime}\right)\right)\right) \\ z_{l}=\operatorname{DeConv}_{l}\left(\hat{z}_{l}\right), \quad y_{l}=\operatorname{Conv}_{l}\left(\hat{z}_{l}\right) \end{array} z^l=Resl(2)(Cat(zl+1,Resl(1)(zl′)))zl=DeConvl(z^l),yl=Convl(z^l)
其中， Cat ( ⋅ ) \operatorname{Cat}(\cdot) Cat(⋅) 表示特征拼接操作， DeConv l ( ⋅ ) \operatorname{DeConv}_{l}(\cdot) DeConvl(⋅) 表示第 l l l阶段的转置卷积， Conv l ( ⋅ ) \operatorname{Conv}_{l}(\cdot) Convl(⋅) 表示带有 1 × 1 1 \times 1 1×1卷积的分割头，用于将特征从维度 d l d_{l} dl投影到类别数 K K K。 h l h_{l} hl 和 w l w_{l} wl 分别表示第 l l l阶段特征图的高度和宽度。 Res l ( 1 ) ( ⋅ ) \operatorname{Res}_{l}^{(1)}(\cdot) Resl(1)(⋅) 和 Res l ( 2 ) ( ⋅ ) \operatorname{Res}_{l}^{(2)}(\cdot) Resl(2)(⋅) 是第 l l l阶段带有残差连接的两个卷积块，每个 Res ( ⋅ ) \operatorname{Res}(\cdot) Res(⋅)都由两个带有LeakyReLU激活的卷积层组成。除了编码阶段和上采样块之间的跳跃连接外，我们还从输入添加了一个额外的跳跃连接，使用了 Res ( ⋅ ) − Cat ( ⋅ ) − Res ( ⋅ ) \operatorname{Res}(\cdot)-\operatorname{Cat}(\cdot)-\operatorname{Res}(\cdot) Res(⋅)−Cat(⋅)−Res(⋅)操作。最后，我们使用一个 1 × 1 1 \times 1 1×1的卷积得到最终的分割输出。
在Markdown中，LaTeX公式通常会被原样保留，并且Markdown本身支持LaTeX公式的嵌入。上述翻译保留了LaTeX公式的格式，并且适应了Markdown的语法。如果需要在具体的Markdown编辑器中显示这些公式，可能还需要根据编辑器的配置进行适当的调整。
2.4、Swin-UMamba†：基于Mamba的解码器的Swin-UMamba
为了进一步探索Mamba在医学语义分割中的潜力，我们提出了一种基于Mamba解码器的Swin-UMamba的变体，即Swin-UMamba†。我们将展示Swin-UMamba†可以在与Swin-UMamba相比具有竞争力的结果的同时，使用更少的网络参数并施加更低的计算负担。此外，我们的研究结果表明，无论解码器结构如何，大规模预训练在医学图像分割任务中都起着重要作用。我们对SwinUMamba†进行了几项修改。首先，Swin-UMamba中的上采样块被替换为patch expanding [34]和两个VSS块。我们发现基于CNN的解码器造成了许多参数和计算负担。其次，我们使用一个4 × 4的patch嵌入层，该层直接将输入图像从H × W × C投影到形状为H/4 × H/4 × 96的特征图，遵循VMamba [24]。值得注意的是，Swin-UMamba†中的最后一个patch扩展块是一个4 ×上采样操作，与4 × patch嵌入层镜像对称。残差patch扩展层是2 ×上采样操作。来自输入图像和Swin-UMamba中2 ×下采样特征的跳跃连接被移除，因为它们没有对应的解码块。此外，在1 ×、1/4 ×、1/8 ×和1/16 ×的分辨率上应用了深度监督，每个尺度都添加了额外的分割头（即，1 × 1卷积将高维特征映射到K），以进行额外的分割。通过结合所有这些修改，我们在AbdomenMRI数据集上将网络参数数量从40M减少到27M，并将FLOPs从58.4G降低到15.0G。有关网络参数数量和FLOPs的更多统计信息，请参见表1、表2和表3。Swin-UMamba†的结构如图2所示。
3、实验
3.1、数据集
我们在三个不同的医学图像分割数据集上评估了Swin-UMamba的性能和可扩展性，这些数据集涵盖了器官分割、仪器分割和细胞分割。这些数据集选择了各种分辨率和图像模式，提供了模型在各种医学成像场景中的有效性和适应性的洞察。
腹部MRI（AbdomenMRI）：这个数据集专注于从MRI扫描中分割出13个腹部器官，包括肝脏、脾脏、胰腺、右肾、左肾、胃、胆囊、食道、主动脉、下腔静脉、右肾上腺、左肾上腺和十二指肠。它最初来自MICCAI 2022年的AMOS挑战赛[35]。我们遵循了[23]中的设置，并额外添加了50个MRI扫描用于测试。有60个MRI扫描，包含5615个切片用于训练，50个MRI扫描，包含3357个切片用于测试。我们使用nnUNet框架[11]将图像裁剪成大小为(320,320)的块进行训练和测试。
内镜图像（Endoscopy）：这个数据集旨在从内镜图像中分割出7种器械，包括大针驱动器、prograsp钳、单极弯曲剪刀、cadiere钳、双极钳、血管封闭器和内置超声探头。它最初来自MICCAI 2017年的EndoVis挑战赛[36]。该数据集包含1800个图像帧用于训练，1200个图像帧用于测试。根据nnU-Net中的数据处理程序，训练和测试图像都被裁剪为(384,640)的大小。值得注意的是，与其他数据集相比，该数据集中的图像具有独特的长宽比。
显微镜图像（Microscopy）：这个数据集专注于NeurIPS 2022年细胞分割挑战赛[37]中的各种显微镜图像的细胞分割。它包含1000张训练图像和101张评估图像。Microscopy中的图像在训练和测试时被裁剪为(512,512)的大小。默认情况下，它是一个实例分割数据集。我们为这个数据集采用了与[23]中描述的相同的数据处理策略。
3.2、实现细节
我们在广受欢迎的nnU-Net框架[11]之上实现了Swin-UMamba。其自配置功能使我们能够专注于网络设计，而不是其他琐碎的细节。损失函数是Dice损失和交叉熵损失的和，我们在每个尺度上执行深度监督[33]。我们使用了带有权重衰减=0.05的AdamW优化器，遵循[24]。采用了余弦学习率衰减，初始学习率等于0.0001。我们使用预训练的VMamba-Tiny模型^7来初始化Swin-UMamba的所有数据集。在训练过程中，我们冻结了预训练模型的前10个周期参数，以对齐其他模块。超参数在所有三个数据集上保持一致，除了训练周期数和特定于数据的设置（例如，图像块大小）。Swin-UMamba在AbdomenMRI数据集上训练了100个周期，在内镜数据集上训练了250个周期，在显微镜数据集上训练了350个周期。遵循[23]，我们禁用了测试时间参数化，以实现更流畅和高效的评估。值得注意的是，通过额外的训练和适当的超参数调整，可以进一步改进，我们将其留待未来的工作。我们的主要目标是评估预训练模型对医学图像分割的影响，而不仅仅是追求最先进（SOTA）的性能。更多详细信息，请参阅我们的代码实现。
3.3、基线方法和评估指标
为了进行全面评估，我们选择了三种类型的方法作为基线方法，包括基于CNN的方法（nnU-Net [11]、SegResNet [13]）、基于Transformer的方法（UNETR [38]、Swin-UNETR [39]、nnFormer [9]）以及最新的基于Mamba的分割网络U-Mamba [23]。具体来说，U-Mamba有两种变体：U-Mamba_Bot和U-Mamba_Enc。U-Mamba_Bot仅在瓶颈部分采用Mamba块，而U-Mamba_Enc在每个编码器阶段都采用Mamba块。我们将Swin-UMamba与U-Mamba_Bot和U-Mamba_Enc进行了比较。值得注意的是，由于预训练模型与U-Mamba的结构差异，将预训练模型直接应用于U-Mamba并不简单[24]。
我们使用Dice相似系数（DSC）和标准化表面距离（NSD）来评估AbdomenMRI和内镜数据集上的分割性能。对于显微镜数据集，我们使用F1分数进行评估，因为它是一个实例分割任务。此外，我们使用fvcore包计算了网络参数的数量（#param）和浮点运算次数（FLOPs），以评估每个模型的规模和计算负担。基线结果中，除了nnFormer的DSC、NSD和F1分数参考自[23]外，其余结果均基于官方实现进行报告。
3.4、AbdomenMRI数据集上的比较
表1展示了在AbdomenMRI数据集上的分割性能。Swin-UMamba和Swin-UMamba†均超过了所有基线方法，包括基于CNN的网络、基于Transformer的网络和基于Mamba的网络。这一卓越的结果证明了Mamba网络在医学图像分割方面的巨大潜力。值得注意的是，所有基于Mamba的网络在DSC和NSD上至少比基于CNN和基于Transformer的基线方法高出1%。Swin-UMamba在DSC上比U-Mamba_Enc提高了1.43%，而U-Mamba_Enc是该数据集上之前的最佳模型。如图3a所示，Swin-UMamba能够准确识别目标器官的形状和类型，而U-Mamba_Enc则无法识别所有器官的类型。
在我们的实验中，基于ImageNet的预训练发挥了至关重要的作用，为Swin-UMamba带来了3.04%的DSC显著提升和4.19%的NSD显著改进。此外，利用基于ImageNet的预训练有助于更快、更稳定的训练，与基线方法相比，仅需要十分之一的训练迭代次数。对于Swin-UMamba†，我们观察到了一个戏剧性的现象。在不使用基于ImageNet的预训练的情况下，Swin-UMamba†无法在此数据集上使用默认设置正常收敛。为了解决这个问题，我们禁用了Swin-UMambat的深度监督，并将其训练周期延长至200个周期。尽管如此，当使用ImageNet预训练权重时，Swin-UMamba†仍然超过了所有基线方法。考虑到Swin-UMambat的网络参数不到之前最佳模型U-Mamba的一半，且FLOPs大约只有三分之一，这一改进尤为值得注意。
我们还观察到Swin-UMambat和Swin-UMamba在参数数量和FLOPs之间存在差异。这种差异主要归因于基于CNN的解码器，因为Swin-UMamba†和Swin-UMamba在编码器部分具有几乎相同的结构。我们选择保留基于CNN的解码器，因为它在大多数情况下都能获得更好的结果。
3.5、内窥镜数据集上的比较
表2展示了各模型在内窥镜数据集上的分割性能。Swin-UMamba在DSC上比U-Mamba_Bot高出2.24%，在NSD上高出2.35%。值得注意的是，与AbdomenMRI数据集相比，使用预训练模型时的性能差距要大得多。图3b展示了Swin-UMamba在内窥镜数据集上的可视化结果。此外，我们发现Swin-UMamba在使用预训练模型后，在DSC上获得了13.08%的显著性能提升，在NSD上获得了13.21%的显著性能提升。一个可能的解释是，内窥镜数据集比AbdomenMRI数据集小，模型更容易对训练数据过拟合。利用预训练模型的力量是缓解此类小数据集中过拟合问题的有效策略。此外，我们还发现，当从头开始训练时，Swin-UMamba†的性能优于Swin-UMamba，这可能得益于其较少的参数，从而有助于避免过拟合。
3.6、显微镜数据集上的比较
表3展示了在显微镜数据集上的分割性能。SwinUMamba和Swin-UMamba†继续以2.42%至18.83%的幅度超越所有基线方法。与先前提到的数据集相比，显微镜数据集具有更高的图像分辨率、更少的样本和更大的视觉差异。这对模型的长距离信息建模能力和数据效率提出了更高的要求。如图3c所示，Swin-UMamba能够准确地分割目标细胞，而U-Mamba_Enc则遗漏了一些。有些令人惊讶的是，我们观察到Swin-UMambat在这个数据集上的表现优于Swin-UMamba，这表明了纯基于Mamba的网络在图像分辨率较高且样本有限的情况下具有潜力。
4、结论
本研究旨在揭示基于ImageNet的预训练对Mambabased模型在二维医学图像分割中的影响。我们提出了一种新型的Mambabased模型Swin-UMamba及其变体Swin-UMambat，两者都能够利用预训练模型的力量进行分割任务。我们在各种医学图像分割数据集上的实验表明，对于基于Mamba的模型，基于ImageNet的预训练具有多个优势，包括更高的分割精度、稳定的收敛性、减轻过拟合问题、数据效率和较低的计算资源消耗。我们认为，我们的研究结果强调了预训练在提高Mamba-based模型在视觉任务中的性能和效率方面的重要性。