Tokenised Flow Matching for Hierarchical Simulation Based Inference

NeurIPS 2025 Workshop FPI Submission45 Authors

Track: Main Track

Keywords: Flow Matching, Transformers, Simulation Based Inference, Bayesian Inference

Abstract:

Large simulation costs are a persistent challenge in Simulation Based Inference
(SBI). Thus we strive for methods which are more sample efficient, requiring fewer
simulations to achieve accurate parameter estimates. Across many scientific do-
mains, SBI can be made more sample efficient by exploiting hierarchical structure.
Rather than directly learning the full posterior, we learn single-site likelihoods, or
posterior factors, that can be combined into a full hierarchical posterior. Current
approaches assume conditional independence of local posteriors given the global
parameters, which is often not suitable, and require training separate estimators
for each level of the hierarchy which adds complexity to training. We present a
tokenised flow matching estimator for posterior estimation (TFMPE), along with
a sample efficient algorithm (Bottom-up sampling) for hierarchical parameter in-
ference which makes no structural assumptions on the dependence of local and
global parameters. We find that our method exhibits stable inference with improved
sample efficiency compared to non-hierarchical methods for hierarchical inference
tasks. We also examine posterior estimates for an infectious disease model and
find that they are as reliable as MCMC approaches despite reduced computational
demands.

Submission Number: 45