Few-Shot Distribution-Aligned Flow Matching for Data Synthesis in Medical Image Segmentation

Abstract

Data heterogeneity hinders clinical deployment of medical image analysis models, and generative data augmentation helps mitigate this issue. However, recent diffusion-based methods that synthesize image-mask pairs often ignore distribution shifts between generated and real images across scenarios, and such mismatches can markedly degrade downstream performance. To address this issue, we propose AlignFlow, a flow matching model that aligns with the target reference image distribution via differentiable reward fine-tuning, and remains effective even when only a small number of reference images are provided. Specifically, we divide the training of the flow matching model into two stages: in the first stage, the model fits the training data to generate plausible images; Then, we introduce a distribution alignment mechanism and employ differentiable reward to steer the generated images toward the distribution of the given samples from the target domain. In addition, to enhance the diversity of generated masks, we also design a flow matching based mask generation to complement the diversity in regions of interest. Extensive experiments demonstrate the effectiveness of our approach, i.e., performance improvement by 3.5-4.0% in mDice and 3.5-5.6% in mIoU across a variety of datasets and scenarios.

Figures (7)

• Figure 1: Illustration of the data distribution of images generated by AlignFlow that align with the target domain distribution. The green dots represent the data points of images generated by AlignFlow, which are accurately scattered at the center of the target-domain data points.
• Figure 2: (a) Illustration of the AlignFlow architecture. In stage 1, we optimize the denoising loss to enable the model to generate reasonable images; in stage 2, we simultaneously optimize the denoising loss and the alignment loss, allowing the model to align the generated images with reference images from the target domain while maintaining its original image generation capability. (b) Illustration of the mask synthesis pipeline.
• Figure 3: Qualitative comparison on REFUGE2 dataset. The source domain is Canon, and the target domains are annotated on the right side of each row.
• Figure 4: t-SNE visualization of the data distributions of images generated by different methods.
• Figure 5: Qualitative comparison on FedPolyp dataset. The source domain is Canon, and the target domains are annotated on the right side of each row.