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.

• The paper introduces AlignFlow, a framework that employs differentiable reward fine-tuning and flow matching to generate target-aligned synthetic image-mask pairs.
• It leverages few-shot learning with as few as 5–10 reference images to improve segmentation metrics by up to 4–5.6% across diverse medical datasets.
• The method stabilizes training using ControlNet tuning and MMD-based alignment, offering a robust solution for domain adaptation in clinical imaging.

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

Motivation and Problem Statement

Medical image segmentation models are fundamentally constrained by limited and heterogeneous data, which impedes clinical deployment and generalization across sites and devices. Traditional and GAN-based augmentation schemes inadequately address domain shifts, and diffusion-based approaches, though effective in quality, fail to explicitly align synthetic data distributions with scarce target-domain samples. This limitation leads to performance drops when models are deployed in real-world scenarios characterized by device/setting-specific distributional variations.

AlignFlow Architecture and Methodology

The proposed AlignFlow framework introduces distribution-aligned, few-shot augmentation by leveraging flow matching generative models and differentiable reward fine-tuning for explicit data distribution alignment. AlignFlow is structured around a two-stage training regimen:

• Stage 1: Standard flow matching with denoising loss to ensure plausible image-mask pair generation aligned with the source training distribution.
• Stage 2: Joint optimization of denoising loss and an alignment loss, employing a differentiable reward computed via Maximum Mean Discrepancy (MMD) between DINOv3 feature space representations of generated and reference target-domain images.

The reward-driven alignment utilizes a partially-tunable inference strategy to enable stable gradient propagation during reward-based fine-tuning, mitigating memory and instability issues prevalent in multi-step generative RL optimization. Mask diversity is further enhanced by a complementary flow-matching-based mask synthesis pipeline, including post-processing with blurring, morphological operations, and DINOv3-based filtering for quality selection.

Technical Contributions

• Few-Shot Distribution Alignment: AlignFlow employs only a small set of reference images to steer the synthetic image-mask pair distribution, using differentiable MMD-based reward learning to align with target domains, supporting practical deployment in low-data settings.
• Reward Function Design: Two reward metrics were considered—MMD, leveraging feature statistics in kernel space, and symmetric KL divergence over pooled feature means. Empirically, MMD outperforms SKL, and is adopted as the default criterion.
• Mask Synthesis for Enhanced Segmentation: The method introduces robust mask generation and selection strategies to maximize anatomical and textural diversity, directly improving downstream segmentation generalization.
• Efficient and Stable Fine-Tuning: By freezing most generative model parameters and only tuning ControlNet modules, AlignFlow offers efficient transfer and domain alignment, retaining large-scale generative priors.

Experimental Results

AlignFlow's performance was extensively validated across six diverse medical image segmentation datasets (gastrointestinal polyp and retinal fundus images) and with three segmentation architectures (UNet, SegFormer, DPT). Results demonstrate consistently superior performance to all tested generative baselines (T2I-Adapter, ControlNet variants, Siamese-Diffusion):

• Segmentation Performance: AlignFlow improves mDice and mIoU by 3.5–4.0% and 3.5–5.6%, respectively, over baselines, with marked robustness across architectures and domains.
• Image Quality Metrics: On CVC-ClinicDB, TOPCON, and Zeiss test sets, AlignFlow achieves the lowest FID, KID, and LPIPS, and the highest SSIM and PSNR, confirming distributional alignment with target domains and perceptual realism.
• Ablation on Reference Count: Using as few as 5–10 reference images suffices for close-to-optimal distribution adaptation and substantial downstream segmentation gain; further scaling reference count yields diminishing returns.
• Mask Generation Ablation: Inclusion of synthetically diversified masks further boosts segmentation metric gains on most datasets, underscoring the benefit of anatomical variance in synthetic augmentation.
• Reward Function Ablation: MMD-based alignment provides clear superiority to SKL in aligning implicit generative distributions for downstream segmentation.
• Hyperparameter Robustness: AlignFlow exhibits stable performance across a range of alignment loss weights and training step schedule parameters.

Implications and Theoretical Impact

AlignFlow establishes a new paradigm for few-shot domain alignment in data-driven medical imaging workflows. Its capacity to leverage small target-specific reference sets addresses a critical bottleneck in multi-site/model generalization, facilitating cross-institutional AI deployment and robust real-world performance. The explicit reward-driven alignment in feature space permits precise, differentiable control over generative domain transfer, extending the utility of flow-matching generative models beyond natural image domains.

Theoretical implications extend to RL-based fine-tuning of diffusion/flow models for arbitrary statistical objectives, and the robust adaptation of frozen large-scale priors with lightweight, data-efficient modules (e.g., ControlNet), offering a practical path for foundation model adaptation in resource-limited domains.

Future Directions

Potential future research avenues include:

• Automated Reference Selection: Active learning for optimal reference set selection could mitigate performance variance from sparse target sampling.
• Broader Modalities and 3D Data: Extension to non-2D or multi-modal imaging and exploration of 3D latent flows could expand applicability.
• Fully Unsupervised Domain Adaptation: Investigating reward proxies not requiring annotation for even more restrictive data regimes.
• Generalized Reward Functions: Task-aware or adversarial reward mechanisms could further increase target fidelity in highly variable clinical settings.

Conclusion

This work introduces AlignFlow, a few-shot distribution-aligned flow matching framework that advances synthetic data generation for medical image segmentation through explicit distribution alignment using differentiable rewards. It demonstrates state-of-the-art data augmentation performance in challenging, heterogeneous domain conditions, representing a substantial methodological advancement in domain-adaptive synthesis and downstream model robustness in medical imaging (2604.02868).