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

Manuel Gloeckler; J. P. Manzano-Patrón; Stamatios N. Sotiropoulos; Cornelius Schröder; Jakob H. Macke

stat.MLMar 16, 2026

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

Authors: Manuel Gloeckler, J. P. Manzano-Patrón, Stamatios N. Sotiropoulos, Cornelius Schröder, Jakob H. Macke

Organizations: Machine Learning in Science, University of Tübingen and Tübingen AI Center, Tübingen, Germany · Sir Peter Mansfield Imaging Centre, School of Medicine, University of Nottingham, UK · National Institute for Health Research (NIHR) Nottingham Biomedical Research Centre, Queens Medical Centre, Nottingham, United Kingdom · Max Planck Institute for Intelligent Systems, Department Empirical Inference, Tübingen, Germany

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.

Figures & tables

Figure 1: PRISM overview. (a) During training, we sample from a model family ${\mathcal{M}}$ with a hierarchical prior $p(\mathcal{M}\mid\lambda)$ , where $\lambda$ controls a chosen notion of model complexity and approximate jointly the model posterior and the conditional parameter posterior $p(\mathcal{M},\theta\mid{\bm{x}},\lambda)$ by optimizing the loss $\mathcal{L}=\mathcal{L}({\color[rgb]{1,0.3984,0}{{\mathcal{M}}}},p({\mathcal{M}}\mid{\bm{x}},{\color[rgb]{1,0.3984,0}{\lambda}}))+\mathcal{L}({\color[rgb]{0,0.6992,1}{\theta_{\mathcal{M}}}},p(\theta\mid{\color[rgb]{1,0.3984,0}{{\mathcal{M}}_{\lambda}}},{\bm{x}}))$ , for model evaluations ${\bm{x}}\sim p_{\color[rgb]{1,0.3984,0}{{\mathcal{M}}}}({\bm{x}}\mid{\color[rgb]{0,0.6992,1}{\theta_{\mathcal{M}}}})$ . (b) At inference time, we set $\lambda$ to tune parsimony, and select or explore models in the combinatorial space, infer model parameters, and deploy the resulting posteriors in downstream analyses.

Figure 2: Illustration on symbolic regression task. Left: Ground-truth function and noisy observations (black,

x<10

), compared to posterior predictives samples for two model complexities

\lambda

(95% credible interval and one noiseless sample each, chosen as the simplest posterior draw). Right: Median number of model components as function of

\lambda

. Below: Equations of true function (black, 2 components), and two sampled equations for low and high complexity priors (2 and 9 components).

Figure 3: Architecture overview. PRISM is based on two parallel transformers to infer (i) the model posterior from tokenized model masks and (ii) the parameter posterior via a diffusion

v

-prediction network (

t

denotes diffusion time). Observations

{\bm{x}}

enter through cross-attention;

(\lambda,t)

are injected via adaptive layer normalization (skip connections omitted).

Figure 4: PRISM on symbolic regression task. (a) Model posterior across $\lambda$ and parameter posterior for the example in Fig. 2 (for $\lambda=0.1$ ). (b) Comparison to ( Schröder and Macke, 2024 ) for a fixed prior. (c) Scaling to large model spaces for a fixed computational training budget; Red line indicates regimes beyond which not all models can be sampled during training.

Figure 6: PRISM on dMRI. (a) Two example voxels (V1, V2) with acquired dMRI signal (black), posterior predictives (color; 95% quantile) and complex multimodal parameter posteriors (bottom, right, all marginals in Fig. A.6 ). Model posteriors (bottom, left) closely match unbiased Monte Carlo estimates (black). (b) Left: Simulation-based calibration over $10^{5}$ simulations averaged over all $\lambda$ , ${\mathcal{M}}$ and ${\bm{\theta}}_{\mathcal{M}}$ . Right: Calibration error for individual models across the model space for two acquisition schemes (gray: expected error), indicating excellent calibration. (c) rKSD distance to the posterior on synthetic and experimental data for B3S-family models, for a ’fixed’ or a ’selected’ model on UKB data (see Fig. A.5 for HCP data). (Models with same label refer to differences in prior.)

Figure 7: Downstream Tractography. Comparisons against MCMC (BedpostX) and previous SBI approaches. (a) 3-way fiber crossings from the Centrum Semiovale area. (b) Uncertainty of the third stick (fiber) orientation in terms of dispersion. High uncertainty (dark) indicates the third fibre is not supported by the data. Last row: Tractography of brain regions based on inferred posteriors. (c) Tractwise correlation with MCMC, compared to SBI ${}_{\text{joint}}$ (tract labels from a).

Figure 8: PRISM on an extended model family for dMRI. (a) Calibration error for the BSZT model on average (left) and across model classes for the parameter posterior (right). (b) Reconstruction and generalization RMSE for reference models as well as for B3S, B3T, BSZT sel. (i.e. the Bayesian model average) (generalization estimated by leave-one-out cross-validation, LOOCV). (c) Examples of predicted fiber orientation distributions (Sec. B.2.2 ) for a local patch.

Appendix figures & tables20 assets

Supplementary material from the paper’s appendix.

Appendix

Figure A.1: (a) Illustration of symbolic regression given on the same ground truth function and noise model as in Fig. 2 but with 1000 observations. (b) Associated model and parameter posterior of the simple function (orange), which becomes more constrained compared to Fig. 2 with 10 observations.

Figure A.2: Evaluation metrics across symbolic regression tasks. The rRMSE ( a ), rKSD ( b ) and calibration error ( c ) across all considered model spaces (x-axis) and model sizes (colors). The gray shaded area indicates the expected Monte Carlo error if perfectly calibrated. Model sizes refer to model_dim

\in\{32,64,128\}

Figure A.3: Confusion matrix for model posterior network. For a subset of 200 randomly selected models (indicated by their binary mask, x/y axis, grouped by similarity) we show the confusion matrix averaged over random $\lambda$ for the symbolic regression tasks with $K=30$ , $K=50$ , $K=100$

Metric	15	30	50	80	100
Accuracy (Top-1)	0.796	0.730	0.690	0.560	0.503
Accuracy (Top-5)	0.981	0.995	0.995	0.940	0.905
Macro precision	0.813	0.776	0.741	0.647	0.597
Macro recall	0.797	0.742	0.688	0.548	0.503
Macro F1	0.794	0.733	0.674	0.525	0.463

Appendix

Table 1: Classification metrics for symbolic classification across different settings as derived from the confusion matrices in Fig. A.3 . Please note that even a Bayes optimal classifier will spread mass across multiple models.

Figure A.4: ESS and run times as a function of the number of diffusion steps for 1000 posterior samples for V2 B3S posterior in Fig. 6 . Evaluated on an Nvidia H100.

Figure A.5: Calibration for dMRI data. (a) Model posterior calibration in terms of SBC across

\lambda

for UKB- and HCP-like data. (b) rKSD on the UKB dataset for synthetic data (true model) and real data, comparing a fixed model to the model selected by the model-posterior network (same as Fig. 6 c). (c) Same as (b) for the HCP dataset.

	Synthetic		Real (UKB)
Metric	UKB-like	HCP-like	B1S	B2S	B3S
ESS	$0.615\,[0.098,\,0.962]$	$0.578\,[0.084,\,0.979]$	$0.177\,[0.021,\,0.695]$	$0.384\,[0.053,\,0.685]$	$0.276\,[0.070,\,0.621]$
KSD	$0.304\,[0.000,\,11.753]$	$0.141\,[0.000,\,11.867]$	$0.306\,[0.000,\,1.546]$	$0.389\,[0.000,\,1.390]$	$0.062\,[0.000,\,0.758]$
RMSE	$0.020\,[0.012,\,0.076]$	$0.020\,[0.012,\,0.076]$	$0.030\,[0.019,\,0.054]$	$0.023\,[0.015,\,0.034]$	$0.020\,[0.014,\,0.032]$

Appendix

Table 2: Parameter inference performance metrics. Median and 10th/90th quantiles of performance metrics on synthetic and real data for the effective sample size (ESS), the Kernel Stein discrepancy (KSD) and the predictive rooted mean square error (RMSE). Obtained via the tailored B3S model. Synthetic is over the whole model space, where real is restricted to specific models.

Figure A.6: Parameter posterior for Ball-and-Stick model. Full parameter posterior distribution conditioned on the B3S model (a) for V1 (b) for V2 as shown in Fig. 6 .

Figure A.7: B3S model selection evaluation. (a) Confusion matrix for all models. (b) Alignment between two Monte Carlo (MC) approximations of the evidence. (c) Alignment between the approximated evidence obtained from the PRISM model posterior probabilities compared to MC approximations. (d) Total variation (TV) distance across different model complexities $\lambda$ .

Figure A.8: SBC evaluation for extended dMRI model spaces. Calibration plot and error (over

\lambda

and models) for the model trained on (a ) the BSZT and (b) BSZT + conv space.

	Synthetic		Real (UKB)
Metric	UKB-like	HCP-like	B3S	B3Z	B3T
ESS	$0.109\,[0.026,\,0.962]$	$0.081\,[0.028,\,0.986]$	$0.133\,[0.045,\,0.302]$	$0.131\,[0.040,\,0.398]$	$0.165\,[0.061,\,0.366]$
KSD	$1.144\,[0.000,\,7.086]$	$1.776\,[0.000,\,10.411]$	$0.004\,[0.000,\,0.107]$	$0.000\,[0.000,\,0.014]$	$0.000\,[0.000,\,0.000]$
RMSE	$0.019\,[0.011,\,0.072]$	$0.019\,[0.011,\,0.072]$	$0.020\,[0.012,\,0.032]$	$0.020\,[0.012,\,0.032]$	$0.018\,[0.012,\,0.030]$

Appendix

Table 3: Parameter inference performance metrics. Median and 10th/90th quantiles of performance metrics on synthetic and real data for the effective sample size (ESS), the Kernel Stein discrepancy (KSD) and the predictive rooted mean square error (RMSE) for the BSZT space model. Synthetic is over the whole model space, where real is restricted to specific models.

	Synthetic		Real (UKB)
Metric	UKB-like	HCP-like	BS W S B	No B	2TS W S B No W
ESS	$0.073\,[0.025,\,0.453]$	$0.083\,[0.021,\,0.433]$	$0.109\,[0.028,\,0.464]$	$0.053\,[0.020,\,0.335]$	$0.135\,[0.043,\,0.410]$
KSD	$3.873\,[0.671,22.23]$	$6.597\,[0.922,35.49]$	$0.001\,[0.000,\,0.016]$	$1.000\,[0.089,\,4.306]$	$0.000\,[0.000,\,0.000]$
RMSE	$0.021\,[0.013,\,0.072]$	$0.019\,[0.012\,0.073]$	$0.022\,[0.013,\,0.049]$	$0.028\,[0.014,\,0.065]$	$0.018\,[0.012,\,0.031]$

Appendix

Table 4: Parameter inference performance metrics. Median and 10th/90th quantiles of performance metrics on synthetic and real data for the effective sample size (ESS), the Kernel Stein discrepancy (KSD) and the predictive rooted mean square error (RMSE) for the BSZT + conv model. Synthetic is over the whole model space, where real is restricted to specific models.

Figure A.9: Confusion matrices for PRISM on dMRI data with full model family. (a) on a 50 model subspace for the BSZT space model (b) on a 50 model subspace for the BSZT+conv space model

Figure A.10: Parameter posterior for dMRI. Parameter posterior for the two voxels shown in Fig. 6 as estimated by the BSZT-space model. (a) Posterior for the B3S model (different from previous due to unshared diffusivity and uninformative fraction prior). (b) Posterior for the B3T model showing large degeneracy due to symmetries and nonidentifiability on current datasets.

Figure A.11: Posterior predictive fiber orientation distributions for a classical spherical deconvolution based method (Rumba, ( Canales-Rodríguez et al., 2015 ) , as implemented in dipy ( Garyfallidis et al., 2014 ) ), presuming the B3S model, the B3T or by considering the global posterior predictive FOD averaged over all possible tensor-based models.

Metric	B3S space	BSZT space	BSZT + conv space
Accuracy (Top-1, Macro)	0.351	0.428	0.422
Accuracy (Top-5, Macro)	0.933	0.940	0.980
Macro precision	0.362	0.404	0.420
Macro recall	0.351	0.428	0.422
Macro F1	0.337	0.403	0.408

Appendix

Table 5: Classification metrics across BSZT spaces (macro-averaged).

Figure A.12: Model selection results on extended model spaces a Model selection frequencies if constrained to a set of B(1/2/3)(S/Z/T) models for the UKB and HCP dataset. b) Top 10 ranking of the whole BSZT space across lambda for the UKB dataset. c) Top 10 ranking of the BSZT+conv space across lambda for the UKB dataset.


Setting	Basis functions
$K{=}15$	Linear, Linear, Quadratic, ShiftedSquare, Cubic, Sinusoidal, Cosinusoidal, ConstantWide, ConstantPositive, TanhRight, TanhLeft, GaussianBump, GaussianWide, RampUp, RampDown
$K{=}30$	Linear, Quadratic, ShiftedSquare, Cubic, Sinusoidal, Cosinusoidal, ConstantWide, ConstantPositive, TanhRight, TanhLeft, TanhCentered, GaussianBump, GaussianWide, RampUp, RampDown, QuarticScaled, QuinticScaled, SinusoidalPhase, CosinusoidalPhase, ExponentialDecay, SaturatingExponential, Logarithmic, SquareRoot, Reciprocal, AbsoluteValue, InverseQuadratic, Lorentzian, Sigmoid, DampedSinusoidal, DampedCosinusoidal
$K{=}50$	Linear, Quadratic, ShiftedSquare, Cubic, Sinusoidal, Cosinusoidal, ConstantWide, ConstantPositive, TanhRight, TanhLeft, TanhCentered, GaussianBump, GaussianWide, RampUp, RampDown, QuarticScaled, QuinticScaled, SinusoidalPhase, CosinusoidalPhase, ExponentialDecay, SaturatingExponential, Logarithmic, SquareRoot, Reciprocal, AbsoluteValue, InverseQuadratic, Lorentzian, Sigmoid, DampedSinusoidal, DampedCosinusoidal, ExponentialGrowth, PowerLawDecay, ArctangentStep, HyperbolicSecant, SincDecay, AbsoluteSinusoidal, RectifiedLinear, Softplus, Gompertz, LinearFractional, SineSquared, TriangularBump, Linear, Quadratic, ShiftedSquare, Cubic, Sinusoidal, Cosinusoidal, ConstantWide, ConstantPositive
$K{=}80$	Linear, Quadratic, ShiftedSquare, Cubic, Sinusoidal, Cosinusoidal, ConstantWide, ConstantPositive, TanhRight, TanhLeft, TanhCentered, GaussianBump, GaussianWide, RampUp, RampDown, QuarticScaled, QuinticScaled, SinusoidalPhase, CosinusoidalPhase, ExponentialDecay, SaturatingExponential, Logarithmic, SquareRoot, Reciprocal, AbsoluteValue, InverseQuadratic, Lorentzian, Sigmoid, DampedSinusoidal, DampedCosinusoidal, ExponentialGrowth, PowerLawDecay, ArctangentStep, HyperbolicSecant, SincDecay, AbsoluteSinusoidal, RectifiedLinear, Softplus, Gompertz, LinearFractional, SineSquared, TriangularBump, Linear, Quadratic, ShiftedSquare, Cubic, Sinusoidal, Cosinusoidal, ConstantWide, ConstantPositive, TanhRight, TanhLeft, TanhCentered, GaussianBump, GaussianWide, RampUp, RampDown, QuarticScaled, QuinticScaled, SinusoidalPhase, CosinusoidalPhase, ExponentialDecay, SaturatingExponential, Logarithmic, SquareRoot, Reciprocal, AbsoluteValue, InverseQuadratic, Lorentzian, Sigmoid, DampedSinusoidal, DampedCosinusoidal, ExponentialGrowth, PowerLawDecay, ArctangentStep, HyperbolicSecant, SincDecay, AbsoluteSinusoidal, RectifiedLinear, Softplus
$K{=}100$	Linear, Quadratic, ShiftedSquare, Cubic, Sinusoidal, Cosinusoidal, ConstantWide, ConstantPositive, TanhRight, TanhLeft, TanhCentered, GaussianBump, GaussianWide, RampUp, RampDown, QuarticScaled, QuinticScaled, SinusoidalPhase, CosinusoidalPhase, ExponentialDecay, SaturatingExponential, Logarithmic, SquareRoot, Reciprocal, AbsoluteValue, InverseQuadratic, Lorentzian, Sigmoid, DampedSinusoidal, DampedCosinusoidal, ExponentialGrowth, PowerLawDecay, ArctangentStep, HyperbolicSecant, SincDecay, AbsoluteSinusoidal, RectifiedLinear, Softplus, Gompertz, LinearFractional, SineSquared, TriangularBump, Linear, Quadratic, ShiftedSquare, Cubic, Sinusoidal, Cosinusoidal, ConstantWide, ConstantPositive, TanhRight, TanhLeft, TanhCentered, GaussianBump, GaussianWide, RampUp, RampDown, QuarticScaled, QuinticScaled, SinusoidalPhase, CosinusoidalPhase, ExponentialDecay, SaturatingExponential, Logarithmic, SquareRoot, Reciprocal, AbsoluteValue, InverseQuadratic, Lorentzian, Sigmoid, DampedSinusoidal, DampedCosinusoidal, ExponentialGrowth, PowerLawDecay, ArctangentStep, HyperbolicSecant, SincDecay, AbsoluteSinusoidal, RectifiedLinear, Softplus, Gompertz, LinearFractional, SineSquared, TriangularBump, Linear, Quadratic, ShiftedSquare, Cubic, Sinusoidal, Cosinusoidal, ConstantWide, ConstantPositive, TanhRight, TanhLeft, TanhCentered, GaussianBump, GaussianWide, RampUp, RampDown, QuarticScaled

Appendix

Table 6: Explicit listing of base function (see Tab. 7 for definitions) in order i.e. mask index

i

corresponds to the i’th base function listed.

ID	Component	Token	Expression (implementation string)	Parameter prior
0	Linear	lin	c_1*x	$\;c_{1}\sim\mathcal{U}(-2,\,2)$
1	Quadratic	quad	c_1xx	$\;c_{1}\sim\mathcal{U}(-0.5,\,0.5)$
2	ShiftedSquare	shift2	(c_1+x)*(c_1+x)	$\;c_{1}\sim\mathcal{U}(-5,\,0)$
3	Cubic	cub	c_1xx*x	$\;c_{1}\sim\mathcal{U}(-0.1,\,0.1)$
4	Sinusoidal	sin	c_1sin(c_2x)	$c_{1}\sim\mathcal{U}(0,\,5)$ $c_{2}\sim\mathcal{U}(0.5,\,5)$
5	Cosinusoidal	cos	c_1cos(c_2x)	$c_{1}\sim\mathcal{U}(0,\,5)$ $c_{2}\sim\mathcal{U}(0.5,\,5)$

Appendix

Table 7: Extended component and noise catalogue (as implemented in code).

ID	Name	Symbol	Signal model	${\bm{\theta}}$	Prior
0	Ball	B	$\displaystyle S_{B}(b)=\exp\!\big(-b\,d_{\mathrm{iso}}\big)$	$d_{\mathrm{iso}}$	$\displaystyle d_{\mathrm{iso}}\sim\mathcal{U}\!\big(0,0.01\big)$
1	Stick	S	$\displaystyle S_{S}(b,\vec{\mathbf{b}})=\exp\!\left(-b\,d_{\parallel}(\vec{\mathbf{b}}\cdot\mathbf{n})^{2}\right)$	$d_{\parallel}$	$\displaystyle d_{\parallel}\sim\mathcal{U}\!\big(0,\;0.01\big)\$
1	Stick	S		$\mathbf{n}$	$\displaystyle\mathbf{n}\sim\mathcal{U}(\mathbb{S}^{2}_{+})$
2	Zeppelin	Z	$\displaystyle S_{Z}(b,\vec{\mathbf{b}})=\exp\!\left(-b\left[d_{\parallel}(\vec{\mathbf{b}}\cdot\mathbf{n})^{2}+d_{\perp}\big(1-(\vec{\mathbf{b}}\cdot\mathbf{n})^{2}\big)\right]\right)$	$d_{\parallel}$	$\displaystyle d_{\parallel}\sim\mathcal{U}\!\big(0,\;0.01\big)$
				$d_{\perp}$	$\displaystyle d_{\perp}\sim\mathcal{U}(0,0.01)$
				$\mathbf{n}$	$\displaystyle\mathbf{n}\sim\mathcal{U}(\mathbb{S}^{2}_{+})$

Appendix

Table 8: Diffusion MRI signal models ( DMIPY -style) with per-parameter priors.

Model Selection

Model selection chooses models, features, or configurations that balance predictive fit, generalization, complexity, and deployment constraints.

See all 96 papers

Program Synthesis for Simulation-Based Inference: Joint Model Selection and Parameter Estimation
Jul 20, 2026Siddharth Mishra-SharmaSimulation-Based Inference Parameter Estimation
Faculty of Computing & Data Sciences Boston University · Boston University
Diffusion Models in Simulation-Based Inference: A Tutorial Review
Dec 22, 2025Jonas Arruda, Niels Bracher, Ullrich Köthe +2Simulation-Based Inference Parameter Estimation
Bonn Center for Mathematical Life Sciences, Life & Medical Sciences Institute, University of Bonn, Germany · Center for Modeling, Simulation, & Imaging in Medicine (CeMSIM), Rensselaer Polytechnic Institute, NY, USA · Computer Vision and Learning Lab, Heidelberg University, Germany
Spectra: Exact Component Transport for Test-Time Prior Adaptation in Simulation-Based Inference
Oct 6, 2026Xin Zhao, Nico Scherf, Robert Trampel +2Simulation-Based Inference Test-Time Adaptation
Department of Neurophysics, Max Planck Institute for Human Cognitive and Brain Sciences, Leipzig, Germany · International Max Planck Research School on Cognitive NeuroImaging, Leipzig, Germany · Methods and Development Group Neural Data Science and Statistical Computing, Max Planck Institute for Human Cognitive and Brain Sciences, Leipzig, Germany +3

Jul 20, 2026J/K move · Enter open · S save

Program Synthesis for Simulation-Based Inference: Joint Model Selection and Parameter Estimation

Siddharth Mishra-Sharma

Faculty of Computing & Data Sciences Boston University · Boston University

Simulation-Based Inference Parameter Estimation Bayesian Inference LLM-Based Program Synthesis

Open paper AskPDF arXiv

Neural simulation-based inference enables parameter estimation for complex models, but typically requires the user to specify a simulator encoding a fixed model structure. We present a framework for joint model selection and parameter estimation that combines large language models for program synthesis with neural simulation-based inference. Given a natural language description of the system and data under investigation, an LLM proposes candidate simulator programs which are iteratively refined via feedback-driven mutation and evaluated using neural density estimation. The approach enables simulation-based inference over a pool of models, not just parameters within a fixed model. On benchmarks spanning deterministic dynamics, stochastic epidemic models, and dark matter substructure inference from gravitational-lensing images, the method identifies plausible model families from open-ended prompts, with accuracy that reflects the information content of the data and identifiability of candidate models.

Abstract

Model Selection

Explore similar work

Program Synthesis for Simulation-Based Inference: Joint Model Selection and Parameter Estimation

Diffusion Models in Simulation-Based Inference: A Tutorial Review

Spectra: Exact Component Transport for Test-Time Prior Adaptation in Simulation-Based Inference