[논문 리뷰] Large-Vocabulary Segmentation for Medical Images with Text Prompts

Ziheng Zhao, Yao Zhang외 6명|arXiv (Cornell University)|2023. 12. 28.

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 및 분할 데이터셋으로부터 다중 모달 의료 지식 트리를 구축한다.
• 해부학적 텍스트와 도해 기반 시각 개념을 맞추기 위해 지식 강화 대조 학습으로 텍스트 및 시각 인코더를 사전 학습한다.
• 텍스트 프롬프트에 의해 안내되는 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 스캔을 사용하며 신체 부위 간의 데이터 간 일반화가 효과적으로 나타난다.
• 지식 주입이 포함된 두 단계의 시각-언어 학습은 텍스트의 의료 개념과 도해 기반 시각 특징 간의 정합성을 향상시킨다.