Latent Diffusion for Medical Image Segmentation: End to end learning for fast sampling and accuracy

TL;DR

This paper introduces LDSeg, a latent diffusion model for medical image segmentation that achieves faster inference, lower memory use, and higher accuracy by operating in a learned latent shape space.

Contribution

The work presents a novel end-to-end latent diffusion framework for medical segmentation, addressing efficiency and noise robustness issues of traditional diffusion models.

Findings

Achieved state-of-the-art accuracy on three medical datasets.

Significantly faster sampling compared to traditional diffusion models.

Enhanced robustness to noise in segmentation results.

Abstract

Diffusion Probabilistic Models (DPMs) suffer from inefficient inference due to their slow sampling and high memory consumption, which limits their applicability to various medical imaging applications. In this work, we propose a novel conditional diffusion modeling framework (LDSeg) for medical image segmentation, utilizing the learned inherent low-dimensional latent shape manifolds of the target objects and the embeddings of the source image with an end-to-end framework. Conditional diffusion in latent space not only ensures accurate image segmentation for multiple interacting objects, but also tackles the fundamental issues of traditional DPM-based segmentation methods: (1) high memory consumption, (2) time-consuming sampling process, and (3) unnatural noise injection in the forward and reverse processes. The end-to-end training strategy enables robust representation learning in the latent space related to segmentation features, ensuring significantly faster sampling from the posterior distribution for segmentation generation in the inference phase. Our experiments demonstrate that LDSeg achieved state-of-the-art segmentation accuracy on three medical image datasets with different imaging modalities. In addition, we showed that our proposed model was significantly more robust to noise compared to traditional deterministic segmentation models. The code is available at https://github.com/FahimZaman/LDSeg.git.