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

• URL: http://arxiv.org/abs/2604.02868v1
• Date: Fri, 03 Apr 2026 08:35:42 GMT
• Title: Few-Shot Distribution-Aligned Flow Matching for Data Synthesis in Medical Image Segmentation
• Abstract summary: We propose AlignFlow, a flow matching model that aligns with the target reference image distribution via differentiable reward fine-tuning.<n>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.
• Score: 7.127259067323343
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• LatentFM: A Latent Flow Matching Approach for Generative Medical Image Segmentation [0.0]
  We propose LatentFM, a flow-based model operating in the latent space for medical image segmentation.<n>By sampling multiple latent representations, our method synthesizes diverse segmentation outputs.<n>Our approach achieves superior segmentation accuracy while remaining highly efficient in the latent space.
  arXiv  Detail & Related papers  (2025-12-04T14:06:08Z)
• LEAF: Latent Diffusion with Efficient Encoder Distillation for Aligned Features in Medical Image Segmentation [2.529281336118734]
  We propose LEAF, a medical image segmentation model grounded in latent diffusion models.<n>During the fine-tuning process, we replace the original noise prediction pattern with a direct prediction of the segmentation map.<n>We also employ a feature distillation method to align the hidden states of the convolutional layers with the features from a transformer-based vision encoder.
  arXiv  Detail & Related papers  (2025-07-24T09:08:04Z)
• Latent Drifting in Diffusion Models for Counterfactual Medical Image Synthesis [55.959002385347645]
  Latent Drifting enables diffusion models to be conditioned for medical images fitted for the complex task of counterfactual image generation.<n>We evaluate our method on three public longitudinal benchmark datasets of brain MRI and chest X-rays for counterfactual image generation.
  arXiv  Detail & Related papers  (2024-12-30T01:59:34Z)
• Dataset Augmentation by Mixing Visual Concepts [3.5420134832331334]
  This paper proposes a dataset augmentation method by fine-tuning pre-trained diffusion models.<n>We adapt the diffusion model by conditioning it with real images and novel text embeddings.<n>Our approach outperforms state-of-the-art augmentation techniques on benchmark classification tasks.
  arXiv  Detail & Related papers  (2024-12-19T19:42:22Z)
• FlowSDF: Flow Matching for Medical Image Segmentation Using Distance Transforms [60.195642571004804]
  We introduce FlowSDF, an image-guided conditional flow matching framework, to represent an implicit distribution of segmentation masks.<n>Our framework enables accurate sampling of segmentation masks and the computation of relevant statistical measures.
  arXiv  Detail & Related papers  (2024-05-28T11:47:12Z)
• Conditional Distribution Modelling for Few-Shot Image Synthesis with Diffusion Models [29.821909424996015]
  Few-shot image synthesis entails generating diverse and realistic images of novel categories using only a few example images. We propose Conditional Distribution Modelling (CDM) -- a framework which effectively utilizes Diffusion models for few-shot image generation.
  arXiv  Detail & Related papers  (2024-04-25T12:11:28Z)
• Training Class-Imbalanced Diffusion Model Via Overlap Optimization [55.96820607533968]
  Diffusion models trained on real-world datasets often yield inferior fidelity for tail classes. Deep generative models, including diffusion models, are biased towards classes with abundant training images. We propose a method based on contrastive learning to minimize the overlap between distributions of synthetic images for different classes.
  arXiv  Detail & Related papers  (2024-02-16T16:47:21Z)
• The Journey, Not the Destination: How Data Guides Diffusion Models [75.19694584942623]
  Diffusion models trained on large datasets can synthesize photo-realistic images of remarkable quality and diversity. We propose a framework that: (i) provides a formal notion of data attribution in the context of diffusion models, and (ii) allows us to counterfactually validate such attributions.
  arXiv  Detail & Related papers  (2023-12-11T08:39:43Z)
• Real-World Image Variation by Aligning Diffusion Inversion Chain [53.772004619296794]
  A domain gap exists between generated images and real-world images, which poses a challenge in generating high-quality variations of real-world images. We propose a novel inference pipeline called Real-world Image Variation by ALignment (RIVAL) Our pipeline enhances the generation quality of image variations by aligning the image generation process to the source image's inversion chain.
  arXiv  Detail & Related papers  (2023-05-30T04:09:47Z)
• Training on Thin Air: Improve Image Classification with Generated Data [28.96941414724037]
  Diffusion Inversion is a simple yet effective method to generate diverse, high-quality training data for image classification. Our approach captures the original data distribution and ensures data coverage by inverting images to the latent space of Stable Diffusion. We identify three key components that allow our generated images to successfully supplant the original dataset.
  arXiv  Detail & Related papers  (2023-05-24T16:33:02Z)
• SinDiffusion: Learning a Diffusion Model from a Single Natural Image [159.4285444680301]
  We present SinDiffusion, leveraging denoising diffusion models to capture internal distribution of patches from a single natural image. It is based on two core designs. First, SinDiffusion is trained with a single model at a single scale instead of multiple models with progressive growing of scales. Second, we identify that a patch-level receptive field of the diffusion network is crucial and effective for capturing the image's patch statistics.
  arXiv  Detail & Related papers  (2022-11-22T18:00:03Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.