# Flow Matching with Optimized Subclass Priors for Medical Image Augmentation

arXiv: https://arxiv.org/abs/2605.16469 · HTML: https://mlsanity.com/paper/2605.16469/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Medical imaging and segmentation, one of 14 directions in Diffusion models.
- Builds on 27 of its 36 references that are in the map.
- First posted on arXiv 15 May 2026.

## Placed in

- Diffusion models / [Medical imaging and segmentation](https://mlsanity.com/diffusion/medical-imaging-and-segmentation/)

## Builds on (most cited first, 27 of 27)

- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [ChestX-ray8: Hospital-scale Chest X-ray Database and Benchmarks on Weakly-Supervised Classification and Localization of Common Thorax Diseases](https://mlsanity.com/paper/1705.02315/) (2017)
- [Making the Most of Text Semantics to Improve Biomedical Vision--Language Processing](https://mlsanity.com/paper/2204.09817/) (2022)
- [Multisample Flow Matching: Straightening Flows with Minibatch Couplings](https://mlsanity.com/paper/2304.14772/) (2023)
- [Generative models improve fairness of medical classifiers under distribution shifts](https://mlsanity.com/paper/2304.09218/) (2023)
- [PriorGrad: Improving Conditional Denoising Diffusion Models with Data-Dependent Adaptive Prior](https://mlsanity.com/paper/2106.06406/) (2021)
- [Generating High Fidelity Data from Low-density Regions using Diffusion Models](https://mlsanity.com/paper/2203.17260/) (2022)
- [Class-Balancing Diffusion Models](https://mlsanity.com/paper/2305.00562/) (2023)
- [Long-Tailed Classification of Thorax Diseases on Chest X-Ray: A New Benchmark Study](https://mlsanity.com/paper/2208.13365/) (2022)
- [Don't Play Favorites: Minority Guidance for Diffusion Models](https://mlsanity.com/paper/2301.12334/) (2023)
- [Why Are Conditional Generative Models Better Than Unconditional Ones?](https://mlsanity.com/paper/2212.00362/) (2022)
- [Subclass-balancing Contrastive Learning for Long-tailed Recognition](https://mlsanity.com/paper/2306.15925/) (2023)
- [Image Generation Diversity Issues and How to Tame Them](https://mlsanity.com/paper/2411.16171/) (2024)
- [Self-Guided Generation of Minority Samples Using Diffusion Models](https://mlsanity.com/paper/2407.11555/) (2024)
- [DiverseFlow: Sample-Efficient Diverse Mode Coverage in Flows](https://mlsanity.com/paper/2504.07894/) (2025)
- [Label-Noise Robust Diffusion Models](https://mlsanity.com/paper/2402.17517/) (2024)
- [LeFusion: Controllable Pathology Synthesis via Lesion-Focused Diffusion Models](https://mlsanity.com/paper/2403.14066/) (2024)
- [Designing a Conditional Prior Distribution for Flow-Based Generative Models](https://mlsanity.com/paper/2502.09611/) (2025)
- [Better Source, Better Flow: Learning Condition-Dependent Source Distribution for Flow Matching](https://mlsanity.com/paper/2602.05951/) (2026)
- [Rethinking cluster-conditioned diffusion models for label-free image synthesis](https://mlsanity.com/paper/2403.00570/) (2024)
- [Is There a Better Source Distribution than Gaussian? Exploring Source Distributions for Image Flow Matching](https://mlsanity.com/paper/2512.18184/) (2025)
- [LCMem: A Universal Model for Robust Image Memorization Detection](https://mlsanity.com/paper/2512.14421/) (2025)
- [TreeDiffusion: Hierarchical Generative Clustering for Conditional Diffusion](https://mlsanity.com/paper/2410.16910/) (2024)
- [Causal Disentanglement for Robust Long-tail Medical Image Generation](https://mlsanity.com/paper/2504.14450/) (2025)
- [Enabling PSO-Secure Synthetic Data Sharing Using Diversity-Aware Diffusion Models](https://mlsanity.com/paper/2506.17975/) (2025)
- [GRASP: Guided Residual Adapters with Sample-wise Partitioning](https://mlsanity.com/paper/2512.01675/) (2025)
- [Subtyping Breast Lesions via Generative Augmentation based Long-tailed Recognition in Ultrasound](https://mlsanity.com/paper/2507.22568/) (2025)

## Abstract (arXiv)

Rare diseases dominate the diagnostic challenge in medical imaging yet are severely underrepresented in clinical datasets, causing classifiers to fail on exactly the conditions where reliable detection matters most. Generative augmentation can supply the missing tail-class coverage, but coarse disease labels aggregate diverse subtypes and acquisition settings into multi-modal conditionals that bias generators toward dominant submodes, while a shared Gaussian source forces rare subpopulations through disproportionately long transport paths. We propose an offline strategy that introduces informative priors at two levels: first, we partition each coarse label into coherent submodes via Gaussian mixture modeling in the generative model's latent space; second, we learn subclass-conditioned source distributions that re-center and re-scale the starting distribution per submode, shortening trajectories and reducing within-subclass dispersion. To prevent degenerate solutions we impose explicit geometric control, moderately concentrating normalized displacement directions around learnable prototypes while capping path-length outliers. On long-tailed chest X-ray (MIMIC-LT, NIH-LT) and CT slice (CT-RATE) benchmarks the proposed method consistently improves tail-class generation fidelity and diversity (FID, IRS) and is a promising augmentation strategy that reliably improves downstream balanced accuracy and macro-F1 over a non-augmented baseline across modalities.