当前位置：首页 > news >正文

【论文解读】MSVM-UNet: Multi-Scale Vision Mamba UNet for Medical Image Segmentation

news 来源：原创 2025/4/29 15:05:05

MSVM-UNet: Multi-Scale Vision Mamba UNet for Medical Image Segmentation

论文链接： https://arxiv.org/abs/2408.13735

Code： https://github.com/gndlwch2w/msvm-unet

来源： 2024 IEEE International Conference on Bioinformatics and Biomedicine

1. 论文信息

本文介绍了一种名为Multi-Scale Vision Mamba UNet（MSVM-UNet）用于医学图像分割，旨在解决Mamba用于2D图像数据方向敏感以及如何有效的学习多尺度详细特征表示和全局上下文依赖性来提高医学图像分割的准确性的问题。通过在VSS块中引入多尺度卷积，使其能更有效的捕获和聚合从VMamba编码器的层次特征中提取多尺度特征表示，更好的处理2D图像数据；此外Large Kernel Patch Expanding（LKPE）层通过同时聚合特征图的空间信息和通道信息来进行更高效的上采样。

2. 论文概要

2.1 研究背景

研究问题： 精确高效地进行医学图像分割是医学图像分析领域的一项基础且具有挑战性的任务。在这个领域的研究中，利用诸如深度学习等技术分析各种类型的医学图像并产生特定器官或病变区域的分割图，以协助医生和研究行人分析疾病和做出诊断。
研究难点： 尽管使用卷积神经网络（CNNs）和视觉 Transformer （ViTs）进行医学图像分割取得了显著的成功，但是CNNs依赖于局部卷积核进行特征提取，虽然对于捕捉局部特征模式有效，但限制了其描述全局和几何特征的能力，基于 Transformer 的方法在建模长程依赖方面表现良好，但是自注意力机制相对于序列长度具有平方的计算复杂度。U-Mamba 提出了在SSMs中嵌入卷积操作，以将卷积层局部特征提取的能力与SSMs的长序列依赖捕获能力相结合。Swin-UMamba表明将VMMamba在ImageNet-1k上预训练的模型转移到医学图像分割领域，可以有效地解决有限的数据资源问题。
文献综述：U-Mamba和Swin-UMamba 采用CNN和SSMs的混合结构，但没有具体解决多尺度特征学习问题，导致在分析目标大小时存在缺陷。为了解决这些问题，作者提出了一种多尺度视觉状态空间（MSVSS）块，它使用一套具有不同核大小的并行卷积操作来捕获和聚合多尺度特征表示，不仅可以模拟原始四个方向的依赖关系，还可以使用卷积操作聚合剩下的四个对角线方向的信息。Swin-UNet 和 VM-Net 中，都使用了扩展 Patch 化的卷积层进行特征上采样。然而，由于扩展 Patch 化卷积层只考虑了通道信息，而没有考虑上采样过程中的空间关系，导致其判别力不足。为解决这个问题，作者提出了一种大核 Patch 扩展（LKPE）层进行上采样。该层通过在扩展通道维度的深度卷积来集成空间信息，从而实现更具判别力的特征表示

2.2 本文贡献

提出了一种新的多尺度视觉状态空间（MSVSS）块，结合了 CSM 与多尺度卷积操作，不仅能够有效地模拟像素之间的长序列依赖关系，而且还能捕捉多尺度特征表示；
提出了一种大核 Patch 扩展（LKPE）层进行上采样。该层通过在扩展通道维度的深度卷积来集成空间和通道信息，从而实现更具判别力的特征表示；

3. 创新方法

MSVM-UNet采用了一个U型分层次编码器-解码器架构，并在编码器与解码器之间采用了short-cut连接。

编码器采用由ImageNet-1k数据集预训练的VMamba V2，其包含4个stage，第一个stage由Patch Embedding模块和VSS块组成，剩下3个stage均由负责下采样的Patch Merging模块和VSS块组成；
解码器包含3个stage和FLKPE输出层构成，其中每个stage均由负责上采样的LKPE模块和MSVSS块组成，MSVSS模块的输入来自编码器VSS块的输出和来自负责上采样的LKPE模块的输出的concat结果；

3.1 Multi-Scale Vision State Space (MSVSS) Block

MSVSS块通过在VSS块中引入多尺度前馈网络（MS-FFN）来解决同时捕获多尺度详细特征和在二维视觉数据中有效地解决方向敏感性问题。

二维选择扫描块（SS2DBlock）实现了每个特征在四个方向上的长程依赖性；
MS-FFN中的卷积操作从四个剩余的对角方向聚合信息，以增强特征表示；
为了有效地捕获和聚合多尺度特征表示，MSVSS 采用了一组具有不同核大小并行的卷积操作来实现这一目标；

2D-Selective-Scan Block (SS2DBlock)

SS2D首先将2D输入特征图沿着四个不同的扫描路径进行 flatten，得到四个一维序列。这些序列随后被输入S6块进行选择性扫描，以模拟长程依赖关系。最后，将四个一维序列恢复为原始2D形式并将它们相加以产生输出。

Multi-Scale Feed-Forward Neural Network (MS-FFN)

MSFNN中引入了卷积操作来聚合这四个对角方向的信息，并采用了一组使用不同 Kernel 大小的卷积操作用于有效地捕捉层叠特征的多细节信息和高分辨率特征表示

3.2 Large Kernel Patch Expanding (LKPE) Layer

Patch Expanding层仅考虑特征的通道信息，而忽略了相邻特征之间的空间关系，为解决此问题，LKPE首先应用一个卷积将通道维数翻倍，然后进行批量归一化和ReLU激活函数，接着使用有效卷积聚合空间信息，并最终通过扩展包含空间和通道信息的特征表示进行上采样

4. 实验分析

4.1 实验细节

在NVIDIA GeForce RTX 3090基于Pytorch框架实现
数据增强：将图像Resize至224大小、水平翻转、垂直翻转、随机旋转、高斯噪声、高斯模糊和对比增强；
Batch Size：32
优化器：AdamW
迭代次数：300
学习率：初始学习率为5e-4，并采用余弦衰减策略
损失函数：Dice Loss和CE Loss的组合

4.2 定量比较

在腹部多器官数据集（Synapse）上训练的各种医学图像分割方法的定量比较中，MSVM-UNet在多个指标均超过其他基线方法。

4.3 定性比较

在腹部多器官数据集（Synapse）上训练的各种医学图像分割方法的定性比较中，MSVM-UNet不仅能有效处理形状和大小各异的器官，还能更好地定位器官边界。

4.4 消融研究

在腹部多器官数据集（Synapse）上消融实验表明：MSVSS块与LKPE能提高分割Dice指标；
对不同上采样方式进行消融实验表明：LKPE效果最优；
对不同大小卷积核进行消融实验表明：卷积核为1、3、5时，多尺度特征提取模块效果最优；

5. 结论

本文提出了一种新颖的多尺度视觉Mamba UNet，旨在解决医学图像分割面临的挑战。由于多尺度深度卷积的设计，MSVM-UNet不仅能捕获不同尺度下的信息并建模所有方向上的长序列依赖性，还能保持计算效率和可接受的参数数量；
通过有效集成通道和空间信息进行上采样，MSVM-UNet实现了更具有判别性的特征表示，从而使得医学图像分割的结果更加准确；