分享自:

Weak-Mamba-UNet:视觉Mamba使CNN和ViT在基于涂鸦的医学图像分割中表现更佳

期刊:IEEE Transactions on Biomedical EngineeringDOI:10.1109/TBME.2026.3668882

本文介绍了一项发表于IEEE Transactions on Biomedical Engineering(已录用,待最终出版)的原创性研究。该研究由来自多所国际知名高校的研究人员共同完成,主要作者包括Ziyang Wang(阿斯顿大学,通讯作者)、Tianli Tao(伦敦国王学院)、Yiyuan Ge(华南理工大学)、Zhihao Chen(北京邮电大学)、Tianxiang Chen(中国科学技术大学)、Zi Ye(都柏林圣三一大学)和Yongxiang Lei(华威大学)。这项研究旨在解决医学图像分割领域的一个关键挑战:在标注成本高昂的情况下,如何利用弱监督学习(Weakly-Supervised Learning, WSL)技术,特别是基于涂鸦(Scribble)的稀疏标注,来实现高效且准确的分割。

学术背景与目标 医学图像分割是医学影像分析的核心任务,对于疾病诊断和治疗规划至关重要。以U-Net为代表的深度学习模型在该领域取得了巨大成功,但其高性能通常依赖于大量精确的像素级标注,这在临床实践中耗时耗力且成本高昂。为了降低标注负担,半监督学习(SSL)和弱监督学习(WSL)方法应运而生。其中,基于涂鸦的标注方式因其便捷性而受到关注,专家只需快速勾画目标区域的大致轮廓,无需进行精细的像素级标注,这大大减轻了标注工作量。

近年来,除了传统的卷积神经网络(CNN)和后来兴起的视觉变换器(Vision Transformer, ViT),一种基于状态空间模型(State Space Model, SSM)的新型架构——Mamba(特别是其视觉版本Visual Mamba, VMamba)在计算机视觉任务中展现出巨大潜力。Mamba以其线性计算复杂度和强大的长距离依赖建模能力著称,为处理医学图像中的全局上下文信息提供了新的工具。然而,在弱监督学习,尤其是涂鸦标注的设定下,如何有效整合CNN、ViT和Mamba这三种具有不同归纳偏置(CNN擅长局部特征,ViT和Mamba擅长全局建模)的架构,以发挥各自优势并实现协同学习,仍是一个未被充分探索的问题。

因此,本研究的目标是提出一个新颖的弱监督学习框架,名为Weak-Mamba-UNet。该框架的核心创新在于构建了一个包含三种异构编码器-解码器网络(CNN-based U-Net, Swin Transformer-based Swin-UNet, 以及VMamba-based Mamba-UNet)的多视角交叉监督学习系统,旨在利用涂鸦标注进行训练,并通过网络间的协作与相互监督来生成高质量的伪标签,从而迭代提升分割性能。研究旨在验证该框架在公开数据集上的有效性,并探索Mamba架构在弱监督医学图像分割场景下的潜力。

详细工作流程 本研究的工作流程是一个精心设计的、端到端的弱监督训练框架,主要包括数据准备、网络架构构建、训练策略与损失函数设计、以及评估验证几个关键环节。

首先,在数据准备阶段,研究使用了两个公开的医学图像分割数据集:ACDC(心脏MRI)和GlaS(结肠组织学图像)。为了模拟弱监督场景,研究者从这些数据集的原始密集标注(像素级真实标签)中人工生成了涂鸦标注。具体而言,对于每个需要分割的解剖结构(如右心室、左心室、心肌等),仅保留其轮廓线上的一部分稀疏点或短线,其余区域标注为“未标记”。所有图像被统一缩放至224×224像素的分辨率以供输入。

其次,构建核心网络架构。Weak-Mamba-UNet框架并行集成了三个独立的、具有对称编码器-解码器结构的U型分割网络,但它们的核心构建块不同: 1. CNN-based U-Net:采用经典的两层3x3卷积块作为基础模块,进行4级下采样和上采样。它主要负责提取精细的局部特征。 2. Swin Transformer-based Swin-UNet:使用两个Swin Transformer块作为基础模块,进行3级下采样和上采样,并在ImageNet上进行了预训练。它擅长建模分层的全局上下文信息。 3. VMamba-based Mamba-UNet:这是本研究引入的关键组件,使用两个视觉Mamba块作为基础模块,同样进行3级下采样和上采样。其核心是2D选择性扫描(SS2D)模块,该模块将2D图像特征沿四个对角线方向展开成序列,分别送入相同的S6(Mamba)算子进行处理,以线性复杂度高效捕获长距离的空间依赖关系,最后将四个方向的输出合并。这种设计使其能够在不显著增加计算负担的情况下,获得全局感受野。

第三,设计训练与监督机制。这是本研究方法学的精髓。三个网络使用相同的涂鸦标注进行独立初始化,并在训练中协同工作。训练损失由两部分组成: 1. 部分交叉熵损失(Partial Cross-Entropy Loss, L_pce):这是弱监督的基础损失。它仅计算在涂鸦标注像素点上的交叉熵损失,迫使网络在已标注的稀疏位置上做出正确预测。 2. 基于伪标签的Dice损失(L_dice):这是实现交叉监督的关键。在每次训练迭代中,框架会动态生成一个密集的软伪标签(Soft Pseudo Label)。生成方式并非简单平均,而是采用了狄利克雷采样混合(Dirichlet-sampled Mixing)策略:从狄利克雷分布Dir(1,1,1)中随机采样三个权重(α, β, γ),满足α+β+γ=1,然后用这些权重对三个网络当前迭代产生的类别概率图进行加权求和,得到混合伪标签 y_pseudo = α * p_cnn + β * p_vit + γ * p_mamba。这个软伪标签为网络提供了全图范围的、稠密的监督信号。每个网络除了最小化自身的L_pce,还需最小化其预测与这个混合伪标签之间的软Dice损失。

这种设计形成了一个多视角交叉监督(Multi-view Cross-supervised Learning)的闭环:每个网络从涂鸦获得直接监督,同时又从其他两个网络的“共识”(以随机混合的形式)中获得间接的、密集的监督。这种随机混合机制作为一种模型侧的扰动,起到了正则化的作用,鼓励网络学习更鲁棒的特征表示,避免了过拟合到单一架构的偏见。

第四,进行实验与评估。实验在配备NVIDIA GeForce RTX 3090 GPU的服务器上进行。所有方法均使用随机梯度下降(SGD)优化器训练30,000次迭代。为了公平比较,研究者将提出的框架与多种先进的弱监督分割基线方法进行了对比,这些方法均使用相同的U-Net或Swin-UNet作为骨干网络,包括:部分交叉熵(PCE)、不确定性感知自集成和变换一致性均值教师模型(USTM)、Mumford-Shah损失、Gated CRF,以及专门为涂鸦分割设计的ScribFormer和ScribbleVC。评估指标涵盖了相似性度量(Dice系数、准确率、精确率、灵敏度、特异性,值越高越好)和差异性度量(95%豪斯多夫距离HD95、平均表面距离ASD,值越低越好)。

主要结果 研究在两个数据集上进行了全面的定量和定性评估,结果有力地支持了Weak-Mamba-UNet框架的有效性。

ACDC心脏MRI数据集上(表I),Weak-Mamba-UNet在Dice系数上达到了0.9171,显著优于所有基线方法。例如,它比性能次优的Gated CRF + U-Net(0.9046)高出约1.25个百分点,比ScribbleVC(0.8676)高出近5个百分点。在衡量边界准确度的HD95指标上,Weak-Mamba-UNet仅为3.96,远低于其他方法(多数在6-150之间),表明其分割结果与真实边界更为贴合。统计检验(双尾配对t检验)显示,这些性能提升具有高度统计学显著性(p值 < 1e-6)。

GlaS组织学图像数据集上(表II),Weak-Mamba-UNet同样取得了最佳性能,Dice系数为0.9360,HD95为17.44,进一步证明了其在不同模态医学图像上的泛化能力。

消融实验(Ablation Study) 提供了更深入的见解(表III、IV)。关键发现包括: 1. 架构异构性的重要性:当框架中使用三个相同的网络时(如3×U-Net, 3×Swin-UNet, 3×Mamba-UNet),性能均不如异构组合(U-Net + Swin-UNet + Mamba-UNet)。特别是3×Swin-UNet组合性能大幅下降,表明单纯堆叠相同架构缺乏足够的特征多样性,可能导致模型收敛到相似的表示,无法有效互补。 2. Mamba-Unet的贡献:在异构组合中,包含Mamba-UNet的组合(如2×U-Net+Mamba-UNet)通常优于不包含它的组合。这证实了Mamba架构在捕获长距离依赖、提供独特视角方面的价值。 3. 伪标签混合策略:实验比较了固定权重混合与狄利克雷随机采样混合(表V)。固定权重(如各1/3)的Dice为0.8736,而狄利克雷采样策略将性能提升至0.9171。这验证了随机混合作为一种有效的正则化手段,能促进模型学习更鲁棒的特征。 4. 计算效率:表VI显示,Mamba-UNet在计算成本上相比Swin-UNet有显著优势:FLOPs减少59%,参数量减少63%,GPU内存占用减少21%,训练和推理时间也更短。这证明了Mamba在效率和性能间的良好平衡。

定性结果(图5)直观展示了Weak-Mamba-UNet的优越性。相比于基线方法(如PCE、USTM)在复杂边界或低对比度区域产生的错误预测或粗糙分割,Weak-Mamba-UNet的分割结果更完整、边界更清晰,更接近真实标签。注意力热图可视化(图6)也显示,网络能有效地聚焦于相关的解剖结构区域。

结论与意义 本研究成功提出并验证了Weak-Mamba-UNet,这是一个创新的、基于涂鸦标注的弱监督医学图像分割框架。其核心贡献在于首次将CNN、ViT和新兴的Visual Mamba三种异构架构整合到一个多视角交叉监督学习范式中。通过狄利克雷采样混合伪标签的机制,框架能够将稀疏的涂鸦监督有效地转化为密集的监督信号,促使三个具有不同归纳偏置的网络相互学习、互补不足。

该研究的科学价值在于:1)为弱监督学习,特别是涂鸦标注下的分割,提供了一种新的、有效的多模型协同学习范式;2)率先探索并证实了Visual Mamba架构在弱监督医学图像分割任务中的有效性和高效性;3)通过系统的消融实验,揭示了在弱监督框架中引入架构多样性(而非容量堆叠)对于提升性能的关键作用。

应用价值非常明确:能够显著降低医学图像分割对精细标注的依赖,减轻临床专家的标注负担,从而加速深度学习模型在医疗领域的实际部署。框架的代码、数据和基线方法均已公开,有助于推动相关领域的研究。

研究亮点 1. 首创性框架:首次提出并实现了融合CNN、ViT和Mamba三种异构架构的弱监督分割框架。 2. 创新的监督机制:设计了基于狄利克雷随机采样的软伪标签混合策略,实现了有效的多视角交叉监督与正则化。 3. 验证了Mamba的潜力:在弱监督医学图像分割任务中,系统性地验证了Visual Mamba在性能和效率上的双重优势。 4. 全面的实验验证:在多个数据集上与众多强基线进行对比,并通过详尽的消融实验分析了各个组件的作用,结论坚实可靠。 5. 实用性:为解决医学图像分析中的标注瓶颈问题提供了一个高效、可行的解决方案,具有直接的临床转化潜力。

其他有价值的内容 研究还简要讨论了失败案例(图7),指出模型在边界模糊、对比度低或结构重叠的区域仍可能出现错误,注意力机制有时会聚焦到不相关区域。这为未来研究指明了改进方向,例如引入不确定性建模或多尺度上下文细化。最后,作者展望了未来工作,包括将框架扩展到更多成像模态(如CT、超声)、处理3D体积数据,乃至应用于遥感、工业检测等非医学领域,显示了该方法的通用性前景。

上述解读依据用户上传的学术文献,如有不准确或可能侵权之处请联系本站站长:admin@fmread.com