[论文解读] MambaU-Lite: A Lightweight Model based on Mamba and Integrated Channel-Spatial Attention for Skin Lesion Segmentation

Thi‐Nhu‐Quynh Nguyen, Quang Tri Ho+4 more|arXiv (Cornell University)|Dec 2, 2024

Cutaneous Melanoma Detection and Management被引用 4

一句话总结

MambaU-Lite 是一种轻量级医学图像分割模型，结合了 Mamba 的高效序列建模能力与 CNN，并引入了一种新型 P-Mamba 模块，该模块集成了通道-空间注意力机制。该模型在 ISIC2018 和 PH2 数据集上实现了最先进性能，仅使用 420K 参数和 1.25G FLOPs，优于现有轻量级模型，在 Dice 和 IoU 分数上表现更优，同时将计算和内存成本降至最低。

ABSTRACT

Early detection of skin abnormalities plays a crucial role in diagnosing and treating skin cancer. Segmentation of affected skin regions using AI-powered devices is relatively common and supports the diagnostic process. However, achieving high performance remains a significant challenge due to the need for high-resolution images and the often unclear boundaries of individual lesions. At the same time, medical devices require segmentation models to have a small memory foot-print and low computational cost. Based on these requirements, we introduce a novel lightweight model called MambaU-Lite, which combines the strengths of Mamba and CNN architectures, featuring just over 400K parameters and a computational cost of more than 1G flops. To enhance both global context and local feature extraction, we propose the P-Mamba block, a novel component that incorporates VSS blocks along-side multiple pooling layers, enabling the model to effectively learn multiscale features and enhance segmentation performance. We evaluate the model's performance on two skin datasets, ISIC2018 and PH2, yielding promising results. Our source code will be made publicly available at: https://github.com/nqnguyen812/MambaU-Lite.

研究动机与目标

• 开发一种轻量级深度学习模型用于皮肤病变分割，以在保持高精度的同时最小化计算和内存成本。
• 解决在高分辨率输入图像下，边界模糊的皮肤病变分割挑战，采用高效架构。
• 将 Mamba 的长距离建模能力与 CNN 相结合，以增强全局上下文建模和局部特征提取。
• 减小模型大小和推理成本，以实现资源受限医疗设备上的部署。
• 通过一种新型混合架构，在 ISIC2018 和 PH2 等基准数据集上提升分割性能。

提出的方法

• 提出一种混合 U-Net 架构，结合 Mamba 和 CNN 组件，以 P-Mamba 模块作为核心模块。
• 引入 P-Mamba 模块，该模块将 VSS（向量选择性扫描）模块与多个池化层结合，以捕获多尺度特征。
• 在跳跃连接中使用 CBAM（卷积块注意力模块），以增强空间和通道维度的特征表示。
• 采用双损失函数，结合 Dice 损失和 Tversky 损失，超参数 γ₁=0.7 和 γ₂=0.3，以优化分割性能。
• 采用受 U-Lite 启发的轻量化设计，将参数量减少至 0.42M，内存使用量降低至 1.67MB。
• 使用 Adam 优化器进行训练，配合余弦退火学习率调度策略，并应用数据增强技术。

Figure 1 : The architecture of the proposed MambaU-Lite model

实验结果

研究问题

• RQ1 混合 Mamba-CNN 架构是否能在极低参数量和 FLOPs 下实现高分割精度，适用于皮肤病变分析？
• RQ2 所提出的 P-Mamba 模块相较于标准 Mamba 或 CNN 模块，在特征学习方面有何改进？
• RQ3 在跳跃连接中集成通道-空间注意力机制是否能提升在具有挑战性的皮肤病变数据集上的分割性能？
• RQ4 与现有轻量级模型相比，MambaU-Lite 在精度、参数量和推理成本方面表现如何？
• RQ5 该模型是否能在 ISIC2018 和 PH2 等多样化的皮肤病变数据集上实现良好泛化，且架构复杂度极低？

主要发现

• 在 ISIC2018 数据集上，MambaU-Lite 达到 0.9057 的 Dice 相似系数（DSC）和 0.8361 的 IoU，优于 U-Net、Attention U-Net、UNeXt、DCSAU-Net 和 U-Lite。
• 在 PH2 数据集上，MambaU-Lite 达到 0.9572 的 DSC 和 0.9189 的 IoU，成为轻量级模型中的新最先进性能。
• 模型仅包含 0.42 百万个参数和 1.67MB 内存占用，显著小于 U-Lite（0.88M 参数，3.51MB）及其他模型。
• 尽管 FLOP 数量（1.25G）高于 U-Lite（0.69G），MambaU-Lite 仍实现了更优的精度，表明其具有更优的效率-精度权衡。
• P-Mamba 模块有效捕获了多尺度特征，显著提升了对边界不清晰病变的分割性能。
• 在跳跃连接中集成 CBAM 有效增强了特征精炼，从而生成更精确的分割掩码。

Figure 2 : The main components’ architectures of the proposed MambaU-Lite model