[논문 리뷰] VIS-MAE: An Efficient Self-supervised Learning Approach on Medical Image Segmentation and Classification

Zelong Liu, Andrew Tieu외 10명|arXiv (Cornell University)|2024. 02. 01.

Brain Tumor Detection and ClassificationNeuroscience인용 수 3

한 줄 요약

VIS-MAE는 CT, MRI, PET, X-ray, 초음파 모odal리티에서 250만 장의 레이블이 없는 의료 영상으로 사전 훈련된 자기지도 기반 시각 트랜스포머 기반 기초 모델이다. 이 모델은 레이블 효율성이 높아 50~80%의 레이블 데이터로 미세조정된 경우에도 벤치마크를 초월하는 최신 기술 성능을 달성하여 의료 영상 분할 및 분류에서 최고의 성능을 보였다.

ABSTRACT

Artificial Intelligence (AI) has the potential to revolutionize diagnosis and segmentation in medical imaging. However, development and clinical implementation face multiple challenges including limited data availability, lack of generalizability, and the necessity to incorporate multi-modal data effectively. A foundation model, which is a large-scale pre-trained AI model, offers a versatile base that can be adapted to a variety of specific tasks and contexts. Here, we present VIsualization and Segmentation Masked AutoEncoder (VIS-MAE), novel model weights specifically designed for medical imaging. Specifically, VIS-MAE is trained on a dataset of 2.5 million unlabeled images from various modalities (CT, MR, PET,X-rays, and ultrasound), using self-supervised learning techniques. It is then adapted to classification and segmentation tasks using explicit labels. VIS-MAE has high label efficiency, outperforming several benchmark models in both in-domain and out-of-domain applications. In addition, VIS-MAE has improved label efficiency as it can achieve similar performance to other models with a reduced amount of labeled training data (50% or 80%) compared to other pre-trained weights. VIS-MAE represents a significant advancement in medical imaging AI, offering a generalizable and robust solution for improving segmentation and classification tasks while reducing the data annotation workload. The source code of this work is available at https://github.com/lzl199704/VIS-MAE.

연구 동기 및 목표

• 레이블이 부족한 의료 영상 데이터 문제를 해결하기 위해 일반화 가능한 기초 모델을 개발하기 위해.
• CT, MRI, PET, X-ray, 초음파 등 다양한 의료 영상 모달리티 간의 모델 일반화 능력을 향상시키기 위해.
• 하류 분류 및 분할 작업에서 레이블 효율성을 향상시키기 위해.
• 사전 훈련 중 수동 주석이 필요 없이 효과적인 다중 모달 표현 학습을 가능하게 하기 위해.
• 임상 적용을 위한 강력하고 이식 가능한 의료 영상 AI 기초 모델을 제공하기 위해.

제안 방법

• VIS-MAE는 입력 의료 영상의 무작위 패치를 마스킹하고, 가시 패치들로부터 이를 복원하는 마스킹 자동에코딩(MAE) 프레임워크를 사용한다.
• 모델은 다양한 영상 모달리티에서 온 250만 장의 레이블이 없는 대규모 다중 모달 데이터셋으로 사전 훈련된다.
• 의료 영상의 장거리 상관관계와 공간 계층을 포착하기 위해 시각 트랜스포머 백본이 사용된다.
• 미세조정 단계에서 모델은 레이블 데이터의 일부만 사용하여 분류 및 분할 작업에 적응된다.
• 모달리티별 정규화 및 데이터 증강 전략을 아키텍처에 통합하여 다양한 영상 유형 간의 강건성을 향상시킨다.
• 자기지도 사전 훈련을 통해 하류 작업에 대한 미세조정 이전에 일반화 가능한 특징을 학습할 수 있다.

실험 결과

연구 질문

• RQ1 다양하고 레이블이 없는 의료 영상으로 사전 훈련된 자기지도 기반 기초 모델이 하류 분할 및 분류 성능을 향상시킬 수 있는가?
• RQ2 제한된 레이블 데이터로 미세조정된 경우 VIS-MAE의 성능이 기존 사전 훈련 모델과 비교해 어떻게 되는가?
• RQ3 VIS-MAE는 다양한 영상 모달리티와 임상 작업 간에 얼마나 잘 일반화되는가?
• RQ4 의료 영상에 마스킹 자동에코딩을 적용하면 더 강건하고 이식 가능한 표현을 얻을 수 있는가?
• RQ5 VIS-MAE는 정확도를 저하시키지 않으면서도 의료 영상 AI의 주석 부담을 줄일 수 있는가?

주요 결과

• VIS-MAE는 다양한 의료 영상 분할 및 분류 벤치마크에서 최신 기술 성능을 달성하였으며, 도메인 내 및 도메인 외 설정 모두에서 기존 모델을 능가하였다.
• 모델은 레이블 데이터의 50%만으로도 높은 성능을 유지하여 뛰어난 레이블 효율성을 입증하였다.
• VIS-MAE는 CT, MRI, PET, X-ray, 초음파를 포함한 다양한 영상 모달리티 간에 강력한 일반화 능력을 보였다.
• 250만 장의 레이블이 없는 영상으로 자기지도 사전 훈련을 수행함으로써, 제한된 데이터로 감독 사전 훈련을 수행한 경우보다 하류 작업 정확도가 크게 향상되었다.
• 모델은 분류 및 인스턴스 분할 작업 모두에서 경쟁 가능한 성능을 달성하여 광범위한 적용 가능성을 보였다.
• 소스 코드는 공개되어 있어 재현성과 의료 AI 분야의 추가 연구를 지원한다.