AI Research Portal

This paper introduces Ambiguity-aware Truncated Flow Matching (ATFM), a novel method for ambiguous medical image segmentation (AMIS) that aims to simultaneously improve prediction accuracy and diversity. ATFM addresses the inherent trade-off in existing methods by proposing a Data-Hierarchical Inference paradigm, Gaussian Truncation Representation (GTR), and Segmentation Flow Matching (SFM). Empirical evaluations on LIDC and ISIC3 datasets demonstrate that ATFM outperforms state-of-the-art methods and achieves more efficient inference.

The paper addresses the following challenges in ambiguous medical image segmentation (AMIS):

Entangled Accuracy and Diversity: Existing Truncated Diffusion Probabilistic Models (TDPMs) suffer from an entangled accuracy and diversity of predictions, making it difficult to improve both simultaneously.
Insufficient Fidelity and Plausibility: Traditional TDPMs can have sub-optimal approximations of the underlying distribution at the truncation point, compromising prediction fidelity. Furthermore, the absence of semantic guidance after truncation adversely affects the plausibility of generated predictions.

Introduction of Ambiguity-aware Truncated Flow Matching (ATFM), a novel inference paradigm and dedicated model components for AMIS.
Proposal of Data-Hierarchical Inference, an AMIS-specific inference paradigm that disentangles accuracy and diversity by marginalizing stochasticity during diffusion.
Development of Gaussian Truncation Representation (GTR) to explicitly model the Gaussian latent distribution at the truncation point, enhancing prediction fidelity and reliability.
Introduction of Segmentation Flow Matching (SFM) to enhance prediction plausibility and diversity by extending semantic-aware flow transformation.
Demonstration of ATFM's superior performance and efficiency on LIDC and ISIC3 datasets compared to state-of-the-art methods.

ATFM employs a Data-Hierarchical Inference paradigm to achieve synergistic optimization of accuracy and diversity. This paradigm consists of two main stages:

Stage 1: Gaussian Truncation Representation (GTR): This stage focuses on improving prediction fidelity and truncation distribution reliability. GTR explicitly models the distribution at the truncation point (Ttrunc) as a Gaussian distribution, parameterizing its mean and covariance. This explicit modeling, supervised by the ground truth, preserves low-frequency modes and ensures sample consistency, thereby enhancing accuracy. The loss term used is Lprior = -log p(Y|XTtrunc) P(XTtrunc|X), which minimizes the negative log-likelihood to optimize the explicit Gaussian distribution.
Stage 2: Segmentation Flow Matching (SFM): This stage enhances prediction plausibility and diversity. SFM extends semantic-aware flow transformation by modeling semantic consistency among labels, predictions, and intermediate states at each timestep. It uses a Semantic-aware Transformation Network (ST-Net) and leverages Flow Matching (FM) to avoid Gaussian constraints that can disturb fine-grained predictions. The loss for SFM is LSF = LFM + Σ α × L'Dice, where LFM is the Flow Matching loss and L'Dice is the Dice loss for semantic consistency.

The overall inference process first improves accuracy at the distribution level using GTR and then enhances diversity and plausibility at the sample level using SFM.

Architecture:
- GTR uses a segmentation backbone and separate convolutional layers for estimating the mean and covariance of the Gaussian distribution at the truncation point.
- SFM incorporates a Semantic-aware Transformation Network (ST-Net).
Key Equations:
- Equation 1: {s1, s2, ..., sn} ∈ P ~ N(μ, Σ) (Marginal distribution at Ttrunc)
- Equation 2: pred; {pred;}=1 supervise {gt;}=1 (Sample level for diversity)
- Equation 4: Z = fe(X), μ = gφ(Z), Σ = hψ(Z), XTrunc ~ N(μ, Σ) (Gaussian distribution at truncation point)
- Equation 5: Lprior = -log p(Y|XTtrunc) P(XTtrunc|X) ≈ (1/M) Σ [ -log p(Y|XTtrunc) ] (Lprior loss for GTR)
Loss/Objective:
- GTR uses Lprior for accuracy and fidelity.
- SFM uses LSF = LFM + Σ α × L'Dice for plausibility and diversity.
Optimization:
- Adam optimizer with a learning rate of 10-4.
- Hyper-parameter M set to 20.
- Hyper-parameter α set to 10-3 and 10-4 for LIDC and ISIC3 respectively.
- SFM trained for 200 epochs on LIDC and 120 epochs on ISIC3.
- GTR pretrained for 1000 epochs on LIDC and 400 epochs on ISIC3.
Augmentation: Not stated in the paper.
Data Handling:
- LIDC-IDRI dataset: 15,096 slices, each with four corresponding segmentation labels.
- ISIC3 subset: 300 images, each with three annotations.
Implementation Details:
- Single RTX 3090 GPU with 24GB memory.
- λ = 10-3 (T = 1000) with a linear schedule.
- Inference steps for SFM: Optimal performance achieved within 25 steps.

Datasets: LIDC-IDRI and ISIC3 subset.
Metrics:
- Generalised Energy Distance (GED) for diversity (GED16, GED32, GED100).
- Hungarian-Matching Intersection-over-Union (HM-IoU) for sample fidelity (HM-IoU32).
- Maximum Dice Matching (MDM) for individual segmentation accuracy (MDM32).
Baselines: Prob. Unet, PHiSeg, SSN, MoSE, P2SAM, CIMD, AB, CCDM, c-Prob. Unet, c-SSN.
Numeric Results:
- On LIDC: ATFM achieved 11.5% improvement on GED100 and 7.3% improvement on HM-IoU32 compared to advanced methods.
- On ISIC3: ATFM outperformed the runner-up method by 12% in GED and achieved top-tier results in HM-IoU and MDM.
- Ablation studies showed ATFM outperformed Act. GTR and SFM by at least 10% and 6% respectively.
Runtime/Training Cost:
- Inference time for 100 samples on LIDC: ATFM takes approximately 113s with Ttrunc = 25 steps, significantly faster than other diffusion-based methods.

No explicit paper-reported limitations were found.

Method evidence: "We propose Ambiguity-aware Truncated Flow Matching (ATFM) to achieve a synergistic enhancement of accuracy and diversity of predictions in AMIS tasks, supported by three designs (as Fig. 1(b)): (1) We propose Data-Hierarchical Inference, an innovative AMIS-specific inference paradigm redefinition inspired by TDPMs, where stochasticity during diffusion is marginalized for an effective disentanglement of accuracy and diversity. Specifically, ATFM performs truncated steps at the data distribution level to prioritize accuracy with diversity marginalized, whereas the final diffusion stage operates at the data sample level to enhance diversity without sacrificing accuracy. (2) We propose Gaussian Truncation Representation (GTR), which explicitly models the Gaussian latent distribution at truncation point to enhance prediction fidelity. While traditional TDPMs approximate the implicit distribution via adversarial sampling, GTR encodes image-level semantic features to logit distribution, thereby directly modeling the reliable distribution at the truncation point. This design preserves low-frequency modes and improves consistency, leading to predictions with higher fidelity in AMIS tasks. (3) We propose Segmentation Flow Matching (SFM), which introduces semantic-aware flow transformation to increase diversity and enhance plausibility simultaneously." (section: Abstract)
Result evidence: "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." (section: Abstract)
Limitation evidence: Not stated in the paper.

Ambiguity-aware Truncated Flow Matching for Ambiguous Medical Image Segmentation

Summary

Problem

Key Contributions

Method

Technical Details

Results

Weaknesses / Limitations

Evidence

Additional Figures