Flow Matching with Optimized Subclass Priors for Medical Image Augmentation

2026-05-15

AI Summary

This work addresses the degraded classification performance for rare disease categories in medical imaging, stemming from data scarcity and long-tailed distributions. The authors propose a generative approach based on latent-space subclass modeling and geometrically constrained flow matching. Specifically, a Gaussian mixture model is employed to perform fine-grained clustering of coarse disease labels, thereby constructing subclass-conditional source distributions. To refine the transport trajectories and avoid degenerate solutions, a geometric control mechanism incorporating directional concentration and path-length constraints is introduced. Evaluated on long-tailed benchmarks including MIMIC-LT, NIH-LT, and CT-RATE, the method significantly improves generation quality—measured by FID and IRS—and enhances diversity. Moreover, it effectively boosts balanced accuracy and macro-F1 scores in downstream diagnostic tasks.

Abstract

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.

Problem

• rare diseases
• medical image augmentation
• long-tailed classification
• subclass modeling
• generative modeling

Innovation

• Flow Matching
• Subclass Priors
• Latent Space Clustering
• Tail-Class Augmentation
• Geometric Regularization