LightM-UNet：Mamba 助力轻量级 UNet 用于医学图像分割
图像与视频处理
2024-03-12 v2 计算机视觉与模式识别
摘要
UNet 及其变体已广泛应用于医学图像分割。然而，这些模型，尤其是基于 Transformer 架构的模型，由于参数量大和计算负载高而面临挑战，使其不适用于移动健康应用。最近，以 Mamba 为代表的状态空间模型（State Space Models, SSMs）已成为 CNN 和 Transformer 架构的有力替代方案。在此基础上，我们采用 Mamba 作为 UNet 中 CNN 和 Transformer 的轻量级替代品，旨在解决真实医疗场景中计算资源受限带来的挑战。为此，我们引入了轻量级 Mamba UNet（LightM-UNet），该模型在一个轻量级框架中集成了 Mamba 和 UNet。具体而言，LightM-UNet 利用纯 Mamba 方式的残差视觉 Mamba 层（Residual Vision Mamba Layer）提取深层语义特征并建模长程空间依赖，且具有线性计算复杂度。在两个真实世界的 2D/3D 数据集上进行的大量实验表明，LightM-UNet 超越了现有的最先进文献。值得注意的是，与著名的 nnU-Net 相比，LightM-UNet 实现了更优越的分割性能，同时将参数量和计算成本分别大幅降低了 116 倍和 21 倍。这突显了 Mamba 在促进模型轻量化方面的潜力。我们的代码实现已公开在 https://github.com/MrBlankness/LightM-UNet。
引用
@article{arxiv.2403.05246,
title = {LightM-UNet: Mamba Assists in Lightweight UNet for Medical Image Segmentation},
author = {Weibin Liao and Yinghao Zhu and Xinyuan Wang and Chengwei Pan and Yasha Wang and Liantao Ma},
journal= {arXiv preprint arXiv:2403.05246},
year = {2024}
}