Many segmentation tasks, such as medical image segmentation or future state prediction, are inherently ambiguous, meaning that multiple predictions are equally correct. Current methods typically rely on generative models to capture this uncertainty. However, identifying the underlying modes of the distribution with these methods is computationally expensive, requiring large numbers of samples and post-hoc clustering. In this paper, we shift the focus from stochastic sampling to the direct generation of likely outcomes. We introduce mode proposal models, a deterministic framework that efficiently produces a fixed-size set of proposal masks in a single forward pass. To handle superfluous proposals, we adapt a confidence mechanism, traditionally used in object detection, to the high-dimensional space of segmentation masks. Our approach significantly reduces inference time while achieving higher ground-truth coverage than existing generative models. Furthermore, we demonstrate that our model can be trained without knowing the full distribution of outcomes, making it applicable to real-world datasets. Finally, we show that by decomposing the velocity field of a pre-trained flow model, we can efficiently estimate prior mode probabilities for our proposals.
Figures & tables
Figure 1: Method overview: In ambiguous segmentation, multiple segmentation masks are valid solutions for a given input, each associated with a prior probability. Our mode proposal model deterministically produces a fixed-size set of proposals and corresponding scores d∈[0,1] . To select the most promising candidates, we threshold the selection scores. Finally, we use a separate flow model, trained on the same data, to estimate the prior probabilities of each proposed mode.
Figure 2: MMFire proposals: MMFire contains multiple simulated wildfire spread outcomes, differing by wind direction. Our mode proposal network successfully produces proposals covering all modes. This model was trained in the single-label scenario, only seeing a single label yi per input at each training step, sampled according to a highly-skewed distribution with probabilities ranging from 0.4% for 0° to 50.2% for 315°.
Figure 3: LIDC proposals: LIDC consists of expert annotations of lung CT images. Ground truths are usually white round shapes in the center of the image, making it very difficult for the model to discard implausible proposals. This model was trained in the single-label scenario. The bottom-right proposal likely did not receive many gradient updates, resulting in a bad value range, making it easy to discard.
Dataset
Method
Proposals
HM IoU*
Runtime
MMFire
ProbU-Net [ 9 ]
16
0.637
0h12m
ProbU-Net [ 9 ]
32
0.687
0h27m
Diffusion + SPELL [ 9 ]
16
0.784
0h55m
Diffusion + SPELL [ 9 ]
32
0.830
2h17m
ModeProposalNet (Single)
16
0.909
0h2m
ModeProposalNet (Full)
16
0.923
0h2m
Table 1: Main results across datasets: Runtimes should be taken as a rough measure, since they represent the total runtime, measured on shared infrastructure.
Method
Proposals
HM IoU
HM IoUmulti
CIMD [ 22 ] , reported in [ 26 ]
32
0.592
BerDiff [ 3 ]
16
0.596
CCDM [ 26 ]
32
0.631
MoSE [ 8 ]
100
0.638
JProbU-Net [ 27 ]
32
0.647
P2SAM [ 11 ]
16
0.679
Table 2: Comparison to LIDC literature: We compare performances on the LIDC test set. HM IoU is the standard metric, i.e . without deduplication of ground truths. In HM IoUmulti , duplicate ground truths are matched to the same proposal, removing the need to produce duplicate proposals. This is a more appropriate measure for our approach, since our training does not encourage duplicate proposals in any way.
Scenario
Filter duplicates
Selection F1
HM IoU*
Proposals
MMFire - 8 modes
Full
✗
0.994
0.923
16.0
Full
✓
0.920
7.9
Single
✗
0.692
0.910
16.0
Single
✓
0.894
14.0
Cityscapes - 16 modes
Table 3: MMFire: Filtering superfluous proposals by selection scores.
Table 7
Appendix figures & tables8 assets
Supplementary material from the paper’s appendix.
Appendix
Dataset
Epochs
Proposals
Scenario
Loss function
Learning rate
λ
MMFire
100
16
Single
MSE
1e-3
1e-2
MMFire
100
16
Full
DiceFocal
1e-3
1e-2
Cityscapes
400
32
Single
MSE
1e-4
1e-1
Cityscapes
400
32
Full
MSE
1e-4
1e-1
LIDC
100
8
Single
DiceFocal
1e-3
1e-3
LIDC
100
8
Full
DiceFocal
1e-3
1e-1
Appendix
Table 6: Training hyperparameters , mostly determined via hyperparameter search. The number of proposals is instead set to be twice the number of ground truth masks that each dataset provides.
Training regime
Pretrained
Frozen weights
HM IoU*
Selection F1
From scratch
✗
✗
0.927
0.948
All weights
✓
✗
0.927
0.952
Decoder only
✓
✓
0.923
0.994
Last conv. only
✓
✓
0.752
0.933
Appendix
Table 7: Training with frozen weights - fully-labeled: The proposal model is initialized with the pretrained diffusion model weights. Varying parts of the model are kept frozen during training. We find that preserving the diffusion-pretrained weights in the encoder leads to greatly increased selection performance .
Figure 4: Cityscapes proposals: The binary ambiguous Cityscapes dataset creates diverse modes for a given input image by flipping each of four classes independently either to the positive or negative class, following fixed probabilities. The model used here was trained in the single-label scenario. Proposals displayed are thresholded. The model can confidently discard many redundant proposals, because the non-thresholded value ranges of those proposals are unusual, indicating that they have not been trained well. Given that only a single proposal per input is updated at each training step, this is unsurprising. Some modes correctly specialize and will often be the best-matching ones and thus receive further gradient updates, while others are too far from the correct outputs to be chosen in Hungarian matching, and thus barely receive gradient updates.
λ
Proposals
HM IoU
HM IoU multi
HM IoU*
0.0
8
0.624
0.834
0.799
0.0
16
0.690
0.852
0.818
0.0
32
0.797
0.863
0.832
0.0
64
0.868
0.872
0.844
0.0
128
0.869
0.869
0.840
0.1
8
0.622
0.831
0.794
Appendix
Table 8: LIDC validation: Selection loss reduces HM IoU. We vary the number of proposals, training with selection loss ( λ>0 ) or without ( λ=0 ). HM IoU is the standard metric, matching against the full set of ground truth masks. HM IoU multi modifies this by matching all duplicate ground truth masks to the same proposal. HM IoU* removes duplicate ground truth masks, matching proposals with the set of unique ground truths. Models were trained using only 10 epochs instead of 100.
Figure 5: Cityscapes: Estimating class templates from proposals First, proposals are filtered according to their selection scores (for this picture: from 32 to 22, not pictured). Second, principal components are computed based on these proposals (first column, one row per component) and then thresholded into binary masks that are supposed to represent one class each (second and third column). Otsu’s method does not separate the vegetation and car class in the first component (marked by the dashed ellipse), making it impossible to correctly estimate the underlying probabilities, though this is possible with a manually chosen threshold.
Method
proad
psidewalk
pvegetation
pcar
BSS
Target distribution
0.050
0.250
0.750
0.950
-
Proposals (manual)
0.050
0.142
0.762
0.910
0.983
GT masks
0.053
0.145
0.706
0.921
0.982
Appendix
Table 9: Cityscapes: Estimated probabilities for the image in Fig. 5 . Both the manually thresholded principal components, computed from proposals, as well as the ground truth masks, lead to very good estimates of the target distribution, though both underestimate the sidewalk class by about 0.11.
Training targets
Proposals
HM IoU*
Selection F1
Full
8
0.923
0.997
Full
16
0.923
0.994
Full
32
0.924
0.978
Single
8
0.883
0.995
Single
16
0.910
0.692
Single
32
0.914
0.306
Appendix
Table 10: Ablation - Redundant proposals, fully labeled: We investigate how different choices for the number of proposals K^ influences the model performance. Note that the selection scores for eight proposals are trivial (marked in italic), since the number of proposals equals the number of ground truth masks, unless duplicate ground truth masks occur. Dataset: MMFire validation
Samples
Brier Skill Score
8
0.913
16
0.910
32
0.916
64
0.919
Appendix
Table 11: Ablation: Vary number of Monte Carlo samples in mode probability estimation. We vary the number of initial samples yt to approximate the expectation in Eq. 22 .
Generative models have achieved remarkable progress with the emergence of flow matching (FM). It has demonstrated strong generative capabilities and attracted significant attention as a simulation-free flow-based framework capable of learning exact data densities. Motivated by these advances, we propose LatentFM, a flow-based model operating in the latent space for medical image segmentation. To model the data distribution, we first design two variational autoencoders (VAEs) to encode both medical images and their corresponding masks into a lower-dimensional latent space. We then estimate a conditional velocity field that guides the flow based on the input image. By sampling multiple latent representations, our method synthesizes diverse segmentation outputs whose pixel-wise variance reliably captures the underlying data distribution, enabling both highly accurate and uncertainty-aware predictions. Furthermore, we generate confidence maps that quantify the model certainty, providing clinicians with richer information for deeper analysis. We conduct experiments on two datasets, ISIC-2018 and CVC-Clinic, and compare our method with several prior baselines, including both deterministic and generative approach models. Through comprehensive evaluations, both qualitative and quantitative results show that our approach achieves superior segmentation accuracy while remaining highly efficient in the latent space.
Ngoc Huynh Trinh, Hoang Anh Nguyen Kim, Hai Toan Nguyen +1
Institute for Artificial Intelligence, University of Engineering and Technology, Vietnam National University, Hanoi, Vietnam
Predicting future states in uncertain environments, such as wildfire spread, medical diagnosis, or autonomous driving, requires models that can consider multiple plausible outcomes. While diffusion models can effectively learn such multi-modal distributions, naively sampling from these models is computationally inefficient, potentially requiring hundreds of samples to find low-probability modes that may still be operationally relevant. In this work, we address the challenge of sample-efficient ambiguous segmentation by evaluating several training-free sampling methods that encourage diverse predictions. We adapt two techniques, particle guidance and SPELL, originally designed for the generation of diverse natural images, to discrete segmentation tasks, and additionally propose a simple clustering-based technique. We validate these approaches on the LIDC medical dataset, a modified version of the Cityscapes dataset, and MMFire, a new simulation-based wildfire spread dataset introduced in this paper. Compared to naive sampling, these approaches increase the HM IoU* metric by up to 7.5% on MMFire and 16.4% on Cityscapes, demonstrating that training-free methods can be used to efficiently increase the sample diversity of segmentation diffusion models with little cost to image quality and runtime. Code and dataset: https://github.com/SebastianGer/wildfire-spread-scenarios
Sebastian Gerard, Josephine Sullivan
KTH Royal Institute of Technology, Stockholm, Sweden
Flow matching has recently emerged as a principled framework for learning continuous-time transport maps, enabling efficient ODE-based sampling without relying on stochastic diffusion processes. While generative modeling has shown promise for medical image segmentation, particularly in capturing uncertainty and complex anatomical variability, existing approaches are predominantly based on diffusion models, which require iterative sampling and incur substantial computational overhead. In this work, we propose MedFlowSeg, a conditional flow matching framework that formulates medical image segmentation as learning a time-dependent vector field that transports a simple prior distribution to the target segmentation distribution. Compared to diffusion-based methods, our formulation enables more efficient inference through solving an ordinary differential equation, while preserving the flexibility of generative modeling. To effectively incorporate conditional information, we introduce a dual-conditioning mechanism. Specifically, we propose a Dual-Branch Spatial Attention (DB-SA) module to inject multi-frequency structural priors, and a Frequency-Aware Attention (FA-Attention) module to model interactions between spatial and spectral representations via discrepancy-aware fusion and time-dependent modulation. These components improve the alignment between noisy intermediate states and clean semantic features, leading to better structural consistency and boundary delineation. We conduct extensive experiments across multiple medical imaging modalities, where MedFlowSeg consistently outperforms prior state-of-the-art (SOTA) baselines, including diffusion-based and flow-based methods.
Zhi Chen, Runze Hu, Le Zhang
School of Engineering, University of Birmingham · School of Public Health, Peking University