[논문 리뷰] BenchMD: A Benchmark for Unified Learning on Medical Images and Sensors

Kathryn Wantlin, Chenwei Wu외 13명|arXiv (Cornell University)|2023. 04. 17.

COVID-19 diagnosis using AI인용 수 10

한 줄 요약

BenchMD는 7개 모달리티에 걸친 19개의 실제 의료 데이터셋에 걸친 통합적이고 모달리티-무관 학습을 평가하며, few-shot 및 SSL/사전학습 시나리오에서 ID 및 OOD 성능을 평가합니다.

ABSTRACT

Medical data poses a daunting challenge for AI algorithms: it exists in many different modalities, experiences frequent distribution shifts, and suffers from a scarcity of examples and labels. Recent advances, including transformers and self-supervised learning, promise a more universal approach that can be applied flexibly across these diverse conditions. To measure and drive progress in this direction, we present BenchMD: a benchmark that tests how well unified, modality-agnostic methods, including architectures and training techniques (e.g. self-supervised learning, ImageNet pretraining),perform on a diverse array of clinically-relevant medical tasks. BenchMD combines 19 publicly available datasets for 7 medical modalities, including 1D sensor data, 2D images, and 3D volumetric scans. Our benchmark reflects real-world data constraints by evaluating methods across a range of dataset sizes, including challenging few-shot settings that incentivize the use of pretraining. Finally, we evaluate performance on out-of-distribution data collected at different hospitals than the training data, representing naturally-occurring distribution shifts that frequently degrade the performance of medical AI models. Our baseline results demonstrate that no unified learning technique achieves strong performance across all modalities, leaving ample room for improvement on the benchmark. Code is released at https://github.com/rajpurkarlab/BenchMD.

연구 동기 및 목표

• 의료 AI를 위한 보편적이고 모달리티-무관 학습 방법의 개발을 촉진한다.
• 다양한 의료 모달리티에서 통합 아키텍처와 학습 기법의 성능을 평가한다.
• 레이블 부족 및 분포 변화가 실제 현장 제약을 반영하도록 성능을 평가한다.
• 실제 세계 작업과 전문가가 검증한 데이터셋을 포함한 표준화되고 공개적으로 접근 가능한 벤치마크를 제공한다.

제안 방법

• 입력을 공유 256 차원 공간으로 투영하는 독립적인 1D, 2D, 3D 임베딩 모듈을 갖춘 모달리티-무관 트랜스포머 아키텍처를 사용한다.
• 세 가지 SSL 목표를 사용한 사전학습: e-Mix (대조적 임베딩 믹스업), ShED (섞인 임베딩 예측), MAE (마스킹된 자동복원)
• 모달리즘 간 이미지넷(pretraining) 대비 SSL 방법을 비교하고 스크래치 기본 모델과의 차이를 비교한다.
• 단일 라벨 작업에서 클래스당 8, 64, 256의 라벨 가용성을 바꿔 선형 평가 및 미세조정을 통해 학습한다.
• 분포 변화에 따른 일반화를 측정하기 위해 분포 외 대상 데이터셋으로의 제로샷 전이가를 평가한다.
• 전처리 및 검증 지표를 표준화하고 평가 지표로 AUROC를 사용한다.

실험 결과

연구 질문

• RQ1 통합 학습 기법이 여러 의료 모달리티에서 높은 성능을 제공하는가?
• RQ2 모달리티 간 라벨 가용성 및 분포 변화에 따라 성능이 어떻게 달라지는가?
• RQ3 어떤 사전학습 전략(SSL vs ImageNet 대비 Scratch)이 모달리티 간 견고한 OOD 일반화를 제공하는가?
• RQ4 단일 아키텍처/접근법이 1D, 2D, 3D 의료 데이터 전반에서 강력한 성과를 달성할 수 있는가?

주요 결과

• 모든 모달리티에서 단일 기법이 우위를 보이지 않으며, 성능은 모달리티 및 데이터셋에 따라 달라진다.
• ImageNet 사전학습은 2D 모달리티(CXR, 유방촬영, 망막촬영)에서 SSL보다 종종 우수하지만, OOD 피부과 영상에서는 SSL이 우수할 수 있다.
• MAE는 EEG 및 일부 피부과 데이터셋에서 강한 성능을 보이지만 ECG와 같은 다른 모달리티에서는 하회한다; e-Mix와 ShED는 모달리티별로 강점이 다르다.
• 처음부터 학습하는 것은 거의 최상위 성능에 도달하지 못하지만 일부 설정에서 여전히 경쟁력이 있다; SSL 베이스라인이 전반적으로 완만한 이득을 제공한다.
• 2단계 사전학습(ImageNet 다음으로 MAE 같은 SSL이 따라오는 경우)은 일부 2D 모달리티에서 이익을 낼 수 있다.
• 레이블 가용성은 일반적으로 OOD 성능을 향상시키나, 일부 경우 전체 데이터에서 과적합이 발생할 수 있다.