Ambiguity-aware Truncated Flow Matching for Ambiguous Medical Image Segmentation

AI summary

Overview

Research area: Computer vision, specifically ambiguous medical image segmentation (AMIS), combining truncated diffusion probabilistic models (TDPMs) with Flow Matching.

Technical level: Intermediate. The paper assumes familiarity with diffusion models, denoising timesteps, and segmentation metrics, though the three proposed components are described at a conceptual level.

Scope: The paper proposes Ambiguity-aware Truncated Flow Matching (ATFM), a framework that redefines the inference paradigm of truncated diffusion for medical segmentation so that prediction accuracy and prediction diversity can be improved at the same time, and validates it on two public medical datasets.

What This Paper Is About

In ambiguous medical image segmentation, different experts label the same image differently, so a good model must output a set of plausible predictions that is both accurate and diverse. Existing methods treat these two goals as a trade-off: pushing diversity usually costs accuracy, and vice versa. The paper's goal is to disentangle the two objectives by assigning them to different stages and different data levels of a truncated diffusion-style inference process, while also fixing the fidelity and plausibility problems that make vanilla truncated diffusion models underperform on this task.

Key Contributions

1. Data-Hierarchical Inference for AMIS. A redefinition of the AMIS inference paradigm that marginalizes stochasticity during truncation: the stage before the truncation point is supervised at the data-distribution level for accuracy, while the stage after the truncation point operates at the data-sample level to promote diversity without sacrificing accuracy.

2. Gaussian Truncation Representation (GTR). Instead of approximating the distribution at the truncation point by sampling (as vanilla TDPMs do with adversarial sampling), GTR explicitly parameterizes the truncated latent as a Gaussian with a learned mean and covariance, which the paper argues preserves low-frequency (rare but clinically plausible) modes and improves consistency across predictions.

3. Segmentation Flow Matching (SFM). A semantic-aware flow transformation that adds supervision of semantic consistency among intermediate states, predictions, and ground truths at each timestep after truncation, using a Semantic-aware Transformation Network (ST-Net), and that avoids the Gaussian constraints of DDPM that the paper says disturb fine-grained predictions.

4. Comprehensive evaluation. Experiments on the LIDC-IDRI and ISIC3 subset datasets showing that ATFM outperforms state-of-the-art methods across the reported metrics while using a more efficient inference process.

Main Findings

• LIDC results (Table 1). ATFM records GED_16 of 0.206 ± 0.002, GED_32 of 0.188 ± 0.001, GED_100 of 0.162 ± 0.002, HM-IoU_32 of 0.667 ± 0.002, and MDM_32 of 0.948 ± 0.001. The paper reports this as an 11.5% improvement on GED_100 on LIDC over the runner-up method and a minimum of 7.3% improvement in HM-IoU_32.

• ISIC3 results (Table 2). ATFM records GED_16 of 0.183 ± 0.001, GED_32 of 0.152 ± 0.002, GED_100 of 0.147 ± 0.003, HM-IoU_32 of 0.732 ± 0.003, and MDM_32 of 0.942 ± 0.002. The paper states ATFM outperforms the runner-up by 12% in GED on this dataset.

• Abstract-level claim. The abstract states that ATFM improves GED and HM-IoU by up to 12% and 7.3% compared to advanced methods.

• Inference efficiency (Table 3). For generating 100 samples on LIDC, ATFM uses GTR plus a truncation of 25 steps and takes approximately 113s. For comparison, the paper reports CIMD with T = 100 steps at approximately 420s, AB with T = 250 steps at approximately 1050s, and CCDM with T = 250 steps at approximately 1100s.

• Ablation study (Table 4, LIDC). Activation-based GTR gives GED_100 of 0.230 ± 0.001 and HM-IoU_32 of 0.550 ± 0.010; SFM without L_SF gives 0.185 ± 0.002 and 0.624 ± 0.002; SFM gives 0.176 ± 0.002 and 0.631 ± 0.002; full ATFM without L_SF gives 0.249 ± 0.001 and 0.597 ± 0.003; full ATFM gives 0.162 ± 0.002 and 0.667 ± 0.002. The paper reports that ATFM outperformed both Act. GTR and SFM by a minimum of 10% and 6% on both metrics, and that the average gap between models with and without L_SF is 11%.

• Hyper-parameter behavior. Five values between 1 and 50 were tested for the SFM inference step count, with optimal performance achieved within 25 steps; α in L_SF was tested between 0 and 5 × 10⁻³, with too-small values limiting L_SF's impact and too-large values diminishing the effect of L_FM.

• Qualitative findings. Figures 4 and 5 report that ATFM predictions better reflect the range of plausible annotations and preserve fine-grained structures compared with two advanced methods on LIDC and ISIC3 respectively.

• Not reported. The provided content does not report dataset-level uncertainty beyond mean ± standard deviation over five independent runs, and it does not report a dedicated limitations section or a quantitative failure analysis.

Methodology in Plain English

The starting point is the truncated diffusion probabilistic model: rather than running a full denoising chain from noise, a TDPM stops at a truncation timestep T_trunc ≪ T and uses a separate module to estimate the distribution at that point, which shortens inference. ATFM keeps that acceleration but changes what each stage is asked to do.

• Stage structure (Data-Hierarchical Inference). Before the truncation point, the model is not asked to produce varied samples; stochasticity is marginalized so the stage is supervised purely on how well the modeled distribution matches the ground-truth distribution. After the truncation point, samples are drawn from that distribution, and diversity is encouraged at the individual-sample level. The two goals therefore sit in different stages instead of competing inside one objective.

• Modeling the truncation point (GTR). A segmentation backbone f_θ produces a feature Z from the image X; two separate convolutional heads g_φ and h_ψ turn Z into a mean μ and covariance Σ, and the truncation latent is drawn from a Gaussian N(μ, Σ). The paper proves (Theorems 1 and 2) that a Gaussian of this form Σ = DD^T + L can represent the marginal distribution at any timestep, and that any Gaussian can be matched at some specific timestep τ*, which justifies using a Gaussian as the truncation distribution. This component is trained with a prior loss L_Prior, approximated by a Monte Carlo average with M samples, and then frozen for later inference.

• Generating the predictions (SFM). After truncation, the model uses Flow Matching with an Optimal Transformation schedule, so the trajectory between the truncation latent and the target is treated as a straight line in latent space. At each timestep t, the intermediate state is obtained by linear interpolation between the truncation endpoint and the target endpoint, and a predicted segmentation x_1t is derived by projecting along the trajectory. The Dice loss is computed between x_1t and every ground-truth annotation and added to the flow-matching loss (L_SF = L_FM + averaged α-weighted Dice terms), so the transformation is pushed to stay semantically consistent at every timestep.

• Training details. Everything runs on a single RTX 3090 GPU with 24GB. GTR is pretrained for 1000 epochs and SFM for 200 epochs on LIDC; on ISIC3, GTR is trained for 400 epochs and SFM for 120 epochs. λ = 10⁻³ (i.e., T = 1000) with a linear schedule; Adam optimizer with learning rate 10⁻⁴; M = 20; α is 10⁻³ for LIDC and 10⁻⁴ for ISIC3.

• Evaluation. GED measures alignment between the distribution of predictions and ground truths, HM-IoU measures fidelity of the samples under Hungarian matching, and MDM measures the best Dice score between each prediction and each ground truth. Metrics are written with a subscript n denoting the number of samples used, and results are averaged over five independent runs on LIDC-IDRI (15,096 slices, four labels each) and the ISIC3 subset (300 images, three annotations each).

Why This Matters

Impact on research. The paper reframes the accuracy-versus-diversity trade-off as an inference-paradigm problem rather than a loss-weighting problem, and it shows that the truncation point of a diffusion model can be explicitly modeled as a distribution rather than approximated by sampling. It also extends Flow Matching into a segmentation setting with per-timestep semantic supervision, and it reports these gains together with substantially reduced sampling time (about 113s versus approximately 420s, 1050s and 1100s for the compared diffusion baselines on LIDC for 100 samples), which makes the approach practical for further study.

Real-world applications

• Lung CT lesion analysis, where the model would produce a set of plausible lesion boundaries consistent with multiple expert readings rather than a single possibly misleading mask.
• Dermoscopic skin lesion segmentation, where clinician boundaries differ and capturing the range of plausible contours supports downstream assessment.
• Uncertainty-aware clinical decision support, where the range of plausible masks reflects genuine diagnostic ambiguity.
• Dataset annotation and review workflows, where plausible alternative segmentations can be surfaced to annotators or reviewers.

Industry relevance. Any pipeline that depends on segmentation masks for high-stakes decisions benefits from predictions that are not only accurate on average but also honestly represent variation. The reported inference speed matters for deployment because it lowers the compute cost of generating many samples per image.

Future Directions

• Evaluating ATFM beyond the two datasets used here (LIDC-IDRI and the ISIC3 subset) to see whether the Data-Hierarchical Inference paradigm transfers to other modalities and other ambiguity-heavy segmentation tasks.
• Understanding and tuning the balance carried by α in L_SF and by the number of inference steps more systematically, since the paper reports an optimum near 25 steps and an α range of 0 to 5 × 10⁻³.
• Investigating whether the Gaussian parameterization in GTR is sufficient for annotation distributions that are strongly non-Gaussian or multi-modal, given that the proof only guarantees Gaussian matchability at some timestep.
• Extending the analysis of where accuracy and diversity still interact after disentanglement, since the paper reports the two objectives as separated by design but does not provide a dedicated limitations study.

Target Audience

Researchers and graduate students working on medical image segmentation, probabilistic or diffusion-based generative models, and uncertainty-aware prediction, as well as applied engineers who need to generate multiple plausible segmentations at practical inference cost. Readers with a background in diffusion models or Flow Matching will get the most from the methodological sections; the experimental tables are readable without specialist knowledge if the three metrics and the meaning of the GED, HM-IoU and MDM subscript notation are understood.

Authors’ 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.