UniSegDiff: Boosting Unified Lesion Segmentation via a Staged Diffusion Model

TL;DR

UniSegDiff introduces a staged diffusion model framework that enhances lesion segmentation across multiple organs and modalities by maintaining high attention throughout training, outperforming previous methods.

Contribution

It proposes a novel staged training and inference strategy for diffusion models to improve unified lesion segmentation across diverse medical imaging modalities.

Findings

Significantly outperforms previous state-of-the-art methods.

Effective across six different organs and multiple imaging modalities.

Maintains high attention across all timesteps during training.

Abstract

The Diffusion Probabilistic Model (DPM) has demonstrated remarkable performance across a variety of generative tasks. The inherent randomness in diffusion models helps address issues such as blurring at the edges of medical images and labels, positioning Diffusion Probabilistic Models (DPMs) as a promising approach for lesion segmentation. However, we find that the current training and inference strategies of diffusion models result in an uneven distribution of attention across different timesteps, leading to longer training times and suboptimal solutions. To this end, we propose UniSegDiff, a novel diffusion model framework designed to address lesion segmentation in a unified manner across multiple modalities and organs. This framework introduces a staged training and inference approach, dynamically adjusting the prediction targets at different stages, forcing the model to maintain high attention across all timesteps, and achieves unified lesion segmentation through pre-training the feature extraction network for segmentation. We evaluate performance on six different organs across various imaging modalities. Comprehensive experimental results demonstrate that UniSegDiff significantly outperforms previous state-of-the-art (SOTA) approaches. The code is available at https://github.com/HUYILONG-Z/UniSegDiff.