[Paper Review] Diff-UNet: A Diffusion Embedded Network for Volumetric Segmentation

Zhaohu Xing, Liang Wan+3 more|arXiv (Cornell University)|Mar 18, 2023

Radiomics and Machine Learning in Medical Imaging41 citations

TL;DR

Diff-UNet integrates a diffusion model into a 3D U-Net for volumetric medical image segmentation and uses Step-Uncertainty based Fusion (SUF) to combine multi-step predictions for robust results across MRI and CT datasets.

ABSTRACT

In recent years, Denoising Diffusion Models have demonstrated remarkable success in generating semantically valuable pixel-wise representations for image generative modeling. In this study, we propose a novel end-to-end framework, called Diff-UNet, for medical volumetric segmentation. Our approach integrates the diffusion model into a standard U-shaped architecture to extract semantic information from the input volume effectively, resulting in excellent pixel-level representations for medical volumetric segmentation. To enhance the robustness of the diffusion model's prediction results, we also introduce a Step-Uncertainty based Fusion (SUF) module during inference to combine the outputs of the diffusion models at each step. We evaluate our method on three datasets, including multimodal brain tumors in MRI, liver tumors, and multi-organ CT volumes, and demonstrate that Diff-UNet outperforms other state-of-the-art methods significantly. Our experimental results also indicate the universality and effectiveness of the proposed model. The proposed framework has the potential to facilitate the accurate diagnosis and treatment of medical conditions by enabling more precise segmentation of anatomical structures. The codes of Diff-UNet are available at https://github.com/ge-xing/Diff-UNet

Motivation & Objective

• Motivate robust 3D medical volumetric segmentation leveraging denoising diffusion models.
• Propose a diffusion-based end-to-end framework that handles multi-label segmentation via Label Embedding.
• Introduce a Denoising Module with a Denoising-UNet and a separate Feature Encoder to extract semantic information.
• Develop a Step-Uncertainty based Fusion (SUF) module to fuse predictions across diffusion steps during testing.
• Demonstrate superiority over state-of-the-art methods on BraTS2020, MSD Liver, and BTCV datasets.

Proposed method

• Convert multi-class segmentation into a one-hot encoded multi-channel label map to enable diffusion-based prediction.
• Use a Denoising Module consisting of a Feature Encoder and a Denoising-UNet that fuses volume features with the noisy one-hot label for x0 prediction.
• Train with a composite loss combining Dice, BCE, and MSE losses.
• During testing, generate multiple DDIM steps and fuse predictions using a Step-Uncertainty based Fusion (SUF) weighting scheme that accounts for step index and predicted uncertainty.
• Uncertainty is estimated via multiple forward passes per step (MC-like sampling) and used to weight step predictions in the final segmentation output.

Experimental results

Research questions

• RQ1 Can a diffusion model be effectively embedded into a 3D U-Net architecture to perform multi-label medical volumetric segmentation?
• RQ2 Does a Step-Uncertainty based Fusion mechanism improve robustness of diffusion-based segmentation across challenging datasets?
• RQ3 How does Diff-UNet perform on multi-organ CT and multimodal MRI segmentation benchmarks compared to state-of-the-art methods?

Key findings

• Diff-UNet achieves superior average Dice scores on BraTS2020 (85.35% average) and improved TC/WT/ET regions compared to several SOTA methods.
• On MSD Liver, Diff-UNet achieves 95.72% Dice on Liver with 0.222 HD95, and 51.65% Dice on Liver Tumor with 17.280 HD95, averaging 73.69 Dice and 8.751 HD95 across regions.
• On BTCV, Diff-UNet achieves the best average Dice (83.75) with HD95 of 8.115 across eight organs.
• Ablation studies show that incorporating a separate Feature Encoder (FE) and the SUF module improves performance, with optimal results at S=4 for uncertainty sampling.
• Qualitative results indicate Diff-UNet produces more accurate segmentation, especially for small targets, than competing methods.