Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Ambiguity-aware Truncated Flow Matching for Ambiguous Medical Image Segmentation

Fanding Li, Xiangyu Li, Xianghe Su, Xingyu Qiu, Suyu Dong, Wei Wang, Kuanquan Wang, Gongning Luo, Shuo Li

TL;DR

This work tackles ambiguous medical image segmentation by addressing the persistent trade-off between accuracy and diversity in predictions. It introduces Ambiguity-aware Truncated Flow Matching (ATFM), a TDPM-inspired framework with three innovations: Data-Hierarchical Inference to separate distribution-level accuracy from sample-level diversity, Gaussian Truncation Representation for faithful truncation modeling, and Segmentation Flow Matching to enforce semantic plausibility during flow-based prediction. Across LIDC-IDRI and ISIC3, ATFM yields superior GED, HM-IoU, and MDM scores while maintaining efficient inference. The results establish ATFM as a practical, high-performance AMIS solution with strong theoretical and architectural grounding.

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_{\text{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.

Ambiguity-aware Truncated Flow Matching for Ambiguous Medical Image Segmentation

TL;DR

This work tackles ambiguous medical image segmentation by addressing the persistent trade-off between accuracy and diversity in predictions. It introduces Ambiguity-aware Truncated Flow Matching (ATFM), a TDPM-inspired framework with three innovations: Data-Hierarchical Inference to separate distribution-level accuracy from sample-level diversity, Gaussian Truncation Representation for faithful truncation modeling, and Segmentation Flow Matching to enforce semantic plausibility during flow-based prediction. Across LIDC-IDRI and ISIC3, ATFM yields superior GED, HM-IoU, and MDM scores while maintaining efficient inference. The results establish ATFM as a practical, high-performance AMIS solution with strong theoretical and architectural grounding.

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 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 and compared to advanced methods.

Paper Structure

This paper contains 37 sections, 2 theorems, 16 equations, 10 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Ambiguous Medical Image Segmentation
  4. Truncated Diffusion Probabilistic Models
  5. Methods
  6. Data-Hierarchical Inference Enables Principled Disentanglement and Joint Enhancement of Accuracy and Diversity
  7. Gaussian Truncation Representation Improves Fidelity via Explicit Gaussian Modeling
  8. Segmentation Flow Matching Enhances Plausibility via Semantic Consistency Modeling
  9. Experiments
  10. Experimental Setup
  11. Datasets.
  12. Implementation Details.
  13. Evaluation Metrics.
  14. Experimental Results
  15. Comparison with SOTAs for Performance Superiority.
  16. ...and 22 more sections

Key Result

Theorem 1

The marginal distribution of the latent variable at any diffusion timestep $\tau$ can be parameterized as

Figures (10)

  • Figure 1: (a) Traditional TDPMs face the challenges of low fidelity and plausibility of predictions by improving diversity at the expense of accuracy. (b) The proposed ATFM enhances fidelity and plausibility of predictions by assigning distinct inference goals into two stages.
  • Figure 2: The proposed ATFM addresses the challenge of a synergistic optimization by boosting accuracy at distribution level and diversity at sample level within Data-Hierarchical Inference (Sec. 3.1) while enhancing fidelity and plausibility with GTR (Sec. 3.2) and SFM (Sec. 3.3), respectively.
  • Figure 3: Data-Hierarchical Inference disentangles accuracy and diversity by marginalizing stochasticity during diffusion with a data-distribution level supervision.
  • Figure 4: Comparative qualitative results on LIDC dataset among ground truths, two advanced methods and the proposed ATFM demonstrate both better alignment with ground truths and higher per-sample accuracy.
  • Figure 5: Comparative qualitative results on ISIC3 subset dataset among ground truths, two advanced methods and the proposed ATFM demonstrate both better alignment with ground truths and higher per-sample accuracy.
  • ...and 5 more figures

Theorems & Definitions (5)

  • Theorem 1
  • Theorem 2
  • proof
  • proof
  • proof