Ambiguous Medical Image Segmentation Using Diffusion Models

Collective insights from a group of experts have always proven to outperform an individual's best diagnostic for clinical tasks. For the task of medical image segmentation, existing research on AI-based alternatives focuses more on developing models that can imitate the best individual rather than harnessing the power of expert groups. In this paper, we introduce a single diffusion model-based approach that produces multiple plausible outputs by learning a distribution over group insights. Our proposed model generates a distribution of segmentation masks by leveraging the inherent stochastic sampling process of diffusion using only minimal additional learning. We demonstrate on three different medical image modalities- CT, ultrasound, and MRI that our model is capable of producing several possible variants while capturing the frequencies of their occurrences. Comprehensive results show that our proposed approach outperforms existing state-of-the-art ambiguous segmentation networks in terms of accuracy while preserving naturally occurring variation. We also propose a new metric to evaluate the diversity as well as the accuracy of segmentation predictions that aligns with the interest of clinical practice of collective insights. Implementation code: https://github.com/aimansnigdha/Ambiguous-Medical-Image-Segmentation-using-Diffusion-Models.

1. Introduction

Diagnosis is the central part of medicine, which heavily relies on the individual practitioner assessment strategy. Re-cent studies suggest that misdiagnosis with potential mor-tality and morbidity is widespread for even the most com-mon health conditions [32], [49]. Hence, reducing the frequency of misdiagnosis is a crucial step towards improving healthcare. Medical image segmentation, which is a cen-tral part of diagnosis, plays a crucial role in clinical out-comes. Deep learning-based networks for segmentation are now getting traction for assisting in clinical settings, how-ever, most of the leading segmentation networks in the lit-erature are deterministic [17], [23], [34], [36], [41], [42], [44], meaning they predict a single segmentation mask for each input image. Unlike natural images, ground truths are not deterministic in medical images as different diagnosticians can have different opinions on the type and extent of an anomaly [1], [1]5,[37], [39]. Due to this, the diagnosis from medical images is quite challenging and often results in a low inter-rater agreement [22], [24], [56]. Depending on only pixel-wise probabilities and ignoring co-variances between the pixels might lead to misdiagnosis. In clinical practice, aggregating interpretations of multiple experts have shown to improve diagnosis and generate fewer false negatives [57]. Figure 1.

A) deterministic networks produce a single output for an input image. B) c-vae-based methods encode prior information about the input image in a separate network and sample latent variables from there and inject it into the deterministic segmentation network to produce stochastic segmentation masks. C) in our method the diffusion model learns the latent structure of the seg-mentation as well as the ambiguity of the dataset by modeling the way input images are diffused through the latent space. Hence our method does not need an additional prior encoder to provide latent variables for multiple plausible annotations.