Supercharging Simulation-Based Inference for Bayesian Optimal Experimental Design

TL;DR

This work addresses BOED under intractable likelihoods by leveraging SBI to compute EIG bounds and optimize experimental designs. It introduces a unified framework linking neural posterior, likelihood, and ratio estimators to variational EIG bounds, plus a novel neural likelihood-based EIG estimator and a robust multi-start gradient ascent (MPR-GA) optimization strategy. The key contributions include (i) explicit SBI-BOED connections across NPE, NLE, and NRE, (ii) a direct EIG estimator via neural likelihoods, and (iii) the MPR-GA method with a diversity penalty that substantially improves per-trajectory optimization, achieving gains up to 22% over prior state-of-the-art on benchmarks. The practical impact is a more flexible, scalable approach to data acquisition in expensive scientific settings, capable of matching or surpassing policy-based methods in several regimes, while also highlighting the potential of static design baselines in some scenarios.

Abstract

Bayesian optimal experimental design (BOED) seeks to maximize the expected information gain (EIG) of experiments. This requires a likelihood estimate, which in many settings is intractable. Simulation-based inference (SBI) provides powerful tools for this regime. However, existing work explicitly connecting SBI and BOED is restricted to a single contrastive EIG bound. We show that the EIG admits multiple formulations which can directly leverage modern SBI density estimators, encompassing neural posterior, likelihood, and ratio estimation. Building on this perspective, we define a novel EIG estimator using neural likelihood estimation. Further, we identify optimization as a key bottleneck of gradient based EIG maximization and show that a simple multi-start parallel gradient ascent procedure can substantially improve reliability and performance. With these innovations, our SBI-based BOED methods are able to match or outperform by up to $22\%$ existing state-of-the-art approaches across standard BOED benchmarks.

Figures (3)

• Figure 1: Posterior samples for the first (left) and second (center) sources after making observations at two design points. The observed designs $\xi$, the source locations and the distribution of candidate designs after running the MPR-GA acquisition optimization with 256 restarts for 12000 steps for the NRE method (right).
• Figure 2: Ablation on the number of restarts used in MPR-GA for different SBI methods. Using multiple restarts significantly improves performance across all SBI methods considered. Error bars show standard error across 50 random seeds. Different initialization strategies are considered: random designs (left) and sampling from the current posterior (middle). The timing of each method is also shown (right).
• Figure 3: Average sPCE lower bound after 10 measurements for CES experiments. Errors show standard error on the mean over different runs. For NPE-NRE 100 runs were performed, for all other methods 4096 runs were used.