SDSeg

Abstract

Diffusion models have demonstrated their effectiveness across various generative tasks. However, when applied to medical image segmentation, these models encounter several challenges, including significant resource and time requirements. They also necessitate a multi-step reverse process and multiple samples to produce reliable predictions. To address these challenges, we introduce the first latent diffusion segmentation model, named SDSeg, built upon stable diffusion (SD). SDSeg incorporates a straightforward latent estimation strategy to facilitate a single-step reverse process and utilizes latent fusion concatenation to remove the necessity for multiple samples. Extensive experiments indicate that SDSeg surpasses existing state-of-the-art methods on five benchmark datasets featuring diverse imaging modalities. Remarkably, SDSeg is capable of generating stable predictions with a solitary reverse step and sample, epitomizing the model's stability as implied by its name. The code has been anonymized and is available at Anonymous Code.

Methods

Comparison of DDIM convergence speed with and without latent estimation loss on BTCV.

Illustration of our Stability Evaluation on REF. To prepare for the evaluation, we first conduct M times of inference process. Then, Dataset-level Stability (DS)(in red) is evaluated on every two sets of the inference results; Instance-level Stability (IS)(in blue) is estimated on every two segmentation maps of each image conditioning. In practice, we choose the first inference results to conduct the DS evaluation with the other M-1 sets, and the first instance of segmentation result to conduct IS evaluation with the other M-1 segmentation maps.

Visualization of the latent representations of medical images from the trainable vision encoder, on CVC. At iteration 0, the encoder pre-trained on natural images couldn't capture enough meaningful semantic features for segmentation. During training, the conditioning encoder gradually learns to focus on segmentation targets.

Visualization of the predicted probability maps in reverse process on STS. The horizontal axis denotes DDIM sampling steps. Upper two rows: SDSeg trained for only 5,000 steps, sub-optimal results appear in early reverse process; Lower two rows: SDSeg trained for 50,000 steps, DDIM sampler generates fine and stable results during the entire reverse process. This demonstrates that SDSeg can generate great results under limited steps of the reverse process.

Other Results

Comparison of DDIM convergence speed with and without latent estimation loss on BTCV.

Illustration of our Stability Evaluation on REF. To prepare for the evaluation, we first conduct M times of inference process. Then, Dataset-level Stability (DS)(in red) is evaluated on every two sets of the inference results; Instance-level Stability (IS)(in blue) is estimated on every two segmentation maps of each image conditioning. In practice, we choose the first inference results to conduct the DS evaluation with the other M-1 sets, and the first instance of segmentation result to conduct IS evaluation with the other M-1 segmentation maps.

Visualization of the latent representations of medical images from the trainable vision encoder, on CVC. At iteration 0, the encoder pre-trained on natural images couldn't capture enough meaningful semantic features for segmentation. During training, the conditioning encoder gradually learns to focus on segmentation targets.

Visualization of the predicted probability maps in reverse process on STS. The horizontal axis denotes DDIM sampling steps. Upper two rows: SDSeg trained for only 5,000 steps, sub-optimal results appear in early reverse process; Lower two rows: SDSeg trained for 50,000 steps, DDIM sampler generates fine and stable results during the entire reverse process. This demonstrates that SDSeg can generate great results under limited steps of the reverse process.

Stable Diffusion Segmentation for Biomedical Images with Single-step Reverse Process

The overview of SDSeg.

Abstract

Methods

Comparison of DDIM convergence speed with and without latent estimation loss on BTCV.

Other Results

Comparison of DDIM convergence speed with and without latent estimation loss on BTCV.

BibTex