[Paper Review] All-in-one simulation-based inference

Manuel Gloeckler, Michael Deistler+3 more|arXiv (Cornell University)|Apr 15, 2024

Simulation Techniques and Applications5 citations

TL;DR

The Simformer is a diffusion-model–based amortized SBI method that uses a transformer over joint parameter-data inputs to sample arbitrary conditionals (posterior, likelihood, or others) and handle unstructured data, function-valued parameters, and observation intervals.

ABSTRACT

Amortized Bayesian inference trains neural networks to solve stochastic inference problems using model simulations, thereby making it possible to rapidly perform Bayesian inference for any newly observed data. However, current simulation-based amortized inference methods are simulation-hungry and inflexible: They require the specification of a fixed parametric prior, simulator, and inference tasks ahead of time. Here, we present a new amortized inference method -- the Simformer -- which overcomes these limitations. By training a probabilistic diffusion model with transformer architectures, the Simformer outperforms current state-of-the-art amortized inference approaches on benchmark tasks and is substantially more flexible: It can be applied to models with function-valued parameters, it can handle inference scenarios with missing or unstructured data, and it can sample arbitrary conditionals of the joint distribution of parameters and data, including both posterior and likelihood. We showcase the performance and flexibility of the Simformer on simulators from ecology, epidemiology, and neuroscience, and demonstrate that it opens up new possibilities and application domains for amortized Bayesian inference on simulation-based models.

Motivation & Objective

• Motivate the need for more flexible, efficient amortized SBI that works with irregular/structured data and function-valued parameters.
• Develop a single model that can sample arbitrary conditionals of the joint distribution of parameters and data.
• Leverage transformers with diffusion models and attention masking to encode dependencies and allow conditioned/inferred outputs.
• Demonstrate the method on diverse simulators from ecology, epidemiology, and neuroscience to show broad applicability.

Proposed method

• Construct a token-based representation for all variables (parameters and data) as input to a transformer.
• Train a score-based diffusion model on the joint distribution p(theta, x) using denoising score matching.
• Introduce an attention mask M_E to encode dependency structures and allow sampling of arbitrary conditionals.
• Enable conditioning on intervals and constraints via guided diffusion to sample from constrained distributions.
• Use randomization of conditioning masks to train the model to output any conditional of the joint distribution.
• Allow sampling of posterior, likelihood, and other conditionals from a single trained network.

Experimental results

Research questions

• RQ1 Can a single diffusion-transformer model accurately sample arbitrary conditionals of the joint distribution p(theta, x)?
• RQ2 Does incorporating dependency-structure masks improve accuracy and efficiency for SBI tasks?
• RQ3 Can the method handle unstructured data, missing values, and function-valued (infinite-dimensional) parameters?
• RQ4 How does guided diffusion enable conditioning on observation intervals or additional constraints?
• RQ5 Is the approach effective across diverse simulators in ecology, epidemiology, and neuroscience?

Key findings

• Simformer outperforms neural posterior estimation on benchmark SBI tasks in terms of accuracy for a given simulation budget.
• Using domain-informed attention masks improves accuracy, especially for sparser dependency structures.
• On average, Simformer requires about 10x fewer simulations than NPE across tasks and observations.
• The model can sample arbitrary conditionals, including posterior and likelihood, and handle unstructured data and function-valued parameters.
• Guided diffusion enables conditioning on observation intervals and other constraints, influencing posterior and posterior-predictive samples.
• The method achieves accurate inference for Lotka-Volterra unstructured observations, SIRD with infinite-dimensional parameter spaces, and Hodgkin-Huxley with energy constraints.