Mediffusion: Joint Diffusion for Self-Explainable Semi-Supervised Classification and Medical Image Generation

TL;DR

This work introduces Mediffusion -- a new method for semi-supervised learning with explainable classification based on a joint diffusion model that combines standard classification with a diffusion-based generative task in a single shared parametrisation.

Abstract

We introduce Mediffusion -- a new method for semi-supervised learning with explainable classification based on a joint diffusion model. The medical imaging domain faces unique challenges due to scarce data labelling -- insufficient for standard training, and critical nature of the applications that require high performance, confidence, and explainability of the models. In this work, we propose to tackle those challenges with a single model that combines standard classification with a diffusion-based generative task in a single shared parametrisation. By sharing representations, our model effectively learns from both labeled and unlabeled data while at the same time providing accurate explanations through counterfactual examples. In our experiments, we show that our Mediffusion achieves results comparable to recent semi-supervised methods while providing more reliable and precise explanations.

Figures (9)

• Figure 1: Mediffusion training and capabilities. Our proposed method utilizes both labeled and unlabeled data samples to build joint representations suitable for generative and discriminative tasks. We evaluate our method in 3 tasks (1) semi-supervised classification (2) inherent visual explainability of classifier decision (3) synthesising new pseudo-labeled data samples.
• Figure 2: Joint training of diffusion model and classifier. Data representation $h_t$ extracted in the UNet-based architecture of our joint latent diffusion model is utilized as input to the classifier component.
• Figure 3: Counterfactual examples. We modify the images to reduce the probability of disease prediction across various classes. From top to bottom: Atelectasis, Cardiomegaly, Mass. The bounding boxes annotated by experts highlight regions that indicate respective diseases.
• Figure 4: Mean pixel differences (0-1) inside and outside the ground-truth bounding boxes when generating counterfactual examples. For all diseases, the majority of changes of our counterfactual examples generation method occur within the ground-truth boxes (green) assigned by the trained physicians.
• Figure 5: Examples of enforcing disease indicators. Visualisation of counterfactual images generated for healthy data samples that are modified towards cardiomegaly disease characterized by an enlarged heart.