Diffusion Model in Latent Space for Medical Image Segmentation Task

TL;DR

Medical image segmentation often yields a single mask, which fails to reflect clinical uncertainty. MedSegLatDiff introduces a latent-space diffusion approach operating on VQ-VAE encodings of images and masks to produce multiple, image-conditioned segmentation masks accompanied by confidence maps. The method achieves state-of-the-art or competitive Dice and IoU on ISIC-2018, CVC-Clinic, and LIDC-IDRI while providing richer uncertainty information through an explicit multi-hypothesis framework. This latent-space diffusion design offers improved interpretability and potential for robust clinical deployment due to efficiency and principled uncertainty representation.

Abstract

Medical image segmentation is crucial for clinical diagnosis and treatment planning. Traditional methods typically produce a single segmentation mask, failing to capture inherent uncertainty. Recent generative models enable the creation of multiple plausible masks per image, mimicking the collaborative interpretation of several clinicians. However, these approaches remain computationally heavy. We propose MedSegLatDiff, a diffusion based framework that combines a variational autoencoder (VAE) with a latent diffusion model for efficient medical image segmentation. The VAE compresses the input into a low dimensional latent space, reducing noise and accelerating training, while the diffusion process operates directly in this compact representation. We further replace the conventional MSE loss with weighted cross entropy in the VAE mask reconstruction path to better preserve tiny structures such as small nodules. MedSegLatDiff is evaluated on ISIC-2018 (skin lesions), CVC-Clinic (polyps), and LIDC-IDRI (lung nodules). It achieves state of the art or highly competitive Dice and IoU scores while simultaneously generating diverse segmentation hypotheses and confidence maps. This provides enhanced interpretability and reliability compared to deterministic baselines, making the model particularly suitable for clinical deployment.

Figures (4)

• Figure 1: The latent diffusion forward process, where noise is progressively added across diffusion steps.
• Figure 2: The latent diffusion reverse process, where multiple segmentation masks are generated for a single input medical image.
• Figure 3: Illustration of our pipeline: two initial mask samples, their averaged output, and the confidence map.
• Figure 4: Dice results as the number of sampling steps increase.