点击下方卡片,关注
「AI视界引擎」
公众号
( 添加时备注:方向+学校/公司+昵称/姓名 )
近期的研究突出了Mamba框架,这是一种以线性计算复杂度高效捕捉长程依赖关系的状态空间模型。尽管Mamba在医学图像分割中表现出竞争性的性能,但由于传统基于位置的扫描方法的间歇性特性以及医学图像中常见的复杂、模糊的边界,它在建模局部特征时遇到了困难。
为了克服这些挑战,作者提出了不确定性驱动的Mamba(UD-Mamba),它通过将通道不确定性纳入扫描机制来重新定义像素顺序扫描过程。
UD-Mamba引入了两种关键的扫描技术:
1)顺序扫描,通过逐行扫描优先处理高不确定性区域;
2)跳过扫描,垂直处理列,以固定间隔从高不确定性到低不确定性或从低不确定性到高不确定性。顺序扫描高效地将高不确定性区域,如边界和前景目标,聚类,以提高分割精度,而跳过扫描增强了背景和前景区域之间的交互,允许及时整合背景信息以支持更准确的前景推理。
认识到从确定区域到不确定区域扫描的优势,作者引入了四个可学习参数以平衡不同扫描方法提取的特征的重要性。此外,采用余弦一致性损失来减轻扫描过程中在不确定和确定区域之间转换的缺点。
作者的方法在涉及病理学、皮肤病学病变和心脏任务的三种不同的医学图像数据集上展示了鲁棒的分割性能。
代码可在https://github.com/piooip/UD-Mamba获取。
unset
unset
- 引言
unset
unset
本文旨在探讨人工智能(AI)领域的研究进展及其在各个行业的应用。随着计算能力的提升和大数据技术的普及,AI技术取得了显著的进步,并在众多领域展现出巨大的潜力。本文将概述AI的基本概念、发展历程以及当前的研究热点,同时分析其在工业、医疗、教育等领域的应用现状和未来发展趋势。
Transformer在图像处理领域展现出巨大的潜力,得益于其建模长距离依赖的能力。然而,它们与序列长度成二次关系的计算复杂度给计算带来了巨大负担,尤其是在医学图像分割等高分辨率任务中。最近,状态空间模型(SSM)作为一种更高效的计算替代方案出现,在保持建模长距离依赖能力的同时,实现了线性复杂度[15]。在这些模型中,Mamba架构[14,8]尤为突出,它采用选择性扫描技术和硬件优化设计,在各种视觉任务中取得了令人印象深刻的成果:基于像素级通道不确定性的扫描机制。
(a)输入图像;
(b)真实情况;
(c)通过基于通道不确定性的计算得到的图像,其中最高不确定性的像素为红色,最低为蓝色;
(d)按不确定程度排序的特征图像;(e)使用SS2D[32]扫描机制的前一种方法;(f)作者的UD-SSM扫描机制,包括顺序扫描和跳过扫描。
在医学图像分割领域,主要目标是准确划定与目标器官或病理组织相对应的区域,为临床诊断提供关键支持[39,5,22,16,29,44]。由于Mamba框架能够捕捉长距离依赖关系并高效处理高分辨率图像,它在医学影像领域的应用越来越广泛[55,54]。然而,Mamba的传统基于位置的顺序扫描方法往往导致不同语义区域的间歇性扫描(见图1(e)),在处理医学图像中的复杂背景和模糊边界时尤为成问题。这阻碍了Mamba准确建模对有效分割至关重要的局部特征[11,45]。
基于先前不确定性方法在医学图像分割中的有效性[53,41,36,24,2],作者提出了不确定性驱动Mamba(UDMamba)以克服这一局限性。UD-Mamba的核心是不确定性驱动选择性扫描模型(UD-SSM),该模型利用通道不确定性作为指导性指标来重新定义像素级扫描过程。如图1(c)所示,具有较高中值通道不确定性的像素通常与关键区域相关,如前景和边界。相反,不确定性较低的区域通常与背景相关。通过计算不确定性图并按不确定性 Level 对像素进行排序,如图1(d)所示,作者确保将不确定(因而关键)的区域与更确定(通常代表背景信息)的区域区分开来。
在UD-SSM中提出的扫描策略,如图1(f)所示,包含两种关键方法:
1)顺序扫描:该方法严格根据像素的不确定性水平顺序处理像素,有效聚类高不确定性区域,如边界和前景区域。通过专注于这些关键区域,顺序扫描确保模型能够捕捉到准确分割所需区域的细微细节。
2)跳过扫描:该技术以恒定的不确定性间隔垂直穿过图像,增强背景和前景信息之间的交互。它补充了模型对背景区域的理解,同时确保前景分割的精确性。
通过结合顺序扫描和跳过扫描,UD-Mamba能够在关注关键区域的细微结构的同时,维持对更广泛背景的理解。这种双扫描方法使得分割性能更加平衡和有效。受Vision Mamba[67]的启发,作者在顺序扫描和跳过扫描中进一步引入了正向和反向扫描机制。
然而,作者发现从低不确定性区域扫描到高不确定性区域通常比反向扫描产生更好的结果,如图2所示。为了优化这一过程,作者引入了四个可学习的参数,这些参数调整了从不同扫描技术收集到的特征的重要性。此外,作者应用余弦一致性损失,以确保从高不确定性到低不确定性区域扫描得到的特征与从低不确定性到高不确定性区域扫描得到的特征相一致,从而进一步提高分割精度。
作者的贡献可以概括如下:
作者提出了一种新型的基于信道不确定性的像素级选择性扫描方法,旨在解决传统基于位置的顺序扫描方法的局限性。作者引入可学习参数以平衡不同扫描方向上的特征重要性,并采用余弦一致性损失来对齐正向和反向扫描结果,从而提高特征一致性。
在三个医学影像数据集上进行的广泛实验表明,UD-Mamba能够有效地识别模糊区域,与现有的基于Mamba的方法相比,其分割结果更为可靠。
unset
unset
- 相关工作
unset
unset
2.1 医学图像分割
医学图像分割是指将医学图像分割成对应病变或器官的密集像素预测的过程,这一过程基于如CT[65,64,49,13]和MRI[23,57]等成像方法。其中,卷积神经网络(CNN)和Transformer是主导的框架。基于CNN的分割技术取得的重要进展是由UNet[39]引入的,它采用了一种对称的编码器-解码器架构,并配有 Shortcut 。这些 Shortcut 有效地整合了编码器中的局部特征和解码器中的语义信息,为后续许多改进奠定了基础。尽管CNN方法取得了成功,但它们受限于其局部感受野,这阻碍了对密集预测任务中至关重要的长距离依赖关系的捕捉。
图3:展示了UD-Mamba架构的示意图,该架构包括一个 Patch 嵌入层、一个带有不确定性驱动(UD)块的编码器-解码器以及一个分割头。每个UD块都具备不确定性驱动选择性扫描模型(UD-SSM)以处理输入。
受到视觉Transformer(ViTs)[1o]的启发,越来越多的研究兴趣集中在将Transformer应用于医学图像分割[16,63,30,21,47,61]中。TransUNet[5]作为开创性工作之一,引入了一种混合模型,在编码器中使用Transformer来建模全局上下文,同时保留了UNet的整体结构。Swin-UNet[4]进一步探索了基于Transformer的全框架分割任务。
虽然Transformer在建模长距离依赖关系方面表现出色,但其自注意力机制相对于输入大小引入了二次复杂度,这给可扩展性带来了挑战,尤其是在像素级任务如医学图像分割中。
2.2 状态空间模型在分割中的应用
状态空间模型(SSMs)最近已成为视觉任务的有力工具,Mamba [14,8] 通过以线性复杂度高效地建模全局上下文,展示了有希望的结果。基于Mamba的模型已在多种应用中展示了其多功能性。U-Mamba [35] 提出了一种结合卷积神经网络(CNNs)和SSMs的混合框架,有效地捕捉了局部和全局特征。Swin-UMamba [31] 将基于ImageNet的预训练整合到基于Mamba的UNet中,以增强医学图像分割的性能。P-Mamba [56] 将Perona-Malik扩散与Mamba结合,以改善儿科心脏病学中的超声心动图左心室分割。此外,王等 [48] 引入了LMa-UNet,这是一种基于Mamba的网络,采用大窗口设计以改善全局上下文建模。
尽管取得了这些进步,由于背景复杂和类别边界模糊,准确分割复杂医学图像仍然是一个挑战。此外,传统的扫描机制,它间歇性地扫描不同的语义区域,限制了模型在图像中持续捕捉完整上下文信息的能力。
2.3 分割中的不确定性估计
在最近关于医学图像分割不确定性估计的研究进展中,各种方法都强调了将不确定性纳入模型以增强可靠性和性能的重要性[34,52,37,46,62,12]。王等人[50]提出了一种领域自适应分割框架,通过引入不确定性 Aware 来优化伪标签,减少了噪声标签的影响。类似地,张等人[59]引入了一种不确定性引导的互一致性学习框架,利用估计的不确定性来选择半监督分割中的可靠预测。
李等人[27]采用了基于证据的深度学习(EDL),专注于基于区域的EDL框架,该框架利用德普斯特-沙弗理论来提供具有可量化不确定性的鲁棒分割结果。总体而言,这些方法突出了不确定性估计在提升医学图像分割模型的可信度和临床应用性方面的重要作用。
在本节中,作者将详细介绍所采用的研究方法和实验设计。首先,作者将阐述数据集的选取标准,随后解释算法的选择及其背后的理论依据。接下来,作者将描述实验的具体步骤,包括数据预处理、模型训练和性能评估。最后,作者将讨论实验结果的分析方法,以及如何确保实验结果的可靠性和有效性。
在本节中,作者首先对选择性SSM[14]的预备知识进行描述,这是Mamba框架中的一个核心概念。接着,作者提供了作者提出的UD-Mamba架构的全面概述,整体框架如图3所示。最后,作者详细阐述了UD-Mamba的关键组件,包括图4所示的Uncertainty-Driven Selective Scanning Model(UD-SSM)的操作流程以及由此派生的优化策略。
3.1 前言
曼巴块利用一种专门设计的选择性扫描机制(SSM),通过保持紧凑的内存表示来高效处理长距离依赖关系。受卡尔曼滤波的启发,该SSM作为一个线性时不变(LTI)系统[25],通过结构化的状态空间公式将输入参数(Δ,A,B,C)转换为。连续时间动力学由以下公式给出:
在应对传统LTI SSMs在捕捉上下文信息方面的局限性时,Mamba模块采用了一种输入依赖的选择机制(称为S6),该机制能够根据输入变化动态调整。其递归关系表达如下:
在本文中,
表示时间索引,
代表在时间
的输入序列,
表示捕捉时间上下文的隐藏状态,而
是输出。通过采用具有线性复杂度的关联扫描算法,Mamba模块有效地计算响应,允许在时间步长间传播上下文,同时最小化计算负担。
3.2 UD-Mamba:这是一种基于深度学习的文本分类模型,它结合了上下文依赖性(UD)和Mamba架构。该模型旨在提高文本分类任务的准确性和效率,通过利用上下文信息来增强分类器的性能。
UD-Mamba架构利用了一种简洁而强大的UNet框架[39],如图3所示。它包括三个关键组件:一个将输入图像转换为一系列 Patch 以供后续处理的 Patch 嵌入层,由不确定性驱动(UD)块组成的编码器-解码器结构,这些块能够捕捉并整合不同尺度上的局部和全局特征,以及一个基于解码特征的输出最终像素级分割结果的分割头。通过 Shortcut 增强了编码器-解码器配置,这些连接促进了多尺度特征表示的整合。这种架构选择增强了不同层次间的信息传播,最终提高了分割精度。
每个UD块通过一系列操作处理输入。首先,输入通过层归一化(LN)转换,然后通过一个线性层。接下来,应用深度卷积(DW-Conv),随后是SiLU激活函数。之后,数据由UD-SSM模块处理,该模块利用提出的顺序和跳过扫描策略,以线性复杂度捕捉长距离依赖关系。通过残差连接将UD-SSM输出与早期特征结合,之后通过最后的线性层进行进一步的精炼。
3.3 用户驱动的序列到序列模型(UD-SSM)
传统自底向上的语义分割模型在有效建模局部特征方面存在困难,这是因为它们对目标区域的扫描是间歇性的。作者提出了一种基于像素级不确定性的选择性扫描方法,即UD-SSM。这种方法与传统的像素顺序扫描机制不同,因为它利用像素级的不确定性来指导扫描顺序。如图4所示,作者的UD-SSM引入了以下关键组件:通道不确定性计算、扫描扩展操作、S6模块[14]和恢复处理。
图4:UD-SSM的详细描述。I. 介绍了基于UD-SSM信道不确定性的不确定性计算过程。II. 解释了扫描扩展操作,主要包括两种策略:顺序扫描(Scan)和跳过扫描(
),以及后续的重加权操作。III. 展示了S6块和恢复处理,同时介绍了一致性约束的计算过程。
通道不确定性计算。给定一个输入特征张量
,其中
和
分别表示通道数、高度和宽度。作者首先为所有通道中的每个空间位置计算一个不确定性图
,其公式为:
在本文中,作者采用标准差作为不确定性度量指标,这一选择在第四章第4.4.2节中呈现的结果得到了验证。具体来说,对于输入特征图
,作者计算每个空间位置
在所有通道
上的标准差。
表示所有通道在该空间位置的平均值。这一计算捕捉了通道间的像素级标准差,其中更高的不确定性通常对应关键区域,如物体边界或前景区域,而较低的不确定性则表明背景的一致性。通过关注像素级不确定性,作者可以更精确地识别出医学图像分割中的关键区域,这在识别病理区域或器官边界时往往至关重要。
不确定性地图U随后按降序排列,得到
,它将空间位置从高不确定性区域(前景和边界)排序到低不确定性区域(背景),这些背景区域被定义为:
这使得模型能够在后续操作中优先考虑具有更高复杂度或重要性的区域。
使用排序后的索引Idc,作者对原始特征图
进行重新排列,以创建
,其中对高不确定性区域进行强化处理,如图所示。
本次重组为高效扫描扫描扩展操作准备了特征图。作者在重新排列的特征图
上引入了两种主要的扫描策略,以针对不同不确定性水平的区域进行聚焦,如图
所示:
- 顺序扫描(Sequential scanning
):该方法按像素不确定性升序或降序处理空间位置,允许高不确定性区域,如前景目标和边界,以更高的密度进行建模。顺序扫描确保这些关键区域得到彻底捕捉,使模型能够聚焦于对准确分割至关重要的区域。
跳扫(Skipscanning):该方法在确定性谱上以固定间隔采样空间位置,促进了高不确定性区域与低不确定性区域之间的交互。
将这两种方法结合,UD-SSM利用四种独特的扫描序列全面捕捉局部和全局特征:从高到低不确定性的顺序扫描和跳过扫描(
和
),以及从低到高不确定性的顺序扫描和跳过扫描(
和
),具体定义为:
为了提升从低不确定性到高不确定性扫描的优势,并调整每个单独扫描序列的贡献,作者引入可学习的权重(
,
,
,
),以根据其捕获关键区域的有效性动态调整每条扫描路径的权重。
S6块与恢复处理。四个扫描序列的重新加权特征经过S6块[14]的处理。最后,一个恢复步骤巩固这些方向性特征,并将它们恢复到原始的空间配置,以保留精确分割所必需的位置精度,这被定义为:
UD-SSM的最终输出,记作yuD-ssM,是这些恢复特征的累加。
3.4 目标函数
为了提高解码阶段的从高到低不确定性扫描的性能并提升整体分割结果,作者在解码器末尾引入了一种余弦一致性约束。该约束应用于双向(从高到低和从低到高不确定性)进行的顺序扫描和跳过扫描。通过将低到高不确定性扫描的结果与高到低不确定性扫描的结果对齐,作者确保了不同扫描方向上特征表示的一致性。余弦一致性损失定义为:
表示正向和反向序列以及跳过扫描的平均余弦相似度。通过最大化这一相似度,作者旨在减少两种扫描方向之间的差异,从而加强最终分割输出的连贯性。
最终,整体损失函数将监督损失
与余弦一致性损失
相结合。
代表结合的交叉熵损失和Dice损失,而
是一个超参数,用于平衡这两个组成部分。
4.1 数据集
作者评估了所提出的UD-Mamba在三个医学图像数据集上的性能:
DigestPath数据集[7]包含了结肠镜肿瘤病变分割的全切片图像(WSIs)。作者将130个恶性WSIs随机分为三个子集:100个用于训练,10个用于验证,20个用于测试。在模型训练过程中,WSIs被进一步划分为256×256像素的图像块,从而得到29,412个图像块的训练集。作者的模型评估是在WSI Level 进行的。
ISIC 2018数据集[6]是一个公开的皮肤病变分割数据集,包含2,694张镜检图像以及对应的像素级标签。作者遵循VM-UNet[40]的方法将数据集划分为训练集和测试集。
表1:不同网络在DigestPath和ISIC 2018数据集上的性能比较。
ACDC数据集[3]包含了100名患者的心脏 cine MRI 扫描图像,用于三个心脏亚结构的分割:左心室(LV)、右心室(RV)和心肌(MYO)。作者将数据集分为70%用于训练,10%用于验证,20%用于测试。在训练前,所有切片都被调整到统一的256×256像素分辨率。
对于DigestPath和ISIC 2018数据集,作者采用了包括平均交并比(mIoU)、Dice相似系数(DsC)、准确率(Acc)、灵敏度(Sen)和特异性(Spe)在内的性能指标进行了详细评估。对于ACDC数据集,评估性能时使用了DsC、mIoU以及95% Hausdorff距离(HD95)。鉴于ACDC数据集中固定的解剖结构,加入HD95可以提供对边界精度的更稳健评估。
4.2 实施细节
所有实验均在配备NVIDIA RTX A6000 GPU的Ubuntu桌面电脑上,使用PyTorch框架进行。训练采用随机梯度下降(SGD)算法,并采用多步学习率策略,初始设置学习率为0.01。总训练轮数固定为300轮。对于UD-Mamba,编码器和解码器的每一层对应两个UD块。作者使用在ImageNet-1K[9]上预训练的权重来初始化编码器。λ的值为0.3。更多实现细节请参考补充材料。
4.3 与先进方法的比较
作者对作者的UD-Mamba与最先进的医学图像分割方法进行了比较,包括基于CNN的方法(UNet [39]、UNet++ [66]和Att-UNet [38])、基于Transformer的模型(TransUNet [5]、SwinUNet [4]和H2Former [17]),以及基于Mamba的模型(Mamba-UNet [51]和SwinUmamba [31])。
表1展示了在ISIC 2018和DigestPath数据集上的定量结果。作者的UD-Mamba在性能上显著优于基于CNN的方法。具体来说,UD-Mamba在ISIC 2018和DigestPath数据集上分别实现了相对于最佳CNN方法的
和
的提升。此外,mloU得分分别提升了
和
。与基于Transformer的模型,如TransUNet[5]相比,作者的方法在mIoU上展现出显著优势,ISIC 2018和DigestPath分别提升了
和
。此外,与代表性的基于Mamba的模型Mamba-UNet[51]相比,UD-Mamba在这两个数据集上分别将mIoU提升了
和
。
表2:ACDC数据集上不同网络的比较。在这里,“RV”、“MYO”和“*LV”分别代表用于分割不同心脏亚结构的深度学习分割模型(DsC)。
表3:ACDC数据集上不同方法测量模型复杂度和效率的比较。
对于ACDC数据集,表2展示了与其他方法的对比结果。与表现最佳的Mamba-UNet[51]相比,作者的方法在DSC和mIoU上分别提高了0.91%和1.45%,同时将HD_{95}指标降至1.31毫米。表3中的结果也显示,与基于Transformer和Mamba的网络相比,作者的模型显著减少了参数数量,并在所有网络中实现了最低的FLOPs。这突显了它在提供强大的分割性能的同时,优化计算效率的能力。三个数据集上的可视化结果展示在图5中。
4.4 削减研究
作者在消融研究部分对DigestPath数据集[7]进行了实验验证。
图5:在三个不同的数据集上,对UD-Mamba与其他多种方法的分割结果进行了视觉比较。
不确定性扫描策略 作者评估了在UD-SSM框架下,单个像素级不确定性驱动的扫描组件(
)以及优化策略(重新加权与
)的影响。表4总结了不同配置下的分割性能,突出了每个组件和策略的有效性。
研究结果表明,结合扫描策略比单独使用各个组件能带来更好的性能。例如,采用策略
实现了
的DSC(区分度)。然而,结合正向和反向双向扫描策略,特别是跳过扫描组件
,将DSC提升至
。类似地,顺序扫描产生了稳定的结果(
,DSC为
对比
的
DSC)。整合所有四个组件
实现了更高的DSC值,达到
,超越了单独的顺序扫描和跳过扫描,并在普通Mamba模型(第一行)的位置基础顺序扫描中表现更优。这证明了整合顺序扫描和跳过扫描的优势,能够同时捕捉局部和全局特征。
表4:比较不确定性扫描组件(y1、V2、
、
)和优化策略的效果。
表5:不同方法测量信道不确定性的比较。
此外,加入优化策略进一步提升了性能。仅通过重新加权,DsC从80.32%提升至80.72%,而引入余弦相似度约束(L_COS)后,DsC达到80.89%,同时mIoU为68.64%,准确率为94.98%。这种组合优化了扫描策略的贡献,并在不同方向上对齐了特征,从而实现了更高的分割精度。
总结来说,表4展示了将扫描组件与重新加权及一致性约束完全集成,最大化了模型处理不确定区域的能力,验证了基于不确定性的分割方法UDSSM的有效性。
为了评估测量信道不确定性的各种标准,作者进行了消融实验。这些标准包括平均绝对偏差(MAD)、标准差(STD)、方差、熵以及两个最高值之间的范围。如表5所示,使用STD来衡量数据分散提供了稳定的度量。这种稳定性使模型能够更可靠地识别真实的不确定区域,而不是被噪声或异常值误导。因此,采用STD计算不确定性的方法取得了最佳结果,达到了最高的DsC
,最高的mloU
和最高的ACC
。
不确定计算区域为评估在缺乏明确空间特征的场景中像素级不确定性驱动的扫描的有效性,作者进行了比较实验,重点关注用于不确定性计算的区域大小。作者不仅依赖于单个像素的不确定性,还将计算扩展到更大的区域,以保留一定程度的空间信息。这些区域被定义为不确定性块,尺寸为
。
表6:计算不确定性区域的不同方法的比较
作者的实验设计探讨了对于参数
的固定值和动态调整值。对于固定大小的区域,作者将
从1调整到
。在动态调整区域的情况下,采用了两种策略:(1)比例缩放,其中
,允许
与特征向量大小
成比例增加;(2)反比例缩放,其中
,导致
随着特征向量大小
的增加而减少。在此,
指的是进入UD-SSM前每个阶段的特征向量大小,
代表第一次进入UD-SSM时的特征向量大小,而
表示 Bottleneck 层的特征向量大小。在UD-Mamba中,
和
分别设置为64和8。计算每个区域的不确定性平均值后,这些值被用于对区域进行排名,以便后续扫描。如表6所示,像素级不确定性驱动的扫描始终优于动态和静态区域方法。这一结果突出了在确定像医学图像分割这类细粒度任务的不确定性时,像素级粒度的优势。与更广泛的基于区域的确定性方法相比,像素级不确定性专注于捕捉局部变化,为识别关键分割目标提供了一种更精确的方法。
分析重加权值如图6所示,展示了四个可学习参数
、
、
和
在整个训练过程中的演变情况。这四个参数均呈现下降趋势,其中
和
的下降幅度相较于
和
较为缓和。这一模式表明,在训练过程中,从高不确定性水平到低不确定性水平的扫描过程(对应于
和
)相较于从低不确定性水平到高不确定性水平的扫描过程(与
和
相关)贡献更为显著。这一观察结果间接证实了图2中提出的结论。
图6:分析四个可学习重加权参数的记录值。
unset
unset
5 结论
unset
unset
本研究通过深入分析,探讨了人工智能在各个领域的应用及其对人类社会的影响。研究发现,人工智能技术不仅极大地推动了生产力的发展,而且改变了人们的生活方式。
然而,随着人工智能技术的快速发展,作者也面临着诸多挑战,如数据安全、隐私保护、伦理道德等问题。因此,未来应加强对人工智能技术的监管,确保其在造福人类的同时,也能规避潜在的风险。本研究旨在为我国人工智能产业的发展提供有益的参考,并为相关政策制定提供理论依据。
unset
unset
参考
unset
unset
[1]. UD-Mamba: A pixel-level uncertainty-driven Mamba model for medical image segmentation .
点击上方卡片,关注
「AI视界引擎」
公众号