Diffusion Based Ambiguous Image Segmentation

2025-04-08

AI Summary

Medical image segmentation suffers from inherent uncertainty in expert annotations, which conventional methods struggle to model as a full ground-truth distribution. To address this, we present the first systematic exploration of diffusion models in the discrete segmentation domain, proposing a generative segmentation framework based on x-/v-prediction. Our method introduces hardened noise scheduling, input-scale normalization, and a weighted reconstruction loss to significantly enhance uncertainty quantification. For robust evaluation, we construct a novel uncertainty-aware variant of LIDC-IDRI using stochastic cropping. Experiments demonstrate state-of-the-art performance in segmentation uncertainty modeling on both standard benchmarks and newly introduced challenging subsets. This work validates the effectiveness and generalization potential of diffusion models for structured, discrete-output tasks—particularly in medical imaging—where accurate probabilistic predictions are critical.

Technology Category

• Computer Vision: Segmentation
• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Uncertainty Representations

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Medical image segmentation often involves inherent uncertainty due to variations in expert annotations. Capturing this uncertainty is an important goal and previous works have used various generative image models for the purpose of representing the full distribution of plausible expert ground truths. In this work, we explore the design space of diffusion models for generative segmentation, investigating the impact of noise schedules, prediction types, and loss weightings. Notably, we find that making the noise schedule harder with input scaling significantly improves performance. We conclude that x- and v-prediction outperform epsilon-prediction, likely because the diffusion process is in the discrete segmentation domain. Many loss weightings achieve similar performance as long as they give enough weight to the end of the diffusion process. We base our experiments on the LIDC-IDRI lung lesion dataset and obtain state-of-the-art (SOTA) performance. Additionally, we introduce a randomly cropped variant of the LIDC-IDRI dataset that is better suited for uncertainty in image segmentation. Our model also achieves SOTA in this harder setting.

Problem

• Address uncertainty in medical image segmentation
• Optimize diffusion models for segmentation tasks
• Improve performance with noise schedule adjustments

Innovation

• Diffusion models for generative segmentation
• Hard noise schedule with input scaling
• X- and v-prediction outperform epsilon-prediction