论文链接：https://arxiv.org/pdf/2402.03302
Code： https://github.com/JiarunLiu/Swin-UMamba
来源： Medical Image Computing and Computer Assisted Intervention -- MICCAI 2024
摘要：
医学图像分割的挑战在于需要同时整合多尺度信息，从局部细节到全局依赖关系。但现有方法难以有效建模长距离的全局信息：
- 卷积神经网络（CNN）由于受限于局部感受野，难以捕获远距离的全局信息。
- 视觉变换器（ViTs）虽然能处理全局信息，但其注意力机制在高分辨率图像上具有高二次复杂度，计算成本较大。
最近，基于Mamba的模型因为其在长序列建模方面表现出色，越来越受到关注。多项研究已显示，这些模型在多种任务中的表现优于流行的视觉模型，具有更高的准确率、较低的内存需求和较少的计算资源。
目前大多数Mamba模型都是从头训练（随机初始化）而不是利用预训练，而预训练已被证明能显著提升数据有限情况下的医学图像分析效果。
为此，本文提出了一个新颖的Mamba基础模型------Swin-UMamba，专为医学图像分割设计，结合了基于ImageNet的预训练优势。实验显示，ImageNet预训练在提升Mamba模型性能方面起到了关键作用。
Swin-UMamba在多个数据集上（如腹部MRI、内窥镜、显微镜）都表现优异，明显优于CNN、ViTs及其他最新的Mamba模型，平均提升了2.72%的性能。
1. 引言
医学图像分割的作用与挑战
- 作用：医学图像分割在诊断、治疗方案制定和疗法实施中发挥关键作用。它可以帮助医生自动识别和划分医学影像中的不同结构，提升诊断效率和准确性。
- 现有方法的局限：目前大部分分割依赖经验丰富的医生进行手工操作，这不仅耗费大量人力和时间，而且不同医生之间的判读可能存在差异，影响结果的一致性。
深度学习在医学图像分割中的应用
- 进步：近年来，深度学习已极大推动医学图像分割的发展，但仍然存在一些难题。
- 需求：高效准确的分割需要结合局部特征（细节信息）和全局依赖（整体结构、关系），这是一个复杂的任务。
模型的优势与局限
- 卷积神经网络（CNNs）：如SegResNet、U-Net、nnU-Net，擅长提取局部特征（如边缘、纹理），但在捕捉长距离全局信息方面有限，因为它们的感受野有限。
- 视觉Transformers（ViTs）：能处理全局上下文，适合捕获长程依赖，但因为基于注意力机制，计算复杂度很高，随着输入分辨率增加会呈指数增长（quadratic complexity），在高分辨率医疗图像（如全-slide病理切片或高分辨率MRI/CT）中面临挑战。此外，ViTs对数据量也很敏感，容易过拟合，尤其是在数据有限的医学领域。
结构化状态空间模型（SSMs）与Mamba模型
- 近期研究表明，结构化状态空间模型（比如Mamba）在长序列建模方面十分高效，能够线性或接近线性扩展、同时捕获长距离依赖。
- 在视觉任务中，相关研究（如Vim、VMamba）展示了这些模型的潜力。它们优于传统模型，具备更高的性能、更低的存储和计算需求。
- 具体应用在医学图像分割上，包括基于nnU-Net和Swin-UNETR的专用架构（U-Mamba、SegMamba），已取得不错的效果。
关键问题与挑战
- 预训练的缺乏：目前大多数Mamba模型是从零开始训练的，缺乏借助ImageNet等大规模预训练模型的经验，导致模型可能难以充分利用已有的知识。
- 迁移学习的重要性：预训练已在CNN和ViT中被证明能提升数据有限情况下的效果，探索其在Mamba模型中的作用尤为关键。
- 模型结构调整的需求：由于Mamba块在视觉领域的应用尚属新颖，现有模型需重新设计以更好地利用预训练模型。
- 实用性与效率：在资源有限的临床环境中，模型的规模和运算效率至关重要，需要确保模型具有良好的可扩展性和部署能力。
模型设计与目的：
- Swin-UMamba：这是一个结合了通用预训练视觉模型（如通过ImageNet预训练的模型）与专门为医学图像分割设计的解码器（decoder）的网络。它利用预训练模型的强大特征提取能力，从而提高医学图像分割的性能。
- Swin-UMamba†：这是Swin-UMamba的变体，采用了基于Mamba的解码器，具有更少的参数和较低的计算量（FLOPs），从而实现更高效的应用。
创新点和贡献：
- 首创探索预训练Mamba网络在医学图像分割中的作用：这是首次尝试评估预训练的Mamba模型对医学图像分割任务的影响，实验结果显示，基于ImageNet的预训练在提高模型性能方面起到关键作用，有时候甚至是至关重要的。
- 提出新型网络结构 ：
- 设计了Swin-UMamba，旨在充分发挥预训练模型的优势，同时保持模型架构的统一性。
- 还提出了参数更少、计算更低（低FLOPs）的Swin-UMamba†变体，以适应资源有限的应用场景，同时保持了良好的性能。
性能表现：
- 结果显示，Swin-UMamba和Swin-UMamba†均超过了之前的多种模型，包括传统的卷积神经网络（CNNs）、视觉变换器（ViTs）以及最新的基于Mamba的模型，在多个医学图像数据集上都表现出明显优势。
- 这一表现充分验证了预训练（特别是ImageNet预训练）和所提出架构在医学图像分割中的有效性。
2. 方法
整体架构：由三个主要部分组成：
- Mamba基础的编码器：在大规模数据集（如ImageNet）预训练，用以多尺度提取特征。
- 解码器：利用多个上采样块逐步恢复图像的空间信息，以生成分割结果。
- 跨级跳跃连接：连接不同尺度的特征，弥补低层细节与高层语义之间的差距。
2.1 Mamba基础的VSS块
Mamba模型的背景：
- Mamba在自然语言处理中的应用利用空间状态序列模型（SSMs）将注意机制的复杂度从二次（quadratic）降低到线性（linear），提升了模型的效率和性能。
- 在计算机视觉中的应用优势包括更高的准确性、较低的计算成本和较少的内存消耗。
挑战：
- 2D视觉数据与1D语言序列的差异：视觉上的空间关系至关重要，但直接将Mamba应用到扁平化的图像中会导致"感受野"受限，难以捕捉未扫描区域的关系。
VSS块的设计：
- 目的：解决二维图像数据的处理难题，确保全面捕获空间关系。
- 方法 ：
- 将输入图像块沿四个不同的方向展开（如上下左右），形成四个不同的序列（ z v z_v zv，其中 v v v表示方向）。
- 每个序列经过空间状态序列模型（S6）进行处理，获取长距离依赖关系。
- 最后，将四个方向的输出序列合并，形成完整的二维特征图。
- S6（空间状态序列模型）
- 核心作用：允许序列中的每个元素都能与之前扫描过的样本进行交互，这通过压缩的隐状态实现，从而高效捕获长距离依赖关系。
- 原理：模仿文本处理中的序列关系，从而在图像或特征序列中建模空间依赖。
2.2 基于ImageNet预训练的集成
Swin-UMamba的编码器由五个阶段组成：
- 第一个阶段是"干线（stem）"阶段，采用一个7×7卷积层进行2×下采样，配合padding和步长为2，之后加上2D实例归一化。不同于
VMamba的设计，作者在此阶段采用了渐进式的2×下采样策略，以保留低层细节。 - 第二个阶段采用2×2的
patch embedding层，保持特征分辨率为原始图像的1/4，与VMamba中的特征保持一致。 - 后续阶段则遵循
VMamba-Tiny的设计格式，每个阶段由patch merging层进行2×下采样，配合若干VSS块（空间状态序列模型块）提取高层次特征。VSS块的数量在各阶段分别为 2 , 2 , 9 , 2 {2, 2, 9, 2} 2,2,9,2。这些阶段的特征维度随着阶段递增，最终达到 D = { 48 , 96 , 192 , 384 , 768 } D = \{48, 96, 192, 384, 768\} D={48,96,192,384,768}。
预训练方面，作者将从ImageNet预训练所得的权重初始化VSS块和patch merging层，以提升模型性能。然而，由于patch embedding层的patch大小和输入通道的差异，其未采用预训练的权重。
2.3 主干解码器
- 引入了带残差连接的额外卷积块，用于处理跳跃连接的特征。
- 在每个尺度层加入深层监督的分割头，以提升训练效果。
- 在上采样阶段
- 利用拼接操作（Cat）结合上一尺度和当前尺度的特征
- 通过两个残差卷积块 （Res）处理后，经过转置卷积 （DeConv）和1×1卷积生成最终的分割图。
2.4 Swin-UMamba†
目标与背景
- 目的：探索
Mamba模型在医学语义分割中的潜力。 - 方法：通过将
Swin-UMamba的解码器替换为基于Mamba的解码器 ，得到Swin-UMamba†，希望其在性能上具有竞争力，同时拥有更低的参数量和计算成本。 - 发现：大规模预训练在医学图像分割中具有重要作用，无论解码器结构如何。
主要改动
- 替换上采样块 ：原
Swin-UMamba中的上采样块被替换为patch expanding和两个VSS块。Patch expanding是直接将输入图像从原始尺寸 （ H × W × C ） （H×W×C） （H×W×C）通过4×4的补丁嵌入层投影到特征图 （ H / 4 × H / 4 × 96 ） （H/4 × H/4 × 96） （H/4×H/4×96），从而简化了结构。 - 减少参数：原模型参数为60M，FLOPs为68.0G；改造后，参数减至28M，FLOPs降低至18.9G。
- 移除部分连接：去除了源自输入图像和较低分辨率特征的跳跃连接，因为这些在解码过程中没有对应的解码块。
- 深度监督：在尺度为1×、1/4×、1/8×和1/16×的特征层上加入了深度监督（通过1×1卷积映射到K类），增强模型的监督信号。
结构示意 ：Swin-UMamba†的结构由多次4×上采样和VSS块组成，整体架构简化，更加高效，结构细节如图所示。
3. 实验
3.1 数据集：
Abdomen MRI（腹部MRI）：该数据集由MICCAI 2022的AMOS挑战提供，目标是分割13个腹部器官（如肝脏、脾脏、胰腺、左右肾、胃、胆囊、食管、主动脉、下腔静脉、左右肾上腺和十二指肠）。训练集包含60个MRI扫描（共5615个切片），测试集包括50个扫描（共3357个切片）。图像被裁剪成尺寸为（320, 320）的块，利用nnU-Net框架进行训练和测试。
Endoscopy（内窥镜）：该数据集来自MICCAI 2017的EndoVis挑战，目标是分割7种内窥镜仪器（如针驱动器、夹钳、弯曲剪刀、卡迪尔钳、双极钳、血管封堵器和超声探头）。训练集有1800帧图像，测试集1200帧。图像裁剪尺寸为（384, 640），且该数据集的图像具有特殊的长宽比例，展示出不同的影像特性。
Microscopy（显微镜）：来自NeurIPS 2022细胞分割挑战，用于细胞实例分割。训练集包含1000张图像，评估集有101张。图像裁剪为（512, 512）尺寸。
3.2 实现细节：
基础框架 ：模型在成熟的nnU-Net框架上构建。nnU-Net具有自适应配置的特性、。
损失函数 ：结合Dice损失 和交叉熵损失 的加和，用于训练。并在不同尺度上引入深层监督，以提升模型的梯度传播效果。
优化器 ：采用AdamW优化器，权重衰减（weight decay）设置为0.05，以改善正则化效果。
学习率策略 ：使用 余弦退火（cosine decay） 的学习率调整策略，初始学习率设为0.0001。
预训练模型的使用 ：利用预训练的VMamba-Tiny模型来初始化Swin-UMamba的参数，确保模型在训练 少量数据和加快收敛。
参数冻结：在训练的前10个epoch中，冻结预训练模型中的参数，以使其他模块（例如新加的层）更好地适应特定任务。
训练周期差异：
- AbdomenMRI：训练100个epoch。
- Endoscopy：训练350个epoch。
- Microscopy：训练450个epoch。
3.3 基线方法和评估指标：
- 对比模型 ：使用三类不同的模型作为基线，以确保评估的全面性：
- 指标 ：
- Dice 相似系数（DSC）：用于评价AbdomenMRI和Endoscopy数据集的全局分割效果。
- 归一化表面距离（NSD）：衡量边界的精确性。
- F1分数：对于显微镜数据集（实例分割任务）。
- 模型复杂度 ：
- 计算网络参数数（#param）；
- 浮点运算量（FLOPs），均用fvcore包衡量。
3.4 在AbdomenMRI数据集上的表现：
Swin-UMamba 和 Swin-UMamba† 都优于所有基线方法（包括CNN、Transformer和其他Mamba模型），证明了它们在医学图像分割中的潜力。
这些模型在Dice相似系数（DSC）和归一化表面距离（NSD）指标上都比CNN和Transformer方法高至少1%，且Swin-UMamba在DSC上比前一状态最优模型（U-Mamba_Enc）提升了1.34%。
ImageNet预训练 在性能提升中起到了关键作用。例如，使用预训练后，Swin-UMamba的DSC和NSD分别提升了7.06%和7.74%，且预训练模型使训练更快、更稳定，只需训练的迭代次数的十分之一。
Swin-UMamba† 在没有预训练时，未能良好收敛，即使延长训练时间也如此，但在使用预训练权重后，表现明显优于所有基线，且参数量和计算量都少于前一状态最优模型U-Mamba。
图像显示Swin-UMamba能更准确识别目标器官的形状和类型，而基线模型则未能全部正确识别。
3.5 在Endoscopy（内窥镜）数据集上的表现：
- Swin-UMamba†在DSC和NSD指标上比U-Mamba_Bot分别提升2.43%和2.41%。
- 预训练带来的性能提升非常明显（分别提升了12.84%和12.90%），尤其是在数据较少、易过拟合的小数据集（如该内窥镜数据集）上，预训练显得尤为重要，有助于缓解过拟合。
- Swin-UMamba†的参数较少，更易避免过拟合，表现优于未加入预训练的模型。
3.6 在Microscopy（显微镜）细胞图像数据集上的表现：
- Swin-UMamba 和Swin-UMamba†继续超越所有基线模型，提升范围从1.99%到20.15%。
- 由于Microscopy数据集具有更高的图像分辨率、样本更少且差异更大，因此对模型的长距离依赖建模能力和数据效能提出了更高要求。
- 预训练再一次表现出显著优势（提升12.45%和7.96%），即使在这种更复杂的高分辨率、有限样本的场景下，预训练依然非常关键。
4. 结论
本研究旨在探讨基于ImageNet预训练对Mamba模型 在二维医学图像分割中的影响。研究团队提出了一种新颖的基于Mamba的模型，叫做Swin-UMamba ，以及它的变体Swin-UMamba† ，这两种模型都能充分利用预训练模型的能力来改善分割任务。在多个医学图像分割数据集上的实验结果显示，采用ImageNet预训练的方法带来了多方面的优势，包括更高的分割准确率、训练过程更稳定、减少了过拟合问题、提高了数据利用效率，以及降低了模型的计算资源需求。