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

Jie Yang; Ziqi Ye; Aihua Ke; Jian Luo; Bo Cai; Xiaosong Wang

eess.IVApr 3, 2026

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

Authors: Jie Yang, Ziqi Ye, Aihua Ke, Jian Luo, Bo Cai, Xiaosong Wang

Organizations: Key Laboratory of Aerospace Information Security and Trusted Computing, Ministry of Education, School of Cyber Science and Engineering, Wuhan University, Wuhan, China · Shanghai Innovation Institute, Shanghai, China · Shanghai AI Laboratory, Shanghai, China

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 & tables

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.

Method	FID( $\downarrow$ )	KID( $\downarrow$ )	LPIPS( $\downarrow$ )	SSIM( $\uparrow$ )	PSNR( $\uparrow$ )
T2I-Adapter	226.13	0.2374	0.4426	0.4519	18.07
ControlNet-Diff	162.87	0.1378	0.4335	0.4109	19.61
ControlNet-FM	166.20	0.1449	0.4396	0.4394	20.31
Siamese-Diffusion	156.82	0.1344	0.4497	0.3982	19.33
AlignFlow(Ours)	149.27	0.1301	0.3931	0.5099	21.23

Table 1: Image quality evaluation on CVC-ClinicDB dataset. ”ControlNet-Diff” and ”ControlNet-FM” respectively denote ControlNet models that use diffusion and flow matching as their foundation models.

Methods	FedPolyp ( Chen et al., 2024 )						REFUGE2 ( Fang et al., 2022 )
	CVC-ClinicDB		CVC-ColonDB		ETIS		KOWA		TOPCON		Zeiss
	mDice	mIoU	mDice	mIoU	mDice	mIoU	mDice	mIoU	mDice	mIoU	mDice	mIoU
UNet ( Ronneberger et al., 2015 )	88.33	83.11	79.07	74.16	75.46	71.32	89.29	82.14	84.64	76.96	87.05	79.23
+T2I-Adapter ( Mou et al., 2024 )	88.35	83.45	77.48	73.11	75.60	72.08	88.44	81.02	89.13	82.22	89.20	82.25
+ControlNet-Diff ( Zhang et al., 2023 )	88.06	83.25	80.33	75.63	74.17	70.67	72.96	66.73	77.33	70.72	87.35	80.23
+ControlNet-FM	88.11	83.32	80.67	75.75	75.75	72.15	88.75	81.30	87.53	80.28	90.01	83.26

Table 2: Comparison of medical image segmentation performance on UNet, SegFormer, and DPT. The evaluation is performed using the mDice(%) and mIoU(%) metrics. ”ControlNet-Diff” and ”ControlNet-FM” respectively denote ControlNet models that use diffusion and flow matching as their foundation models. The ”+” symbol indicates that synthetic data were added to the original training set to form the complete training dataset. The best results are highlighted in bold, while the second-best results are underlined.

Methods	CVC-ClinicDB		CVC-ColonDB		ETIS
Methods	mDice	mIoU	mDice	mIoU	mDice	mIoU
UNet	88.33	83.11	79.07	74.16	75.46	71.32
+AlignFlow	88.68	83.77	82.41	77.38	76.45	72.83
+AlignFlow(Mask)	87.23	82.20	82.52	77.47	77.99	73.60
DPT	90.03	85.13	81.66	76.61	73.39	69.76
+AlignFlow	90.51	85.67	84.65	79.67	77.39	73.22

Table 3: Validation of the effectiveness of the mask generation method, evaluated using mDice (%) and mIoU (%) metrics.

Figure 4: t-SNE visualization of the data distributions of images generated by different methods.

Methods	FedPolyp ( Chen et al., 2024 )						REFUGE2 ( Fang et al., 2022 )
	CVC-ClinicDB		CVC-ColonDB		ETIS		KOWA		TOPCON		Zeiss
	mDice	mIoU	mDice	mIoU	mDice	mIoU	mDice	mIoU	mDice	mIoU	mDice	mIoU
UNet	88.33	83.11	79.07	74.16	75.46	71.32	89.29	82.14	84.64	76.96	87.05	79.23
+AlignFlow(1)	88.14	82.92	79.11	74.10	76.10	72.14	87.87	80.22	88.85	81.78	89.65	82.61
+AlignFlow(5)	88.56	83.75	81.41	76.24	75.82	72.18	88.72	81.36	87.60	80.89	89.44	82.52
+AlignFlow(10)	88.68	83.77	82.41	77.38	76.45	72.83	89.77	82.77	89.14	82.54	89.94	83.22

Table 4: Comparison of the impact of different number of reference images on medical image segmentation performance, evaluated using mDice (%) and mIoU (%) metrics. The numbers in parentheses indicate the number of reference images used.

$t_{s}^{i}$	KOWA				Zeiss
	UNet		SegFormer		UNet		SegFormer
	mDice	mIoU	mDice	mIoU	mDice	mIoU	mDice	mIoU
0	85.52	76.95	86.30	78.10	87.71	79.98	88.43	81.31
10	89.14	81.96	89.70	83.25	88.62	81.31	89.29	82.69
15	89.49	82.48	89.24	82.34	89.78	82.88	90.32	84.00
16	87.40	79.61	89.86	83.20	88.20	80.90	88.34	81.40

Table 5: Comparison of the impact of different values of

t_{s}^{i}

on medical image segmentation performance, evaluated using mDice (%) and mIoU (%) metrics.

$w_{align}$	KOWA				Zeiss
	UNet		SegFormer		UNet		SegFormer
	mDice	mIoU	mDice	mIoU	mDice	mIoU	mDice	mIoU
0.25	89.07	81.87	89.97	83.42	89.03	82.06	90.05	83.35
0.5	88.62	81.22	89.71	83.07	89.24	82.17	89.78	83.32
1.0	89.77	82.77	90.36	83.96	89.94	83.22	90.19	83.85
2.0	88.69	81.37	89.19	82.27	88.72	81.70	88.81	81.89

Table 6: Comparison of the impact of different values of

w_{align}

on medical image segmentation performance, evaluated using mDice (%) and mIoU (%) metrics.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Method	FID( $\downarrow$ )	KID( $\downarrow$ )	LPIPS( $\downarrow$ )	SSIM( $\uparrow$ )	PSNR( $\uparrow$ )
T2I-Adapter ( Mou et al., 2024 )	176.46	0.1982	0.3245	0.3933	19.07
ControlNet-Diff ( Zhang et al., 2023 )	64.96	0.0637	0.2005	0.5812	21.92
ControlNet-FM	54.64	0.0578	0.1787	0.5530	23.25
Siamese-Diffusion ( Qiu et al., 2025 )	70.73	0.0646	0.2060	0.5154	21.26
AlignFlow(Ours)	40.26	0.0356	0.1644	0.5993	25.96

Appendix

Table 7: Image quality evaluation on TOPCON dataset.

Method	FID( $\downarrow$ )	KID( $\downarrow$ )	LPIPS( $\downarrow$ )	SSIM( $\uparrow$ )	PSNR( $\uparrow$ )
T2I-Adapter ( Mou et al., 2024 )	152.02	0.2135	0.3299	0.3783	20.31
ControlNet-Diff ( Zhang et al., 2023 )	64.78	0.0763	0.1887	0.6441	24.61
ControlNet-FM	72.52	0.0908	0.2020	0.5152	22.14
Siamese-Diffusion ( Qiu et al., 2025 )	66.67	0.0776	0.2046	0.5741	23.08
AlignFlow(Ours)	44.88	0.0471	0.1821	0.5871	27.64

Appendix

Table 8: Image quality evaluation on Zeiss dataset.

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.

Figure 6: Qualitative comparison on FedPolyp dataset using the same mask.

Figure 7: Qualitative comparison on REFUGE2 dataset using the same mask.

Synthetic Data Augmentation

Synthetic data augmentation creates or modifies training examples to improve learning when real data are scarce, imbalanced, costly, or sensitive.

See all 448 papers

FMS $^2$ : Unified Flow Matching for Segmentation and Synthesis of Thin Structures
Mar 14, 2026Babak Asadi, Peiyang Wu, Mani Golparvar-Fard +2Thin-Structure Segmentation Topology-Preserving Image Segmentation
University of Illinois at Urbana-Champaign
SymmAdapt: Symmetrical Flow Matching for Source-Free Domain Adaptation in Medical Image Segmentation
Aug 23, 2026Tal Grossman, Noa Cahan, Hayit GreenspanSource-Free Domain Adaptation Medical Image Segmentation
Faculty of Engineering, Tel Aviv University, Tel Aviv, Israel
Diagnose-to-Generate: Eliminating Structural Distortions in Flow Models via Patch-Level Grounded Alignment
Oct 13, 2026Yifan Li, Lishuai Gao, Tong Zhang +8Image Generation Flow-Based Generative Modeling
University of Science and Technology of China, Hefei, AnHui Province, China · Meituan Inc., Beijing, China

Mar 14, 2026J/K move · Enter open · S save

FMS $^2$ : Unified Flow Matching for Segmentation and Synthesis of Thin Structures

Babak Asadi, Peiyang Wu, Mani Golparvar-Fard, Viraj Shah+1

University of Illinois at Urbana-Champaign

Thin-Structure Segmentation Topology-Preserving Image Segmentation Medical Image Segmentation Crack Segmentation Flow Matching Conditional Image Generation

Open paper AskPDF arXiv

Segmenting thin structures like infrastructure cracks and anatomical vessels is a task hampered by topology-sensitive geometry, high annotation costs, and poor generalization across domains. Existing methods address these challenges in isolation. We propose FMS

^2

, a flow-matching framework with two modules. (1) SegFlow is a 2.96M-parameter segmentation model built on a standard encoder-decoder backbone that recasts prediction as continuous image

\rightarrow

mask transport. It learns a time-indexed velocity field with a flow-matching regression loss and outputs the mask via ODE integration, rather than supervising only end-state logits. This trajectory-level supervision improves thin-structure continuity and sharpness, compared with tuned topology-aware loss baselines, without auxiliary topology heads, post-processing, or multi-term loss engineering. (2) SynFlow is a mask-conditioned mask

\rightarrow

image generator that produces pixel-aligned synthetic image-mask pairs. It injects mask geometry at multiple scales and emphasizes boundary bands via edge-aware gating, while a controllable mask generator expands sparsity, width, and branching regimes. On five crack and vessel benchmarks, SegFlow alone outperforms strong CNN, Transformer, Mamba, and generative baselines, improving the volumetric metric (mean IoU) from 0.511 to 0.599 (+17.2%) and reducing the topological metric (Betti matching error) from 82.145 to 51.524 (-37.3%). When training with limited labels, augmenting SegFlow with SynFlow-generated pairs recovers near-full performance using 25% of real annotations and improves cross-domain IoU by 0.11 on average. Unlike classical data augmentation that promotes invariance via label-preserving transforms, SynFlow provides pixel-aligned paired supervision with controllable structural shifts (e.g., sparsity, width, branching), which is particularly effective under domain shift.

Abstract

Synthetic Data Augmentation

Explore similar work

FMS2^22: Unified Flow Matching for Segmentation and Synthesis of Thin Structures

SymmAdapt: Symmetrical Flow Matching for Source-Free Domain Adaptation in Medical Image Segmentation

Diagnose-to-Generate: Eliminating Structural Distortions in Flow Models via Patch-Level Grounded Alignment

FMS $^2$ : Unified Flow Matching for Segmentation and Synthesis of Thin Structures