Improving Robustness and Reliability in Medical Image Classification with Latent-Guided Diffusion and Nested-Ensembles

Xing Shen,Hengguan Huang,Brennan Nichyporuk,Tal Arbel
2024-09-25
Abstract:Ensemble deep learning has been shown to achieve high predictive accuracy and uncertainty estimation in a wide variety of medical imaging contexts. However, perturbations in the input images at test time (e.g. noise, domain shifts) can still lead to significant performance degradation, posing challenges for trustworthy clinical deployment. In order to address this, we propose LaDiNE, a novel and robust probabilistic method that is capable of inferring informative and invariant latent variables from the input images. These latent variables are then used to recover the robust predictive distribution without relying on a predefined functional-form. This results in improved (i) generalization capabilities and (ii) calibration of prediction confidence. Extensive experiments were performed on the task of disease classification based on the Tuberculosis chest X-ray and the ISIC Melanoma skin cancer datasets. Here the performance of LaDiNE was analysed under a range of challenging covariate shift conditions, where training was based on "clean" images, and unseen noisy inputs and adversarial perturbations were presented at test time. Results show that LaDiNE outperforms existing state-of-the-art baseline methods in terms of accuracy and confidence calibration. This increases the feasibility of deploying reliable medical machine learning models in real clinical settings, where accurate and trustworthy predictions are crucial for patient care and clinical decision support.
Machine Learning,Computer Vision and Pattern Recognition
What problem does this paper attempt to address?
The paper aims to address the issues of robustness and reliability in medical image classification. Specifically, it focuses on the significant performance degradation of existing deep learning models when input images are perturbed during testing (e.g., noise, domain shifts, etc.). This performance drop poses a challenge for clinical applications, as accurate and reliable predictions are crucial for patient care and clinical decision support. To tackle this problem, the authors propose LaDiNE (Latent-guided Diffusion Nested-Ensembles), a novel and robust probabilistic approach that infers informative and invariant latent variables from input images. These latent variables are then used to recover robust predictive distributions without relying on predefined functional forms. This approach enhances the model's generalization ability and the calibration of predictive confidence. The paper demonstrates the performance of LaDiNE under various challenging covariate shift conditions through extensive experiments on tuberculosis chest X-rays and ISIC melanoma skin cancer datasets. The results show that LaDiNE outperforms existing state-of-the-art baseline methods in terms of accuracy and confidence calibration, thereby increasing the feasibility of deploying reliable medical machine learning models in real clinical settings.