Ambiguity-aware Truncated Flow Matching for Ambiguous Medical Image Segmentation

TL;DR

This paper introduces ATFM, a novel method for ambiguous medical image segmentation that improves accuracy, diversity, and plausibility of predictions through a new inference paradigm and model components, outperforming state-of-the-art methods.

Contribution

The paper proposes ATFM with Data-Hierarchical Inference, Gaussian Truncation Representation, and Segmentation Flow Matching, addressing accuracy-diversity entanglement and enhancing prediction fidelity and plausibility.

Findings

Outperforms SOTA methods on LIDC and ISIC3 datasets.

Improves GED and HM-IoU by up to 12% and 7.3%.

Achieves more efficient inference in ambiguous medical image segmentation.

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.