LatentFM: A Latent Flow Matching Approach for Generative Medical Image Segmentation

TL;DR

LatentFM introduces a flow-based latent space model for medical image segmentation that produces accurate, diverse, and uncertainty-aware predictions with confidence maps, outperforming prior methods.

Contribution

It is the first to combine flow matching with VAEs for latent space medical segmentation, enabling diversity and uncertainty quantification.

Findings

Achieves superior segmentation accuracy on ISIC-2018 and CVC-Clinic datasets.

Generates diverse segmentation outputs reflecting data distribution.

Provides confidence maps for uncertainty estimation.

Abstract

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.