Swin-UMamba：基于Mamba和ImageNet预训练的医学图像分割模型

点击下方卡片，关注“CVer”公众号

AI/CV重磅干货，第一时间送达

最后一天特惠活动！CVer学术星球新人和续费优惠券来了！为了感谢大家的支持，限时赠送10张50元新用户优惠券(下图左边二维码)，10张20元老用户续费券(下图右边二维码)，龙年坚持分享更多更好的工作，寒假不停更！助力你的科研和工作！顺利毕业升学，升职加薪！

代码&模型：https://github.com/JiarunLiu/Swin-UMamba

论文：https://arxiv.org/abs/2402.03302

Highlight

1. 相比较于U-Mamba，Swin-UMamba在三种医学图像分割任务指标上可以取得平均3.58%的提升。

2. 该篇工作通过实验验证了ImageNet预训练对基于Mamba的医学图像分割模型起到非常重要的作用，在迭代次数不变的情况下最高可为Swin-UMamba带来13.08%的DSC提升。

3. 提出了一种变体网络Swin-UMamba，其仅需要相比于U-Mamba不到1/2的网络参数量和约1/3的FLOPs就能够实现与Swin-UMamba相近的性能。

研究背景介绍

对精准的医学图像分割来说，局部信息建模和全局信息建模都很重要，目前两种主流的网络结构CNNs和ViTs在全局信息建模上都有一定的缺陷：CNNs在局部特征提取上非常有效，但由于其固有的局部归纳偏置导致感受野受限；ViTs具备处理全局上下文信息的能力，然而ViT的注意力机制具有较高的二次计算复杂度，这不利于高分辨率图像的建模，高分辨率图像在医学领域并不罕见。

最近的一些工作尝试将Mamba应用到视觉任务中，例如U-Mamba、VMamba和Vim。基于Mamba的模型优势十分鲜明：准确率更高、显存占用更少、计算开销更低。然而，这些工作主要关注于视觉模型结构设计，模型本身往往是从零开始训练的，没有进一步探讨预训练Mamba模型在医学图像分割中的作用——预训练已经被证明在CNNs和ViTs中是非常有效的，尤其是考虑到许多医学图像数据集数据量较少，使用预训练模型可以有效缓解模型过拟合问题并提升模型泛化能力。因此该工作提出了Swin-UMamba，希望借助ImageNet预训练模型的力量进一步提升基于Mamba的模型在医学图像分割任务中的性能。

方法

如上图所示，Swin-UMamba使用类似于UNet的结构，包括编码器、解码器以及两者之间的跳跃连接。作者在设计Swin-UMamba时首先考虑的问题是如何更好地利用ImageNet预训练模型学习到的多尺度特征进行医学图像分割任务，为此Swin-UMamba的编码器部分基本遵循VMamba-Tiny的网络结构设计以便加载预训练模型参数，主要由patch merging模块和VSS模块构成。与VMamba不同的是，Swin-UMamba的第一个stage使用一个卷积层进行2倍下采样，这是为了使每一个stage的下采样倍数都能保持一致，对应地，Swin-UMamba将patch embedding的图像块尺寸也由调整为。对于解码器部分，Swin-UMamba在上采样模块中加入了一个额外的残差卷积模块用来处理跳跃连接的特征，并在每个尺度上都加入了一个额外的分割头用于深度监督，上采样模块如图1右上角所示。

为了进一步探索Mamba在医学图像分割中的潜力，并验证预训练模型在不同网络结构下的影响，作者还提出了一种变体网络Swin-UMamba 。相比于Swin-UMamba，Swin-UMamba 的网络参数更少、FLOPs更低，且仍能保持与Swin-UMamba相近的结果。为此做了几点修改：首先上采样模块被替换为patch expanding和2个VSS Block，这可以有效地降低网络参数量和计算开销；其次，Swin-UMamba解码器的最后使用了一个4倍上采样的patch expanding操作；对应地，作者去除了Swin-UMamba stage 1中的卷积层，并将编码器中patch embedding的图像块尺寸恢复为。通过上述操作，Swin-UMamba的参数量从40M降至27M、FLOPs从58.4G降至15G，Swin-UMamba的网络结构如图2所示。

实验结果

实验结果如表1、表2和表3所示，在AbdomenMRI（腹部MRI器官分割）、Endoscopy（内窥镜手术器械分割）和Microscopy（细胞分割）三个数据集上，Swin-UMamba和Swin-UMamba相比于baseline方法均有显著的提升。例如，Swin-UMamba相比较于U-Mamba_Enc的平均指标提升3.8%，Swin-UMamba相较于U-Mamba_Enc的平均指标提升2.4%。这里需要注意的一点是Swin-UMamba仅需U-Mamba_Enc不到1/2的网络参数量和约1/3的FLOPs，因此Swin-UMamba的性能提升也是非常可观的。

作者发现ImageNet预训练对Swin-UMamba和Swin-UMamba都起到了非常重要的作用，包括更高的分割精度、更稳定的收敛、减轻过拟合问题、数据效率和较低的计算开销。例如，相较于不使用ImageNet预训练，使用ImageNet预训练能够为Swin-UMamba在Endoscopy上的DSC带来13.08%的显著提升，这可能是因为Endoscopy数据集较小，导致模型更容易过拟合。此外，Swin-UMamba仅需baseline算法1/10的训练迭代次数就能够在AbdomenMRI数据上收敛，这有利于节省训练期间的计算开销。实验中还观察到了一个非常有趣的现象，ImageNet预训练有时对模型收敛稳定性起到至关重要的作用，Swin-UMamba在没有使用预训练模型时很难在AbdomenMRI数据上正常收敛，而当使用预训练模型后，Swin-UMamba又能够以较少的参数量和FLOPs超过所有baseline算法，这也进一步佐证了预训练对Mamba-based模型的重要性。

最后提供一些可视化案例供大家参考：

有关Swin-UMamba的更多细节可以参考arXiv上的文章，模型和代码都已在GitHub上开源，欢迎大家一起交流讨论。

Reference

• J. Liu et al., “Swin-UMamba: Mamba-based UNet with ImageNet-based pretraining.” arXiv, Feb. 05, 2024.

• Y. Liu et al., “VMamba: Visual State Space Model.” arXiv, Jan. 18, 2024.

• J. Ma, F. Li, and B. Wang, “U-Mamba: Enhancing Long-range Dependency for Biomedical Image Segmentation.” arXiv, Jan. 09, 2024.

• L. Zhu, B. Liao, Q. Zhang, X. Wang, W. Liu, and X. Wang, “Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model.” arXiv, Jan. 17, 2024.

最后一天特惠活动！CVer学术星球春节优惠券来了！为了感谢大家的支持，现在赠送10张50元新用户优惠券(左图领取)，10张20元老用户续费券(右图领取)，龙年坚持分享更多更好的工作，寒假不停更！共同进步！

医学影像交流群成立

扫描下方二维码，或者添加微信：CVer444，即可添加CVer小助手微信，便可申请加入CVer-医学影像微信交流群。另外其他垂直方向已涵盖：目标检测、图像分割、目标跟踪、人脸检测&识别、OCR、姿态估计、超分辨率、SLAM、医疗影像、Re-ID、GAN、NAS、深度估计、自动驾驶、强化学习、车道线检测、模型剪枝&压缩、去噪、去雾、去雨、风格迁移、遥感图像、行为识别、视频理解、图像融合、图像检索、论文投稿&交流、PyTorch、TensorFlow和Transformer、NeRF等。
一定要备注：研究方向+地点+学校/公司+昵称（如医学影像+上海+上交+卡卡），根据格式备注，可更快被通过且邀请进群

▲扫码或加微信号: CVer444，进交流群
CVer计算机视觉（知识星球）来了！想要了解最新最快最好的CV/DL/AI论文速递、优质实战项目、AI行业前沿、从入门到精通学习教程等资料，欢迎扫描下方二维码，加入CVer计算机视觉（知识星球），已汇集近万人！

▲扫码加入星球学习

▲点击上方卡片，关注CVer公众号
整理不易，请点赞和在看