SelectiveKD: A semi-supervised framework for cancer detection in DBT through Knowledge Distillation and Pseudo-labeling

TL;DR

SelectiveKD is a semi-supervised learning framework that leverages unlabeled DBT slices through knowledge distillation and pseudo-labeling to improve cancer detection performance with limited annotations.

Contribution

The paper introduces SelectiveKD, a novel semi-supervised framework that enhances DBT cancer detection by effectively utilizing unlabeled data and mitigating noise via selective pseudo-labeling.

Findings

Significant improvement in cancer classification AUC.

Effective utilization of unlabeled slices from DBT stacks.

Enhanced generalization across multiple datasets.

Abstract

When developing Computer Aided Detection (CAD) systems for Digital Breast Tomosynthesis (DBT), the complexity arising from the volumetric nature of the modality poses significant technical challenges for obtaining large-scale accurate annotations. Without access to large-scale annotations, the resulting model may not generalize to different domains. Given the costly nature of obtaining DBT annotations, how to effectively increase the amount of data used for training DBT CAD systems remains an open challenge. In this paper, we present SelectiveKD, a semi-supervised learning framework for building cancer detection models for DBT, which only requires a limited number of annotated slices to reach high performance. We achieve this by utilizing unlabeled slices available in a DBT stack through a knowledge distillation framework in which the teacher model provides a supervisory signal to the student model for all slices in the DBT volume. Our framework mitigates the potential noise in the supervisory signal from a sub-optimal teacher by implementing a selective dataset expansion strategy using pseudo labels. We evaluate our approach with a large-scale real-world dataset of over 10,000 DBT exams collected from multiple device manufacturers and locations. The resulting SelectiveKD process effectively utilizes unannotated slices from a DBT stack, leading to significantly improved cancer classification performance (AUC) and generalization performance.