点击下方卡片，关注 「AI视界引擎」 公众号
在最近医学图像分析技术的进展中，卷积神经网络（CNN）和视觉 Transformer （ViT）已经设定了重要的基准。前者通过其卷积操作在捕捉局部特征方面表现出色，而后者则通过利用自注意力机制实现了对全局上下文的显著理解。然而，这两种架构在高效地对医学图像内的长距离依赖进行建模方面存在局限，这对于精确分割是一个关键方面。
受到擅长处理长序列和全局上下文信息且计算效率增强的状态空间模型（SSM）的Mamba架构启发，作者提出了Mamba-UNet，这是一种将U-Net在医学图像分割中的能力与Mamba能力相结合的新型架构。Mamba-UNet采用基于纯视觉Mamba（VMamba）的编码器-解码器结构，并融入跳跃连接以保留网络不同尺度上的空间信息。
这种设计促进了全面的特征学习过程，捕捉医学图像中的复杂细节和更广泛的语义上下文。作者在VMamba块中引入了一种新颖的集成机制，以确保编码器和解码器路径之间无间断的连接和信息流动，提升了分割性能。作者在公开可获得的MRI心脏多结构分割数据集上进行了实验。
结果显示，在相同的超参数设置下，Mamba-UNet在医学图像分割方面优于UNet、Swin-UNet。
源代码和 Baseline 实现：https://github.com/ziyangwang007/Mamba-UNet
1 Introduction
医学图像分割对于诊断和治疗至关重要，基于深度学习的网络已在这一领域显示出主导性能。U-Net是最重要的架构之一，以其对称的编码器-解码器风格架构和跳跃连接而闻名，其中不同的编码器和解码器在不同层次提取特征信息，跳跃连接使得特征信息的有效转换成为可能。大多数研究进一步探索了使用先进的网络块技术（如密集连接，残差块，注意力机制，深度卷积，以及孔洞卷积）的U-Net，从而在CT、MRI、超声医学图像分割中产生了各种改进的UNet。
受到自然语言处理中自注意力机制成功的启发，ViT首次使用了纯粹的多头自注意力机制来进行图像识别任务，并达到了最先进的表现。这展示了其在建模长距离依赖方面的强大能力。像移位窗口这样的技术进一步定制了ViT，从而产生了Swin-Transformer，它增强了它们在计算机视觉中的密集预测任务中的应用性，例如图像分割和检测。
在医学图像分割中，将ViT与U-Net架构结合，受到传统CNN设计的启发，也导致了各种基于ViT的混合和纯U-Net。例如，TransUNet是首个利用ViT在UNet编码器中的特征学习能力的工作。UNETR将ViT与UNet结合用于3D分割，而Swin-UNet和DCSUnet进一步探索了基于U-Net结构的纯Swin Vision Transformer网络块。
虽然Transformer在捕捉长距离依赖方面表现出色，但由于自注意力机制与输入大小成二次方扩展，其计算成本高昂，这对于高分辨率生物医学图像尤其是一个挑战。状态空间模型（SSMs），尤其是结构化SSM（S4）的最新发展，由于在处理长序列时的高效性能，提供了一个有前景的解决方案。Mamba模型通过选择机制和硬件优化增强了S4，在密集数据领域展现了卓越的性能。
在视觉状态空间模型（VMamba）中引入的跨扫描模块（CSM）进一步增强了Mamba在计算机视觉任务中的应用性，通过允许遍历空间域并将非因果视觉图像转换为有序的块序列。受到这些能力的启发，作者提出在U-Net架构内利用视觉Mamba块（VSS）来改善医学图像分析中的长距离依赖建模，从而得到Mamba-UNet。图1简要介绍了带有各种网络块的U-Net的演变以及作者提出的Mamba-UNet的位置。
2 Approach
Architecture Overview
所提出的Mamba-UNet架构如图2所示，其灵感来源于UNet和Swin-UNet。首先将尺寸为的输入2D灰度图像以类似于ViT和V Mamba的方式分割成块，然后转化为1-D序列，其维度为。一个初始的线性嵌入层将特征维度调整到任意大小，记作。这些块标记随后通过多个VSS块和块合并层进行处理，以创建分层特征。
块合并层负责下采样和维度增加，而VSS块专注于学习特征表示。编码器每个阶段的输出具有，，，以及的分辨率。解码器由VSS块和块扩展层组成，遵循编码器风格，实现相同特征尺寸的输出，从而通过跳跃连接增强在下采样中丢失的空间细节。在编码器和解码器中，各使用2个VSS块，并在编码器中加载预训练的V Mamba-Tiny，这个过程与Swin-UNet加载预训练的SwinViT-Tiny相同。以下小节将讨论VSS块，编码器的块合并以及解码器的块扩展的细节。
VSS Block
VSS网络块如图3所示，主要基于Visual Mamba。在VSS块中，输入特征首先遇到一个线性嵌入层，然后分为两条路径。一个分支进行深度卷积 和SiLU激活，进入SS2D模块，并在层归一化后，与另一条流经过SiLU激活后的路径合并。这个VSS块与典型的视觉 Transformer 不同，它省略了位置嵌入，选择了一个简化的结构，不包括MLP阶段，这使得在相同的深度预算内能够堆叠更多的块。
Encoder
在编码器中，降低分辨率的维标记输入经过两个连续的VSS块进行特征学习，同时保持维度和分辨率不变。在Mamba-UNet的编码器中，将降采样过程——即切块合并——应用了三次，每次将输入分割成四分之一，连接后通过层归一化处理，使得标记数量减少到原来的，同时将特征维度翻倍。
Decoder
类似于编码器，解码器也使用两个连续的VSS块进行特征重建，它采用 Patch 扩展层而不是合并层来放大深层特征。这些层在提高分辨率（放大）的同时将特征维度减少一半，即减少到原来的，一个初始层先将特征维度加倍，然后在重排和减少特征以增强分辨率之前作为示例。
Bottleneck & Skip Connections
两个VSS模块用于Mamba-UNet的 Neck 分。编码器和解码器的每个层级都采用跳跃连接，将多尺度特征与上采样的输出混合，通过融合浅层和深层的信息增强空间细节。随后的线性层保持了这一整合特征集的维度，确保与上采样分辨率的一致性。
3 Experiments and Results
Data Sets
作者使用了公开可获得的ACDC MRI心脏分割数据集进行实验，该数据集来源于MICCAI 2017挑战赛。这个数据集包含了100名患者的MRI扫描，对多个心脏结构进行了标注，如右心室以及左心室的内膜和心外膜。它涵盖了一系列多样的病理状况，分为五个亚组：正常、心肌梗死、扩张型心肌病、肥厚型心肌病和异常右心室，确保了特征信息的广泛分布。为了符合ViT分割 Backbone 网络的输入要求，所有图像都被重新调整到224×224的大小。数据集被分割，以便20%的图像分配给测试集，其余的用于训练（包括验证）。
Implementation Details
实施工作在Ubuntu 20.04系统上进行，使用了Python 3.8.8，PyTorch 1.10和CUDA 11.3。硬件配置包括Nvidia GeForce RTX 3090 GPU和Intel Core i9-10900K CPU。平均运行时间约为5小时，包括数据传输、模型训练和推理过程。该数据集专门为2D图像分割处理。Mamba-UNet模型用批量大小为24进行了10,000次迭代训练。采用了随机梯度下降（SGD）优化器，学习率为0.01，动量为0.9，权重衰减设置为0.0001。网络性能在每次200次迭代后在验证集上进行评估，只有在验证集上达到新的最佳性能时，才会保存模型权重。
Baseline Methods
为了比较目的，UNet和Swin-UNet也在相同的超参数配置下进行了训练。Mamba-UNet与其他 Baseline 方法，包括UNet和Swin-UNet直接进行了比较。
Evaluation Metrics
对Mamba-UNet与基准方法的评估使用了广泛的评价标准。相似性度量指标越 高越好，包括：Dice系数、交并比（IoU）、准确度、精确度、灵敏度和特异性，用向上箭头（）表示，表明更高的值代表更好的性能。
相反，差异度量如Hausdorff距离（HD）95%和平均表面距离（ASD），用向下箭头（）标记，越低越好，表示预测与真实分割之间更接近的相似度。
在文中， 代表真正例的数量， 表示真负例的数量， 指代假正例的数量，而 代表假负例的数量。
豪斯多夫距离（HD） (6)
平均表面距离
（ASD）= (7)
在哪里， 和 分别代表预测表面和真实表面上的点集。 表示两点之间的欧几里得距离。"95%" 是豪斯多夫距离的一个修正版本，它关注距离的95个百分点以减少异常值的影响。
Qualitative Results
图4展示了三张随机选取的原始样本图像及相应推理结果，与包括Mamba-UNet在内的所有 Baseline 方法的公开发布的 GT 值对比，其中不同颜色表示 GT 值的边界。
Quantitative Results
表1报告了Mamba-UNet与其他分割网络的直接比较，包括相似性度量和差异性度量。最佳性能用粗体表示，Mamba-UNet的第二佳性能用下划线表示。定量结果显示Mamba-UNet更有可能预测出精确的分割 Mask 。
为了进一步在测试集上验证Mamba-UNet，作者也以逐图像的方式验证，根据Dice系数绘制的分割预测分布如图5所示，其中X轴是Dice系数，Y轴是预测的数量。这个直方图进一步证明Mamba-UNet更有可能提供具有高Dice系数性能的预测。
4 Conclusion
在本文中，作者引入了Mamba-UNet，这是一个纯粹的基于视觉Mamba块式的UNet风格网络，用于医学图像分割。性能表明，Mamba-UNet相较于经典的类似网络如UNet和Swin-UNet具有更优越的性能。
未来，作者计划对不同模态和目标的更多医学图像分割任务进行更深入的探索，并将其与更多的分割 Backbone 网络进行比较。此外，作者旨在将Mamba-UNet扩展到3D医学图像，并结合半/弱监督学习[14]以进一步促进医学成像领域的发展。
参考
[1].Mamba-UNet: UNet-Like Pure Visual Mamba for Medical Image Segmentation.
点击上方卡片，关注 「AI视界引擎」 公众号