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

TL;DR

Mediffusion is a joint diffusion model for semi-supervised medical image classification that combines explainability and generative capabilities, effectively learning from limited labeled data while providing reliable explanations.

Contribution

We introduce Mediffusion, a unified model that integrates classification and generative tasks for explainable semi-supervised learning in medical imaging.

Findings

Achieves competitive classification performance with semi-supervised data.

Provides accurate and reliable explanations via counterfactual examples.

Effectively learns from both labeled and unlabeled data.

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.