顶刊 MIA'24 | MA-SAM: 用于医学图像分割的模态无关SAM适应框架

科技   2024-10-18 10:06   中国香港  

点击下方“ReadingPapers”卡片,每天获取顶刊论文解读

论文信息

题目:MA-SAM: Modality-agnostic SAM adaptation for 3D medical image segmentation

MA-SAM: 用于3D医学图像分割的模态无关SAM适应框架

作者:Cheng Chen, Juzheng Miao, Dufan Wu, Aoxiao Zhong, Zhiling Yan, Sekeun Kim, Jiang Hu, Zhengliang Liu, Lichao Sun, Xiang Li, Tianming Liu, Pheng-Ann Heng, Quanzheng Li

源码链接:https://github.com/cchen-cc/MA-SAM

关键字:

  • 基础模型
  • 分割任何东西
  • 医学图像分割

摘要

“Segment Anything Model”(SAM)是一个用于通用图像分割的基础模型,已经在众多自然图像分割任务中展示了令人印象深刻的零样本性能。然而,当应用于医学图像时,SAM的性能显著下降,主要是由于自然和医学图像领域之间的巨大差异。为了有效适应医学图像,作者提出在微调过程中需要结合关键的第三维度信息,即体积或时间知识。同时,作者的目标是充分利用SAM在其原始2D骨干网络中的预训练权重。在本文中,作者介绍了一个名为MA-SAM的模态无关SAM适应框架,适用于各种体积和视频医学数据。作者的方法基于参数高效的微调策略,只更新一小部分权重增量,同时保留大部分SAM的预训练权重。通过在图像编码器的变换器块中注入一系列3D适配器,作者的方法使得预训练的2D骨干网络能够从输入数据中提取第三维度信息。作者在五个医学图像分割任务上,使用11个公共数据集,涵盖了CT、MRI和手术视频数据,全面评估了我们的方法。值得注意的是,不使用任何提示,我们的方法在Dice指标上一致性地超越了各种最先进的3D方法,分别在CT多器官分割、MRI前列腺分割和手术场景分割上超过了nnU-Net 0.9%、2.6%和9.9%。我们的模型还展示了强大的泛化能力,在挑战性的肿瘤分割任务中表现优异,当使用提示时,Dice得分超过了nnU-Net 38.7%。

1. 引言

基础模型(Bommasani等人,2021)的兴起,这些模型在庞大且多样化的数据集上进行训练,催化了智能模型开发的范式转变。由于它们显著的泛化能力和少样本学习能力,适应一个预训练的大型模型以适应多种下游任务,而不是传统的从头开始构建和训练不同的特定任务模型的方法,变得越来越有吸引力。“分割任何东西模型”(SAM)(Kirillov等人,2023)是最近开发的视觉基础模型,用于可提示的图像分割,在1亿个掩码上预训练了1100万个自然图像。由于其大规模训练数据和通用模型架构,SAM在各种自然图像任务中展示了令人印象深刻的零样本性能。鉴于这些优点,一个自然的问题出现了:是否可以将SAM直接扩展到解决关键的医学图像分割任务,这是一个领域,一直在努力获得训练深度模型所需的高质量图像和标签?然而,由于自然图像和医学图像之间的显著领域差距,最新的评估SAM在医学图像上的研究表明,无论是否使用提示,SAM的零样本能力都不足以直接部署在医学图像上(Huang等人,2023b;He等人,2023;Wald等人,2023)。在这些评估中,与特定于领域的医学图像训练的模型相比,SAM的性能较差,甚至在一些具有挑战性的任务中完全失败。
基于这些评估,很明显微调是将SAM应用于医学图像的一个重要步骤。但是,我们为什么倾向于适应SAM进行医学图像任务呢?这可以归因于与SAM相关的三个潜在优势。首先,SAM的训练数据集包含了大量图像。在医学应用的背景下获取类似的大规模训练数据集极其具有挑战性。尽管SAM的训练数据只包括自然图像,但它不限制于任何特定的医学成像方式。如果SAM的微调对某种医学成像有效,那么很有可能相同的方法也适用于其他方式。其次,微调后,作为预训练的大型模型,SAM可能具有强大的泛化潜力,这对于在关键医疗应用中有效部署智能模型非常重要。第三,SAM的提示设计为处理困难任务(如肿瘤分割)提供了便捷的半自动分割解决方案。在这些方面,SAM提供了一个通用的基础模型,有潜力适应多种医学成像方式,为全自动和半自动分割提供了良好的泛化能力。
适应SAM进行医学应用的努力正在迅速增长,这些方法大多依赖于SAM的提示设计(Cheng等人,2023b;Wu等人,2023a;Deng等人,2023b;Dai等人,2023)。然而,为医学数据中的每个对象提供合适的提示并非易事。例如,考虑一个包含多个器官的腹部CT体积,即使为每个器官的每个切片提供一个基本的点提示,也需要大量的努力。此外,在分割对象呈现相对规则的形状和位置的情况下,自动分割方法已经获得了令人鼓舞的结果,从而避免了半自动分割中提示的需求。在自动医学图像分割的背景下,一些最近的研究采用了参数高效迁移学习(PETL)技术,如LoRA(Hu等人,2021)或Adapters(Houlsby等人,2019),在自动分割中展示了有希望的性能(Zhang和Liu,2023;Wang等人,2023a)。然而,这些方法专注于纯2D适应,忽视了医学图像中固有的有价值的第三维度信息。这包括医学体积数据中的关键3D空间信息和医学视频数据中的时间信息。
在本文中,我们提出了一种用于医学图像分割的模态无关SAM适应方法,称为MA-SAM,它有效地且有效地捕获医学数据中的体积或时间信息。对于图像编码器的微调,我们采用了基于张量化-分解的PETL技术,称为FacT(Jie和Deng,2023),以增强调整效率。这种微调方法在很大程度上保留了预训练的权重,并且只更新轻量级的权重增量,确保了对象分割所需的一般知识,并减少了需要调整的参数数量。为了弥合2D自然图像和体积或视频医学数据之间的差距,我们进一步将一组3D适配器整合到每个图像编码器的变换器块中,以提取宝贵的第三维度信息。对于轻量级掩码解码器的适应,我们采用了全微调,并修改了其原始架构,通过一个简单而有效的渐进式上采样机制来恢复预测分辨率。我们展示了我们的SAM适应框架在多种医学成像方式上处理各种分割任务的有效性。通过与多种SOTA方法比较,我们的自动分割展示了优越的性能和显著的泛化能力。我们的主要贡献如下:
  • 我们提出了一种参数高效的微调方法,将SAM适应于体积和视频医学数据。通过创新一系列3D适配器,我们的方法有效地将医学图像中的基本第三维度信息纳入了SAM的2D网络骨干。
  • 我们证明了我们的SAM适应可以应用于各种医学成像方式,包括CT、MRI和手术视频数据,用于解剖、手术场景和肿瘤分割。不使用任何提示,我们的自动分割一致性地超越了竞争的SOTA方法。
  • 我们验证了在医学图像上微调后获得的模型呈现出卓越的泛化能力,甚至比SOTA领域泛化方法表现更好。
  • 我们展示了通过进一步利用提示,我们的方法在挑战性的肿瘤分割任务中取得了令人印象深刻的结果,Dice得分超过了nnU-Net 38.7%。

3. 方法

在本节中,我们首先简要介绍SAM架构的概述,然后分别介绍图像编码器的参数高效微调、体积或时间信息的整合以及掩码解码器的适应。我们框架的有效SAM适应概述如图1所示。
图像编码器的参数高效微调
为了有效提取图像特征,SAM的图像编码器包含了网络参数的相当大一部分。微调所有这些权重在计算上是密集的。先前的研究表明,PETL技术可以实现与完整微调相似的适应性能,但更新的网络参数数量显著减少(Hu等人,2021;Pan等人,2022)。因此,我们采用了FacT(Jie和Deng,2023),这是一种SOTA PETL技术,它在引入较少数量的可训练参数的同时,可以获得与其他PETL方法相当或更优越的性能。基于变换器模型在秩上往往是冗余的共同观察,FacT假设用于微调的密集权重增量矩阵可以由一组具有跨层权重共享的低秩因子近似。按照FacT中的张量分解,我们将每层的权重增量分解为三个因子,和,其中表示ViT中的特征维度,表示这些因子的秩,且。值得注意的是,两个因子在所有层之间共享,而因子对每层都是唯一的。然后可以使用以下等式计算权重增量:
其中表示用于调整因子学习率的超参数。我们在实验中将固定为1,并使用优化器调整整体学习率以实现类似的缩放效果。FacT权重增量应用于每个变换器块中的查询和值变换,而从SAM初始化的所有其他权重保持冻结状态,因为经验上在FacT应用于其他层时没有观察到明显的改进。有了FacT权重增量,查询和值变换变为:
其中表示微调后的查询或值变换,代表SAM预训练权重。

整合体积或时间信息

SAM最初是针对2D图像进行预训练的,但医学成像通常涉及超过两个维度。例如,体积CT和MRI数据包含了描述解剖结构或病变的关键3D空间信息,手术视频数据具有宝贵的帧间时间关系。整合这些医学成像数据中固有的体积或时间知识对于成功迁移学习SAM在医学应用中至关重要。为了解决这个关键挑战,我们提出了将一系列3D适配器整合到SAM架构中的2D变换器块中。这些适配器的目的是提取医学图像分析所需的基本体积或时间洞察。通过整合这些适配器,我们弥合了医学成像数据的固有复杂性与SAM预训练的2D骨干网络之间的差距,使其能够有效处理多维医学数据。具体来说,如图1所示,每个3D适配器由一个归一化层、一个线性下投影层、一个3D卷积层后跟一个激活层和一个线性上投影层组成。体积或时间信息的核心提取主要位于3D卷积层。下投影层的目的是将原始的维特征降低到一个更紧凑的维表示,以便控制新引入参数的数量。相反,上投影层恢复了特征维度。用表示特征图,3D适配器可以表示为:
其中表示层归一化,表示激活函数,分别表示线性下投影和上投影层,表示3D卷积层,核大小为,专门用于提取第三维度信息。为了使3D适配器与2D SAM骨干网络兼容,对于网络输入,我们提取一组相邻切片。这里,表示批量大小,表示相邻切片的数量,表示切片的尺寸。在输入被送入SAM骨干网络之前,应用重塑操作将转换为,通过合并相邻切片到批量维度。然后对于特征图,在送入3D适配器的3D卷积层之前,它们被重塑从。这里表示由于变换器中的补丁嵌入过程而下采样16倍的空间尺寸。在3D卷积操作之后,特征图的形状再次改变。通过这种方式,可以在2D网络骨干内有效提取体积或时间信息。对于每个变换器块,我们在注意力层之前和之后分别整合两个3D适配器,因为经验上可以以这种设计获得更优越的性能。

4. 实验

声明

本文内容为论文学习收获分享,受限于知识能力,本文对原文的理解可能存在偏差,最终内容以原论文为准。本文信息旨在传播和学术交流,其内容由作者负责,不代表本号观点。文中作品文字、图片等如涉及内容、版权和其他问题,请及时与我们联系,我们将在第一时间回复并处理。

深度学习爱好者
分享机器学习、深度学习和Python等知识与原理,每天分享深度学习与计算机视觉领域的经典和最新的论文总结,带读者一起跟踪前言科技!
 最新文章