[논문 리뷰] Training Like a Medical Resident: Context-Prior Learning Toward Universal Medical Image Segmentation

Yunhe Gao, Zhuowei Li외 4명|arXiv (Cornell University)|2023. 06. 04.

Radiomics and Machine Learning in Medical Imaging인용 수 9

한 줄 요약

본 논문은 다양한 데이터셋, 모달리티, 및 과제에 걸친 범용 의료 영상 분할을 위한 오라클-가이드 컨텍스트-프라이어 학습 방식인 Hermes를 제안하여, 기존의 작업별 모델에 비해 성능과 확장성을 향상시킨다.

ABSTRACT

A major focus of clinical imaging workflow is disease diagnosis and management, leading to medical imaging datasets strongly tied to specific clinical objectives. This scenario has led to the prevailing practice of developing task-specific segmentation models, without gaining insights from widespread imaging cohorts. Inspired by the training program of medical radiology residents, we propose a shift towards universal medical image segmentation, a paradigm aiming to build medical image understanding foundation models by leveraging the diversity and commonality across clinical targets, body regions, and imaging modalities. Towards this goal, we develop Hermes, a novel context-prior learning approach to address the challenges of data heterogeneity and annotation differences in medical image segmentation. In a large collection of eleven diverse datasets (2,438 3D images) across five modalities (CT, PET, T1, T2 and cine MRI) and multiple body regions, we demonstrate the merit of the universal paradigm over the traditional paradigm on addressing multiple tasks within a single model. By exploiting the synergy across tasks, Hermes achieves state-of-the-art performance on all testing datasets and shows superior model scalability. Results on two additional datasets reveals Hermes' strong performance for transfer learning, incremental learning, and generalization to downstream tasks. Hermes's learned priors demonstrate an appealing trait to reflect the intricate relations among tasks and modalities, which aligns with the established anatomical and imaging principles in radiology. The code is available: https://github.com/yhygao/universal-medical-image-segmentation.

연구 동기 및 목표

• 다양하고 부분적으로 주석이 달린 의료 데이터셋을 활용하여 범용 의료 영상 분할을 고무한다.
• 분할을 안내하는 학습 가능한 토큰으로 작업(prior)과 모달리티 프라이어를 주입하는 Hermes를 도입한다.
• 작업별 모델과 비교했을 때 단일 범용 모델의 성능향상, 전이성 및 확장성을 입증한다.
• 신체 부위, 모달리티, 임상 표적을 포괄하는 일곱 개의 데이터세트와 추가적으로 다운스트림 작업에서 Hermes를 평가한다.
• 성능 향상에 기여하는 핵심 구성 요소를 식별하기 위한 차감 분석(ablations)을 분석한다.

제안 방법

• 다른 ROI 및 모달리티를 가진 여러 데이터셋에서 학습하는 것을 범용 분할로 정의한다.
• 학습 가능한 토큰(오라클)으로 작업 프라이어 풀과 모달리티 프라이어 풀을 갖춘 Hermes를 도입한다.
• 양방향 주의 기반 프라이어 융합 모듈을 통해 프라이어를 백본 특징과 융합한다.
• 계층적이고 다중 스케일 컨텍스트 프라이어를 사용하고 데이터셋 간 혼합 배치로 학습한다.
• BCE 및 Dice 손실을 사용하여 분할 백본과 컨텍스트 프라이어를 함께 최적화한다.
• 일곱 개의 업스트림 데이터셋과 두 개의 다운스트림 작업(전이/증분/일반화)을 평가한다.

실험 결과

연구 질문

• RQ1 다른 ROI와 모달리티의 다양한 작업이 보편적 분할 모델에 상호 이익이 될 수 있는가?
• RQ2 해부학적 차이에도 불구하고 신체 부위 간에 전이 가능한 표현이 달성될 수 있는가?
• RQ3 모달리티의 공통점을 활용하여 CT와 MRI 간의 성능 향상을 이끌어낼 수 있는가?
• RQ4 부분 주석 및 이질적인 데이터 하에서 작업 및 모달리티 프라이어가 분할을 향상시키는가?
• RQ5 보편적 학습이 새로운 작업으로의 전이, 증분 학습 및 일반화에 어떻게 영향을 미치는가?

주요 결과

패러다임	모델	BCV	StructSeg	LiTS 종양	KiTS 종양	AMOS CT	AMOS MR	CHAOS	평균
Traditional	nnUNet	84.23	88.53	64.91	81.72	88.79	85.49	91.34	83.57
Traditional	ResUNet	84.36	88.59	64.87	81.89	88.97	85.43	91.41	83.64
Universal	Hermes-R	85.78	89.43	67.66	85.26	89.59	86.70	92.02	85.20
Traditional	SwinUNETR	83.32	88.36	64.82	74.32	88.29	83.97	88.34	81.63
Traditional	MedFormer	84.61	89.04	66.24	82.09	89.45	85.58	91.85	84.12
Universal	Hermes-M	86.23	89.57	68.55	85.47	89.72	87.15	92.14	85.54

• Hermes는 일곱 개 데이터셋과 두 가지 백본 유형에 걸쳐 전통적인 작업별 백본보다 Dice 점수를 일관되게 향상시킨다.
• 서로 다른 목표 ROI가 Hermes 하에서 서로의 성능을 상호 향상시킨다.
• 신체 부위 간 표현의 전이 가능성이 관찰된다(복부에서 흉부로).
• 교차 모달리티 이점이 관찰된다(CT와 MRI에서 단일 모델 하 개선).
• Hermes와 함께하는 범용 학습은 더 큰 백본에서 전통적 패러다임에 비해 규모 확장성이 더 잘 나타난다.
• 다운스트림 작업에서 Hermes는 강한 전이, 증분 학습 및 일반화 성능 향상을 달성한다.