Knowledge Distillation to Ensemble Global and Interpretable Prototype-Based Mammogram Classification Models
This addresses the need for interpretable AI in clinical mammography, though it is incremental by combining existing methods.
The paper tackles the problem of limited interpretability in state-of-the-art deep learning mammogram classifiers by proposing BRAIxProtoPNet++, which ensembles a global model with a prototype-based model to improve accuracy and interpretability. Experiments on weakly-labelled datasets show it achieves higher classification accuracy than SOTA models and better interpretability and prototype diversity.
State-of-the-art (SOTA) deep learning mammogram classifiers, trained with weakly-labelled images, often rely on global models that produce predictions with limited interpretability, which is a key barrier to their successful translation into clinical practice. On the other hand, prototype-based models improve interpretability by associating predictions with training image prototypes, but they are less accurate than global models and their prototypes tend to have poor diversity. We address these two issues with the proposal of BRAIxProtoPNet++, which adds interpretability to a global model by ensembling it with a prototype-based model. BRAIxProtoPNet++ distills the knowledge of the global model when training the prototype-based model with the goal of increasing the classification accuracy of the ensemble. Moreover, we propose an approach to increase prototype diversity by guaranteeing that all prototypes are associated with different training images. Experiments on weakly-labelled private and public datasets show that BRAIxProtoPNet++ has higher classification accuracy than SOTA global and prototype-based models. Using lesion localisation to assess model interpretability, we show BRAIxProtoPNet++ is more effective than other prototype-based models and post-hoc explanation of global models. Finally, we show that the diversity of the prototypes learned by BRAIxProtoPNet++ is superior to SOTA prototype-based approaches.