Calibrating Generative Flows through Regional Probability Matching
Generative models offer a powerful route to scientific simulation and discovery by enabling the study of complex populations and rare phenomena. Flow matching can generate realistic samples, yet inaccurate occurrence frequencies can distort population statistics and compromise inferences about rare events. Here we developed generalized mass-preserving coupled flow matching (G-MPCF), which matches the probability masses of Gaussian source regions to those of weighted data regions. This construction encodes the desired regional frequencies in the correspondence between noise and observations, shaping the learned dynamics while preserving the prescribed target distribution. On a rare-mode benchmark, G-MPCF reduces relative score error by 95.0% and numerical integration error by 89.9% compared with independent flow matching. On CIFAR-10 with a 100:1 class-probability ratio, it reduces corrected class-mass error by 44.6% compared with minibatch optimal transport, while improving mean image fidelity with matched computational budgets. We envision that incorporating regional probability calibration into generative modelling will support more reliable scientific simulation and inference, enabling the study of systems in which rare events and their occurrence frequencies are central to understanding the underlying mechanisms.