# Ambiguity-aware Truncated Flow Matching for Ambiguous Medical Image Segmentation

arXiv: https://arxiv.org/abs/2511.06857 · HTML: https://mlsanity.com/paper/2511.06857/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Medical imaging and segmentation, one of 14 directions in Diffusion models.
- Builds on 15 of its 27 references that are in the map.
- First posted on arXiv 10 Nov 2025.

## Placed in

- Diffusion models / [Medical imaging and segmentation](https://mlsanity.com/diffusion/medical-imaging-and-segmentation/)

## Builds on (most cited first, 15 of 15)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Improved Techniques for Training Score-Based Generative Models](https://mlsanity.com/paper/2006.09011/) (2020)
- [Skin Lesion Analysis Toward Melanoma Detection 2018: A Challenge Hosted by the International Skin Imaging Collaboration (ISIC)](https://mlsanity.com/paper/1902.03368/) (2019)
- [Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning](https://mlsanity.com/paper/2208.04202/) (2022)
- [A Probabilistic U-Net for Segmentation of Ambiguous Images](https://mlsanity.com/paper/1806.05034/) (2018)
- [Simplifying, Stabilizing and Scaling Continuous-Time Consistency Models](https://mlsanity.com/paper/2410.11081/) (2024)
- [The Cramer Distance as a Solution to Biased Wasserstein Gradients](https://mlsanity.com/paper/1705.10743/) (2017)
- [Stochastic Segmentation Networks: Modelling Spatially Correlated Aleatoric Uncertainty](https://mlsanity.com/paper/2006.06015/) (2020)
- [A Hierarchical Probabilistic U-Net for Modeling Multi-Scale Ambiguities](https://mlsanity.com/paper/1905.13077/) (2019)
- [Stochastic Segmentation with Conditional Categorical Diffusion Models](https://mlsanity.com/paper/2303.08888/) (2023)
- [That Label's Got Style: Handling Label Style Bias for Uncertain Image Segmentation](https://mlsanity.com/paper/2303.15850/) (2023)
- [Finding Local Diffusion Schr"odinger Bridge using Kolmogorov-Arnold Network](https://mlsanity.com/paper/2502.19754/) (2025)
- [Hypernet-Ensemble Learning of Segmentation Probability for Medical Image Segmentation with Ambiguous Labels](https://mlsanity.com/paper/2112.06693/) (2021)

## Built on by (most cited first, 1 of 1)

- [KANResDiff: Learning Local Residual Diffusion via Kolmogorov-Arnold Network for Ambiguous Medical Image Segmentation](https://mlsanity.com/paper/2608.11617/) (2026)

## Abstract (arXiv)

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 Ttrunc 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.