【魔改UNet系列】LightM-UNet：Mamba与UNet结合实现轻量级图像分割

小白学视觉 · 公众号 · · 2024-10-28 10:44

正文

点击上方




    
“小白学视觉”，选择加"星标"或“置顶”
重磅干货，第一时间送达

论文信息

题目：LightM-UNet: Mamba Assists in Lightweight UNet for Medical Image Segmentation

LightM-UNet：Mamba助力轻量级UNet进行医学图像分割

作者：Weibin Liao, Yinghao Zhu, Xinyuan Wang, Chengwei Pan, Yasha Wang, Liantao Ma

源码：https://github.com/MrBlankness/LightM-UNet

论文创新点

LightM-UNet 是一种基于Mamba的轻量级网络，用于医学图像分割，具有以下几个创新点：

轻量级架构 ：作者提出了 LightM-UNet ，这是一个轻量级的UNet和Mamba的融合，仅拥有1M的参数数量。这是通过在UNet架构中使用Mamba来实现的，旨在解决实际医疗环境中计算资源限制所带来的挑战。
残差视觉Mamba层（RVM层） ：作者提出了 残差视觉Mamba层 来增强原始的状态空间模型块，用于深度语义特征提取。这一改进几乎不引入新的参数和计算复杂度，同时通过使用残差连接和调整因子，增强了模型对长距离空间依赖性的建模能力。
性能提升 ：在与现有的最先进模型进行比较时， LightM-UNet 在参数和计算成本上实现了显著的减少。特别是与著名的nnU-Net相比， LightM-UNet 在减少参数和计算成本116倍和21倍的同时，实现了更优越的分割性能。

摘要

UNet及其变体已广泛用于医学图像分割。然而，这些模型，尤其是基于Transformer架构的模型，由于参数众多和计算负载大，给移动健康应用带来了挑战。最近，状态空间模型（SSMs），如Mamba，作为CNN和Transformer架构的有力替代品出现。基于此，作者将Mamba作为轻量级替代CNN和Transformer在UNet中的使用，旨在解决实际医疗环境中计算资源限制所带来的挑战。为此，我们介绍了轻量级Mamba UNet（LightM-UNet），它将Mamba和UNet集成在一个轻量级框架中。具体来说，LightM-UNet利用纯Mamba方式的残差视觉Mamba层来提取深层语义特征和模拟长距离空间依赖性，计算复杂度为线性。在两个真实世界的2D/3D数据集上进行的广泛实验表明，LightM-UNet超越了现有的最先进文献。特别是与著名的nnU-Net相比，LightM-UNet在大幅降低参数和计算成本116倍和21倍的同时，实现了更优越的分割性能。这突出了Mamba在促进模型轻量化方面的潜力。

关键字

医学图像分割 · 轻量级模型 · 状态空间模型

2 方法论

尽管LightM-UNet支持医学图像分割的2D和3D版本，但为了方便起见，本文以LightM-UNet的3D版本来描述方法论。

2.1 架构概述

所提出的LightM-UNet的整体架构如图2所示。给定一个输入图像，其中、、和分别表示3D医学图像的通道数、高度、宽度和切片数。LightM-UNet首先使用深度卷积（DWConv）层进行浅层特征提取，生成浅层特征图，其中32表示固定数量的滤波器。随后，LightM-UNet结合三个连续的编码器块来提取图像的深度特征。每个编码器块之后，特征图中的通道数翻倍，而分辨率减半。因此，LightM-UNet在第个编码器块提取深度特征

，其中。之后，LightM-UNet使用瓶颈块来模拟长距离空间依赖性，同时保持特征图的大小不变。随后，LightM-UNet集成三个连续的解码器块进行特征解码和图像分辨率恢复。每个解码器块之后，特征图中的通道数减半，分辨率翻倍。最后，最后一个解码器块的输出达到与原始图像相同的分辨率，包含32个特征通道。LightM-UNet使用DWConv层将通道数映射到分割目标的数量，并应用SoftMax激活函数生成图像掩码。与UNet的设计一致，LightM-UNet还使用跳跃连接为解码器提供多级特征图。

2.2 编码器块

为了最小化参数数量和计算成本，LightM-UNet采用仅包含Mamba结构的编码器块来提取图像的深度特征。具体来说，给定一个特征图，其中，，，，，编码器块首先将特征图展平并转置成的形状，其中。随后，编码器块使用个连续的RVM层来捕获全局信息，在最后一个RVM层中增加通道数。之后，编码器块将特征图重塑并转置成的形状，接着进行最大池化操作以降低特征图的分辨率。最终，第个编码器块输出新的特征图，形状为。

2.3 残差视觉Mamba层（RVM层）

LightM-UNet提出了RVM层来增强原始SSM块，以进行图像深度语义特征提取。具体来说，LightM-UNet利用先进的残差连接和调整因子来进一步增强SSM对长距离空间建模的能力，几乎不引入新的参数和计算复杂度。如图2（a）所示，给定输入深度特征，RVM层最初使用LayerNorm，然后是VSSM来捕获空间长距离依赖性。随后，它在残差连接中使用调整因子，以获得更好的性能。该过程可以用以下数学公式表示：

之后，RVM层使用另一个LayerNorm对