HiDiff: Hybrid Diffusion Framework for Medical Image Segmentation

Tao Chen,Chenhui Wang,Zhihao Chen,Yiming Lei,Hongming Shan
2024-07-04
Abstract:Medical image segmentation has been significantly advanced with the rapid development of deep learning (DL) techniques. Existing DL-based segmentation models are typically discriminative; i.e., they aim to learn a mapping from the input image to segmentation masks. However, these discriminative methods neglect the underlying data distribution and intrinsic class characteristics, suffering from unstable feature space. In this work, we propose to complement discriminative segmentation methods with the knowledge of underlying data distribution from generative models. To that end, we propose a novel hybrid diffusion framework for medical image segmentation, termed HiDiff, which can synergize the strengths of existing discriminative segmentation models and new generative diffusion models. HiDiff comprises two key components: discriminative segmentor and diffusion refiner. First, we utilize any conventional trained segmentation models as discriminative segmentor, which can provide a segmentation mask prior for diffusion refiner. Second, we propose a novel binary Bernoulli diffusion model (BBDM) as the diffusion refiner, which can effectively, efficiently, and interactively refine the segmentation mask by modeling the underlying data distribution. Third, we train the segmentor and BBDM in an alternate-collaborative manner to mutually boost each other. Extensive experimental results on abdomen organ, brain tumor, polyps, and retinal vessels segmentation datasets, covering four widely-used modalities, demonstrate the superior performance of HiDiff over existing medical segmentation algorithms, including the state-of-the-art transformer- and diffusion-based ones. In addition, HiDiff excels at segmenting small objects and generalizing to new datasets. Source codes are made available at <a class="link-external link-https" href="https://github.com/takimailto/HiDiff" rel="external noopener nofollow">this https URL</a>.
Computer Vision and Pattern Recognition
What problem does this paper attempt to address?
The paper attempts to address the limitations of existing deep learning methods in medical image segmentation. Specifically, current deep learning-based segmentation models typically adopt a discriminative approach, aiming to learn the mapping from input images to segmentation masks. However, these discriminative methods overlook the underlying data distribution and intrinsic class characteristics, leading to an unstable feature space and difficulty in handling fuzzy boundaries and fine objects. To address these issues, the paper proposes a novel hybrid diffusion framework (HiDiff) that supplements discriminative segmentation methods by incorporating knowledge from generative models, thereby improving the effectiveness of medical image segmentation. ### Core Issues of the Paper 1. **Limitations of Existing Discriminative Models**: - **Ignoring Data Distribution**: Existing discriminative models mainly focus on learning the mapping from input images to segmentation masks but ignore the underlying data distribution and class characteristics. - **Unstable Feature Space**: By focusing only on decision boundaries, these models perform poorly when far from the decision boundary, making it difficult to handle fuzzy boundaries and fine objects. 2. **Advantages and Challenges of Generative Models**: - **Advantages**: Generative models can directly model the underlying data distribution, helping to alleviate the limitations of discriminative models. - **Challenges**: Generative models face issues of training instability and slow inference speed. ### Solution The paper proposes a hybrid diffusion framework (HiDiff) that includes two key components: 1. **Discriminative Segmenter**: Utilizes a pre-trained discriminative segmentation model to provide initial segmentation masks. 2. **Diffusion Refiner**: Proposes a novel Binary Bernoulli Diffusion Model (BBDM) to effectively and efficiently refine segmentation masks by modeling the underlying data distribution. ### Main Contributions 1. **Proposes a novel hybrid diffusion framework (HiDiff)** that combines the advantages of discriminative segmentation models and generative diffusion models. 2. **Introduces a novel Binary Bernoulli Diffusion Model (BBDM)** that can effectively and efficiently refine segmentation masks. 3. **Introduces an alternating collaborative training strategy**, enabling the discriminative segmenter and diffusion refiner to mutually enhance each other during training. 4. **Extensive Experimental Results**: Conducted experiments on multiple medical image segmentation datasets, showing that HiDiff outperforms existing medical segmentation algorithms, especially in segmenting small objects and generalizing to new datasets. ### Experimental Validation The paper conducts experiments on multiple benchmark tasks across four widely used modalities (CT, MRI, endoscopic images, and retinal images), including the segmentation of abdominal organs, brain tumors, polyps, and retinal vessels. The experimental results demonstrate that HiDiff significantly outperforms existing medical segmentation algorithms in terms of performance and excels in handling small objects and generalization capabilities. ### Conclusion By proposing the HiDiff framework, the paper effectively addresses the limitations of existing discriminative models in medical image segmentation. By incorporating the advantages of generative models, it improves segmentation accuracy and robustness.