Tokenised Flow Matching for Hierarchical Simulation Based Inference
Abstract
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 neural surrogate of the simulator 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.
Cite
@article{arxiv.2604.20723,
title = {Tokenised Flow Matching for Hierarchical Simulation Based Inference},
author = {Giovanni Charles and Cosmo Santoni and Seth Flaxman and Elizaveta Semenova},
journal= {arXiv preprint arXiv:2604.20723},
year = {2026}
}
Comments
31 pages, 11 figures