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

2026-04-03

AI Summary

This work addresses the challenge of deploying medical image analysis models in heterogeneous clinical settings, where distribution shifts between source and target domains degrade downstream performance. To mitigate this issue with limited target-domain data, we propose AlignFlow—a two-stage few-shot medical image synthesis framework based on flow matching. In the first stage, plausible image–mask pairs are generated; in the second, a differentiable reward-based fine-tuning mechanism leverages a small set of target samples to align synthesized content with the target distribution. A novel flow-matching mask generation module further enhances mask diversity within regions of interest. AlignFlow is the first approach to integrate distribution alignment with differentiable reward fine-tuning, effectively reducing domain shift using only a few target images. Experiments demonstrate consistent and significant improvements over state-of-the-art methods, achieving gains of 3.5–4.0% in mDice and 3.5–5.6% in mIoU across multiple datasets.

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.

Problem

• data heterogeneity
• distribution shift
• medical image segmentation
• data synthesis
• few-shot

Innovation

• flow matching
• distribution alignment
• few-shot data synthesis
• medical image segmentation
• differentiable reward