MedFlowSeg: Flow Matching for Medical Image Segmentation with Frequency-Aware Attention

TL;DR

MedFlowSeg introduces a flow matching approach for medical image segmentation that offers efficient inference and improved structural accuracy by leveraging frequency-aware attention mechanisms.

Contribution

The paper proposes a novel flow matching framework with dual-conditioning and frequency-aware attention modules for more efficient and accurate medical image segmentation.

Findings

MedFlowSeg outperforms state-of-the-art diffusion and flow-based methods.

The dual-conditioning mechanism enhances structural consistency.

Frequency-aware attention improves boundary delineation.

Abstract

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.