Medical SAM Adapter: Adapting Segment Anything Model for Medical Image Segmentation

Medical SAM Adapter:适配 Segment Anything 模型以进行医疗图像分割

本文核心:提出Med-SA适配器,即在SAM基础上添加SD-Trans 和 HyP-Adpt 的参数高效适应,使其不仅可以适应3D医学图像,还可以整合特定领域的医学知识,提高分割任务。

摘要

Segment Anything Model (SAM) 由于其在各种分割任务中令人印象深刻的能力和基于提示的界面,最近在图像分割领域越来越受欢迎。然而,最近的研究和个别实验表明,由于缺乏医学特定知识,SAM 在医学图像分割方面表现不佳。这就提出了一个问题,即如何增强 SAM 的医学图像分割能力。

在本文中,提出了医疗 SAM 适配器 (Med-SA) 而不是微调 SAM 模型,它使用一种轻便而有效的适应技术将特定领域的医学知识整合到分割模型中。在 Med-SA 中,提出了空深转置 (SD-Trans) 来使 2D SAM 适应 3D 医学图像,并提出了超提示适配器 (HyP-ADPT) 来实现提示条件适应。对各种图像模态的 17个医学图像分割任务进行了全面的评估实验,Med-SA 优于几种最先进的 (SOTA) 医学图像分割方法,同时仅更新了 2% 的参数。

1.介绍

背景

最近,Segmentation Anything Model (SAM)(Kirillov et al. 2023) 作为一种功能强大且用途广泛的视觉分割模型而受到广泛关注,它可以根据用户提示生成多样化且详细的分割掩码。尽管它的性能优于自然图像,但最近的许多研究还表明,它在医学图像分割方面的性能不佳。

若采用SAM等技术,使医学图像分割具有交互性,具有巨大的临床价值。交互式系统可以根据临床医生的指示确定感兴趣区域的优先级,为他们提供更加身临其境和个性化的体验。例如,在单个眼底图像中,通常存在重叠和错综复杂的结构,例如血管、视盘、视杯和黄斑。交互式分割可以极大地帮助临床医生有效地区分目标组织与这些复杂结构。考虑到获取大规模注释数据集的难度,采用像 SAM 这样的基础交互式模型进行临床利用变得至关重要。

SAM存在缺陷

SAM 在医学图像上的性能有限是由于缺乏医学特定知识,包括图像对比度低、组织边界模糊和病变区域微小等挑战。

解决这个问题的最先进 (SOTA)方法是专门针对医疗数据对普通 SAM 模型进行全面微调,这在计算和内存占用方面都非常昂贵。 此外,是否有必要进行全面微调是值得怀疑的,因为以前的研究表明,预先训练的视觉模型对医学图像具有很强的可转移性。

解决办法

在本文试图以最小的努力使训练有素的 SAM 适应医学图像分割。从技术上讲,选择使用一种称为 Adaption的参数高效微调 (PEFT) 技术来微调预训练的 SAM。Adaption 一直是自然语言处理 (NLP) 中一种流行且广泛使用的技术,用于为各种下游任务微调基本预训练模型。Adaption 的主要思想是将具有部分参数的 Adapter 模块插入到原始模型中,并且只更新少量额外的 Adapter 参数,同时保持大型预训练模型的冻结。

然而,将 Adaption 技术直接应用于医疗场景并不是那么简单。挑战:(1)图像模态。与自然图像不同,许多医学图像是 3D 的,例如 CT 和 MRI 扫描。目前尚不清楚如何使 2D SAM 模型适应 3D 医学图像分割。(2)虽然 Adaption 在 NLP 方面取得了成功,但将其应用于视觉模型的研究有限,尤其是像 SAM 这样的交互式视觉模型。在交互式视觉模型中,用户提供的视觉提示在最终预测中起着至关重要的作用。如何将 Adaption 与这些重要的视觉提示相结合仍未探索。

为了克服这些挑战,提出了一种称为 Medical SAM Adapter (MedSA) 的新型适应框架。在 Med-SA 中,引入了空深转置 (SD-Trans)技术来实现 2D 到 3D 的适应。在 SD-Trans 中,将输入嵌入的空间维度转置到深度维度,允许相同的自注意力块在给定不同输入的情况下处理不同的维度信息。

然后,提出了 Hyper-Prompting Adapter (HyP-Adpt),以实现基于提示的自适应。在这种方法中,使用视觉提示生成一系列权重,这些权重可以高效地应用于自适应嵌入,从而促进广泛而深入的提示-自适应交互。

实验评估

进行全面的评估实验,涵盖各种图像模式的 17 项 医学图像分割任务,包括 CT、MRI、超声图像、眼底 图像和皮肤镜图像。结果表明,Med-SA 的表现优于 SAM 和完全微调 SAM (MedSAM),并且两者之间存在很大的性能差距。Med-SA 还超越了几种专为医学图像分割量身定制的 SOTA 方法,例如 nnUNet、TransUNet、UNetr 和 Swin-UNetr。更重要的是,Med-SA 通过仅更新总 SAM 参数的 2% 的额外参数来实现这种卓越的性能。

• 本文提出了用于一般医学图像分割的 Adaption 方法,框架 Med-SA 是 SAM 架构的简单而强大的扩展,大大增强了其医疗应用能力,同时仅更新了总参数的2%。

• 建议使用SD-Trans来实现高维 (3D) 医疗数据的分割,解决医学图像模态带来的挑战。

• 提出 HyP-Adpt来促进提示条件适应,承认用户提供的提示在医学领域的重要性。

• 对17个具有各种图像模式的医学图像分割任务进行了广泛的实验,清楚地确立了 Med-SA 优于 SAM 和以前最先进的方法。在广泛使用的腹部多器官分割 BTCV 基准测试中,Med-SA的表现优于 Swin-UNetr 2.9%,比 vanilla SAM 高34.8%,比完全微调 SAM (MedSAM) 高 9.4%。

2.相关工作

交互式分割

强调交互式重要性

交互式分割有着悠久的历史,最初被研究人员视为一种优化技术。DIOS 的开创性工作通过集成深度学习并将正向和负向点击作为距离图,彻底改变了交互式分割。随后的研究则集中于通过预测多个潜在结果来解决不确定性问题,并使选择网络或用户能够在这些结果中进行选择成为可能。CDNet通过结合自我注意来生成更一致的预测,进一步增强了交互式细分。RITM和 AccuracyNet引入了使用以前的掩码作为输入,以提高预测的稳健性和准确性。最近,SAM 证明了交互式分割对零样本分割的重大影响,并强调了其在视觉基础模型中的潜在重要性。

然而尽管交互式医学图像分割在临床实践中发挥着关键作用,但人们对它的关注有限。例如,单个眼底图像可能需要分割多个目标,例如血管、视盘、视杯和黄斑,具体取决于不同的要求和用例。本文的Med-SA 为交互式医学图像分割提供了一个很好的起点,旨在激发该领域的未来研究。

参数高效微调

PEFT 已被证明是一种针对特定用途微调大型基本模型的有效策略。

(1)与完全微调相比,它使大多数参数保持冻结状态,并且学习的参数要少得多,通常不到总数的 5%。这可以通过更快的更新实现高效的学习。

(2)PEFT 方法比完全微调效果更好,因为它们避免了灾难性的遗忘,并且更好地推广到域外场景,尤其是在低数据情况下。在所有 PEFT 策略中,Adaption脱颖而出,成为一种有效的工具,可用于为 NLP 和计算机视觉中的下游任务微调大型基本视觉模型。

(3)最近的研究表明,Adaption 可以很容易地用于各种下游计算机视觉任务。因此, Adaption 是将 SAM 引入医疗领域的最合适技术。预计,这种简单、干净但功能强大的 Med-SA 将为基础医疗模型的开发带来更大的可能性。

3.方法

初步:SAM 架构

(1)概述了 SAM 架构。SAM 由三个主要组件组成:图像编码器、提示编码器和掩码解码器。

(2)图像编码器基于 MAE 预训练的标准 Vision Transformer (ViT)。具体来说,使用 ViT-H/16 变体,它采用 14×14 个窗口式注意力和 4 个等距全局注意力块,如 1 (a) 所示。图像编码器的输出是输入图像的 16× 缩减采样嵌入。

(3)提示编码器可以是稀疏(点、框) 或密集(掩码)。在本文中,只关注稀疏编码器,它将点和框表示为位置编码,并与每种提示类型的学习嵌入相加。

(4)掩码解码器是一个 Transformer 解码器模块,经过修改后包含动态掩码预测头。解码器使用双向交叉注意来学习提示和图像嵌入之间的交互。

之后,SAM 对图像嵌入进行上采样,MLP将输出标记映射到动态线性分类器,该分类器预测给定图像的目标掩码。

Med-SA 架构

目标是通过微调来增强 SAM 架构在医学图像分割任务中的医疗能力。并不完全调整所有参数,而是将预训练的SAM参数保持冻结,设计一个适配器模块并将其集成到指定的位置。(1)适配器用作瓶颈模型,依次由向下投影、ReLU 激活和向上投影组成,如图1(b) 所示。下投影使用简单的 MLP 层将给定的嵌入压缩到较低的维度,而上投影使用另一个 MLP 层将压缩的嵌入扩展回其原始维度。

(2)在 SAM 编码器中,为每个 ViT 块使用两个适配器。对于标准 ViT 块(如 1(a)所示),第一个适配器位于多头注意力之后和残差连接之前(如 1(b)所示)。第二个适配器被放置在多头注意力之后的 MLP 层的残差路径中。在第二个适配器之后,立即使用以下比例因子 s 缩放嵌入(Chen 等人,2022 年)。

(3)在 SAM 解码器中,为每个 ViT 块集成了三个适配器。第一个适配器用于集成提示嵌入,为了实现这一点,引入了一种称为 Hyper-Prompting Adapter (HyPAdpt)的新结构,该结构在解码器中的第二个适配器的部署方式与编码器中的部署方式完全相同,以适应 MLP 增强的嵌入。第三个适配器部署在图像嵌入提示交叉注意的残差连接之后。在适应后连接另一个残差连接和层归一化以输出最终结果。

SD-Trans 系列

由于 2D 图像与 MRI 和 CT 扫描等流行的 3D 模态之间存在尺寸差异,因此使 SAM 适应医学图像分割是一项挑战。在临床使用中,了解切片之间的相关性对于准确决策至关重要。虽然 SAM 可以应用于体积的每个切片以获得最终分割,但它没有考虑 3D 医学图像分割中固有的紧密体积相关性,正如以前的研究中所强调的那样。

为了解决这一限制,提出了SD Trans,其灵感来自图像到视频的适应。具体结构在1(c)中描述。如图所示,在每个块中,将注意力操作分为两个分支:空间分支和深度分支。

对于深度为 D 的给定 3D 样本,将 D × N × L 输入到空间分支的多头注意力中,其中 N 表示嵌入的数量,L 表示嵌入长度。在这里,D 对应于运算的数量,允许在 N × L 上应用交互作用,捕获空间相关性并将其抽象为嵌入。

在深度分支中,转置输入矩阵以获得 N ×D ×L,然后将其馈送到相同的多头注意力中。尽管采用了相同的注意力机制,但现在交互发生在 D × L 上,从而能够学习和抽象深度相关性。最后,将深度分支的结果转置回其原始形状,并将它们添加到空间分支的结果中,并合并深度信息。

HyP-ADPT系列

虽然在以前的一些作品中已经将适应技术应用于视觉模型,但适应的应用对交互式视觉模型的探索在很大程度上仍未得到探索。源任务和下游任务之间的交互行为可能会表现出显著差 异。因此,将视觉提示(在交互式模型中起关键作用)合并到适配器中变得至关重要。

在这方面,提出了一种名为 HyPAdpt 的解决方案,旨在实现快速条件适应。 HyP-Adpt受到 HyperNetworks的启发,它使用一个网络为另一个网络生成权重以进行知识调节。采用了 HyperNetworks 的高级概念,但对其进行了重新设计,以便在功能级别有效地应用它。

具体来说,只利用投影和重塑操作从提示嵌入中生成一系列权重图,然后将这些权重映射直接应用(矩阵乘积)到适配器嵌入。与生成整个网络相比,这种方法实现了广泛而深入的特征级交互,同时还显著减少了所需的参数数量。

具体来说,对适配器 e^down嵌入进行了超提示(hyper prompting)。同时,提示信息(点击位置、点击归因或边界框位置)被连接并降维为提示嵌入e^prompt,然后用 e^prompt生成一个权重序列,取其中的一个来说明,它可以表示为:

其中 Re 表示重塑,M 表示投影为e^prompt∈ R^(N×L)到e^prompt∈ R^[N×(L^in*L^out)] 的 MLP 层,其中 ∗ 是值乘法,L^in 的第一个权重将是 e^down的长度,L^out 最后一个权重将是输出的目标长度。之后,将 e^prompt从 1D 嵌入重塑为 2D 权重 w^prompt∈ R^(N×L^in×L^out),并将其应用于e^down,可以表示为:

其中 ⊗ 是矩阵乘积。沿长度维度对单元进行归一化,然后应用 ReLU 激活。为 hyper-prompting 设置了3个层,每个权重由单独的 MLP 层投射。HyPAdpt 有助于根据提示信息调整参数,并更灵活地适应不同的模态和下游任务。

训练策略

对于交互式分割,在模型训练过程中使用点击提示和边界框 (BBox) 提示。为了生成 BBox 提示,采用了与 SAM 相同的方法。但是,由于原始 SAM 论文提供的有关点击提示生成的细节有限,因此设计了自己的方法。 点击提示生成过程背后的基本概念包括使用正点击来指示前景区域,使用负点击来指示背景区域。

本文结合了随机和迭代点击采样策略,以使用这些提示来训练模型。最初,使用随机抽样进行快速初始化,随后,使用迭代抽样程序进行几次点击。这种迭代采样策略模拟了与真实用户的交互,因为每一次新的点击都会被放置在由网络根据之前的点击集生成的预测的错误区域中。参考 (Lin et al. 2020) 进行随机采样生成,参考 (Mahadevan,Voigtlaender,and Leibe 2018) 来模拟迭代采样过程。

4.实验

数据

五个不同的医学图像分割数据集进行实验,分为两种类型:一般分割效果和不同模态上的泛化能力。

第一种类型侧重于评估一般细分效果,本文选择了腹部多器官分割,因为它代表了医学图像分割中最重要的挑战之一。使用了 BTCV 数据集(Fang and Yan 2020),这是一个广泛使用且公开可用的基准,以 12个解剖结构为基准。

另外四个任务用于验证模型在不同模态上的泛化能力,包括视盘和视杯在眼底图像上的分割、脑肿瘤在脑部MRI图像上的分割、甲状腺结节在超声图像上的分割,以及黑色素瘤或痣在皮肤镜图像上的分割。对于眼底图像分割,在 REFUGE2(Fang et al. 2022)数据集上进行了实验。 对于脑肿瘤分割,在 BraTs 2021 数据集上进行了实验 (Baid et al. 2021)。对于甲状腺结节分割,使用了 TNMIX 基准,这是一个混合数据集,包含来自 TNSCUI 的 4554 张图像(马等人,2017 年)和来自 DDTI 的 637 张图像(Pedraza 等人,2015 年)。最 后,对于黑色素瘤或痣分割,在 ISIC 2019 数据集 (Milton 2019) 上进行了实验。

实现细节

在这项研究中,主要遵循官方 ViT-H SAM GitHub 存储库实施了Med-SA 管道。对于 2D 医学图像训练,遵循 SAM 的默认训练设置。对于 3D 医学图像训练,使用较小的批量,大小为 16。

对于 REFUGE2、 TNMIX 和 ISIC 数据集,训练了40个 epoch 的模型。对于 BTCV 和 BraTs 数据集,将训练扩展到 60个 epoch。与完全微调的训练相比,选择了更小的 epoch 数,因为观察到模型在设置中收敛得更快。

在交互式模型中,试验了四种不同的提示设置。这些包括:(1)一个随机的1个正点,表示为“1 点”,(2)三个正点,表示为“3 点”,(3) 目标重叠 50% 的边界框,表示为“BBox 0.5”,以及 (4) 目标重叠 75% 的边界框,表示为 “BBox 0.75.

所有实验都是使用 PyTorch 平台实现的,并在 4个NVIDIA A100 GPU 上进行了训练/测试。本文利用默认设置来重现比较方法。

与 SOTA 在腹部多器官分割上的比较

为了验证提出的 Med-SA 模型的一般性能,将其与多器官分割数据集 BTCV 上的 SOTA 分割方法进行了比较。定量结果在表1中,将 Med-SA 与公认的医学图像分割方法进行了比较,包括 nnUNet、TransUNet、UNetr、Swin-UNetr、EnsDiff和 SegDiff,以及基础版SAM 和完全微调SAM (MedSAM),使用 Dice 分数评估细分效果。

在表中,可以看到,当仅使用 1点提示时,Med SA 比 SAM 取得了显着改进。并且在 BTCV 数据集上,单点 Med-SA 在所有 12个器官上都实现了 SOTA 性能,在整体性能上超过了其他方法。随着提供更精细的提示,结果不断改善,最终 Dice 达到 89.8%,BBox达到0.75,比之前的 SOTA(Swin-UNetr)高出 2.9%。而且Swin-UNetr 由 138M可旋转参数组成,而本文只更新 13M参数。同时Med-SA在不断改善之后,甚至在所有提示变化下,都超越了完全微调的 MedSAM 模型。通过提出的 SD-Trans 和 HyPAdpt,仅更新了其总可转动参数的 2% (13M vs. 636M),这突出了所提出的技术的有效性。

在比较交互式分割模型(SAM、MedSAM、MedSA)中不同提示的性能时,注意到 3 分提示的性能略高于 1 分提示。BBox 0.75 的性能通常与 3 点提示相当或更好。但是, BBox 0.5 的性能低于标准,这表明准确的边界框注释对于实现性能改进的重要性。所有交互式模型(包括 SAM、MedSAM 和 Med-SA)在不同的提示中都表现出相似的行为,这表明它们对提示的响应是一致的。

考虑到SAM在1点中的性能,可以观察到无论使用何种提示,SAM的零镜头性能在目标医学图像分割任务中通常不如经过充分训练的模型(例如,MedSAM)。虽然这种比较似乎不公平,但因为 SAM 在自然图像数据集中展示了卓越的零镜头性能,所以将 SAM 的零镜头性能与经过充分训练的医学图像模型进行比较。这表明,与自然图像分割相比,SAM 的零镜头可转移性对医学图像的效果较差,这在以前的研究中也观察到。这一发现强调了需要特定的技术来使 SAM 适应医学图像分割。

3 点对 Med-SA 和 SAM 的性能进行了定性比较。从图中可以观察到 Med-SA 准确地分割在人眼难以识别的部位上。相反,SAM 在许多器官边界视觉上清晰的情况下失败。这进一步强调了对医学图像进行微调通用分割模型以实现最佳性能的必要性。

在多模态图像上与 SOTA 的比较

还将 Med-SA 与具有不同图像模态的四个医学图像分割任务中专门优化的分割方法进行了比较。结果如表2中,提出了 ResUnet和 BEAL用于视杯分割,TransBTS和 EnsemDiff用于脑肿瘤分割,MTSeg和 UltraUNet被提议用于甲状腺结节分割,FAT-Net和 BAT被提议用于黑色素瘤分割。SegDiff、 nnUNet、TransUNet、UNetr 和 Swin-UNetr 被提议用于一般医学图像分割。使用 Dice 分数、IoU 和 HD95 指标评估分割性能。

从表中可以看到,这些专门优化的方法通常在各自的领域内表现良好,但在应用于其他领域时性能会下降。例如,UltraUNet 实现了以前的甲状腺结节分割 SOTA,但与其他方法相比,在视盘分割方面的表现最差。

另一方面,通用方法通常在大多数模式下都能取得良好的效果,但在脑肿瘤分割和甲状腺结节分割等特定任务劣于专业方法。 将注意力转向交互式模型 SAM 和 MedSAM,观察到零样本 SAM 与医学图像中边界模糊的器官/组织作斗争,例如视盘/杯分割或甲状腺结节分割。就完全微调的 MedSAM 而言,由于其在 3D 图像处理方面的局限性,它在脑肿瘤分割方面存在不足。

Med-SA 在所有分割任务中都实现了SOTA 性能,展示了它能够推广到各种医疗分割任务和图像模态。在广泛使用的 BraTs 基准测试中,由于其对3D 图像的适应性,MedSA 在 Dice 分数和 HD95 指标方面分别比之前的 SOTA Swin-UNetr 高出 2.1% 和 1.86,而其可转动参数的利用率不到10%。

消融研究

进行了一项全面的消融研究,以验证拟议的 SD Trans 和 HyP-Adpt 的有效性。结果以表3 表示,其中基线(第一行)表示 SAM 和原始 Adaption 方法的简单组合。在基线设置中,3D 图像被视为 2D 图像序列并单独处理,而不涉及 Adaption 过程中的提示。如图所示,与普通的 SAM plus Adaption 设置相比,2D 到 3D 设计在 3D 数据基准(BTCV 和 BrainTumor)上都显著提高了性能。这一改进凸显了本文提出的 2D 到 3D 设计的有效性。

在 Prompt conditional Adaption 中,将 HyP-Adpt 与两个更简单的替代方案进行了比较:加法和连接,用于组合提示嵌入。虽然加法和串联也显示出一些有效性,但所得的改进仍然微乎其微。另一方面,使用提议的HyP-Adpt 会使得性能显著提高,进一步验证了本文提出的 HyP-Adpt 设计的有效性。

5.结论

在本文中,扩展了 SAM(一种强大的通用分割模型)来解决医学图像分割问题,并引入了 Med-SA。利用简单而有效的 SD-Trans 和 HyPAdpt 的参数高效适应,对原始 SAM 模型进行了实质性改进。该方法在跨越 5 种不同图像模态的 17 个医学图像分割任务中实现了 SOTA 性能。预计这项工作将成为推进基础医学图像分割的垫脚石,并激发新型微调技术的发展。

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值