Stochastic Gradient Bayesian Optimal Experimental Designs for
  Simulation-based Inference
Vincent D. Zaballa Elliot E. Hui
Simulation-based inference (SBI) methods tackle complex scientific models
with challenging inverse problems. However, SBI models often face a significant
hurdle due to their non-differentiable nature, which hampers the use of
gradient-based optimization techniques. Bayesian Optimal Experimental Design
(BOED) is a powerful approach that aims to make the most efficient use of
experimental resources for improved inferences. While stochastic gradient BOED
methods have shown promising results in high-dimensional design problems, they
have mostly neglected the integration of BOED with SBI due to the difficult
non-differentiable property of many SBI simulators. In this work, we establish
a crucial connection between ratio-based SBI inference algorithms and
stochastic gradient-based variational inference by leveraging mutual
information bounds. This connection allows us to extend BOED to SBI
applications, enabling the simultaneous optimization of experimental designs
and amortized inference functions. We demonstrate our approach on a simple
linear model and offer implementation details for practitioners.