开源3D医学大模型SAT，支持497类器官，性能超越72个nnU-Nets，上交大团队发布
- 王林原创
- 2024-07-12 10:52:01499浏览
作者 | 上海交通大学、上海人工智能实验室
编辑 | ScienceAI
近日，上海交通大学与上海人工智能实验室联合团队发布3D医学图像分割大模型SAT（Segment Anything in radiology scans, driven by Text prompts），在3D医学图像（CT、MR、PET）上，基于文本提示实现对人体497种器官/病灶的通用分割。所有数据和代码、模型均已开源。
论文链接：https://arxiv.org/abs/2312.17183
代码链接：https://github.com/zhaoziheng/SAT
数据链接：https://github.com/zhaoziheng/SAT-DS/
研究背景
医学图像分割在诊断、手术规划和疾病监测等一系列临床任务中都有重要作用。然而，传统的研究针对每个特定的分割任务训练「专用」模型，导致每个「专用」模型的应用范围都相对有限，无法高效便捷地满足广泛多样的医疗分割需求。
与此同时，大语言模型最近在医疗领域取得了巨大成功，而要进一步推动通用医疗人工智能的发展，构建一个可以连接语言和定位能力的医学分割工具变得十分必要。
为了克服这些挑战，来自上海交通大学和上海人工智能实验室的研究者们提出了第一个基于知识增强，使用文本提示的3D医疗图像通用分割模型，名为SAT（Segment Anything in radiology scans, driven by Text prompts），并做出了以下三点主要贡献：
1. 该研究首次探索将人体解剖学知识注入文本编码器，以精准编码解剖学术语，实现了通过文本提示的放射学影像通用医疗分割模型。
2. 该研究构建了第一个包含6K+人体解剖学概念的多模态医疗知识图谱。同时，构建了当前最大规模的3D医学图像分割数据集，名为SAT-DS，汇集了72个公开数据集，来自CT、MR和PET三种模态的22K+图像，302K+分割标注，涵盖了人体8个主要部位中的497个分割目标。
3. 基于SAT-DS，该研究训练了两款不同大小的模型：SAT-Pro（447M参数）和SAT-Nano（110M参数），并设计实验从多个角度验证了SAT的价值：SAT的性能与72个nnU-Nets专家模型相当（在每个数据集上单独调参和优化，共约2.2B参数），并在域外数据上表现出更强的泛化能力；SAT可以作为一个基于大规模数据预训练的基础分割模型，通过下游微调迁移到特定任务时，可以表现出比nnU-Nets更好的性能；此外，与基于box提示的MedSAM对比，SAT基于文本提示可以实现更精准、更高效的分割；最后，在域外的临床数据上，研究团队展示了SAT可以被用作大语言模型的代理工具，在报告生成等任务中直接赋予后者定位和分割的能力。
接下来将从数据、模型与实验结果三个方面介绍原文细节。
数据构建
多模态知识图谱：为了实现精准编码解剖学术语，研究团队首先收集了一个包含6K+人体解剖学概念的多模态知识图谱，其内容来自于三个来源：
1. Unified Medical Language System（UMLS）是由美国国家医学图书馆构建的生物医学字典。研究团队从中提取了近230K的生物医学概念与定义，以及涵盖1M+条相互关系的知识图谱。
2. 网络上的权威解剖学知识。研究团队筛选了6502个人体解剖学概念，并借助检索增强的大语言模型从网络上检索相关信息，获取了6K+概念与定义，涵盖了38K+解剖结构相互关系的知识图谱。
3. 公开的分割数据集。研究团队收集了大规模的公开3D医学图像分割数据集，将分割区域通过解剖学概念（类别标签）与上述文本知识库中的知识对应连接，提供视觉知识对照。
SAT-DS：为了训练通用分割模型，研究团队构建了领域内最大规模的3D医学图像分割数据集合SAT-DS。特别的，72个多样的公开分割数据集被收集并整理，总计包括22186例3D图像，302033个分割标注，来自CT、MR和PET三种模态，以及涵盖人体8个主要区域的497个分割类别（解剖学结构或病灶）。
为了尽可能降低异构数据集间的差异，研究团队对不同数据集间的方向、体素间距、灰度值等图像属性进行了标准化，用统一的解剖学术语系统命名了不同数据集中的分割类别。
图 3：SAT-DS是一个大规模、多样化的3D医学图像分割数据集合，涵盖人体8个主要区域共497个分割类别。
模型架构
知识注入：为了构建可以精准编码解剖学术语的提示编码器，研究团队首先将多模态解剖学知识用对比学习的方式注入到文本编码器。
如下图a所示，用解剖学概念将多模态知识连接成对，随后使用视觉编码器（visual encoder）和文本编码器（text encoder）分别编码视觉和文本知识，通过对比学习在特征空间中将解剖学结构的视觉特征和文本知识对齐，并构建解剖学结构之间的关系，从而学习到对解剖学概念的更好编码，作为提示引导视觉分割模型的训练。
基于文本提示的通用分割：研究团队进一步设计了基于文本提示的通用分割模型框架，如下图b所示，包含文本编码器、视觉编码器、视觉解码器与提示解码器。
其中，考虑到同一解剖结构在不同图像中存在差异，提示解码器（query decoder）使用视觉编码器输出的图像特征增强解剖学概念特征，即分割提示。最后，在分割提示与视觉解码器输出的像素级特征间计算点积，得到分割预测结果。
模型测评
该研究将SAT与两个代表性的方法进行对比，即「专用」模型nnU-Nets和交互式通用分割模型MedSAM。测评包含了域内数据集测试（综合分割性能）与zero-shot域外数据集测试（跨中心数据迁移能力）两方面，评测结果从数据集、类别和人体区域三个层面进行了整合：
类别：不同数据集之间相同类别的分割结果进行汇总、平均；
区域：基于类别结果，将同一人体解剖区域内的类别结果进行汇总、平均；
数据集：传统的分割模型评估方式，同一数据集内的分割结果进行平均；
与专用模型nnU-Nets的对比实验
为了最大化nnU-Nets的性能，该研究在每个单独的数据集上训练nnU-Nets并与SAT对比，具体设定如下：
1. 在域内测试中，使用SAT-DS中的全部72个数据集进行测试和对比。对于SAT，使用72个训练集的总和进行训练，并在72个测试集上进行测试；对于nnU-Nets，汇总72个nnU-Nets在各自测试集上的结果作为一个整体。
2. 在域外测试中，进一步划分72个数据集，使用其中的49个数据集（命名为SAT-DS-Nano）的训练集训练SAT-Nano，在10个域外的测试集上zero-shot测试；对于nnU-Nets，使用49个nnU-Nets在10个域外测试集上测试并汇总结果。
域内测试结果：从表1可以看到，SAT-Pro在域内测试中表现出与72个nnU-Nets十分接近的性能，并在多个区域上超越nnU-Nets。需要注意的是，SAT可以仅用一个模型完成72个分割任务，并在模型尺寸上远小于nnU-Nets的集合（如下图c所示）。
Fine-tuning迁移测试结果：该研究进一步将SAT-Pro在每个数据集上单独fine-tune后进行测试，命名为SAT-Pro-Ft。从表1中可以看出，SAT-Pro-Ft相比SAT-Pro在所有区域上都有了明显的性能提升，并在总体性能上超过了nnU-Nets。
域外测试结果：如表2所示，SAT-Nano在10个数据集20个指标中的19项上超过了nnU-Nets，表现出整体更强的迁移能力。
表 2：SAT-Nano与nnU-Nets、MedSAM的域外测试对比，结果以数据集为单位呈现。
与交互式分割模型MedSAM的对比实验
该研究直接使用MedSAM的公开checkpoint进行测试和SAT对比，具体设定如下：
1. 在域内测试中，从72个数据集中进一步筛选了32个MedSAM训练中使用过的数据集进行对比。
2. 在域外测试中，筛选了5个MedSAM训练中没有使用过的数据集进行对比。
对于MedSAM，考虑两种不同的Box提示: 使用包含ground truth分割的最小矩形（Oracle Box） ，记为MedSAM（Tight）；在Oracle Box基础上加入随机偏移，记为MedSAM（Loose）。同时测试Oracle Box直接作为预测的效果。对于SAT，直接沿用nnU-Nets对比实验中的模型在这些数据集上测试，不重新训练。
域内测试结果：如表3所示，SAT-Pro几乎在所有的区域上都比MedSAM表现更好，且综合表现来看SAT-Pro和SAT-Nano的性能都优于MedSAM。尽管在病灶上SAT-Pro表现不如MedSAM，Oracle Box本身作为预测在病灶上的表现也足够好，甚至在DSC上超越了MedSAM。这表明MedSAM在病灶的分割上的更优性能很可能来自于Box提示的强先验信息。
表 3：SAT-Pro、SAT-Nano与MedSAM的域内测试对比，结果以区域或病灶为单位整合。
定性比较：图6从域内测试的结果中挑选了两个典型例子进行可视化展示，进一步对比SAT和MedSAM。如图6上所示，在对心肌的分割中，Box提示很难区分心肌和被心肌包裹的心室，因此MedSAM也错误将两者一起分割了出来，这表明Box提示在类似的复杂空间关系中很容易有歧义，导致分割不精准。
相比之下，基于文本提示（直接输入解剖结构的名字）的SAT可以精准地区分心肌和心室。此外，在图6下展示的肠道肿瘤分割中可以看到，Oracle Box对于病灶目标来说已经是很好的预测结果，而MedSAM的分割结果可能不会比所得到的Box提示更好。
图 6：SAT-Pro与MedSAM（Tight）的定性比较。其中MedSAM使用Oracle Box作为提示，Box用蓝色标识。第一行展示了一个心肌分割的例子；第二行展示了一个肠道肿瘤分割的例子。
域外测试结果：如表2所示，与MedSAM（Tight）相比，SAT-Nano在5个数据集共10个指标中的5个上超过了MedSAM。而MedSAM（Loose）在所有指标上都有明显的性能下降，表明MedSAM对于用户输入的Box提示的偏移比较敏感。
消融实验
在设计SAT时，视觉骨干网络和文本编码器是两个关键的部分，该研究尝试在SAT框架中使用不同的视觉网络结构或文本编码器，通用消融实验来探讨它们的影响。
为了节省实验的开销，消融实验中的所有SAT模型训练和测试都在包含49个数据集的SAT-DS-Nano上进行，它包含了13303个3D图像，151461个分割标注，以及429个分割类别。
视觉骨干网络：在SAT-Nano的框架下，该研究选用了三种主流的分割网络结构进行对比，即U-Net（110M参数），SwinUNETR（107M参数）和U-Mamba（114M参数）。为了公平对比，该消融实验中控制它们的参数量大致接近。同时为了计算开销，省略了知识注入的步骤，直接使用MedCPT（MedCPT是基于PubMed文献，使用225M私有用户点击数据训练的文本编码器，在一系列医疗语言任务中取得了最好的性能）作为文本编码器生成提示。三种变体分别被记为U-Net-CPT、SwinUNETR-CPT和U-Mamba-CPT。
从图7中可以看到，使用U-Net与U-Mamba作为视觉骨干网络，最终的分割性能比较接近，其中U-Net略好于U-Mamba；而使用SwinUNETR时的分割性能有明显下降。最后，研究团队选择U-Net作为SAT的视觉骨干网络。
文本编码器：在SAT-Nano的框架下，该研究选择了三种代表性的文本编码器进行对比：使用上文提出的知识注入方式训练的文本编码器（记为Ours），使用先进的医疗文本编码器MedCPT，使用没有针对医疗数据微调的文本编码器BERT-base。
为了公平，该消融实验统一采用U-Net作为视觉网络。三种变体分别被记为U-Net-Ours，U-Net-CPT和U-Net-BB。如图8所示，总体而言，使用MedCPT相比使用BERT-base对于分割性能有微小的提升，表明领域知识对提供好的分割提示有一定的帮助；而使用该研究提出的文本编码器在所有的类别上都取得了最好的性能，表明构建多模态人体解剖学知识库和知识注入对于分割模型有明显的帮助。
长尾分布是分割数据集的一个明显特征。如图9中a图和b图所示，研究团队调研了用于消融实验的SAT-DS-Nano中429个类别的标注数量分布。如果把标注数最多的10个类（前2.33%）定义为头部，标注数最少的150个类别（后34.97%）定义为尾部类，可以发现尾部类的标注数量仅占总体标注数的3.25%。
该研究进一步探究文本编码器对长尾分布中不同类别的分割结果影响。如图9中c图所示，研究团队提出的编码器在头部、尾部和中间类别上都取得了最好的性能，其中尾部类上的提升比头部类更加明显。同时，MedCPT在头部类上表现略低于BERT-base，但在尾部类上效果更好。这些结果都表明，领域知识，特别是多模态人体解剖学知识的注入，对于长尾类别的分割有明显帮助。
与大语言模型的结合
由于SAT可以基于文本提示进行分割，它可以被直接用作大语言模型的代理工具，提供分割能力。为了展示应用场景，研究团队选择了4个多样的真实临床数据，使用GPT4从报告中提取分割目标并调用SAT进行zero-shot分割，结果展示在图10中。
可以看到，GPT-4可以很好地检测到报告中重要的解剖学结构，并且调用SAT，在不需要任何数据fine-tune的情况下，在真实临床图像上很好地分割出这些目标。
研究价值
作为首个基于文本提示的3D医疗图像通用分割大模型，SAT的价值体现在许多方面：
SAT构建了高效灵活的通用分割：SAT-Pro仅用一个模型，在广泛的分割任务上展现出与72个nnU-Nets相当的性能，并且有更少的模型参数量。这表明相比传统的医疗分割方法需要配置、训练和部署一系列专用模型，SAT-Pro作为通用分割模型是更灵活高效的解决方法。同时，研究团队也证明SAT-Pro在域外数据上泛化性能更好，能更好地满足跨中心迁移等临床需求。
SAT是基于大规模分割数据预训练的基础模型：SAT-Pro在大规模的分割数据集上进行训练后，当通过微调迁移到特定数据集上时，表现出了明显的性能提升，并比nnU-Nets总体性能更好。这表明SAT可以被视作强大的基础分割模型，可以通过微调迁移在特定任务上表现更好，从而平衡通用分割与专用分割的临床需求。
SAT实现了基于文本提示的准确、鲁棒分割：相比基于Box提示的交互式分割模型，SAT基于文本提示可以取得更准确和对提示鲁棒的分割结果，并且可以节省使用者圈画Box的大量时间，从而实现自动的、可批量化的通用分割。
SAT可以作为大语言模型的代理工具：研究团队在真实临床数据上展示了SAT可以和大语言模型无缝衔接，通过文本作为桥梁，直接为任何大语言模型提供分割和定位能力。这对进一步推动通用医疗智能（Generalist Medical Artificial Intelligence）的发展有重要价值。
模型尺寸对于分割的影响：通过训练两种不同大小的模型：SAT-Nano和SAT-Pro，该研究观测到在域内测试中SAT-Pro相比SAT-Nano有明显的提升。这暗示着在大规模数据集上训练通用分割模型时，scaling-law依然适用。
领域知识对于分割的影响：研究团队提出了第一个多模态的人体解剖学知识库，并探索用知识增强提升通用分割模型的性能，特别是对长尾类别的分割。考虑到分割标注，特别是长尾类别上的标注，相对稀缺，这一探索对于构建通用的分割模型有重要意义。
以上是开源3D医学大模型SAT，支持497类器官，性能超越72个nnU-Nets，上交大团队发布的详细内容。更多信息请关注PHP中文网其他相关文章！