Aleatoric Uncertainty Medical Image Segmentation Estimation via Flow Matching

TL;DR

This work targets aleatoric uncertainty in medical image segmentation caused by inter-expert variability. It proposes conditional flow matching, a flow-based generative framework that learns an exact density $q(S|X)$ by conditioning a velocity field $u_\theta(t,S,X)$ on the image $X$ and an expert segmentation $S^{(e)}$, and evolving samples from a simple prior $p_0(S)$ along the ODE $\frac{d}{dt}\psi_t(S_0)=u_\theta(t,\psi_t(S_0),X)$. The method aligns sampling with expert annotations through a conditional path $p_t(S|S^{(e)},X)=\mathcal{N}(S; tS^{(e)},(1-t)^2I)$ and optimizes a regression loss $\mathcal{L}_{\mathrm{CFM}}$, enhanced by classifier-free guidance to maintain both fidelity to anatomy and sample diversity. Empirically, it achieves competitive segmentation accuracy and superior uncertainty maps on LIDC-IDRI and MMIS, delivering multiple plausible segmentations that reflect annotator variability; code is released for public use. Overall, the approach avoids diffusion-based stochastic sampling, preserves fine local structure, and provides robust, clinically actionable uncertainty quantification for segmentation tasks.

Abstract

Quantifying aleatoric uncertainty in medical image segmentation is critical since it is a reflection of the natural variability observed among expert annotators. A conventional approach is to model the segmentation distribution using the generative model, but current methods limit the expression ability of generative models. While current diffusion-based approaches have demonstrated impressive performance in approximating the data distribution, their inherent stochastic sampling process and inability to model exact densities limit their effectiveness in accurately capturing uncertainty. In contrast, our proposed method leverages conditional flow matching, a simulation-free flow-based generative model that learns an exact density, to produce highly accurate segmentation results. By guiding the flow model on the input image and sampling multiple data points, our approach synthesizes segmentation samples whose pixel-wise variance reliably reflects the underlying data distribution. This sampling strategy captures uncertainties in regions with ambiguous boundaries, offering robust quantification that mirrors inter-annotator differences. Experimental results demonstrate that our method not only achieves competitive segmentation accuracy but also generates uncertainty maps that provide deeper insights into the reliability of the segmentation outcomes. The code for this paper is freely available at https://github.com/huynhspm/Data-Uncertainty

Figures (3)

• Figure 1: Illustration of our proposed conditional flow matching framework for modeling multi-expert segmentation variability. At $t=0$, samples $S_0$ are drawn from a simple Gaussian prior $p_0(S)$. As $t$ progresses, the flow $\psi_t$ (green lines) evolves these samples toward the target distribution $q(S\mid X)$ at $t=1$. Unlike diffusion methods (red contours) that inject noise and risk blurring fine details, flow matching directly learns a velocity field $u_\theta(t,S,X)$ that preserves local structure. Multiple expert segmentations (blue, purple, and yellow boxes) are incorporated by conditioning on their annotation. The guide image $X$ (bottom) provides anatomical context to ensure samples $S^{(e)}$ align closely with the underlying anatomy.
• Figure 2: Comparative qualitative analysis of our proposed method against three baseline methods is presented. Three examples from the LIDC-IDRI dataset are shown, with the input images in the first row, followed by rows displaying four segmentation masks and uncertainty maps for the ground truth (GT), the three baseline methods, and our proposed method.
• Figure 3: Comparative qualitative analysis of our proposed method against three baseline methods is presented. Three examples from the MMIS dataset are shown, with the input triplet images in the first row, followed by rows displaying four segmentation masks and uncertainty maps for the ground truth (GT), the three baseline methods, and our proposed method.