Tokenised Flow Matching for Hierarchical Simulation Based Inference

Large simulation costs are a challenge in Simulation Based Inference(SBI), especially as the number of parameters and observations increase.Thus we strive for sample efficiency, methods which require fewersimulations to achieve accurate posterior estimates. The sampleefficiency of SBI in many scientific domains can be improved withhierarchical formulationswhere the dependency structure of the parameters and/or observations areexploited to decompose the posterior into more efficiently estimatedtargets, which can later be aggregated into the desired posteriorestimate. At the time of writing, this approach has been limited tocompositional score matching based methods. Outside of hierarchicalmodelling, SBI has seen improvements in robustness and trainingstability with Flow Matching for Posterior Estimation (FMPE), and in flexibility with unified,tokenised estimators which canestimate arbitrary conditionals and function-valued observations. Weintroduce tokenised FMPE, to combine these advances and propose astable, flexible, sample efficient method for SBI on hierarchicalparameter inference problems. Our method combines local and globalestimators into single unified estimator, supports function-valuedparameters and observations, and exhibits stronger sample efficiencythan FMPE on hierarchical inference problems.