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

TL;DR

This paper introduces AlignFlow, a flow matching model that improves medical image data synthesis by aligning generated images with target distributions, enhancing downstream segmentation performance with limited reference data.

Contribution

The paper proposes a novel flow matching approach with distribution alignment and differentiable reward fine-tuning for improved medical image synthesis.

Findings

Performance improved by 3.5-4.0% in mDice.

Performance improved by 3.5-5.6% in mIoU.

Effective even with few reference images.

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.