[논문 리뷰] UniMed-CLIP: Towards a Unified Image-Text Pretraining Paradigm for Diverse Medical Imaging Modalities

Muhammad Uzair Khattak, Shahina Kunhimon외 3명|arXiv (Cornell University)|2024. 12. 13.

Radiomics and Machine Learning in Medical ImagingMedicine인용 수 3

한 줄 요약

이 논문은 6개의 의료 영상 모odalities를 통해 총 530만 개의 이미지-텍스트 쌍으로 구성된 대규모 오픈소스 데이터셋을 기반으로 사전 훈련된 통합 시각-언어 모델인 UniMed-CLIP을 소개한다. 레이블 전용 데이터셋에서 대규모 언어 모델(Large Language Models)을 활용해 고품질의 캡션을 합성함으로써, UniMed-CLIP는 기존의 비공개 데이터로 훈련된 모델들조차도 초월하는 최신 기술 수준의 제로샷 성능을 달성하였으며, 훈련 데이터의 1/3만으로도 +12.61%의 정확도 향상을 기록하였다.

ABSTRACT

Vision-Language Models (VLMs) trained via contrastive learning have achieved notable success in natural image tasks. However, their application in the medical domain remains limited due to the scarcity of openly accessible, large-scale medical image-text datasets. Existing medical VLMs either train on closed-source proprietary or relatively small open-source datasets that do not generalize well. Similarly, most models remain specific to a single or limited number of medical imaging domains, again restricting their applicability to other modalities. To address this gap, we introduce UniMed, a large-scale, open-source multi-modal medical dataset comprising over 5.3 million image-text pairs across six diverse imaging modalities: X-ray, CT, MRI, Ultrasound, Pathology, and Fundus. UniMed is developed using a data-collection framework that leverages Large Language Models (LLMs) to transform modality-specific classification datasets into image-text formats while incorporating existing image-text data from the medical domain, facilitating scalable VLM pretraining. Using UniMed, we trained UniMed-CLIP, a unified VLM for six modalities that significantly outperforms existing generalist VLMs and matches modality-specific medical VLMs, achieving notable gains in zero-shot evaluations. For instance, UniMed-CLIP improves over BiomedCLIP (trained on proprietary data) by an absolute gain of +12.61, averaged over 21 datasets, while using 3x less training data. To facilitate future research, we release UniMed dataset, training codes, and models at https://github.com/mbzuai-oryx/UniMed-CLIP.

연구 동기 및 목표

• 시각-언어 사전 훈련을 위한 대규모 오픈 액세스 의료 영상-텍스트 데이터셋의 부족 문제를 해결하기 위해.
• 모달리티별로 특화된 모델의 한계를 극복하고 통합적인 다중 모달리티 시각-언어 기초 모델을 개발하기 위해.
• 레이블 전용 의료 데이터셋을 고품질의 이미지-텍스트 쌍으로 변환하기 위해 LLM을 활용한 스케일 수 있는 데이터 중심의 사전 훈련을 가능하게 하기 위해.
• 공개 연구를 촉진하기 위해 완전히 오픈소스의 데이터셋, 훈련 코드, 모델을 공개하기 위해.

제안 방법

• 의료 스타일 템플릿을 사용하여 대규모 언어 모델(LLMs)을 활용해 모달리티에 특화된 이미지-레이블 데이터셋을 표준화된 이미지-텍스트 쌍으로 변환하는 데이터 수집 프레임워크를 개발하였다.
• X-레이, CT, MRI, 초음파, 병리, 망막 영상의 6개 의료 영상 모달리티를 아우르는 총 530만 개의 이미지-텍스트 데이터셋인 UniMed를 구축하였다.
• MetaCLIP의 ViT-B/16 시각 인코더와 BioMed-BERT 텍스트 인코더를 기반으로 한 시각-언어 모델을 대비 학습 목표를 사용해 미세조정하였다.
• 16개의 A100 40GB GPU를 사용한 멀티노드 환경에서 훈련하였으며, GPU당 배치 크기 128, 학습률 5e-5, 2000단계의 웜업 스텝을 적용하였다.
• GPT-4o를 활용한 프롬프트 엔지니어링을 통해 질병 레이블에서 일관되고 임상적으로 적절한 캡션을 생성하여 캡션 품질을 보장하였다.
• 합성된 이미지-텍스트 쌍을 기존 오픈 의료 VLM 데이터와 융합하여 다양성과 표현 품질을 향상시켰다.

실험 결과

연구 질문

• RQ1 오픈소스 데이터만을 사용하여도 통합 시각-언어 모델이 다양한 의료 영상 모달리티에서 강력한 제로샷 일반화 성능을 달성할 수 있는가?
• RQ2 LLM에 의해 생성된 합성 이미지-텍스트 쌍으로 훈련된 모델의 성능은 비공개 또는 소규모 데이터셋으로 훈련된 모델과 비교해 어떻게 되는가?
• RQ3 LLM이 레이블 전용 의료 데이터셋을 사전 훈련을 위한 고품질의 임상적으로 관련성이 높은 이미지-텍스트 쌍으로 효과적으로 변환할 수 있는가?
• RQ4 다양한 모달리티에서 훈련된 통합 모델이 제로샷 및 선형 프로빙 평가에서 모달리티별 특화 모델을 능가하는가?
• RQ5 LLM 증강 데이터를 활용한 데이터 중심 접근 방식이 훨씬 적은 데이터 요구량으로도 최신 기술 수준의 성능을 달성할 수 있는가?

주요 결과

• UniMed-CLIP는 훈련 데이터의 1/3만을 사용했음에도 불구하고, 비공개 데이터로 훈련된 BiomedCLIP보다 21개의 다양한 의료 데이터셋 평균에서 +12.61%의 정확도 향상을 기록하였다.
• PCam 병리 영상 데이터셋에서 UniMed-CLIP는 100% 선형 프로빙에서 85.97%의 정확도를 달성하여 BiomedCLIP(83.40%)와 PMC-CLIP(81.03%)를 모두 능가하였다.
• ACL MRI 데이터셋에서 UniMed-CLIP는 100% 선형 프로빙에서 97.28%의 정확도를 기록하여 BiomedCLIP(83.84%)와 CLIP(81.09%)를 초월하였다.
• ODIR-2x300 망막 영상 데이터셋에서 UniMed-CLIP는 100% 선형 프로빙에서 95.00%의 정확도를 기록하여 BiomedCLIP(94.17%)와 CLIP(91.73%)를 능가하였다.
• CT, MRI, 망막 영상과 같은 클래스 불균형 데이터셋에서도 UniMed-CLIP는 강력한 성능을 유지하여 비대칭 데이터 분포 하에서도 뛰어난 일반화 능력을 보였다.
• 모든 6개 모달리티에서 최신 기술 수준의 제로샷 성능을 달성하였으며, 특히 MediMeTA(CT)와 FIVES(망막)와 같은 도전적인 데이터셋에서 이전 모델들보다 뚜렷한 성능 향상을 보였다.