Tokenised Flow Matching for Hierarchical Simulation Based Inference

The cost of simulator evaluations is a key practical bottleneck for Simulation Based Inference (SBI). In hierarchical settings with shared global parameters and exchangeable site-level parameters and observations, this structure can be exploited to improve simulation efficiency. Existing hierarchical SBI approaches factorise the posterior yet still simulate across multiple sites per training sample; We instead explore likelihood factorisation (LF) to train from single-site simulations. In LF sampling we learn a per-site (factorised) likelihood and then assemble synthetic multi-site observations to amortise inference for the full hierarchical posterior. Building on this, we propose Tokenised Flow Matching for Posterior Estimation (TFMPE), a tokenised flow matching approach that supports function-valued observations through likelihood factorisation. To enable systematic evaluation, we introduce a benchmark for hierarchical SBI. We validate TFMPE on this benchmark and on realistic infectious disease and computational fluid dynamics models, finding well-calibrated posteriors while reducing computational cost.

Scientific simulators, such as those of disease spread, blood flow through arteries, or climate, contain unknown parameters that must be fit to real measurements before the simulator can be used for prediction. Fitting is done by running the simulator many times under different parameter guesses and comparing the outputs to data. In modern problems this is the dominant cost: a single fit may involve simulating hundreds of related settings (one simulation per country, one simulation per patient). We show that most of these runs are avoidable. We first teach a neural network to imitate the simulator for a single setting, then combine its outputs as a substitute for the expensive multi-setting runs. A second network learns to recover the unknown parameters from observations. In our cardiovascular experiment, drawing one sample from the neural imitator was roughly 1,800 times cheaper than running the original simulator, and let us efficiently recover the hidden simulator parameters.