Expectation-Maximization as the Engine of Scalable Medical Intelligence

TL;DR

ScaleMAI introduces an EM-based framework that automates and accelerates the creation of large, high-quality annotated medical datasets by iteratively refining annotations and training models, surpassing human performance.

Contribution

The paper presents ScaleMAI, a novel EM-based framework that co-evolves data annotation and model training, significantly scaling up high-quality medical datasets with minimal human intervention.

Findings

Created a dataset of 47,315 CT scans with detailed annotations

Model exceeds human expert performance in tumor diagnosis (+7%)

Achieves significant improvements in tumor detection (+10%) and segmentation (+14%)

Abstract

Large, high-quality, annotated datasets are the foundation of medical AI research, but constructing even a small, moderate-quality, annotated dataset can take years of effort from multidisciplinary teams. Although active learning can prioritize what to annotate, scaling up still requires extensive manual efforts to revise the noisy annotations. We formulate this as a missing-data problem and develop ScaleMAI, a framework that unifies data annotation and model development co-evolution through an Expectation-Maximization (EM) process. In this iterative process, the AI model automatically identifies and corrects the mistakes in annotations (Expectation), while the refined annotated data retrain the model to improve accuracy (Maximization). In addition to the classical EM algorithm, ScaleMAI brings human experts into the loop to review annotations that cannot be adequately addressed by either Expectation or Maximization step (<5%). As a result, ScaleMAI progressively creates an annotated dataset of 47,315 CT scans (4.8x larger than the largest public dataset, PanTS) including 4,163,720 per-voxel annotations for benign/malignant tumors and 88 anatomical structures. ScaleMAI iteratively trains a model that exceeds human expert performance in tumor diagnosis (+7%), and outperforms models developed from smaller, moderate-quality datasets, with statistically significant gains in tumor detection (+10%) and segmentation (+14%) on two prestigious benchmarks.