Conditional flow matching with coordinate-aware patch training for efficient 3D medical image segmentation

Abstract

Medical image segmentation plays a crucial role in clinical diagnosis and treatment planning. In recent years, generative methods such as diffusion models have shown strong potential for medical image segmentation because they can model complex data distributions and capture aleatoric uncertainty. However, these methods usually require hundreds of iterative sampling steps during inference, which results in high computational cost and limits their practical use in clinical environments. In addition, when generative segmentation models are applied to three dimensional medical image segmentation tasks, they often face severe GPU memory limitations. To address these issues, we propose a conditional Flow Matching (CFM) based framework for efficient 3D medical image segmentation. By adopting a straight probability path derived from optimal transport, the proposed method can generate high quality segmentation masks with only a single ODE integration step during inference. To alleviate the memory limitation caused by volumetric data and the loss of spatial location information in conventional patch training, we introduce a coordinate-aware patch training strategy, in which normalized 3D Cartesian coordinates are explicitly concatenated with local image patches, enabling the model to retain global anatomical positional awareness under constrained GPU memory. Experiments on the BraTS2020 brain tumor dataset and a private prostate MRI dataset demonstrate that the proposed method maintains competitive segmentation accuracy while achieving an inference time of only 1.14 seconds per volume, offering a clear efficiency advantage over existing generative segmentation methods and providing a viable solution for efficient 3D medical image segmentation.