Supercharging Simulation-Based Inference for Bayesian Optimal Experimental Design

TL;DR

This paper introduces new methods for Bayesian optimal experimental design using simulation-based inference, leveraging multiple EIG formulations and neural likelihood estimation to improve optimization and performance.

Contribution

It presents a novel EIG estimator compatible with modern SBI density estimators and demonstrates enhanced optimization techniques for better experimental design outcomes.

Findings

Achieves up to 22% improvement over state-of-the-art methods.

Introduces multiple EIG formulations for SBI.

Develops a multi-start gradient ascent procedure for optimization.

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.