Tokenised Flow Matching for Hierarchical Simulation Based Inference

1. Problem Statement

Simulation Based Inference (SBI) is a framework for Bayesian parameter estimation when the likelihood function $p(y|\theta)$ is intractable but a simulator is available. While effective for standard problems, SBI faces a critical bottleneck in hierarchical settings.

In hierarchical models, data arises from multiple sites (e.g., countries, patients) sharing global parameters ($\theta_g$) but having distinct local parameters ($\eta_s$).

• The Bottleneck: Traditional amortized SBI methods require simulating all $n_s$ sites for every training sample to generate a full observation vector $y = (y_1, \dots, y_{n_s})$. As the number of sites increases, the computational cost of training becomes prohibitive.
• Existing Limitations: Current hierarchical approaches often factorize the posterior (e.g., training a global estimator on subsets of data). While this reduces simulation calls per sample, it still requires multiple simulations per training instance or suffers from error propagation when global parameters are inferred from noisy local estimates.
• Complex Data: Many scientific applications involve function-valued observations (e.g., time-series, irregularly sampled data) which are difficult to handle with standard fixed-dimensional vector inputs.

2. Methodology

The authors propose Tokenised Flow Matching for Posterior Estimation (TFMPE), a novel approach combining three key components: Likelihood Factorization, Tokenization, and Flow Matching.

A. Likelihood Factorization (LF) Sampling

Instead of factorizing the posterior, the authors propose factorizing the likelihood.

• Concept: Due to conditional independence, $p(y|\theta_g, \eta) = \prod p(y_s|\theta_g, \eta_s)$.
• Two-Stage Training:
  1. Stage 1 (Surrogate Training): Train a per-site neural surrogate $q_{\phi_l}(y_s | \theta_g, \eta_s)$ using single-site simulations. This requires only one simulator call per training sample, regardless of the total number of sites.
  2. Stage 2 (Posterior Training): Train a posterior estimator $q_{\phi_p}(\theta_g, \eta | y)$ using synthetic multi-site observations generated by the surrogate $q_{\phi_l}$.
• Benefit: This decouples the simulation cost from the number of sites ($n_s$), allowing the model to scale to hundreds or thousands of sites without re-simulating the full system for every training step.

B. Tokenization

To handle variable numbers of sites and function-valued observations (irregular time/space coordinates), the authors embed all variables into a sequence of tokens.

• Token Structure: Each token contains:
  • Value ($v_i$): The parameter or observation value.
  • Variable ID ($\ell_i$): Identifies the variable type (e.g., $\beta$, $A$).
  • Position ($p_i$): Index within a multi-dimensional variable.
  • Group ID ($g_i$): Explicitly indicates which site the token belongs to (crucial for distinguishing local parameters across sites).
  • Functional Input ($\xi_i$): Continuous coordinates (e.g., time) for function-valued data.
• Architecture: An Encoder-only Transformer processes these tokens via self-attention. This allows the model to handle arbitrary numbers of sites and irregular observation schedules naturally.

C. Flow Matching

The inference engine uses Flow Matching (specifically Conditional Flow Matching) to learn the posterior distribution.

• Mechanism: Instead of learning a score function (as in diffusion models), the network learns a vector field $v_t$ that transports samples from a base distribution (e.g., Gaussian) to the target posterior $p(\theta|y)$.
• Advantages: Flow matching offers simpler probability paths (e.g., optimal transport) leading to more stable and faster training compared to score-based methods, while maintaining the theoretical guarantees of normalizing flows.

3. Key Contributions

1. Likelihood Factorization (LF) Sampling: A training strategy that reduces the simulation budget per training sample to a single site by learning a per-site surrogate first. This is particularly advantageous when simulations are expensive.
2. TFMPE Architecture: The first method to combine tokenization (for irregular/functional data) with flow matching for hierarchical SBI. It supports function-valued observations and variable site counts within a single amortized estimator.
3. Hierarchical SBI Benchmark: The authors introduce the first dedicated benchmark for hierarchical SBI, adapting existing tasks (Gaussian Linear, SIR, Two Moons, etc.) to hierarchical structures. This enables systematic evaluation of sample efficiency.
4. Real-World Validation: Demonstrated on realistic, high-cost models:
   • Seasonal SEIR: Infectious disease modeling with irregular time-series data across 100 sites.
   • Computational Fluid Dynamics (CFD): Hemodynamics calibration with 1D blood flow networks, where simulator calls are computationally heavy.

4. Results

The paper evaluates TFMPE against non-hierarchical baselines (NPE, SNPE, Simformer, FMPE) and hierarchical baselines (Posterior Factorization).

• Sample Efficiency: On the hierarchical benchmark, TFMPE with LF sampling significantly outperforms non-hierarchical methods, which fail to scale as the number of sites ($n_s$) increases because they consume the simulation budget on multi-site sampling.
• Comparison with Posterior Factorization (PF):
  • TFMPE generally achieves better sample efficiency than PF methods, especially when the per-site observation model is easier to approximate than the global-to-observation mapping.
  • In tasks with complex nonlinear observation maps (e.g., Two Moons), PF performed slightly better, highlighting a trade-off: LF is superior when the surrogate approximation error is low; PF is safer when the global inference problem is simpler than the local surrogate.
• Scalability:
  • SEIR Model: TFMPE achieved well-calibrated posteriors for 100 sites using single-site simulations, whereas MCMC failed to converge beyond 5 sites due to computational intractability.
  • CFD Model: TFMPE reduced the computational cost of training by 2.0x for 2 patients and showed a speedup scaling linearly with the number of patients. The per-patient sampling speedup via the neural surrogate was ~1850x compared to the raw CFD simulator.
• Ablation Studies:
  • Group Embedding: Removing the explicit group ID token caused performance to collapse, proving that explicitly providing site membership is critical for the transformer to disentangle local parameters.
  • Architecture: While Transformers are powerful, an MLP-based version also performed well, suggesting the primary gain comes from the LF sampling strategy rather than the architecture alone.

5. Significance

This work addresses a major practical bottleneck in scientific computing: the high cost of simulating complex hierarchical systems.

• Tractability: It makes Bayesian inference feasible for problems with hundreds of sites and expensive simulators (like CFD or epidemiological models) where traditional MCMC or standard SBI is impossible.
• Flexibility: By using tokenization, the method naturally handles the irregular, function-valued data common in real-world science (e.g., varying observation times in disease surveillance).
• Efficiency: The Likelihood Factorization approach offers a distinct paradigm shift, trading a small amount of surrogate approximation error for massive gains in simulation efficiency, enabling "amortized" inference at scales previously unattainable.

In summary, TFMPE provides a robust, scalable, and sample-efficient framework for hierarchical Bayesian inference, bridging the gap between complex scientific simulators and modern deep learning inference techniques.