[Paper Review] BenchMD: A Benchmark for Unified Learning on Medical Images and Sensors

Kathryn Wantlin, Chenwei Wu+13 more|arXiv (Cornell University)|Apr 17, 2023

COVID-19 diagnosis using AI10 citations

TL;DR

BenchMD evaluates unified, modality-agnostic learning across 19 real-world medical datasets spanning 7 modalities, assessing ID and OOD performance with few-shot and SSL/pretraining scenarios.

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.

Motivation & Objective

• Motivate development of universal, modality-agnostic learning methods for medical AI.
• Assess how unified architectures and training techniques perform across diverse medical modalities.
• Evaluate performance under label scarcity and distribution shifts to reflect real-world constraints.
• Provide a standardized, publicly accessible benchmark with real-world tasks and expert-validated datasets.

Proposed method

• Use a modality-agnostic transformer architecture with separate 1D, 2D, and 3D embedding modules projecting inputs into a shared 256-dimensional space.
• Pretrain with three SSL objectives: e-Mix (contrastive embedding mixup), ShED (shuffled embedding prediction), and MAE (masked autoencoding).
• Compare SSL methods with ImageNet pretraining (2D modalities) and a scratch baseline across modalities.
• Train with linear evaluation and finetuning, varying label availability (8, 64, 256 per class for single-label tasks).
• Evaluate zero-shot transfer to out-of-distribution target datasets to measure generalization under distribution shifts.
• Standardize preprocessing and validation metrics; use AUROC as the evaluation metric.

Experimental results

Research questions

• RQ1 Do unified learning techniques offer high performance across multiple medical modalities?
• RQ2 How does performance vary with label availability and distribution shifts across modalities?
• RQ3 Which pretraining strategy (SSL vs ImageNet vs training from scratch) yields robust OOD generalization across modalities?
• RQ4 Can a single architecture/approach achieve strong results across 1D, 2D, and 3D medical data?

Key findings

• No single technique dominates across all modalities; performance varies by modality and dataset.
• ImageNet pretraining often outperforms SSL on several 2D modalities (CXRs, mammograms, fundus) but SSL can outperform in OOD dermoscopic images.
• MAE shows strong results on EEG and some dermoscopic datasets but underperforms in other modalities like ECG; e-Mix and ShED show modality-specific strengths.
• Training from scratch is rarely top-performing but remains competitive in some settings; SSL baselines provide modest gains overall.
• Two-stage pretraining (ImageNet followed by SSL like MAE) can yield benefits for some 2D modalities.
• Label availability generally improves OOD performance, though overfitting can occur with full data in some cases.