点击下方卡片，关注 「AI视界引擎」 公众号
( 添加时备注：方向+学校/公司+昵称/姓名 )
早期发现皮肤异常在诊断和治疗皮肤癌中起着至关重要的作用。使用AI驱动设备对受影响的皮肤区域进行分割相对常见，并支持诊断过程。
然而，实现高性能仍然是一个重大挑战，因为需要高分辨率图像，而单个病损的边界往往不够清晰。
同时，医疗设备要求分割模型具有较小的内存占用和较低的计算成本。基于这些需求，作者提出了一种新的轻量级模型MambaU-Lite，该模型结合了Mamba和CNN架构的优点，参数量不足40万，计算成本超过1G FLOPs。
为了增强全局上下文和局部特征提取，作者提出了P-Mamba块，这是一种创新组件，它将VSS块与多个池化层相结合，使模型能够有效地学习多尺度特征并提高分割性能。
作者在ISIC2018和PH2两个皮肤数据集上评估了该模型的表现，取得了令人鼓舞的结果。
1 Introduction
皮肤病变的分割在皮肤癌计算机辅助诊断系统中发挥着重要作用。然而，在自动化技术尚未进入这一医学领域时，手动分割方法被认为既繁琐又不准确，整体而言不可靠且成本高昂。幸运的是，随着深度学习的进步，尤其是U-Net [1] 及其变种 [2]、[3]、[4] 的发展，各种尝试将这些技术应用于分割任务以消除人为错误并提高速度得到了实施。
在机器学习和神经网络的另一端，2017年出现了一种名为Transformer的新模型[5]，其核心机制为“注意力”。该模型在处理NLP任务方面表现出了革命性的突破。理想情况下，Dosovitskiy等人为在NLP任务中使用Transformer[5]与计算机视觉领域的某些先前模型之间架起桥梁，提出了Vision Transformer[6]，其中包含了“ViT”组件，从而开启了基于Transformer的各种图像处理模型的新时代。将这一思路扩展到分割任务中，TransUNet[7]结合了U-Net结构和强大的ViT模块。
UCTransnet[8]、Swin-Unet[9]以及各种后续模型均采用了类似的组合方式，并进行了相应的调整，取得了相对成功的结果。然而，Transformer及其注意力机制[5]存在一个问题，即该机制的计算复杂度随着序列长度的平方而增加，这在某些情况下导致推理速度不够理想。这种问题同样适用于一般的分割和图像处理任务，尤其是高分辨率的图像放大后可能导致序列异常长。
近年来，在2024年，Gu和Dao提出了一种名为Mamba（或称为S6模型）的新方法，通过引入选择机制和硬件感知算法来提升随机结构状态空间模型（S4）的性能。最重要的是，该模型在性能上表现出线性扩展能力，并且与基于Transformer的模型相比，能够取得有竞争力的结果。利用Mamba进行图像处理时，Vision Mamba或Vim [11]引入了双向SSM机制，以选择性地捕捉图像的全局上下文。同年，VMamba [12] 在2D选择扫描机制（SS2D）的基础上被提出，通过从四个方向学习图像，使得Mamba机制更适应图像处理。作者的混合模型MambaU-Lite继承了VSS块的强大功能，并结合了ULite模型 [13] 的优雅设计，在超过40万个参数的情况下取得了潜在的良好效果。以下是作者在研究中的主要贡献：
- 作者提出了一种轻量级模型，即MambaU-Lite，这是一种结合了Mamba和CNN两种方法的混合分割模型，集两者之所长，优化性能同时保持合理的计算成本。
- 建立并实现了一种新的子结构P-Mamba，用于高效地学习不同层次的特征。
- MambaU-Lite模型在ISIC 2018和PH2两个知名的皮肤病变分割数据集上进行了评估，结果显示该模型是一个轻量级模型，表现出色。
2 Related 1 Work
视觉状态空间模型 [12]。受 Mamba [10] 的启发，该模型成功地将控制理论中的状态空间模型（SSM）应用到了自然语言处理（NLP）中，Vision Mamba [12] 作为一种新颖的方法被引入，旨在通过集成基于 SSM 的模块来高效支持视觉表示。此外，该模型不仅有助于全局特征的提取，还能最大限度地减少计算成本和时间消耗。因此，在视觉相关任务中应用 SSM 成为了一个趋势 [14]，医学分割也不例外 [15]。
U-Net网络架构 [1]。U-Net最早由Ronneberger等人在2015年提出，为众多医学图像分割模型奠定了基础。该架构设计简洁，采用对称的编码解码结构并带有 Short-Cut 连接，有效解决了有限标注数据的问题，并在效率上超越了之前的分割模型。后续对U-Net的改进，例如Attention U-Net [3]，进一步证明了该架构在图像分割领域的优越性。
3 The Proposed Model
在本节中，提出了MambaU-Lite模型的架构，并详细展示在图1中。该模型包含三个基本子结构：编码器、 Bottleneck 和解码器，共同形成了类似于经典U-Net [1] 的U形组合。此外，模型经历了四个处理阶段，在这些阶段中，四個CBAM [16] 块辅助解码器从编码器获取丰富的空间信息。
最初，输入图像通过InitConv层调整通道数至16，生成大小为的特征图。随后，图像经过一系列Encoder层。在所 Proposal 的架构中，作者使用了前两个P-Mamba Encoder块（PE块）。在这两个块以及每个Encoder之后的max-pooling层的作用下，特征图的尺寸分别为和。对于接下来的两个Encoder层，输入被分成两部分，有效减半了通道数，然后分别通过PE块和Axial Encoder块（AE块）处理。经过这两个块和max-pooling层，特征图的尺寸均为和。最后两个块的输出被拼接，生成尺寸为的特征图，并传送到 Bottleneck 层，然后与 Shortcut 结合，经过Decoder层处理。经过所有Decoder层和上采样层后，从每个Decoder块输出的结果被插值回原始输入尺寸。这些结果随后被拼接并通过FinalConv层进行处理，以生成输入图像的预测 Mask 。
3.1 The proposed PMamba Block
提出的P-Mamba块，如图2所示，结构设计旨在提高模型学习多样化特征的能力。这通过使输入特征图通过两个独立分支进行处理来实现。
在第一个分支中，输入通过一个带有3x3核的深度可分离卷积层来捕获初始的局部特征。为了在减少参数数量的同时维持稳定的性能，在保持输入通道数减半的情况下，将输入通道分别送入两个VSS块，如图2所示。这些块是由Nguyen等人在AC-MambaSeg [17] 中引入的，旨在使模型能够有效地学习多尺度特征。两个VSS块的输出被连接起来以恢复原来的尺寸，随后使用实例归一化进行标准化处理，并通过ReLU激活函数进一步标准化输出，从而增强模型的稳定性。
在第二分支中，输入通过顺序的平均池化层和最大池化层进行处理。结合使用这两种池化类型可以让模型同时捕捉全局和细节信息，从而丰富特征表示。池化层的输出被连接起来并通过一个3x3卷积层以恢复通道数到原始大小，这也有助于模型关注重要的信息。最后，输出通过一个sigmoid函数，该函数作为一种注意力层来强调重要特征并抑制无关特征。
最后，两个分支的输出被加在一起，使模型能够学习更多的特征。
3.2 The Encoder Block
如第3节所述，编码器由两个主要部分组成：AE块和PE块，如图2d和图2e所示。对于AE块，输出首先通过Dinh等人[13]提出的轴向深度可分离卷积层进行处理，该层使用了7x7的 Kernel ，然后经过批量归一化和ReLU激活函数。在进入用于将通道数翻倍的点积卷积层之前，提取了一个 Shortcut 以避免信息损失，并在解码器中使用。在PE块中，输入首先通过P-Mamba块处理，然后依次经过3x3 Kernel 的轴向深度可分离卷积层、批量归一化、ReLU激活函数以及点积卷积层。与AE块类似，在点积卷积层之前也提取了一个 Shortcut ，以便保留解码过程中所需的重要信息。
3.3 The Decoder Block
解码块的概览如图2f所示。初始时，前一个解码层的输出会被上采样以匹配相应 Short-Cut 的大小。然后，该输出会通过注意力门控（AG）模块，如图2c所示。AG模块的输出会与来自前一个解码层的上采样特征图进行拼接。拼接后的输出随后会经过一系列层处理：首先是一个点wise卷积层以降低维度，接着是BatchNormalization、ReLU激活函数，然后是另一个点wise卷积，最后是一个轴向深度卷积，其核大小为7×7。这些层的组合使得模型能够有效地提取有意义的特征，同时减少参数量和计算开销。
3.4 The Skip Connection and Bottleneck Block
Short-Cut 和 Bottleneck 组件在模型中发挥着关键作用，有助于在处理过程中防止信息丢失。在所提出模型中，作者使用Woo等人[16]引入的CBAM块作为 Short-Cut ，而瓶 Neck 分则采用集成通道-空间注意力（ICSA）块。
ICSA块包含两个连续的优先通道注意(PCA)块，之后是一个优先空间注意(PSA)块，该结构由Le等人[18]提出，并在鱼类分类任务中展示了出色的效果。具体而言，PCA块使用深度可分离卷积来单独增强每个通道上的特征提取，而PSA块则采用逐点卷积来改进空域区域的特征图。ICSA块在瓶 Neck 分有效地捕捉高层特征，然后再将其传递给解码器。
4 Experiment
4.1 Dataset
为了评估所提出方法的有效性，作者在两个皮肤病变数据集中进行了实验：ISIC 2018 和 PH2。ISIC 2018 数据集包含 2,594 张皮肤镜图像以及分割 Mask 。作者将该数据集分为两部分：分配了 2,334 张图像用于训练，剩余的 260 张图像用于测试。对于 PH2 数据集，这是一个规模较小的数据集，共包含 200 张图像，作者也将其分为两部分，分配了 170 张图像用于训练，剩余的 30 张图像用于测试。来自两个数据集的每张图像都被调整为 大小，以便于训练过程。
4.2 Training and Evaluation Metric
作者使用PyTorch框架进行了实验，并采用了Adam优化策略。模型经过了300个epoch的训练，初始学习率为，如果连续10个epoch Dice分数没有改进，则将学习率减半。在训练过程中，作者使用了一个结合Dice损失和Tversky损失的复合损失函数。作者将Tversky损失的超参数设置为和。损失函数的公式如下：
其中， 表示 GT 标签， 表示每个像素 的预测 Mask 值，而 则表示输出分割 Mask 中的总像素数。
为了评估模型的性能，作者采用了语义分割中常用的两大指标：Dice相似系数（DsC）和交并比（IoU）。这些指标有助于确定预测 Mask 与Ground Truth标签之间的相似重叠程度，清楚地展示了模型的有效性。
4.3 Results and Comparison
为了评估所提出模型的效果，作者将其与先前提出的多种方法进行了比较，包括U-Net [1]、Attention U-Net [3]、UNeXt [2]、DCSAU-Net [19] 和 U-Lite [13]。这些模型均在与所提模型相同的条件下训练，并且所有实现都源自作者的开源代码库。MambaU-Lite与其他模型之间的比较结果是在ISIC 2018和PH2数据集上进行的。一些可视化分割结果如图3所示。从该图可以看出，所提出的MambaU-Lite模型生成的输出与原始GT Mask 更加接近，进一步证实了分割模型的准确性与可靠性。
表1中的定量比较显示，在ISIC2018数据集上，提出的MambaU-Lite模型优于其他模型，其DSC值为0.9057，IoU值为0.8361。表现第二好的模型是Ulite，其DSC值为0.9032，IoU值为0.8340。尽管U-Lite的FLOPS比提出的模型低0.69G，但它具有显著更高的参数量和内存大小，参数量为0.88M，内存大小为3.51MB，几乎是MambaU-Lite的一倍，而MambaU-Lite的参数量为0.42M，内存大小仅为1.67MB。表2展示了作者模型在小数据集PH2上的有效性。可以看出，MambaU-Lite模型优于其他模型，其DsC值为0.9572，IoU值为0.9189。此外，作者的模型拥有最低的参数量和内存大小。虽然UNeXt的计算成本低于提出的模型，但其性能相对较低，DsC值仅为0.9409，远低于MambaU-Lite的值。
5 Conclusion
在本研究中，作者提出了轻量级的MambaU-Lite模型，用于皮肤病变分割，旨在最大限度地减少模型参数数量、计算成本和内存使用。
作者提出了一种P-Mamba块，将其集成到一种创新架构中，该架构结合了Mamba和CNN的优势，能够有效捕捉高层和细粒度特征。
尽管作者的模型在皮肤病变数据集上已经展现出了令人鼓舞的结果，但未来的工作将进一步优化并泛化该模型，使其适用于各种医学影像任务，增强其适应性，并使其更适合部署在医学设备中。
参考
[0]. MambaU-Lite: A Lightweight Model based on Mamba and Integrated Channel-Spatial Attention for Skin Lesion Segmentation .
点击上方卡片，关注 「AI视界引擎」 公众号