Improving Flow Matching for Simulation-Based Inference

Flow matching is an incredibly powerful technique that can be used to sample arbitrary distributions. Recently, flow matching posterior estimation (FMPE) has been introduced in simulation-based inference (SBI) as a scalable alternative to standard neural posterior estimation (NPE) using normalizing flows. However, FMPE suffers from a lower sampling efficiency than NPE because it requires multiple network evaluations per sample. In this work, we propose extensions of FMPE based on mini-batch optimal transport that reduce the required number of network evaluations for high-quality samples. We investigate the extensions theoretically and show that they can lead to straight probability flows in the appropriate limit. Finally, we demonstrate the performance of the method on simple toy models and high-dimensional gravitational wave source parameter inference, showing that it is possible to increase the sample efficiency of FMPE while achieving approximately the same performance over different tasks.