Scalable Simulation-Based Model Inference with Test-Time Complexity Control

TL;DR

PRISM is a scalable simulation-based inference method that jointly infers model structures and parameters, allowing test-time control of model complexity, demonstrated on symbolic regression and diffusion MRI data.

Contribution

It introduces PRISM, a novel encoder-decoder framework for joint model and parameter inference with adjustable complexity at test time.

Findings

Scales to billions of model instantiations in synthetic tasks.

Successfully performs model selection in diffusion MRI data.

Enables flexible complexity control during inference.

Abstract

Simulation plays a central role in scientific discovery. In many applications, the bottleneck is no longer running a simulator; it is choosing among large families of plausible simulators, each corresponding to different forward models/hypotheses consistent with observations. Over large model families, classical Bayesian workflows for model selection are impractical. Furthermore, amortized model selection methods typically hard-code a fixed model prior or complexity penalty at training time, requiring users to commit to a particular parsimony assumption before seeing the data. We introduce PRISM, a simulation-based encoder-decoder that infers a joint posterior over both discrete model structures and associated continuous parameters, while enabling test-time control of model complexity via a tunable model prior that the network is conditioned on. We show that PRISM scales to families with combinatorially many (up to billions) of model instantiations on a synthetic symbolic regression task. As a scientific application, we evaluate PRISM on biophysical modeling for diffusion MRI data, showing the ability to perform model selection across several multi-compartment models, on both synthetic and in vivo neuroimaging data.