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 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. Compositional inference. Many posterior estimation workflows exploit the factorisation in Equation 2 to train estimators on datasets of parameter-observation pairs (θ, y s ), enabling inference over arbitrarily sized observation sets (Radev et al., 2023; Boelts et al., 2024; Geffner et al., 2023) . Summary networks. Some compositional methods use permutation-invariant summary networks, such as Deep Sets (Zaheer et al., 2017), to aggregate variable-sized observation sets into fixeddimensional representations for posterior estimation. These approaches typically train by uniformly sampling observation subsets up to size n s , requiring on average ns /2 simulations per training sample. Factorised score estimation. Alternative compositional approaches decompose the posterior score over observation subsets and compose individual scores during sampling via annealed Langevin dynamics (Geffner et al., 2023) . Fully-Factorised Neural Posterior Score Estimation (F-NPSE) trains on single observations, requiring one simulation per sample, while Partially-Factorised variants (PF-NPSE) train on subsets of size n max , requiring n max simulations per sample where 1 < n max < n s . Increasing n max reduces approximation error from score composition at the cost of sample efficiency. Hierarchical SBI. When applied to hierarchical models, recent methods have introduced separate estimators for global and local parameters to reduce the number of simulations required for posterior estimation (