原论文信息如下:

论文标题:
MM-UNet: Meta Mamba UNet for Medical Image Segmentation
发表日期:
2025年03月
作者:
Bin Xie, Yan Yan, Gady Agam
发表单位:
伊利诺伊理工学院, 伊利诺伊大学芝加哥分校
原文链接:
http://arxiv.org/pdf/2503.17540v1


想象一下,医生在诊断疾病时,需要从CT扫描图像中精确分割出肝脏、肾脏等器官。这就像在一张复杂的迷宫中找出特定的路径,传统方法常常力不从心。😮

最近,来自伊利诺伊理工学院和伊利诺伊大学芝加哥分校的研究团队带来了一项突破性研究——MM-UNet,它将自然语言处理中的明星模型Mamba成功引入医学图像分割领域,在多个基准测试中超越了当前最先进的方法!

图5:Mamba注意力图可视化。每个注意力图都能有效捕捉时间维度上的图像模式,即使3D医学图像被展平为1D序列输入MetaSSM块,这突显了在医学图像分割中使用SSMs的动机

医学图像分割的重要性与传统方法的局限性

医学图像分割是医疗影像分析中的核心技术,它帮助医生精确识别器官边界、检测病变区域,并为手术规划提供关键信息。简单来说,就是把CT、MRI等医学图像中的不同组织和器官"描边"出来。

传统方法主要依赖两种技术路线:

卷积神经网络(CNN)特别是U-Net及其变体,在医学图像分割中表现出色。但它们有个致命弱点:局部感受野。就像人眼通过小孔看世界,CNN只能看到图像的一小部分,难以捕捉全局的、长距离的依赖关系。

视觉Transformer(ViT)虽然能建模全局依赖,但计算复杂度呈平方级增长。想象一下处理一张高分辨率CT图像,需要的计算资源会爆炸性增长!

状态空间模型(SSMs)在自然语言处理中的成功及其在医学成像中的挑战

状态空间模型(State Space Models,SSMs)最近在自然语言处理领域大放异彩。它们能够以线性时间复杂度建模长序列,这就像拥有了一种"超能力":处理很长的文本时,计算量不会爆炸性增长。

但把SSMs直接应用到医学图像上,就像让一个习惯处理文字的语言专家去分析CT扫描图,会遇到两个主要挑战

维度不匹配问题SSMs原本是为1D序列设计的,而医学图像是3D的。为了处理图像,研究者通常会把3D图像"展平"成1D序列,但这会引入不连续性——想象把一张纸揉成一团再展开,原来的相邻关系就被破坏了。

高方差数据拟合困难医学图像中不同器官和组织间的强度差异很大,而SSMs使用的高阶多项式投影算子(HiPPO)会平滑预测,难以准确拟合这些高方差点。

图4:使用S4拟合展平的2D医学图像的实验

从图4中可以清楚看到,当图像被展平后,在行与行的交界处(青色线标记),模型的预测出现了明显错误。红色框中的预测结果几乎与真实值相反,这就是不连续性带来的问题。

MM-UNet的设计理念:如何平衡CNN与SSM的优势

面对这些挑战,MM-UNet提出了一个巧妙的解决方案:不是简单地用SSM替换CNN,而是创建一个统一的元架构,让两种技术优势互补。

图1:(a)提出的MM-UNet架构概览。(b)在MM-UNet中用不同模块替换元块的实验,包括纯CNN基、混合和纯SSM基模块。跳跃连接表示残差连接

MM-UNet的核心创新在于其可替换的元块设计。编码器、瓶颈和解码器的每个阶段都由可替换的元块组成,支持纯CNN、纯Mamba或混合配置。

经过大量实验,研究团队发现最优配置是:两个卷积层后接一个Mamba模块,且Mamba模块放置在残差连接内部

图3:来自预训练模型的残差连接内部和外部的特征图强度分布,以及来自两个连续卷积层的特征图

为什么这样设计?图3揭示了关键原因:残差连接内部的特征图方差明显低于外部。由于SSMs难以处理高方差输入,将它们放在残差连接内部正好避开了这个弱点!

另一个巧妙设计是双向扫描策略。为了解决图像展平带来的不连续性问题,MM-UNet不仅按正常顺序扫描图像,还会按相反顺序再扫描一次。

图2:(1)提出的MetaSSM架构概览,其中MetaScan模块被不同的扫描顺序替换。(2)不同MetaScan配置的实验评估

实验设计:从数据集到评估指标的选择

为了全面验证MM-UNet的有效性,研究团队选择了两个具有挑战性的公开数据集:

AMOS2022腹部CT器官分割数据集包含200个腹部CT扫描,手动标注了16个解剖结构,是multi-organ分割任务的重要基准。

Synapse数据集包含30个腹部CT扫描案例,按照既定划分策略,使用24个案例进行训练,4个进行验证。

评估指标采用医学图像分割中广泛使用的Dice相似系数(DSC),它衡量预测分割与真实标注之间的重叠程度,数值越高表示分割越准确。

训练过程中,研究团队采用了nnUNet框架,只修改网络架构而保持其他配置一致。学习率初始设为0.001,使用多项式衰减策略,优化器选择SGD,动量设为0.99。

一个特别巧妙的设计是深度监督:在解码器的最后三个阶段(对应三个最大分辨率)应用辅助损失。这就像有多个老师在不同层次上指导学生,确保每个阶段都能学到有用的特征。

损失函数结合了交叉熵损失和Dice损失,通过5折交叉验证确保结果的可靠性。所有实验的批量大小设为2,这在处理大型3D医学图像时是合理的选择。

SSM在医学图像分割中的应用现状与局限性

状态空间模型(SSMs)在自然语言处理领域大放异彩后,研究者们迫不及待地想把它应用到医学图像分割中。然而,这就像让一位语言学家去解读CT扫描图——专业不对口啊!😅

目前已有的Mamba-based分割模型包括SegMamba、VM-UNet和U-Mamba等。SegMamba采用纯Mamba编码器配CNN解码器,VM-UNet是完全基于Mamba的U-Net,而U-Mamba则是Mamba和CNN的混合体。这些模型都试图将SSMs的强大序列建模能力引入医学图像分割,但效果参差不齐。

SSMs在医学图像分割中面临两个核心挑战:

高方差数据拟合困难医学图像中不同器官和组织的强度差异很大,而SSMs使用的高阶多项式投影算子(HiPPO)会平滑预测,难以准确拟合这些高方差点。这就像用钝刀切精细的蛋糕,总是差那么点意思。

空间不连续性当3D医学图像被展平为1D序列时,原本相邻的像素在序列中可能相隔很远,SSMs很难推断这些"失散兄弟"之间的关系。

图4:使用S4拟合展平的2D医学图像的实验

从图4可以清晰地看到SSMs在处理展平图像时的困境。在行与行的交界处(青色线标记),模型的预测出现了明显错误,红色框中的预测结果几乎与真实值相反。这种不连续性让SSMs在处理空间结构时显得力不从心。

MM-UNet的创新模块设计:残差连接中的SSM融合策略

面对SSMs的局限性,MM-UNet提出了一个精妙的解决方案:不是简单地用SSM替换CNN,而是创建一个智能融合的架构,让两种技术各展所长。

MM-UNet的核心创新在于其可替换的元块设计。编码器、瓶颈和解码器的每个阶段都由可替换的元块组成,支持纯CNN、纯Mamba或混合配置。这种设计就像乐高积木,可以根据需要灵活组合。

经过大量实验验证,研究团队发现最优配置是:两个卷积层后接一个Mamba模块,且Mamba模块放置在残差连接内部。这个设计看似简单,实则蕴含深意。

图3:来自预训练模型的残差连接内部和外部的特征图强度分布,以及来自两个连续卷积层的特征图

图3揭示了这一设计的关键原因:残差连接内部的特征图方差明显低于外部。由于SSMs难以处理高方差输入,将它们放在残差连接内部正好避开了这个弱点!这就像给SSMs提供了一个舒适的"工作环境",让它们能够发挥最佳性能。

实验表明,这种混合配置相比纯CNN模块提升了1.9%,相比纯Mamba模块提升了1.8%。更重要的是,将Mamba模块嵌入残差连接内部比放在外部额外提升了0.5%的Dice分数。这些改进虽然看似微小,但在医学图像分割这种精度要求极高的任务中意义重大。

扫描顺序优化:如何减少空间不连续性的影响

如果说模块设计是MM-UNet的"内功",那么扫描顺序优化就是其"外功"。为了解决图像展平带来的不连续性问题,研究团队进行了深入的扫描策略研究。

医学数据集通常包含3D空间信息,将这些3D图像展平为1D序列时,有多种可能的扫描顺序。例如,Depth-Height-Width(DHW)顺序先沿宽度扫描,然后是高度,最后是深度。同样,还存在DWH、WDH、WHD、HDW和HWD等其他顺序,每种顺序都有其对应的逆扫描顺序。

图2:(1)提出的MetaSSM架构概览,其中MetaScan模块被不同的扫描顺序替换。(2)不同MetaScan配置的实验评估

研究团队测试了多种扫描策略,包括标准1D扫描、2D扫描、3D窗口扫描、Zigzag扫描和倾斜扫描等。结果发现,简单的双向扫描策略效果最佳。

从标准1D DHW扫描(基线模型B1)开始,Dice分数为0.902。引入相反方向的互补扫描顺序形成1D BiScan(B2)后,性能提升了0.008,达到0.910。这种改进突显了双向扫描在减轻空间维度展平带来的不连续性方面的有效性。

有趣的是,添加更多扫描顺序对并不会带来额外的好处。添加第三个独立扫描方向(HWD)到1D BiScan中并没有改善性能,表明未配对的扫描顺序会引入有害的不连续性。同样,扩展到三对(B4)甚至六对(B5)相反方向扫描顺序分别获得0.908和0.907的Dice分数,证明超过一对的添加只会引入冗余。

其他扫描策略如2D扫描(B6)由于未能充分建模时间连续性而表现不佳。3D窗口扫描方法(B7)将整个3D体积划分为多个非重叠窗口,获得了与1D扫描相当的结果(0.902),但这种方法由于窗口边界间的频繁跳跃而固有地增加了不连续性。

这些发现表明,MetaSSM的最佳扫描策略是使用简单的双向方法,仅包含一对相反方向的顺序。这种平衡的方法充分捕捉了空间连续性,并有效减轻了展平过程中由不连续性引起的预测误差。

图6:在推理期间沿轴向、冠状和矢状维度翻转以提高性能。此操作类似于在Mamba块内应用多个扫描顺序

在推理阶段,每个图像块会经历沿轴向、冠状和矢状维度的八次翻转,然后对预测结果进行平均以提高性能。这种操作功能上类似于在Mamba块内应用多个扫描顺序,同时使用3D高斯函数与平均预测相乘,以增强中心区域的权重并减弱外围区域的权重,从而减轻边界不连续性引起的预测误差。

实验结果对比:MM-UNet在AMOS2022和Synapse数据集上的优异表现

理论说再多,不如实验结果有说服力。MM-UNet在两个具有挑战性的医学图像分割数据集上进行了全面测试,结果令人印象深刻!

在AMOS2022数据集上,MM-UNet与多种主流方法进行了对比,包括基于卷积的方法(nnUNet、3D UX-Net)、基于Transformer的方法(UNETR、SwinUNETR、nnFormer)以及基于Mamba的方法(VMUNet、SwinUMamba、UMamba)。

表2:MM-UNet与最先进方法在AMOS测试数据集上的比较,通过Dice分数评估。为公平比较,所有结果均基于5折交叉验证而无任何集成。最佳结果以粗体表示

从表2可以看出,MM-UNet在大多数器官上都优于所有现有方法,实现了最先进的DSC性能。具体而言,它在DSC上分别超过了nnUNet和3D UX-Net 3.2%和1.0%。同时,本方法优于所有基于Mamba的方法。考虑到AMOS2022数据集的复杂性,这些结果证明了本方法的有效性。

在Synapse数据集上,MM-UNet同样表现出色,与多种领先的基于卷积的方法(VNet、nnUNet)、基于Transformer的方法(TransUNet、SwinUNet、nnFormer)以及基于Mamba的方法(VMUNet、SwinUMamba、UMamba)进行了对比。

表3:MM-UNet与最先进方法在Synapse数据集上的比较(DSC以%表示)。最佳结果以粗体突出显示

MM-UNet实现了新的最先进性能,优于所有现有方法。具体而言,在这个竞争激烈的数据集上,它在DSC上分别超过了nnFormer和nnUNet 1.5%和6.9%。值得注意的是,本模型在分割大型器官(如肝脏和脾脏)方面表现出色,这些器官受益于SSMs捕捉长距离依赖关系的能力。

图7展示了与代表性方法的定性比较,证明MM-UNet生成了更准确的分割,特别是对于肝脏和脾脏。这些结果证实了本方法的鲁棒性和有效性。

为了直观深入地理解MetaSSM,研究团队进一步可视化了QK^T中的注意力值,如图5所示。每个注意力图都能有效捕捉时间维度上的图像模式,即使3D医学图像被展平为1D序列作为MetaSSM块的输入。与依赖大量参数在像素级别建立关系的注意力层不同,SSMs仅使用少量参数就实现了优越的注意力图。这些显著特性突显了在医学图像分割中使用SSMs的动机。

在这里呢迪哥也为大家准备了u-net结合mamba的经典案例和代码分享! 

有需要的小伙伴可以添加小助理领取

包含:

二、研究生论文及SCI攻略
三、学术论文写作攻略
四、科研方法分享
五、1v1论文辅导

Logo

魔乐社区(Modelers.cn) 是一个中立、公益的人工智能社区,提供人工智能工具、模型、数据的托管、展示与应用协同服务,为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作,由全产业链共同建设、共同运营、共同享有,推动国产AI生态繁荣发展。

更多推荐