点击下方卡片,关注「AI视界引擎」公众号
( 添加时备注:方向+学校/公司+昵称/姓名 )
医学图像分割在多种临床应用中发挥着重要作用,但现有模型常因计算效率低下以及复杂医学数据带来的挑战而难以应对。状态空间序列模型(SSMs)在以线性计算复杂度建模长距离依赖关系方面展现出潜力,然而其在医学图像分割中的应用仍受限于与图像 Token 的不兼容性及自回归假设。此外,在同时捕捉局部细粒度信息与全局语义依赖关系方面难以取得平衡。
为解决这些挑战,作者提出了SAMA-UNet,一种用于医学图像分割的新型架构。关键创新在于自适应Mamba样聚合注意力(SAMA)模块,该模块将上下文自注意力与动态权重调制相结合,基于局部和全局上下文优先考虑最相关的特征。
这种方法降低了计算复杂度,并提升了跨多尺度复杂图像特征的表征能力。作者还提出了因果共振多尺度模块(CR-MSM),通过因果共振学习增强编码器与解码器之间的信息流。该机制使模型能够自动调整不同尺度的特征分辨率和因果依赖关系,从而在U型架构中实现Low-Level特征与High-Level特征之间的更好语义对齐。
在MRI、CT和内窥镜图像上的实验表明,SAMAUNet在分割精度方面优于当前使用CNN、Transformer和Mamba的现有方法。该实现已在GitHub上公开发布。
- 引言
分割是医学图像分析中的一个基本步骤,它能够将图像划分为有意义的感兴趣区域。该过程在广泛的医学应用中发挥着关键作用,例如疾病诊断、癌症微环境量化、治疗计划和疾病进展跟踪[1]。深度学习等先进方法帮助研究行人检查不同类型的医学图像,以创建精确的分割图,从而在诊断中突出显示特定器官或疾病区域。卷积神经网络(CNN)和Transformer是医学图像分割领域中的先进深度学习方法。CNN使用共享权重,擅长捕捉平移不变性和局部特征提取。UNet[2]是一种突出的架构,它在编码器和解码器之间引入了 Shortcut ,实现了Low-Level和High-Level细节的整合,从而提高了其从高分辨率医学图像中提取层次化图像特征的效率。然而,CNN依赖于局部卷积核,难以捕捉长程依赖关系,这限制了它们在医学图像分割中处理全局上下文的能力。
Transformer模型最初是为自然语言处理(NLP)而开发的,现已成功应用于计算机视觉任务,包括医学图像分割。视觉Transformer(ViT)[3]和Swin Transformer[4]在捕捉全局依赖关系方面表现出色,因为它们将图像解释为一系列 Token 。TransUNet[5]和SwinUNETR[6]等混合模型结合了CNN的局部特征提取能力与Transformer的全局建模优势,从而提升了分割性能。然而,Transformer模型在序列长度方面会产生计算复杂度,在处理需要处理长序列的高分辨率医学图像时,这种复杂度变得难以承受。状态空间序列模型(SSMs)[7, 8]因其能够以线性计算复杂度建模长距离依赖关系而受到广泛关注。这些模型在NLP任务中取得了成功,并被扩展到计算机视觉领域,其中视觉状态空间(VSS)[9]、双向SSM[10]和全向SSM[11]等方法已展现出良好的效果。研究表明,SSM中的宏观设计Mamba模块[12]对基于图像的任务中的特征提取有显著贡献。然而,尽管Mamba在医学图像分割中取得了成功,但它也面临固有的挑战。具体而言,Yu和Wang[13]指出Mamba无法整合图像 Token 之间的因果关系,这限制了其在自回归任务(如分割)中的应用。
为解决这些局限性,作者提出了SAMA-UNet,这是一种新型架构,旨在将受Mamba启发的创新集成到基于transformer的模型中。SAMA-UNet引入了SAMA模块,该模块结合了上下文自注意力机制与动态特征调制,以增强局部和全局特征表示。这种方法提高了分割精度,并降低了传统transformer模型所关联的计算复杂度。此外,作者提出了CR-MSM,该模块通过引入因果共振学习来增强多尺度特征的集成。该模块利用U型网络固有的多尺度特征生成能力,在优化跨尺度特征表示的同时保留因果依赖关系。
作者旨在通过引入受人类视锥系统启发的softmax注意力机制来应对transformer的二次复杂度,从而实现更优的注意力分配,覆盖局部和全局上下文。作者的目标通过差异注意力得到进一步增强,该机制抑制无关的噪声 Token ,而快速注意力机制则提升了注意力机制的计算效率。这些创新使SAMA-UNet能够在保持计算效率的同时实现高性能分割。
在本文中,作者提出以下贡献:
-
作者提出了SAMA模块,该模块结合了动态特征调制和注意力机制,以实现线性复杂度,同时捕捉局部和全局特征依赖。作者建议使用CR-MSM来增强U型网络中的特征融合,并提高多尺度信息流的连贯性。
-
作者提出了一种用于医学图像分割的鲁棒SAMA-UNet架构,该架构在MRI、CT和内窥镜数据集上表现出优于现有基于CNN、Transformer和Mamba模型的性能。
-
在四个多样化的医学影像数据集上的实验结果表明,SAMA-UNet显著优于现有方法,为医学图像分割提供了一种更高效和准确的方法。
-
相关工作
2.1 基于Transformer的方法
Transformer已成为自然语言处理(NLP)、计算机视觉和多模态任务中的主流架构,这得益于其捕捉长距离依赖关系和表示复杂特征的能力。然而,传统Transformer模型在处理高分辨率图像时存在的二次计算复杂度仍然是一个重大挑战。研究行人已提出多种方法来缓解这一问题,同时保持性能[14]。Sparse注意力技术通过缩短token序列来降低计算复杂度。例如,移位窗口注意力[4, 15]在小窗口内计算注意力图,并通过重叠滑动窗口扩展全局注意力覆盖范围。
类似地,扩张注意力[16]采用扩张率,在有限的采样点内捕捉短距离和长距离依赖关系。可变形注意力[17, 18, 19]根据学习到的偏移动态调整采样位置,以实现更灵活的注意力分配,满足特定输入特征的需求。此外,像素焦点注意力模拟人类视觉工作原理,使每个特征点既能看到局部细节也能看到全局细节。此外,线性注意力用非负映射函数替代标准注意力中的softmax运算,将复杂度从二次时间降低到线性时间。然而,这种方法通常会导致性能下降。在此基础上,聚焦线性注意力[20]通过设计高效映射函数改进注意力矩阵的表示。此外,类似Mamba的线性注意力[12]在受Mamba启发的更大设计中结合了SSM和线性注意力,进一步提升了性能。最后,Flash注意力[21]使GPU上softmax注意力的数据传输和计算过程更快更高效,同时仍保持良好性能。
这些进展激励作者开发了SAMA-UNet,该模型通过使用硬件加速的差分注意力机制,将类似Mamba的宏观架构设计融入到像素级注意力中。这种方法使得自注意力机制能够在保持高表示能力的同时实现更高效的计算复杂度,特别适用于高分辨率医学图像分割任务。
2.2. 基于Mamba的方法
状态空间序列模型[7, 8]在自然语言处理和计算机视觉领域中展现出对长距离依赖建模的线性计算复杂度优势。Mamba模块作为一种基于状态空间序列模型的方法,通过线性变换引入选择性机制,实现序列内输入依赖的交互,从而增强基于上下文的推理能力。Mamba的宏观架构设计已被应用于混合CNNSSM结构中,以提升图像处理任务性能,特别是在医学图像分割[22]方面。针对Mamba的若干扩展已被提出以改进分割性能。
例如,U-Mamba[23]将Mamba模块集成到卷积层后的U-Net编码器中,其性能优于传统的CNN和基于transformer的模型。Swin-UMamba[24]与VM-UNet[25]分别用受视觉状态空间启发的组件替换U-Net中的卷积层,以改进局部和全局特征的表征。VM-UNet-V2[26]采用卷积块注意力机制优化 Shortcut ,并引入语义与细节注入模块以增强高低层特征的交互。与此同时,LKM-UNet[27]结合基于块的Mamba模块,实现局部和全局层面的特征提取。此外,MSVM-UNet[28]在视觉状态空间模块中添加多尺度卷积,以有效捕获不同尺度的二维特征。
尽管这些模型成功地将Mamba集成到医学图像分割中,但它们并未充分利用Mamba在分析因果序列数据方面的优势。作者的工作通过关注U-Net中多尺度特征之间的潜在因果关系,强化了Mamba的优势。通过将因果共振学习嵌入编码器-解码器结构中,作者增强了特征表示,并缩小了编码器与解码器输出之间的语义差距,从而在高分辨率医学图像分割中取得了更好的性能。
作者的SAMA-UNet模型通过引入SAMA模块提升了现有混合模型的性能,该模块能够动态调整特征并采用因果共振学习,从而提高了分割精度并减少了计算量。通过融合这些创新技术,SAMA-UNet解决了传统CNN、Transformer和基于Mamba方法的固有挑战,为高分辨率医学图像分割任务提供了一种更高效、更鲁棒的解决方案。
- 方法论
所提出的SAMA-UNet架构旨在提升医学图像分割的效率与准确性,如图1所示。yinqing-SAMA-UNet_2505通过引入多项关键创新来应对医学数据带来的计算低效问题,从而实现对复杂解剖结构的精确分割。具体而言,SAMA-UNet集成了以下组件:(1) SAMA模块取代Mamba模块中的SSM,采用具有线性计算复杂度的softmax注意力机制,有效消除了自回归假设的需求。(2) CR-MSM利用U型网络连续多尺度特征中隐含的因果关系,引入了一种新型扫描扩展方法来增强2D-选择性扫描(SS2D)操作。该设计使SSM的连续长序列建模能力与医学图像分割任务的需求更为契合。
3.1. 主干架构
SAMA-UNet的架构采用U形设计,以实现高效的医学图像分割。块嵌入模块首先处理输入图像,将其转换为特征嵌入。该模块使用重叠块通过卷积操作降低空间维度并增加通道维度来分解图像。这一过程确保了重要空间信息的保留,同时将图像嵌入到高维特征中。嵌入后的特征随后通过编码器中的一系列SAMA模块进行处理,其中它们会经历动态注意力处理。SAMA模块降低了传统注意力机制的计算复杂度,同时增强了模型捕获局部和全局依赖关系的能力。这些特征通过包含因果共振多尺度模块的改进 Shortcut 进行进一步细化,以更好地对齐编码器和解码器之间的信息。多尺度优化特征通过卷积在解码器中逐步融合,并通过转置卷积逐渐恢复到原始图像分辨率。最后,不同尺度的解码头生成分割结果。
3.2. SAMA
SAMA模块是SAMA-UNet中的关键创新,旨在提升Transformer的效率同时保持高表示能力。该模块基于类似Mamba的注意力机制构建,结合了上下文自注意力和动态权重调节的优势。图2(c)展示了所提出的SAMA模块。图2还突出了作者的模块与其他基于Mamba的模块之间的差异。如图2(a)所示,Mamba使用SSM和门控注意力[29]来改变其设计,而MLLA则将线性注意力与Mamba的结构相结合,如图2(b)所示。为清晰和简洁起见,图中未描绘所使用的多个注意力头。
借鉴[13, 12]的研究成果,所提出的SAMA模型将多种Mamba类设计元素整合到注意力机制的大结构中。这些设计包括在注意力层之前引入SiLU激活层、深度卷积层(DepthConv)和线性层,以及绕过激活和加权策略。这些设计增强了模型捕获位置信息和通道信息的能力,从而提升了编码器的性能。SAMA模块内的操作涉及两个分支:局部分支和全局分支。局部分支专注于局部特征提取,而全局分支则聚合全局上下文。局部和全局特征通过聚合注意力并行处理,然后进行拼接以进行进一步处理。这种方法在降低计算成本的同时,确保了局部和全局特征的有效捕获。
SAMA模块的流程可以用公式1描述。
和
分别表示局部分支和全局分支的输入。Split和Concat分别指沿着通道维度的拆分和拼接操作。
此外,作者对基于像素的注意力方法引入了若干改进。首先,作者将特征图沿通道维度分为两部分:一部分用于局部注意力,另一部分用于全局注意力。这两个部分分别处理,然后进行拼接。这种设计解决了两个关键问题:(1) 直接堆叠局部和全局注意力可能会由于下采样导致跨越局部和全局范围的特征引入不一致的偏差;(2) 分割特征图可以降低计算复杂度。其次,作者将局部增强位置编码(PE)融入softmax注意力模块。这一改进通过卷积提供位置信息,同时消除了相对位置偏差。第三,作者移除了一些可学习的嵌入参数和其他配置,以简化注意力模块。
最后,受差分Transformer[30]的启发,作者将差分注意力引入局部和全局注意力的计算中。他们的研究发现,在注意力图中,一些不可忽略的注意力分数被分配给了无关的上下文。因此,无关的上下文稀释了相关信息的权重分布,对整体性能产生了负面影响。通过引入这种结构,作者将差分去噪应用于作者的注意力模块,帮助其消除注意力噪声,更有效地关注相关信息。
作者使用局部分支作为示例,描述所提出的差分聚合注意力模块如下:首先,如方程2所示,局部分支的
和
Token 沿着通道维度分别被分割为
和
。接着,方程3展示了作者如何通过两组Softmax注意力的差值来创建差分注意力,使用可学习的标量
,其初始值为
。
表示方程1中Split操作后特征通道维度。遵循文献[30]中的参数设置,
被设置为固定值0.8。最后,作者对差分注意力应用组归一化(GN)操作,然后进行标量点积生成差分输出。该输出与由PE生成的位置信息相结合,产生差分聚合注意力(DiffAgg)输出。
3.3. CR-MSM
CR-MSM旨在解决多尺度特征融合的挑战,特别是在U型架构中。医学图像通常表现出多尺度信息,必须有效捕捉以确保精确分割。CR-MSM引入了因果共振学习,该学习确保模型在优化编码器和解码器之间信息流的同时,保持不同空间尺度上的因果依赖关系。
与传统仅依赖卷积的多尺度模块不同,CR-MSM通过多视图变换和状态空间建模(SSM)明确捕捉方向性。在此上下文中,SSM指的是一类高效序列建模技术,将空间特征图视为1D序列,并使用结构化线性递归来捕获长距离依赖关系。SSM模型使其适用于对加宽和变换后的特征图上的空间序列进行建模,而无需计算机制带来的负担。
给定编码器层的一组特征图,CR-MSM应用三种空间变换——原始、转置和翻转——以生成多样化的方向视图。这些视图在尺度上进行扩展并连接,以创建三个不同的序列通道。每个方向序列通过一个SSM模块,捕捉跨空间方向的远距离上下文关系。生成的序列随后被分割、逆变换,并重建为原始格式。最后,作者对三个方向结果进行平均,并通过一个线性投影层,得到一个强且对齐的多尺度表示。作者将CR-MSM的输出传递给解码器进行融合,这有助于弥合编码器和解码器之间的语义差距,从而提高分辨率感知的语义一致性。
- 实验与结果
4.1. 实验设置
4.1.1. 数据集
BTCV数据集是Synapse腹部多器官分割任务的一部分,该任务在Multi-Atlas Labeling Beyond the Cranial Vault Workshop and Challenge [31]中被引入。它包含30次腹部CT扫描,共计3779个轴向增强CT切片。每个CT体积包含85到198个切片,分辨率为512×512像素。根据先前工作的数据分割[5],其中18个案例用于训练,12个案例用于测试。默认输入图像大小设置为
。此设置基于nnUNet配置[32]进行了微调。ACDC数据集源自Automated Cardiac Diagnosis Challenge [33],包含100次心脏MRI扫描,每扫描分割为三个子结构:右心室(RV)、心肌(Myo)和左心室(LV)。80个MRI样本随机用于训练,其余20个样本指定用于测试。与nnUNet设置一致,默认输入图像大小设置为
。EndoVis17数据集来自MICCAI 2017 EndoVis Challenge [34],专注于内窥镜图像中七个手术器械的分割。作者采用官方数据集分割,其中训练集包含从八个视频中提取的1800和1200个图像帧。测试集包含来自两个额外未见过视频的帧。输入图像大小设置为
,遵循nnUNet配置。ATLAS23是一个公开可用的分割数据集[35],源自MICCAI 2023 ATLAS Challenge,包括60次CE-MRI T1加权扫描以及60个肝脏和肝脏肿瘤分割 Mask 。标注由放射科医生生成。输入图像大小设置为
,符合nnUNet配置。
4.1.2. 评估指标
遵循Metrics Reloaded [36]和U-Mamba [23]中的建议,作者使用Dice相似度系数(DSC)和归一化表面距离(NSD)来评估这些语义分割任务。DSC是用于评估预测分割与真实分割重叠程度的常用指标。数学上,它表示为:
其中
和
分别表示预测分割区域和真实分割区域。NSD量化了预测分割和真实分割表面之间的差异。它关注边界对齐,因此在评估具有精细结构细节区域的分割质量时特别有用。它计算预测分割边界点中有多少比例位于真实边界预定义的容差距离内,并归一化到总边界点数。数学上,它可以表示为:
这里,
表示预测分割的表面点集,
表示真实分割的表面点集。函数
表示点
到表面
的最短距离。参数
是一个预定义的容差值,
表示集合的基数,即集合中点的数量。
4.1.3. 实现细节
所有实验均使用Python 3.10和PyTorch 2.0.0,在配备24 GB内存的NVIDIA GeForce RTX 3080 GPU上进行。nnUNet框架作为核心,每个训练周期包含500个训练周期,每个训练周期由250次迭代组成。预处理和数据增强方法严格遵循nnUNet为相应数据集实现的方法。为了进行公平比较,作者将SwinUNETR、U-Mamba、LKM-UNet和MLLA-UNet集成到nnUNet框架中。所有网络均使用AdamW优化器从头开始训练,默认初始学习率设置为
或
,尽管某些模型遵循其发布代码中指定的学习率设置。采用一个周期的余弦退火学习率调度器来调整训练过程中的学习率。
4.2. 定量与定性结果
表1 总结了在四个医学图像分割数据集上的定量 2D 分割结果。与基于 CNN、Transformer 和 Mamba 的方法相比,作者提出的 SAMA 模型在 DSC 和 NSD 指标上始终表现出更优的性能。
在BTCV数据集上,yinqing-SAMA-UNet_2505在DSC和NSD指标上分别比第二优模型U-Mamba高出了0.92%和1.15%。作者遵循U-Mamba的评估标准,计算时包括了13个器官类别中的8个(主动脉、胆囊、左肾、右肾、肝脏、胰腺、脾脏和胃)。基于Mamba的方法展现出优异的整体性能,yinqing-SAMA-UNet_2505在DSC和NSD指标上分别比U-Mamba高出0.92%和1.15%。在ACDC数据集上,尽管所有方法都达到了相对较高的准确率,yinqing-SAMA-UNet_2505在DSC和NSD指标上分别比SwineUNETR高出了0.32%和0.20%。在EndoVis17数据集上,基于CNN和Mamba的方法均表现良好,yinqing-SAMA-UNet_2505进一步提升了性能,在DSC和NSD指标上分别比U-Mamba高出了1.69%和1.84%。最后,在ATLAS23数据集上,yinqing-SAMA-UNet_2505在DSC和NSD指标上分别比U-Mamba高出了1.27%和1.68%。
这些结果共同表明,yinqing-SAMA-UNet_2505不仅在不同医学影像模态(包括MRI、CT和内窥镜)上实现了强大的泛化能力,而且有效地提升了基于Transformer的方法,从而显著提高了性能。四个数据集的定性结果如图3所示。在第一行中,许多方法,如nnU-Net、SwinUNETR、U-Mamba和LKMUNet,由于其异质性外观,无法准确分割器官区域的轮廓。在第二行中,U-Mamba和LKM-UNet在右心室产生了分割错误。在第三行中,nnU-Net、SwinUNETR和LKMUNet错误分类了Prograsp Forceps的分割区域,而U-Mamba在Bipolar Forceps和超声波探头中产生了错误分类。在第四行中,包括nnU-Net、SwinUNETR、U-Mamba和LKMUNet在内的几种方法无法准确分割肝脏肿瘤的轮廓。相比之下,SAMA-UNet在这些场景中表现出明显优势,突显了其捕捉局部细节和全局语义信息的能力。
4.3. 单个器官的定量分析
作者对BTCV数据集上的各个类别进行了详细的定量分析,如表2所示。在BTCV数据集的器官分割方面,作者提出的SAMA-UNet在多个器官上实现了最高的DSC分数:主动脉为91.53%,胆囊为68.95%,左肾为87.69%,右肾为86.37%,肝脏为95.73%,胰腺为77.82%,脾脏为92.28%,胃为82.71%。然而,胆囊和胰腺的分割仍有更大的提升空间。
4.4. 消融实验
在本小节中,作者对BTCV数据集进行消融实验,以验证yinqing-SAMA-UNet_2505中关键组件的有效性。首先,作者进行实验以评估yinqing-SAMA-UNet_2505中Mamba-like聚合注意力模块的有效性。此外,作者对具有隐式因果性的多尺度Mamba模块的某些配置进行消融实验。
- SAMA模块的消融研究:为了提升基于Transformer的token混合模块的性能和计算效率,作者使用BTCV数据集对SAMA模块进行了全面的消融研究。SMM模块、MLLA模块以及所提出的SAMA模块的详细设计如图2所示。实验结果聚焦于分割性能和计算复杂度,如图4和图5所示,评估基于DSC、NSD、GFLOPs和参数总数等关键指标。
如图4所示,作者对SwinUNETR的初始修改包括用深度可分离卷积模块替换编码器阶段的patch merging、patch expanding和卷积操作,同时用简单的恒等映射替换 Shortcut 中的卷积操作。这些修改不仅将DSC提高了0.87%,还显著降低了GFLOPs和参数数量,分别从281.69和100.44M减少到124.2和19.8M。基于这些修改,作者进一步用三种替代方案替换了token mixer组件:SSM、MLLA和SAMA模块。
将原始的移位窗口自注意力(SwsA)模块替换为SSM和MLLA模块,分别将DSC提升至82.82%和81.81%,并将NSD提高至84.23%和86.53%。然而,这些改动也略微增加了GFLOPs,从124.2增加到131.9和127.6,参数量分别从19.8M增加到23.0M和21.5M。基于这些基准,作者对所提出的SAMA模块进行了进一步分析。
对于SAMA模块,作者首先用聚合注意力(Agg)模块替换了SwSA模块,该模块集成了通道级分离和改进的局部-全局像素聚焦注意力。这种替换相较于SwSA模块减少了计算量和参数,同时显著提升了DSC,提升了5.3%,如图4所示。这表明作者提出的注意力机制实现了线性复杂度,并有效提取了局部和全局特征。在此基础上,作者将类似Mamba的结构融入Agg模块,包括旁路分支、线性层、卷积层以及在通道分离操作前的SiLU激活函数。这一改进进一步将DSC和NSD提升至84.37%和87.24%,分别优于采用类似Mamba结构的SMM和MLLA模块。尽管这一改动将计算量和参数增至125.3和20.8M(如图5倒数第二列所示),但数值仍低于SMM和MLLA模块。虽然引入了线性层和卷积层增加了参数,但也验证了Mamba模块在增强特征提取能力方面的宏观架构优势。最后,作者将Softmax注意力替换为差分注意力。这一修改未引入额外的计算成本或参数,但进一步将DSC和NSD提升至84.88%和87.43%。这些结果表明差分注意力在抑制注意力噪声、为相关特征分配更高分数并将无关特征推向零方面的有效性。
总体而言,与SMM和基于MLLA的方法相比,所提出的SAMA模块在替代SwSA时表现出更优的性能和更低的计算复杂度。这验证了yinqing-SAMA-UNet_2505的有效性。
- CR-MSM模块的消融研究:为了评估所提出的CR-MSM在捕获长距离、因果对齐的空间依赖关系方面的有效性,作者进行了一系列聚焦的消融实验。这些实验针对CR-MSM的三个核心组件:方向性多视角变换、状态空间建模机制以及因果融合策略。每个实验旨在隔离某一特定设计元素对BTCV数据集整体分割性能的贡献。
以下检查了以下消融设置:
用单个原始替换方向多视图变换
从整体视角出发,评估空间方向多样性对特征的影响
表示。用标准卷积层替换SSM模块,从而
移除用于捕获长距离的结构的递归建模
依赖项。替换原本通过因果关系融合机制实现的部分
方向性平均,结合简单特征拼接,在多尺度整合过程中消除了因果对齐。
表3中的结果清晰地突出了每个组件的重要性。在实验
中,移除多视图变换导致DSC下降了0.33%,NSD下降了0.29%,证实了方向多样性增强了空间理解能力。在实验
中,用卷积层替换SSM导致DSC下降了0.45%,NSD下降了0.34%,强调了结构化递归在长距离上下文建模中的优势。在实验
中,用简单拼接替换因果融合导致DSC性能损失了0.53%,NSD损失了0.31%,突出了跨空间视图的因果平均的有效性。
总体而言,这些发现验证了CR-MSM架构作为一个稳健且高效的模块,能够增强多尺度表示学习。当移除每个组件时,都会导致可测量的性能下降,这证实了其综合设计对分割精度有显著贡献。
4.5. 参数与计算效率
作者通过使用BTCV数据集评估并比较了每个模型的计算复杂度,具体通过计算可训练参数数量和浮点运算次数(FLOPs)来进行。如表4所示,yinqing-SAMA-UNet_2505SAMA-UNet在29.03M参数和145.67 GFLOPs上实现了平衡的权衡,在保持计算效率的同时提供了具有竞争力的性能。与其他最先进模型相比,SAMA-UNet所需的参数和FLOPs显著少于基于重量的Transformer模型,如LKM-UNet(189.55M,993.70G)和SwinUNETR(100.44M,281.69G)。即使与仅使用编码器的模型如U-Mamba(Enc)和nnUNet相比,yinqing-SAMA-UNet_2505仍然是最高效的,这表明其具有更快的推理潜力和更低的内存开销。这些结果突出了SAMA-UNet的计算优势,使其成为实时或资源受限的医疗影像应用中不牺牲分割质量的合适选择。
- 讨论
本研究探讨了状态空间序列模型在图像分割任务中的局限性,特别是其在满足自回归假设方面的困难。具体而言,图像 Token 之间缺乏因果关系,使得基于扫描块 Token 在分割任务中预测下一个 Token 在逻辑上不可行。为应对这些挑战,作者提出了SAMA模块和CR-MSM模型,从两个角度探索基于Mamba设计的潜力以增强医学图像分割。在四个数据集上的定量和定性结果表明,yinqing-SAMA-UNet_2505在分割多个器官和手术器械方面具有有效性。作者将在下文讨论yinqing-SAMA-UNet_2505在技术和临床方面的影响。
5.1. 技术影响
为克服现有状态空间序列模型(SSMs)在建模图像 Token 之间因果关系方面的局限性,作者从两个互补的角度来处理这个问题,这与SAMA-UNet中引入的架构创新相一致。为克服现有状态空间序列模型(SSMs)在建模图像 Token 之间因果关系方面的局限性,作者从两个互补的角度来处理这个问题,这与SAMA-UNet中引入的架构创新相一致。
与自然语言不同,医学分割任务中的图像 Token 并不遵循固有的因果序列。传统自回归模型设计的空间自编码器(SSMs)在这一背景下难以有效捕捉空间依赖关系。为解决这一问题,作者将编码器中的基于SSM的 Token 混合器替换为所提出的SAMA模块。SAMA引入了一种混合 Token 混合策略,该策略结合了像素聚焦softmax注意力和受Mamba启发的宏观架构元素,使模型能够在不强制因果序列的情况下捕捉局部细节和全局语义依赖关系。此外,作者还集成了差分注意力和快速注意力机制。受人类视锥系统启发的这些机制使网络能够聚焦于最相关的区域,同时保持线性计算复杂度。作者的实验(如图4和图5所示)证实,使用SAMA模块代替SSM能够显著提升性能,同时相比基于SSM的设计,还能减少模型参数和FLOPs数量。虽然注意力机制有助于在编码器中混合 Token ,但U形分割网络也高度依赖 Shortcut 来在不同尺度上共享特征。在此,作者引入了CR-MSM来捕捉跨分辨率 Level 的隐式因果关系。在U形网络中,编码器中更高分辨率的特征本质上先于更粗糙的表示,并通常为解码器重建提供基础线索。然而,传统的多尺度融合(例如简单的拼接或基于卷积的聚合)无法保留这种方向性因果性。CR-MSM通过将四个方向性变换(原始、转置、翻转和翻转转置)应用于编码器特征来解决这一问题。一个共享的SSM模块对每个变换视图进行展平和处理,捕捉跨空间方向的长程依赖关系。为在融合过程中保持因果一致性,作者引入了一种因果共振策略,该策略执行方向性平均而不是简单的拼接。这种融合确保模型尊重特征生成的自然顺序。
在消融研究中,表1HI证实了每个组件的重要性。将因果融合替换为拼接会导致DSC下降;移除多视图变换会导致性能下降。这些结果共同验证了作者的假设,即空间因果关系对于在医学图像中保持语义连续性至关重要。
5.2. 临床影响
从临床角度来看,SAMA阻滞和CR-MSM模块显著提高了诊断精度。通过捕获全局上下文信息,yinqing-SAMA-UNet_2505能够准确识别原始图像中的结构和边界,从而生成精确的分割图。这些改进有助于实现更准确的诊断,并促进优化治疗方案的开发。在包括MRI、CT和内窥镜在内的多种成像模式上的大量实验表明,yinqing-SAMA-UNet_2505是一种稳健且通用的医学图像分析工具。该结果突出了集成多模态诊断的潜力,有效利用不同模态的互补信息,以实现更详细和精确的图像分析。此外,它还提供了一种可扩展的解决方案,用于在医学图像中利用大规模多模态数据,符合对大规模模型和数据集日益增长的需求。
此外,在放射科工作流程中,能够快速处理和分析大量高分辨率成像数据的能力可以显著减少诊断时间,从而促进临床决策的加速。yinqing-SAMA-UNet_2505的计算开销降低和内存占用减少,使其非常适合集成到硬件资源有限的临床系统中,进一步扩大了其在实际医疗环境中的应用范围。
5.3. 局限性与未来方向
然而,仍存在若干挑战。首先,将Mamba宏观架构集成到SAMA模块中增加了参数数量,这需要更有效的参数优化策略。其次,扩展模型以处理3D图像分割会消耗大量更多内存,因为它将Vision State Space模块中的扫描序列数量增加了三倍。应对这些挑战需要设计在效率和效果之间取得平衡的模块,尤其针对3D网络应用。未来工作将集中于优化网络架构、提升参数利用率,并探索大规模数据集中的实时分割。
- 结论
在本工作中,作者提出了一种用于医学图像分割的新型SAMA-UNet架构,该架构有效结合了动态注意力、基于Mamba的建模和多尺度特征融合的优势。作者的设计采用自适应Mamba类聚合注意力模块,以提升模型对全局和局部特征的理解能力,同时保持计算效率。
此外,作者提出了因果共振多尺度模块,以促进编码器-解码器阶段中因果对齐的特征集成。通过在四个不同基准数据集-BTCV、ACDC、EndoVis17和ATLAS23上进行广泛评估,SAMAUNet在CNN、Transformer和基于Mamba的模型中均表现优于当前最优水平,并分别达到了85.38%、92.16%、67.14%和84.06%的DSC分数,以及87.82%、96.54%、68.70%和88.47%的NSD分数。
这些结果验证了yinqing-SAMA-UNet_2505在不同模态(包括CT、MRI和内窥镜成像)上的鲁棒性和泛化能力。尽管SAMA-UNet性能优异,但仍存在局限性。Mamba结构的集成有利于序列建模,但会引入额外参数。此外,将架构扩展到3D分割场景会导致因SS2D扫描空间增大而增加内存使用。针对这些挑战,特别是通过轻量级设计优化和高效3D适配,将指导作者未来的研究工作。
点击上方卡片,关注「AI视界引擎」公众号