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

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.
Comment: 14 pages, 6 figures, 5 tables