LightM-UNet论文摘要_mamba设计1.8m参数unet,比u-mamba小近100倍

作者：人工智能uu | 2024-07-18 16:35:45

踩

mamba设计1.8m参数unet,比u-mamba小近100倍

Motivation

（1）作为一种基于卷积神经网络（CNN-based）的模型，UNet面临着卷积操作固有的局部性，这限制了它理解明确的全局和长距离语义信息交互的能力，一些研究尝试通过采用扩张卷积层[5]，自我注意力机制[19]和图像金字塔[25]来缓解这个问题。然而，这些方法在建模长距离依赖方面仍然表现出限制。

（2）当前的Unet变体整合Transformer架构，利用自注意力机制将图像视为一系列连续的 Patch ，从而捕获全局信息。尽管这种方法有效，但基于Transformer的解决方案由于自注意力机制导致了图像尺寸的二次复杂度

Contribution

状态空间模型（SSMs）不仅建立了长距离依赖关系，而且输入规模的线性复杂性也使其成为轻量级UNet道路上CNN和Transformer的强大竞争对手，

引入了LightM-UNet，这是一个基于Mamba的轻量级U形分割模型，是UNet与Mamba的轻量级融合，能不增加额外参数和计算负担的情况下，赋予UNet容纳长距离依赖的能力。

Methodologies

提出了“残差视觉曼巴层（RVM层）”以纯曼巴方式从图像中提取深层特征。在引入的新参数和计算开销最小的情况下，作者通过使用“残差连接”和“调整因子”进一步增强了SSM对视觉图像中长距离空间依赖关系建模的能力

Architecture

技术路线：

首先使用深度可分卷积（DWConv）层进行浅层特征提取，生成浅层特征图

随后，LightM-UNet结合三个连续的编码器块（Encoder Blocks）从图像中提取深层特征。在每个编码器块之后，特征图中的通道数翻倍，而分辨率减半。

在此之后，LightM-UNet使用瓶颈块（Bottleneck Block）来建模长距离空间依赖关系，同时保持特征图的大小不变。之后，LightM-UNet整合三个连续的解码器块（Decoder Blocks）进行特征解码和图像分辨率恢复。在每个解码器块之后，特征图中的通道数减半，分辨率加倍。

最后，最后一个解码器块的输出达到与原始图像相同的分辨率，包含32个特征通道。LightM-UNet使用DWConv层将通道数映射到分割目标数，并应用SoftMax激活函数生成图像 Mask 。与UNet的设计一致，LightM-UNet也采用跳跃连接（skip connections）为解码器提供多 Level 特征图。

Modules

Encoder Block

LightM-UNet采用了仅包含Mamba结构的编码器块来从图像中提取深层特征。

随后，编码器块使用Nl个连续的RVM层来捕捉全局信息，在最后一个RVM层中通道数增加。

RVM层（Residual Vision Mamba Layer ）

LightM-UNet提出了RVM层以增强原始的SSM块，用于图像深层语义特征提取。具体来说，LightM-UNet利用先进的残差连接和调整因子进一步增强了SSM的长距离空间建模能力，几乎不引入新的参数和计算复杂性。

Bottleneck Block

类似于Transformer，当网络深度变得过大时，Mamba也会遇到收敛挑战。因此，LightM-UNet通过结合四个连续的RVM层来构建瓶颈，以进一步建模空间长期依赖关系，从而解决这个问题。在这些瓶颈区域中，特征通道的数量和分辨率保持不变。

Decoder Block

Experiments

数据集和实验设置。为了评估作者模型的性能，作者选择了两个公开可用的医学图像数据集：LiTs数据集[1]，包含3D CT图像；以及Montgomery&Shenzhen数据集[9]，包含2D X光图像。这些数据集在现有的分割研究[12, 24]中被广泛使用，并在本研究中分别用来验证2D和3D版本的LightM-UNet的性能。数据被随机划分为训练集、验证集和测试集，比例分别为7:1:2。

LightM-UNet是使用PyTorch框架实现的，三个编码器块中的RVM层数量分别设置为1、2和2。所有实验都是在单个Quadro RTX 8000 GPU上进行的。采用SGD作为优化器，初始学习率为1e-4。PolyLRScheduler作为调度器，共训练了100个周期。

采用了平均交并比（mIoU）和Dice相似度得分（DSC）作为评估指标。

图3展示了分割结果示例，表明与其他模型相比，LightM-UNet具有更平滑的分割边缘，并且不会对小型目标（如肿瘤）产生错误的识别。

使用卷积或自注意力机制的VSSM会导致性能损失。此外，卷积和自注意力引入了大量参数和计算开销。进一步地，作者观察到基于Transformer和基于VSSM的结果都优于基于卷积的结果，这证明了建模长距离依赖的好处。

作者进一步去除了RVM层中的调整因子和残差连接。实验结果表明，在移除这两个组件后，模型的参数数量和计算开销几乎没有减少，但模型的性能显著下降（mIoU下降了0.44%和0.69%）。这验证了作者在不引入额外参数和计算开销的情况下提升模型性能的基本原则。关于Montgomery&Shenzhen数据集的额外消融分析可以在补充材料中找到。

文章参考

北大提出LightM-UNet | 用Mamba设计1.8M参数UNet，比U-Mamba小近100倍，精度依旧完成超车 (qq.com)

个人笔记

首先使用深度可分卷积（DWConv）层进行浅层特征提取，生成浅层特征图Fs

随后，LightM-UNet结合三个连续的编码器块（Encoder Blocks）从图像中提取深层特征。在每个编码器块之后，特征图中的通道数翻倍，而分辨率减半。

在此之后，LightM-UNet使用瓶颈块（Bottleneck Block）来建模长距离空间依赖关系，同时保持特征图的大小不变。

之后，LightM-UNet整合三个连续的解码器块（Decoder Blocks）进行特征解码和图像分辨率恢复。在每个解码器块之后，特征图中的通道数减半，分辨率加倍。

最后，最后一个解码器块的输出达到与原始图像相同的分辨率，包含32个特征通道。

在解码器（Decoder）部分的结尾，LightM-UNet使用DWConv层将通道数映射到分割目标数，并应用SoftMax激活函数生成图像掩码。与UNet的设计一致，LightM-UNet也采用跳跃连接（skip connections）为解码器提供多级特征图。

解码器块负责将编码器提取的特征进行上采样（恢复分辨率），并通过跳跃连接（skip connections）与编码器的特征进行融合，以保留多尺度的信息。在最后的DWConv层中，模型将特征通道数映射到与分割目标相对应的数量，这通常等于类别数。然后，应用SoftMax激活函数来生成每个像素点的类别概率分布，从而得到最终的图像掩码