点击下方卡片，关注 「AI视界引擎」 公众号
( 添加时备注：方向+学校/公司+昵称/姓名 )
状态空间模型（SSMs），特别是Mamba，在医疗图像分割领域显示出巨大的潜力，因为它们能够以线性计算复杂度来模拟长程依赖性。然而，准确的医疗图像分割需要有效地学习多尺度详细特征表示和全局上下文依赖性。尽管现有工作通过将卷积神经网络（CNNs）与SSMs集成，以发挥各自的优势，试图解决这个问题，但没有设计用于有效捕获和聚合多尺度特征表示的专用模块，也没有充分解决将Mamba应用于2D图像数据时的大小敏感性问题。为了克服这些局限，作者提出了一种名为MSVM-UNet的多尺度视觉Mamba UNet模型，用于医疗图像分割。
Specifically，通过在VSS块中引入多尺度卷积，作者可以更有效地捕获和聚合从VMamba编码器的层次特征中提取的多尺度特征表示，更好地处理2D视觉数据。
此外，LKPE（Large kernel patch expanding）层的扩撒能够同时集成空间和通道信息，实现特征图的有效上采样。
在Synapse和ACDC数据集上的大量实验表明，作者的方法在捕获和聚合多尺度特征表示以及像素之间的长程依赖性建模方面比一些最先进的方法更为有效。
I Introduction
精确高效地进行医学图像分割是医学图像分析领域的一项基础且具有挑战性的任务。在这个领域的研究中，利用诸如深度学习等技术分析各种类型的医学图像并产生特定器官或病变区域的分割图，以协助医生和研究行人分析疾病和做出诊断。
近年来，使用卷积神经网络（CNNs）和视觉 Transformer （ViTs）进行医学图像分割取得了显著的成功。具体来说，UNet [1]，由于其优雅的U形结构以及跳跃连接，在处理高分辨率医学图像方面表现出色，同时将低级细节与高级语义无缝结合，实现了令人印象深刻的分割结果。此外，TransUNet [2]提出了一种CNNs与ViTs的混合结构，以同时利用CNNs的细节提取能力和ViTs的长程依赖建模能力。虽然这些方法在性能上取得了令人称道的表现并产生了高质量的分割结果，但CNNs和 Transformer 固有的特性导致了性能瓶颈 [2, 3]。具体来说，CNNs依赖于局部卷积核进行特征提取，虽然对于捕捉局部特征模式有效，但限制了其描述全局和几何特征的能力 [4]。尽管基于 Transformer 的方法在建模长程依赖方面表现良好，但是自注意力机制相对于序列长度具有平方的计算复杂度 [5]，这使得在高分辨率的分割任务上高效处理具有挑战性。此外，例如Swin Transformer [6]，PVT v2 [7]和BiFormer [8]提出了有效的自注意力计算技术。然而，这些方法在计算复杂性和建模能力之间存在权衡，因此限制了它们在高序列建模方面的能力。
最近，状态空间模型（SSMs） 由于在模拟长序列方面具有巨大的潜力，受到了广泛关注。Mamba[11] 用线性计算复杂度高效建模长序列，在自然语言处理领域取得了显著成功。在这一基础上，VMamba[12] 引入了交叉扫描模块（CSM）和合理设计的层次结构设计，显示出在分析2D图像数据方面具有巨大的潜力。在医学图像分割领域，高效处理高分辨率医学图像仍是一个重大的挑战。受上述工作的启发，U-Mamba[13] 提出了在SSMs中嵌入卷积操作，以将卷积层局部特征提取的力量与SSMs的长程依赖捕获能力相结合。Swin-UMamba[14] 表明将VMMamba在ImageNet-1k上预训练的模型转移到医学图像分割领域，可以有效地解决有限的数据资源问题。与Swin-UNet[15] 相似，VM-UNet[16] 提出使用纯Visual State Space（VSS）块构建医学图像分割框架。
与一维序列不同，2D视觉数据中的像素固有地具有方向依赖性[17]。直接将Mamba中的一维序列处理方法应用到2D数据上，无法有效地捕捉像素之间的长程依赖关系，导致受限于2D数据的响应场，这被称为方向敏感性问题[12]。尽管VMMamba使用了四种扫描策略解决这个问题，但它只关注了四个相邻方向（即，上、下、左、右），导致在与不同大小和形状的物体分析中存在一定局限性。此外，U-Mamba和Swin-UMamba 采用CNN和SSMs的混合结构，但没有具体解决多尺度特征学习问题，导致在分析目标大小时存在缺陷。为了解决这些问题，作者提出了一种多尺度视觉状态空间（MSVSS）块，它使用一套具有不同核大小的并行卷积操作来捕获和聚合多尺度特征表示，不仅可以模拟原始四个方向的依赖关系，还可以使用卷积操作聚合剩下的四个对角线方向的信息。
此外，在 Swin-UNet 和 VM-Net 中，都使用了扩展 Patch 化的卷积层进行特征上采样。然而，由于扩展 Patch 化卷积层只考虑了通道信息，而没有考虑上采样过程中的空间关系，导致其判别力不足。为解决这个问题，作者提出了一种大核 Patch 扩展（LKPE）层进行上采样。该层通过在扩展通道维度的深度卷积来集成空间信息，从而实现更具判别力的特征表示的同时也具备可接受的额外开销增加。
本研究的主要贡献可以总结如下：
- 作者提出了一种新的多尺度视觉状态空间（MSVSS）块，结合了 CSM 与多尺度卷积操作，不仅能够有效地模拟像素之间的长时间依赖关系，而且捕捉多尺度特征表示。
- 作者引入了一种新的大核 Patch 扩展（LKPE）层用于特征图上采样。通过在扩展通道维度前集成大核深度卷积，作者实现了具有可接受额外开销的更具判别力的特征表示。*作者在 Synapse 多器官数据集和 ACDC 数据集上验证了作者提出的 MSVM-Net。具体来说，在 Synapse 多器官数据集上，作者的模型实现了 85.00% 的 DSC 和 14.75mm 的 HD95。在 ACDC 数据集上，作者的模型实现了 92.58% 的 DSC。
II Methods
Overall Architecture of MSVM-UNet
在图1中，作者给出了所提出的MSVM-UNet的整体架构。该架构采用了一个U型分层次编码器-解码器结构，并采用 Short-Cut 。编码器采用由ImageNet-1k数据集预训练的VMMaba V2 [12]，该数据集包含了四个阶段。除了第一个阶段，由纹理嵌入层和VSS块组成，其余三个阶段由融合层和VSS块组成。具体而言，纹理嵌入层将输入划分成尺寸为的非重叠块，并将通道维数映射到维数。VSS块负责学习输入图像的分级特征表示。融合层用于下采特征图。对于输入，作者首先使用编码器的四个阶段依次提取出四级特征表示，分别表示为，，和，然后将它们输入到解码器。具体来说，特征通过扩展路径传递到解码器的最后阶段，而特征，和通过 Short-Cut 传递到解码器相应阶段。解码器包含三个阶段，每个阶段都包括一个Large Kernel Patch Expanding（LKPE）层和一个Multi-Scale Vision State Space（MSVSS）块。与融合层不同，LKPE层负责上采特征图。MSVSS块从压缩路径捕获和聚合细粒度的多尺度信息，并将展开路径的高级语义信息进行聚合。最后，通过最后一个大型 Kernel 块扩展（FLKPE）层得到分割预测。
Multi-Scale Vision State Space (MSVSS) Block
为了同时捕捉分级的详细信息并在二维视觉数据中有效地解决方向敏感性问题，作者提出了多尺度视觉状态空间（MSVSS）块。具体来说，MSVSS 通过在VSS块内引入多尺度前馈网络（MS-FFN）来解决这些问题。首先，二维选择扫描块（SS2DBlock）model了每个特征在四个方向上的长程依赖性，然后，MS-FFN中的卷积操作从四个剩余的对角方向聚合信息，以增强特征表示。此外，为了有效地捕获和聚合多尺度特征表示，MSVSS 采用了一组具有不同核大小并行的卷积操作来实现这一目标。如图1（c）所示，MSVSS块包括两个层规范化层、SS2DBlock和MS-FFN。MSVSS块的定义由方程（1）和（2）给出：
其中，和分别表示第i个阶段的输入和输出特征图。表示SS2DBlock的输出，表示层规范化。
Iii-B1 2D-Selective-Scan Block (SS2DBlock)
2D-Selective-Scan块对输入特征图进行选择性扫描，以捕捉全局上下文信息和长程依赖关系。具体来说，2D输入特征图首先经过一个线性层，一个方向卷积操作和一个激活函数。然后，通过2D-Selective-Scan（SS2D）操作进行更深入的特征提取。最后，在另一层归一化和线性投影之后，获得输出。如图2（a）所示，SS2D首先将2D输入特征图沿着四个不同的扫描路径进行 flatten，得到四个一维序列。这些序列随后被输入S6块 [11]进行选择性扫描，以模拟长程依赖关系。最后，将四个一维序列恢复为原始2D形式并将它们相加以产生输出。SS2D块的定义由方程（3）给出：
其中，和分别表示SS2D块在第i阶段的输入和输出特征图。代表一个线性投影，用于将通道维数翻倍。表示具有核大小为3x3的深度卷积。代表激活函数。表示2D-Selective-Scan操作。代表另一个线性投影，用于将通道维数减半。
Iii-B2 Multi-Scale Feed-Forward Neural Network (MS-FFN)
图2(b)所示，作者在前馈网络中引入了卷积操作来聚合这四个对角方向的信息。此外，为了有效地捕捉层叠特征的多细节信息和高分辨率特征表示，作者采用了一组使用不同 Kernel 大小的卷积操作。为了避免引入过多的参数和计算开销，作者使用了宽深度卷积，它们具有参数和计算效率高、低维张量卷积和高效计算等优点，从而实现了MS-FFN。如图1(d)所示，MS-FFN包括两个线性层和一组并行的宽深度卷积层。MS-FFN的定义由公式(4)和(5)给出：
其中和分别表示MS-FFN在第i个阶段中的输入和输出张量，表示第一线性变换的输出。表示用于通过对数扩展以增加信道维数的线性投影。表示 Kernel 大小为的宽深度卷积。定义了一组并行的卷积核值，即。表示一种激活函数。表示另一个线性投影，它将通道维数减少到输入维数。
Large Kernel Patch Expanding (LKPE) Layer
作者使用LKPE层来将当前阶段的特征图放大，以匹配跳线连接的特征图的尺寸。如图3所示，作者提出了作者提出的LKPE层与作者提出的Patch Expanding层进行比较。与后者所依赖的仅仅使用线性投影（相当于核大小为的卷积）来放大输入特征图的通道维数不同，作者考虑引入大的卷积核。灵感来自其他上采样方法，如转置卷积和上采样[3]，Patch Expanding层仅考虑特征的通道信息，而忽略了相邻特征之间的空间关系，这使这种方法在信息利用方面并非最优。为解决此问题，作者提出了一种大核Patch Expanding层。具体来说，LKPE首先应用一个卷积将通道维数翻倍，然后进行批量归一化和ReLU激活函数。接下来，它使用有效卷积聚合空间信息，并最终通过扩展包含空间和通道信息的特征表示进行上采样。LKPE的定义由公式(6)给出：
其中，和分别表示第i个阶段的特征图在放大之前的和放大后的特征图。表示负责将通道维数翻倍的线性变换。表示批量归一化。表示ReLU(x)=max(0,x)激活函数。表示重新排列操作，表示重形状操作。
Final Large Kernel Patch Expanding (FLKPE) Layer
作者使用FLKPE层来生成分割预测。首先将解码器最后阶段的特征图作为输入，首先使用线性投影来聚合通道维度信息，并将其扩展16倍。接下来，作者使用深度卷积来聚合空间维度信息。随后，作者将得到的结果特征的特征分辨率变为输入图像的大小，同时保持通道维度不变。最后，将转换后的特征图通过卷积映射到分割预测。FLKPE的定义由等式（7）给出：
其中表示解码器的最后一阶段生成的特征图。表示将其上采样至输入图像大小得到的特征图。表示最终的分割预测。表示负责将输入通道维度扩展16倍的线性变换。
III Experiments and Results
在第三部分实验与结果部分的开头。
Datasets
本文将对 proposed MSVM-UNet 模型在 Synapse腹部多器官分割数据集（Synapse）和 Automated Cardiac Diagnosis Challenge 数据集（ACDC）上的性能进行评估。以下是详细的内容：
- Synapse 数据集 [24]
该数据集包含 30 个腹部 CT 扫描及 3779 张轴向对比增强的腹部 CT 图像。每个 CT 卷包含 85 至 198 层，分辨率均为 像素， Voxel 空间分辨率为 。与 TransUNet 类似 [2]，作者将数据集随机划分为 18 个用于训练的案例和 12 个用于测试的案例。本研究只针对 8 种腹部器官进行分割：主动脉、胆囊、左肾、右肾、肝脏、胰腺、脾脏和胃部。
Iii-A2 ACDC dataset [25]
该数据集包含100张心脏MRI扫描图像，每张图像包括三个子器官：右心室（RV）、心肌（Myo）和左心室（LV）。遵循TransUNet [2]的划分，作者将数据集分割为70个病例用于训练，10个病例用于验证，以及20个病例用于测试。
Implementation Details and Evaluation Metrics
开始写 实现细节与评价指标 部分。
Iii-B1 Implementation Details
在作者的实验中，所有模型基于Pytorch 2.0.0框架实现，所有训练在NVIDIA GeForce RTX 3090 GPU上进行。作者使用ImageNet-1k数据集的预训练权重初始化 Backbone 网络。为了减少过拟合并增强模型的泛化能力，作者采用了广泛的数据增强技术，包括将输入图像大小调整为，水平翻转、垂直翻转、随机旋转、高斯噪声、高斯模糊和对比增强。作者将batch大小设置为32，并使用AdamW优化器训练网络至少300个周期。初始学习率设置为5e-4，并在训练过程中使用余弦退火时间表衰减。由于不同数据集的难度 Level 不同，并且为了减少过拟合，作者对不同数据集设置了不同的权重衰减：Synapse数据集为1e-3，ACDC数据集为1e-4。此外，作者使用callflops计算了模型报告的FLOPs和参数数量，输入大小为。作者使用Dice和交叉熵损失函数的组合来训练网络，定义如下：
其中和分别对应Dice损失和交叉熵损失的权重。
Iii-B2 Evaluation Metrics
遵循普遍使用的模型性能评估指标，作者利用Dice SCore（DSC）和95% Hausdorff距离（HD95）评估模型在Synapse和ACDC数据集上的性能。DSC和HD根据方程（9）和（10）计算：
其中，表示真值图和分割图，分别表示分割图中的边界点，表示点和之间的距离。HD95是和边界点的距离的第95个百分位数。
Comparisons with State-of-the-Arts
为了验证所提方法的有效性，作者将其性能与基于CNN、transformer和mamba的最新方法进行了比较。
Iii-C1 Results on Synapse Multi-Organ Segmentation
表1显示，与各种方法相比，作者提出的MSVM-UNet在平均DSC方面取得了85.00%的最佳结果，在HD95方面取得了14.75mm的最佳结果。具体而言，与基于CNN的方法（如2D D-LKA Net）相比，作者的方法在DSC和HD95方面分别提高了0.73%和5.29mm；与基于transformer的方法（如PVT-EMCAD-B2）相比，分别提高了1.37%和0.93mm；与基于mamba的方法（如VM-UNet）相比，分别提高了2.62%和1.47mm。此外，对于小器官，与最佳方法相比，胆囊和胰腺的DSC分别提高了0.09%和1.80%，而大型器官的胃的DSC则提高了1.57%。这是因为MSVM-UNet能同时有效捕捉像素之间的长程依赖关系和局部上下文关系。由于引入了多尺度卷积操作，MSVM-UNet不仅能有效处理形状和大小各异的器官，还能更好地定位器官边界。
Iv-C2 Results on ACDC for Automated Cardiac Segmentation
表2展示了作者的方法与上述几种方法在MRI医学图像上的性能对比，其中作者提出的MSVM-UNet获得了最佳平均DSC为92.58%。此外，在ACDC数据集（RV, Myo, 和LV）的三个类别中，作者的方法分别获得了最佳的DSC结果为91.00%，90.35％，和96.39％，这表明了作者的方法具有很好的泛化能力。这种表现在不同的医学图像数据模式（MRI和CT）上有良好的效果。
Qualitative Analysis
如图4所示，作者在Synapse多器官数据集上对各种方法进行2D视觉比较。观察到，作者的方法在不同的器官上产生了更好的分割结果，并且在一定程度上避免了过度分割（如最后一行胆囊的分割）和欠分割（如第三行胰腺的分割）的问题。这主要是因为作者的MSVM-UNet能更好地捕获具有不同几何形状的特征。与基于mamba的方法相比，作者的方法在器官边缘分割方面显示了更好的性能。此外，在比较第一行肝脏分割时，作者发现采用卷积操作的方法在绘制肝脏边界方面比没有卷积操作的方法取得了更好的结果。这要归因于包括了适当的卷积操作，这些操作有助于模型捕捉局部细节特征和位置信息，从而导致更 discriminative 的特征表示和更好的分割结果。
Ablation Studies
作者在Synapse数据集上进行了一项全面的消融实验，以验证和调查作者提出的法
V-E1 Effect of Different Components of MSVM-UNet
作者在 Synapse 多器官数据集上进行了一系列实验，以了解 MSVM-UNet 解码器中不同组成部分的影响。作者评估了用作者提出的模块替换解码器中特定模块的影响。如表3 所示，用作者的模块替换原始模块后，仅计算开销和参数数量有轻微增加，但性能有显著提高。具体来说，与具有 VSS 块和扩大卷积层的解码器相比，使用作者提出的模块提高了 DSC 和 HD95 分别提高了 2.2% 和 12.54mm，仅增加了额外的 0.42G FLOPs 和 0.25M 参数。这表明了作者提出的 MSVM-UNet 的有效性和效率。
V-E2 Effect of Different Upsampling Methods
为了探究作者提出的LKPE的有效性，作者在Synapse多器官数据集上进行了实验，分别使用了原始解码器、移位卷积、上采样块[3]、 Patch 扩充层和LKPE层作为上采样层，以评估它们的个体性能。如表4所示，作者报告了不同上采样方法对应的表现和开销。为了更清楚地比较不同上采样操作的计算开销和参数数量，作者仅报告解码器的FLOPs和参数。与原始 Patch 扩充层相比，LKPE分别将DSC和HD95提高了1.5%和13.08mm，同时只引入了额外的0.12G FLOPs和0.06M参数。此外，还可以观察到将通道和空间信息聚合的方法通常获得更好的性能，进一步证实了作者提出的LKPE的有效性。
V-E3 Effect of Multi-Scale Kernels in MSVSS Block
作者还在Synapse多器官数据集上进行了额外的实验，以探索MSVSS块中的不同多尺度卷积核的影响。如表5所示，各种多尺度卷积核的表现情况。同样，仅提供解码器的FLOPs和参数数量。为了避免计算开销过大，作者设计了三组核：第一组为和，第二组为和，第三组为三者的并集。随着卷积核数量和规模的增加，作者发现性能下降。基于这些观察，作者在MSVSS块中选择了作为默认的多尺度卷积核。
V-E4 Effect of Encoder Model Scales
为了探究不同的编码器深度对模型性能的影响，作者在Synapse多器官数据集上进行了两组实验，以研究不同的编码器尺度的影响。如表6所示，随着编码器的深度和规模的增加，性能略有降低。由于这两组实验使用了相同的设置，作者假设这种 minor性能下降可能是因为模型复杂度增加导致的轻微过拟合。基于这些观察和考虑计算开销和参数数量，作者选择 tiny 版本的编码器作为默认尺度。
Iii-B5 Effect of Feature Enhancement
原始解码器与作者的MSVM-UNet解码器对应层的特性如图5所示。作者计算特征图的所有通道的平均值，并使用Matplotlib生成 Heatmap 。从图5中可以看出，作者的方法有助于处理大小和形状不同的器官，并获得更具有区分度的特征表示。
IV Conclusion
在本文中，作者提出了一种新颖的多尺度视觉Mamba UNet，旨在解决医学图像分割面临的挑战。由于多尺度深度卷积的设计，MSVM-UNet不仅能捕获不同尺度下的信息并建模所有方向上的长程依赖性，还能保持计算效率和可接受的参数数量。
此外，通过有效集成通道和 spatial信息进行上采样，MSVM-UNet实现了更具有判别性的特征表示，从而使得医学图像分割的结果更加准确。
作者的实验结果显示，MSVM-UNet在两个医学图像数据集上表现尤为出色，在Synapse多器官数据集上的DSC和HD95分别比VM-UNet提高了2.62%和1.47mm。
定性分析也表明，MSVM-UNet可以准确地定位器官，处理大小和形状不同的器官。
参考
[1].MSVM-UNet: Multi-Scale Vision Mamba UNet for Medical Image Segmentation.
点击上方卡片，关注 「AI视界引擎」 公众号