Self-Supervised Multiple Instance Learning for Acute Myeloid Leukemia Classification

TL;DR

This work tackles AML genetic subtype classification from blood smears under weak labeling by pre-training an MIL encoder with self-supervised learning (SSL) methods—SimCLR, SwAV, and DINO—and integrating it into an attention-based MIL classifier. The same unlabeled dataset drives both SSL pre-training and MIL training, enabling data-efficient, cost-effective analysis without extensive single-cell annotations. Results show SSL-pretrained encoders achieve performance comparable to fully supervised pretraining, with SimCLR often delivering the best classification and attention-based interpretability toward malignant cells. The approach demonstrates the practical potential of SSL in MIL for medical image analysis, advancing data-efficient, explainable AI for AML diagnosis.

Abstract

Automated disease diagnosis using medical image analysis relies on deep learning, often requiring large labeled datasets for supervised model training. Diseases like Acute Myeloid Leukemia (AML) pose challenges due to scarce and costly annotations on a single-cell level. Multiple Instance Learning (MIL) addresses weakly labeled scenarios but necessitates powerful encoders typically trained with labeled data. In this study, we explore Self-Supervised Learning (SSL) as a pre-training approach for MIL-based AML subtype classification from blood smears, removing the need for labeled data during encoder training. We investigate the three state-of-the-art SSL methods SimCLR, SwAV, and DINO, and compare their performance against supervised pre-training. Our findings show that SSL-pretrained encoders achieve comparable performance, showcasing the potential of SSL in MIL. This breakthrough offers a cost-effective and data-efficient solution, propelling the field of AI-based disease diagnosis.

Figures (5)

• Figure 1: Overview of our data-efficient MIL model for AML genetic subtypes classification. We pre-train an MIL encoder $f_{\theta}$ with one of the state-of-the-art SSL models, i.e., SimCLR, SwAV, and DINO. Then we embed the trained encoder in the attention MIL architecture. For both SSL and MIL the same training dataset is used.
• Figure 2: The confusion matrix presents the test fold results of the first run for each pre-training method. While different SSL pre-training methods lead to varying class-wise performance, overall SSL pre-training performs correspondingly to fully supervised pre-training
• Figure 3: Similar mean ROC curves for MIL model resulting from different pre-training methods. Differences in the AUC score for genetic subtype prediction is within the margin of error.
• Figure 4: Averaged attention values for all labeled instances per model. Most of cell types with higher attention are malignant, while most lower attention cell types a healthy. SimCLR and DINO pretrained models provide more robust attention scores for malignant cells compared to the supervised pre-training.
• Figure 5: UMAP embedding visualization of encoder features with expert cytologist-assigned cell labels. SSL pre-trained encoder could distinguish cell-types without any label similar to fully supervised pre-trained encoder.