Ambiguity-aware Truncated Flow Matching for Ambiguous Medical Image Segmentation

2025-11-10

AI Summary

In ambiguous medical image segmentation (AMIS), an inherent trade-off exists between prediction accuracy and solution diversity. To address this, we propose a data-hierarchical inference paradigm—first achieving effective decoupling of accuracy and diversity within truncated diffusion probabilistic models (TDPMs). Our method integrates Gaussian truncation representation, semantic-aware flow transformation, and segmentation-flow matching, thereby enhancing the reliability of truncated distributions and the plausibility of multi-solution outputs. Evaluated on LIDC and ISIC-3, our approach improves GED by 12% and HM-IoU by 7.3% over state-of-the-art methods, while reducing inference overhead. Key contributions include: (1) the first data-hierarchical inference framework for AMIS; (2) a segmentation-flow matching formulation enabling joint optimization of fidelity and diversity; and (3) empirical validation that Gaussian truncation representation excels at modeling ambiguous boundaries.

Technology Category

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

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Security and Privacy: Data transparency and provenance
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

Abstract

A simultaneous enhancement of accuracy and diversity of predictions remains a challenge in ambiguous medical image segmentation (AMIS) due to the inherent trade-offs. While truncated diffusion probabilistic models (TDPMs) hold strong potential with a paradigm optimization, existing TDPMs suffer from entangled accuracy and diversity of predictions with insufficient fidelity and plausibility. To address the aforementioned challenges, we propose Ambiguity-aware Truncated Flow Matching (ATFM), which introduces a novel inference paradigm and dedicated model components. Firstly, we propose Data-Hierarchical Inference, a redefinition of AMIS-specific inference paradigm, which enhances accuracy and diversity at data-distribution and data-sample level, respectively, for an effective disentanglement. Secondly, Gaussian Truncation Representation (GTR) is introduced to enhance both fidelity of predictions and reliability of truncation distribution, by explicitly modeling it as a Gaussian distribution at $T_{ ext{trunc}}$ instead of using sampling-based approximations.Thirdly, Segmentation Flow Matching (SFM) is proposed to enhance the plausibility of diverse predictions by extending semantic-aware flow transformation in Flow Matching (FM). Comprehensive evaluations on LIDC and ISIC3 datasets demonstrate that ATFM outperforms SOTA methods and simultaneously achieves a more efficient inference. ATFM improves GED and HM-IoU by up to $12%$ and $7.3%$ compared to advanced methods.

Problem

• Enhancing accuracy and diversity simultaneously in ambiguous medical image segmentation
• Addressing entangled prediction quality and insufficient fidelity in diffusion models
• Improving plausibility of diverse predictions while maintaining efficient inference

Innovation

• Data-Hierarchical Inference enhances accuracy and diversity
• Gaussian Truncation Representation improves fidelity and reliability
• Segmentation Flow Matching enhances plausibility of predictions