English

Tokenised Flow Matching for Hierarchical Simulation Based Inference

Machine Learning 2026-04-23 v1 Artificial Intelligence

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.

Keywords

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

R2 v1 2026-07-01T12:30:44.747Z