[论文解读] Large-Vocabulary Segmentation for Medical Images with Text Prompts

Ziheng Zhao, Yao Zhang+6 more|arXiv (Cornell University)|Dec 28, 2023

Multimodal Machine Learning Applications被引用 15

一句话总结

SAT是一个以文本提示为引导的通用医学影像分割模型，能够在多模态数据上实现对31个数据集、362个类别的3D分割，且参数规模仅107M的SAT-Nano与专用nnU-Nets相当。

ABSTRACT

This paper aims to build a model that can Segment Anything in 3D medical images, driven by medical terminologies as Text prompts, termed as SAT. Our main contributions are three-fold: (i) We construct the first multimodal knowledge tree on human anatomy, including 6502 anatomical terminologies; Then, we build the largest and most comprehensive segmentation dataset for training, collecting over 22K 3D scans from 72 datasets, across 497 classes, with careful standardization on both image and label space; (ii) We propose to inject medical knowledge into a text encoder via contrastive learning and formulate a large-vocabulary segmentation model that can be prompted by medical terminologies in text form; (iii) We train SAT-Nano (110M parameters) and SAT-Pro (447M parameters). SAT-Pro achieves comparable performance to 72 nnU-Nets -- the strongest specialist models trained on each dataset (over 2.2B parameters combined) -- over 497 categories. Compared with the interactive approach MedSAM, SAT-Pro consistently outperforms across all 7 human body regions with +7.1% average Dice Similarity Coefficient (DSC) improvement, while showing enhanced scalability and robustness. On 2 external (cross-center) datasets, SAT-Pro achieves higher performance than all baselines (+3.7% average DSC), demonstrating superior generalization ability.

研究动机与目标

• 构建一个大规模、多数据集的医学分割数据集，在跨数据集上实现统一标注。
• 将多模态医学领域知识融入文本编码器以引导分割。
• 开发一个通用分割模型，利用文本提示对跨模态和区域的多样化目标进行分割。

提出的方法

• 从 e-Anatomy、UMLS 以及分割数据集构建多模态医学知识树。
• 利用知识增强的对比学习对文本与视觉编码器进行预训练，以将解剖文本与基于图谱的视觉概念对齐。
• 使用带有 transformer 查询模块和掩码生成器的3D U-Net骨干网络来训练SAT-Nano，在文本提示的引导下。
• 在后续阶段冻结文本编码器，采用两阶段视觉-语言训练流程。
• 实现数据集预处理和采样策略，以平衡31个数据集和362个类别。

实验结果

研究问题

• RQ1 是否可以仅使用文本提示，让一个单一的通用模型在多模态下分割广泛的解剖结构和病变？
• RQ2 训练语料库需要多大规模，以及何种知识整合能提升跨数据集的泛化？
• RQ3 紧凑的SAT-Nano是否在31个数据集上实现与任务特定的nnU-Nets相当的性能？
• RQ4 知识增强的表示学习对分割提示的文本-图像对齐有何影响？

主要发现

• SAT-Nano（107M 参数）能够使用文本提示对31个数据集的362个类别进行分割。
• 该模型的性能可与为每个数据集/子集训练的36个专门的 nnU-Nets 相提并论。
• 训练使用来自31个数据集的11K个3D扫描，并在不同身体区域显示出有效的跨数据集泛化。
• 带有知识注入的两阶段视觉-语言训练提升了文本医学概念与基于图谱的视觉特征之间的对齐。