SAM-Mamba: 用于广义零射线分割的 Mamba 引导 SAM 架构!

大模型向量数据库机器学习

点击下方卡片,关注

「AI视界引擎」

公众号

( 添加时备注:方向+学校/公司+昵称/姓名 )

picture.image

picture.image

结肠镜下息肉分割对于检测结直肠癌至关重要。然而,由于息肉结构、颜色和大小的变化以及与周围组织缺乏明确边界,这一过程颇具挑战性。基于卷积神经网络(CNN)的传统分割模型难以捕捉到详细的模式和全局上下文,从而限制了其性能。

基于视觉 Transformer (Vision Transformer, ViT)的模型在一定程度上解决了这些问题,但仍然难以捕捉局部上下文,并且在零样本泛化方面表现较弱。

为此,作者提出了一种Mamba引导的Segment Anything Model(SAM-Mamba ),以实现高效的息肉分割。作者的方法引入了Mamba-Prior模块,该模块在编码器中起到了桥梁作用,连接了SAM的一般预训练表示和与息肉相关的微小线索。

通过将息肉图像中的显著线索注入SAM图像编码器作为领域先验知识,同时捕捉各种尺度下的全局依赖关系,从而提高了分割结果的准确性。在五个基准数据集上的广泛实验表明,在定量和定性指标上,SAM-Mamba均优于传统的CNN、ViT和Adapter基模型。

此外,SAM-Mamba展示了卓越的对未见数据集的适应性,使其非常适合用于实时临床应用。

  1. Introduction

结肠直肠癌(CRC)被认为是发病率最高的消化道癌症,并位居全球第三位常见癌症之列。它通常在结肠内壁上以息肉的形式发展,如果未能及时发现和治疗,则会引发CRC。因此,早期检测和及时治疗对于预防CRC、降低死亡率至关重要。结肠镜检查是一种广泛采用的方法,用于检测和切除结肠内的息肉。然而,通过手动检查结肠镜图像来识别和分割息肉是一个耗时的过程,需要具备高技能和丰富经验的临床医生。此外,在手动检查过程中,小尺寸的息肉容易被忽视。因此,迫切需要开发自动化的息肉分割方法,以提高检测性能,并且可能在结肠镜检查过程中协助临床医生。

过去十年见证了深度学习架构在医疗图像分割方法发展中的显著进步,尤其是编码器-解码器卷积神经网络(CNN)[22, 37]的应用。早期的息肉分割方法基于流行的U-Net架构[22],并包含了一些辅助组件,如残差和密集连接以及注意力机制,包括U-Net [37] 和 ResUNet

[15]。这些方法缺乏处理关键边界信息的能力。为了解决这一问题,研究行人设计了多种分割方法,例如FCN [11]、PraNet [10]、CFA-Net [36] 和 MEGANet [3]。 另一方面,也有一些方法如MSNet [34]、

[27] 和

[33] 被提出以应对不同息肉之间的尺度多样性问题。尽管这些方法在分割息肉及其边界方面取得了巨大成功,但它们无法捕捉到对检测复杂且微小息肉至关重要的全局特征关系。此外,由于息肉与周围组织在颜色和纹理上的高度相似性、息肉之间显著的形状和尺寸变化以及不分明的边界,息肉分割仍然具有挑战性。此外,反复的下采样操作使得恢复微小息肉变得困难。进一步地,由于这些模型难以从各种图像获取设备捕获的不同图像特征中学习到一致性的特征,因此它们在未见过的数据上往往表现出较差的一般化能力。

Transformer架构在计算机视觉任务中,尤其是在医学图像分析领域,因其能够通过自注意力机制建模广泛的特征依赖关系而取得了卓越的成功[7, 24]。受到这些成功的启发,人们开始开发基于Transformer的分割方法,例如TransUNet[4]和UNETR[13]。然而,这些方法限制了它们捕捉局部上下文信息的能力。尽管最近有一些研究引入了卷积层来克服上述问题,但它们在息肉分类领域的应用仍未得到探索。最近,提出了一些基于Transformer的方法,如PVT-Cascade[21]和CTNet[31],并且这些方法在已见数据集上的表现令人印象深刻。然而,它们在未见数据集上的泛化性能有限,而这些模型的特征学习能力还有待进一步提升,以满足实时临床需求。

Segment Anything Model (SAM) 最近被引入作为图像分割的基础模型,并以其实现的令人印象深刻的零样本泛化性能而闻名,特别是在未见过的数据集上的表现[16]。然而,当直接应用于医学图像分割时,包括息肉分割,SAM 的性能较低,原因是没有特定领域的知识[35]。

与此同时,对 SAM 进行医学数据微调会导致较高的计算成本和内存需求[30]。最近出现的 Adapter 模块已被用来克服上述限制,并以较少的努力适应目标任务[5,30]。更近的是,利用状态空间模型 (SSM) 的 Mamba 已因其在序列数据中有效建模长距离依赖关系方面展现出卓越的计算速度和内存效率而备受关注[38]。一些基于 Mamba 的分割模型包括 U-Mamba [18] 和 SegMamba [32]。这些模型的泛化性能尚未得到充分探索。受这些技术近期取得的成功启发,提出了一种基于 Mamba 的先验结合 SAM,用于实现对已见过和未见过的数据集的有效息肉分割。具体而言,设计了一个包含多尺度空间分解 (MSD) 和 Mamba 块的 MambaPrior 模块,以全面捕获不同尺度的全局上下文,从而有助于对不同尺度的息肉及其复杂边界进行分割。广泛的数据集评估表明,所提出的 SAM-Mamba 模型的有效性以及其在零样本泛化性能方面优于当前最先进的 CNN、ViT 和基于 Adapter 的模型的能力。

作者的贡献如下:

  1. 作者提出了一个基于Mamba的先验知识在SAM(SAMMamba)中,以增强泛化零样本息肉分割性能。该方法通过有效捕捉息肉图像的多尺度和全局上下文线索,利用传统SAM的学习能力。据作者所知,这是首次尝试在SAM中探索Adapter和Mamba的有效性用于息肉分割。

  2. 作者提出一个Mamba-Prior模块,该模块包含一个MSD块,随后是Mamba块,用于将学习到的特征注入到SAM编码器中。前者块有助于在不同尺度上学习空间特征,而后者块则全面捕捉特征图中的更广泛上下文,丰富了学习到的特征表示,从而有效地分割复杂息肉及其边界。

  3. 作者将在五个不同的基准数据集上评估SAM-Mamba,并将其效果与最新的息肉分割方法进行对比。此外,作者还进行了消融实验以探讨所提出的模块中各个组件的重要性。实验结果表明,SAM-Mamba具备有效的零样本泛化能力。

  4. Related Work


在结肠息肉分割领域的初期努力主要采用了基于编码器-解码器结构的最流行卷积神经网络架构U-Net。例如,周等人[37]开发了UNet++,通过在编码器和解码器子网络之间引入一系列嵌套密集连接来提高结肠息肉的分割精度。贾等人[15]提出了改进的ResUNet模型ResUNet++,通过增加注意力模块和池化层提高了结肠息肉的精确分割。为了建立息肉区域边界关系,方等人[11]提出了一种选择性特征聚合(SFA)网络,该网络采用卷积基础共享编码器、双解码器和边界敏感损失函数。萍等人[10]提出了PraNet,该网络采用并行部分解码器和反向注意力模块以精炼分割边界并增强息肉分割精度。魏等人[29]基于颜色交换操作、浅层注意力模块和概率修正策略设计了SANet,以提高息肉分割精度,解决颜色不一致、小息肉退化和像素不平衡等问题。在[34]中,设计了多尺度减法网络(MSNet),通过在金字塔级联的多个减法单元上进行连接,并引入损失函数以实现从细节到结构的监督,从而解决了息肉尺度多样性的问题。其改进版本

在[33]中被提出,利用了内层和跨层的多尺度减法单元以实现高效的息肉分割。马等人[20]提出使用基于EfficientNetV2主干的UNet和新的位置嵌入特征块,以增强特征传递并提高息肉分割精度和泛化能力。周等人[36]引入了一种跨越层特征聚合网络(CFA-Net),该网络结合了边界感知特征和跨越层特征融合,以解决结肠息肉中的尺度变化和边界模糊问题。最近,在[3]中设计了一种多尺度边缘引导注意力网络MEGANet,该网络在编码器和解码器之间集成边缘引导注意力模块以保留边缘信息,从而改善弱边界息肉的分割效果。

Transformers已经在医学图像分割领域展示了其重要性,因为它们能够捕捉全局上下文细节,因此最近在息肉分割任务中得到了应用。Trinh等人[27]设计了

Unet,该模型采用混合CNN-Transformer编码器,并结合多尺度上采样模块以综合多层次解码器信息,增强息肉的局部和全局特征表示。为了确保息肉分割的稳健特征学习能力,Rahman等人[21]设计了PVT-Cascade,该模型采用分层级联注意力解码器,结合了多尺度特征与注意力门控机制和卷积注意力模块,从而增强了全局和局部上下文。在最近的研究贡献中,设计了一种对比变换网络(CTNet)[31],包含对比 Transformer Backbone 网、自我多尺度交互模块和收集信息模块,旨在获得稳定的息肉分割结果和更好的泛化性能。尽管取得了显著进展,但这些方法仍难以应对息肉及其伪装特性带来的挑战。此外,还存在大量提高泛化能力的空间。

  1. Methodology

在本节中,作者将逐步介绍作者的SAM-Mamba框架用于息肉分割。由于作者的SAM-Mamba大大扩展了SAM主干网络的概念,作者首先详细说明SAM的基本原理,SAM作为一种通用的目标分割模型已经崭露头角。然后,作者讨论SAM-Mamba模块,通过一个新颖的Mamba-Prior模块增强了SAM在息肉分割中的适应性,以应对该任务的关键挑战。

3.1. Preliminaries of SAM

SAM架构由三个主要组件组成:图像编码器、 Prompt 编码器和 Mask 解码器。图像编码器:图像编码器基于已经使用掩蔽自动编码器(Masked Autoencoder, MAE)预训练的标准ViT架构。具体来说,它采用

变体,该变体包括沿14×14窗口化的注意力机制,并且包含四个等间隔的全局注意力模块。编码器输出输入图像的16倍下采样的嵌入表示。 Prompt 编码器: Prompt 编码器可以处理稀疏 Prompt (例如,点、框)或密集 Prompt (例如, Mask )。在此,稀疏编码器使用位置编码与针对每种 Prompt 类型学习到的特定嵌入相结合来编码点和框。 Mask 解码器: Mask 解码器是一个修改过的Transformer解码器块,其中包括一个动态 Mask 预测头,该预测头采用双向交叉注意机制,促进了 Prompt 和图像嵌入之间交互的学习。在处理后,SAM对图像嵌入进行上采样,输出 Token 通过MLP映射到一个动态线性分类器以预测给定图像的目标 Mask 。得益于SAM的这三个组件,它们帮助在通用目标分割任务中取得了令人鼓舞的结果,但在将SAM适应肠息肉分割任务时仍然存在一些挑战。

首先,迁移学习能力较弱:通过采用传统的全微调策略,SAM 可能会导致过拟合、遗忘,甚至特征退化,尤其是在下游数据集不够大且多样化时,对于大规模预训练模型尤为明显。其次,计算成本增加:SAM 处理的输入图像分辨率比经典 ViT 提高了 4 倍,这增加了 patch 的数量,从而使全微调 SAM 的计算成本增加了 4 倍。此外,SAM 对点、框或文本 Prompt 的依赖性:SAM 需要一个 Prompt 或一组 Prompt 来生成分割 Mask ;然而,在大多数息肉分割模型中,输入仅仅是息肉图像,输出是分割 Mask 。因此,直接将 SAM 原样应用于息肉分割任务仍然可行性较低。

3.2.SAM-Mamba

在图1a所示的SAM-Mamba中,作者的目标是通过有效的轻量化微调来增强SAM架构在息肉分割任务中的适应性。与传统的全量微调方法优化所有参数不同,作者保持预训练的SAM参数冻结,并采用基于Adapter的微调方法。Adapter作为 Bottleneck 模型,包含低层参数以适应息肉图像领域。然而,与一般图像不同的是,息肉图像分割任务对几个关键属性的注意力相对简单,即在颜色、形状和难以区分的边界方面区分息肉区域和邻近像素。因此,一个具有有限可学习参数的小型化Adapter可能不足以学习息肉分割任务的关键鉴别特征表示。

picture.image

3.2.1 Mamba-Prior Module

为了弥合SAM的一般先验表示与结肠息肉相关琐碎线索之间的差距,在图1b中提出了一个Mamba-Prior模块。为了将结肠息肉图像中的显著线索注入SAM图像编码器作为领域先验,它采用了三种策略:

(i) 多尺度空间分解:在各种空间尺度上编码结肠息肉区域的低级特征(如大小、形状、边界),从而帮助分析从细粒度到粗粒度的空间语义;

(ii) 通道显著性和背景积累:在每一尺度的空间分解中,分别提取显著性和背景线索,并沿通道深度进行累积;

(iii) Mamba通道交互:利用Mamba层捕捉多尺度空间分解中显著和背景通道中的长程内部变化。这些功能策略与 Adapter 一起学习,以向SAM编码器注入结肠息肉领域的关键线索。每种策略的功能细节如下:

多尺度空间分解(MSD):该方法通过并行卷积层处理具有不同感受野(

)的输入息肉图像

,其中

。任意卷积层的结果图可以用

表示,存在

。值得注意的是,结果图

被填充以匹配原始图像大小,但具有不同的滤波器数量

。为了构建一个空间多尺度特征金字塔

,作者将结果图沿通道深度进行堆叠,并确保粗粒度图

位于金字塔顶部,细粒度图

保留于底部。这样,模型可以有序地分析息肉区域。

信道显著性和上下文积累:这部分专注于从

中提取和积累显著且更广泛的上下文特征。对于显著性和上下文的提取,作者应用标准的最大池化和全局平均池化操作,分别生成

Mamba 信道交互:Mamba 展示了在保持线性计算复杂度的同时处理长序列数据的能力,因此作者利用其优势来捕捉显著性和上下文图即

内部像素之间的交互关系。

首先,Mamba 可以独立地编码

的多尺度信道分布之间的依赖关系。接着,Mamba 使用图1 所展示的门控机制进一步细化特征表示。对于给定的输入

,经过 Mamba 后可获得如下特征图:

其中,

表示一个线性层,

表示 SiLU 激活函数,

表示矩阵乘法。

中密集且有序的通道内交互编码使得能够理解多尺度特征图

中息肉区域的关键线索。然而,通过 Mamba 门控机制,可能会遗忘细粒度和稀疏的线索。

因此,作者通过 Shortcut 将息肉图像的原始多尺度特征图与 Mamba 的结果

相乘,并随后在通道深度方向上进行拼接,以获得域先验嵌入特征图

,其形式表示为:

3.2.2 Adapter

作者框架中集成的Adapter模块借鉴了[5]的工作,主要基于两个顺序的交叉注意力机制,如图la所示。第一个和第二个交叉注意力分别用于增强多尺度特征,并将Mamba-Prior注入到ViT块中。这种注入确保了ViT块的特征分布不会发生剧烈变化,从而更好地利用预训练的ViT。

3.2.3 SAM Decoder

对于 Mask 解码器,作者采用了Kirillov等人的[16]提出的一种架构,该架构利用边界框、 Mask 、点或文本等 Prompt 来进一步丰富由SAM编码器提取的特征以进行分割。

然而,这使得模型在没有这些 Prompt 的情况下无法使用。为此,作者的SAM-Mamba 首先通过训练Mamba-Prior和Adapter从

中提取一个伪 Mask ,并使用

优化模型。随后,从

中获得的伪 Mask 作为 Prompt 被输入到解码器中,通过进一步训练Mamba-Prior、Adapter和SAM-Decoder并使用

进行监督,以此来细化 Mask 。

3.2.4 Objective Function

作者提出的SAM-Mamba及其功能模块,包括Mamba-prior、Adapter、

和SAM解码器,均采用联合可训练的方式,并通过一个由Dice损失和加权二元交叉熵(BCE)损失组成的损失函数进行定义:

。Dice损失通过增加困难像素的权重以突出它们的重要性,而BCE损失则更注重困难像素而不是对所有像素等同对待。由于SAM解码器高度依赖于由

生成的伪 Mask ,因此作者采用了两阶段训练制度。在第一阶段中,图像编码器内的Adapter在深层监督下使用次级输出

进行训练,图1中的几行虚线说明了这一点。第一阶段的损失函数如下所示:

其中,

是来自图像编码器的上采样侧输出,并与 Ground Truth

监督。第二阶段:在随后的阶段中,整个模型,即 Mask 解码器和图像编码器的 Adapter ,将在完全监督下进行训练,如图1 中实线所示。该阶段的总损失计算公式为:

这里,

分别是 Mask 解码器的输出和来自图像编码器的上采样输出,并与 GT 进行比较。

  1. Experiments

4.1.DatasetsandEvaluationMetrics

数据集:为了评估SAM-Mamba的表现,作者在五个具有挑战性的息肉分割数据集上进行了实验:ETIS [25]、CVC-ColonDB [26]、EndoScene [28]、Kvasir-SEG [14] 和 CVC-ClinicDB [2]。为了确保公平比较并展示其零样本泛化能力,作者采用了与PraNet [10] 中相同的经验设置。在指定条件下,选择了1450张图像作为训练集,其中900张图像来自Kvasir-SEG,550张图像来自CVC-ClinicDB数据集。剩下来自Kvasir-SEG的100张图像和来自CVC-ClinicDB的62张图像用于测试集。此外,作者还从CVC-ColonDB中采用了380张图像、从ETIS中采用了196张图像以及从CVC-300(EndoScene的测试集)中采用了60张图像进行测试。这种配置由于不同数据集间分辨率的不同以及图像采集设备的差异而带来了各种挑战。

评价指标:为了进行彻底的评估和比较,作者采用了六种不同的指标:Dice系数、交并比(IoU)、S-测量值

[8]、加权F-测量值

[19]、最大E-测量值

[9] 和平均绝对误差(MAE),这符合现有的前沿方法。值得一提的是,在作者的研究中,Dice系数和交并比的均值分别表示为mDice和mIoU。

4.2. Implementation Details

SAM-Mamba模型使用PyTorch实现,并利用NVIDIA A100 GPU进行加速。所有输入图像大小均调整为

像素。采用多尺度训练策略,尺度设置为

以增强数据。使用Adam优化器,学习率为

来训练模型。模型最多训练200个epoch。

4.3. Quantitative Comparison

为了验证作者提出的SAM-Mamba的鲁棒性,作者对其进行了一系列广泛的对比实验,与包括UNet [22]和UNet++ [37]、SFA [11]、PraNet [10]、SANet [29]、MSNet [34]、Polyp-PVT [6]、PEFNet [20]、

UNet [27]、PVT-Cascade [21]、

SNet [33]、CFANet [36]、CTNet [31]和MEGANet [3]在内的14种当前最先进的分割方法进行了比较。作者在两个基准数据集Kvasir-SEG和CVC-ClinicDB上验证了SAM-Mamba的学习能力。如表1详细所示,SAM-Mamba严格地与最先进的CNN和基于ViT的分割模型进行了对比。在具有挑战性的Kvasir-SEG数据集上,SAM-Mamba超过了竞争性的CTNet [31],在mDice指标中取得了令人印象深刻的92.4%。同样,在CVC-ClinicDB数据集上,它展示了出色的且一致的表现,在关键的评价指标如mIoU、

中都超越了竞争对手,这表明其在息肉分割中的稳健性和优越性。然而,对于

指标,96.1%和95.5%的评分则显示出边缘检测方面还有改进的空间。尽管其MAE值分别为2.5和0.6具有竞争力,但仍然暗示有进一步优化的可能性。总体而言,SAM-Mamba在大多数指标上表现出色,巩固了其领先地位,并为进一步提高边缘敏感性和错误最小化留下了空间。该模型的学习轨迹在图4中可视化地展示出来,通过中间特征图、编码器输出、解码器热力图以及精细化的分割 Mask 与输入和 ground truth 的综合对比,展现了其逐步学习的过程。值得注意的是,PEFNet的结果来自于

,而其他 Baseline 则直接来自它们的原始作品。

picture.image

picture.image

零样本泛化能力验证:模型评价的一个关键方面在于其在零样本场景下有效泛化到未见过的数据的能力,这是实际医疗图像分割应用中的一个关键要求。为了评估这一点,作者对SAM-Mamba进行了基准测试,使用了三个数据集:CVC-300、CVC-ColonDB 和 ETIS,并特别测试了其零样本泛化能力。

如表2和表3所示,SAM-Mamba取得了显著的性能,在CVC-ColonDB和ETIS数据集上分别以

指标高出SOTA模型4%和3.8%,展现了卓越的性能。在其他指标上也观察到可比的性能增益,进一步验证了所提出模型在零样本泛化方面的鲁棒性和适应性。

picture.image

picture.image

4.4. Qualitative Comparison

作者的模型在已见数据集和未见数据集上均进行了定性评估。如图2所示,MSD模块使作者的模型能够在CVC-ClinicDB数据集上准确识别次级息肉。类似地,在Kvasir-SEG数据集上,作者的模型实现了显著低的误检率,这归功于Mamba层在探索全面全局上下文信息方面的有效性。对于未见数据集,例如CVC-300,作者的模型始终维持较低的误检率,与在已见数据集上的表现一致。

这种一致性在图3中通过CVCColonDB和ETIS数据集的热力图可视化进一步得到了体现,展示了作者的模型具有类似的稳健性。图4展示了作者的SAM-Mamba模型编码器和解码器的学习进展,一系列热力图可视化序列地呈现了其稳健的学习能力。

picture.image

picture.image

总结来说,MSD模块增强了模型对不同大小息肉的检测能力,而MambaPrior模块有效减少了假阳性,从而整体提高了分割准确性。

4.5. Ablation Study and Discussion

Mamba-Prior 成分的影响:表4 中呈现的消融研究结果突显了 SAM-Mamba 模型在不同数据集上各组成部分的影响。结果表明,在 SAM 模型中加入 Mamba 组件可以显著提升性能。具体而言,包含 MSD 和 Mamba Adapter 的模型在所有数据集上均表现出更优异的结果,KvasirSEG 的 mDice 得分为 92.4%,CVC-ClinicDB 的得分为 94.2%,CVC-ColonDB 的得分为 85.3%,CVC-300 的得分为 92.0%,ETIS 的得分为 84.8%,同时仅需额外增加 9.5% 的参数量。相比之下,不包含 Mamba 组件的配置则获得较低的性能。

这表明 Mamba 和 Adapter 的加入显著增强了模型捕捉更详细和重要特征的能力,从而提升了分割结果的质量。Mamba 所获得的性能增益进一步验证了其在已见过的数据集上的有效性及其对未见过的数据集的良好鲁棒性。

picture.image

不同核大小的影响:在本实验中,作者验证了MSD不同组件的有效性。表5显示,“多尺度+Mamba”配置在CVC-ColonDB、CVC-300和ETIS等未见过的数据集上表现始终优于单一尺度的变体。

对于Kvasir-SEG (90.6) 和 CVC-ClinicDB (92.5) 这些已见过的数据集,单一尺度配置表现出色,但在未见过的数据集上表现不佳,得分分别为80.9和81.3,这很可能是因为它们固定了核大小。多尺度方法改善了不同数据尺度下的泛化能力,在所有数据集上均取得最佳效果,并提升了分割性能和鲁棒性。

picture.image

  1. Conclusion

本文提出了一种名为SAM-Mamba的新方法,用于泛化零样本息肉分割。该方法的主要创新之处在于整合了Mamba-Prior模块,该模块结合了多尺度空间分解和内尺度特征的依赖建模,以提取形状和大小各异的息肉。

得益于Mamba在建成长距离特征依赖方面的能力增强,SAM-Mamba能够在已见和未见数据集上有效定位复杂的息肉及其边界。

在五个基准数据集上的定量和定性结果表明,SAM-Mamba在特征学习和泛化能力方面优于传统的CNN、ViT和基于Adapter的模型。

参考

[0]. SAM-Mamba: Mamba Guided SAM Architecture for Generalized Zero-Shot PolypSegmentation .

点击上方卡片,关注

「AI视界引擎」

公众号

0
0
0
0
关于作者
关于作者

文章

0

获赞

0

收藏

0

相关资源
CV 技术在视频创作中的应用
本次演讲将介绍在拍摄、编辑等场景,我们如何利用 AI 技术赋能创作者;以及基于这些场景,字节跳动积累的领先技术能力。
相关产品
评论
未登录
看完啦,登录分享一下感受吧~
暂无评论