Information Bridging
Introduction
Simulators and mechanistic models in the sciences are helpful tools for making infrences about a process that may be hard, or impossible, to observe its latent variables. A biolgoical example that motivated me to get in this area is inferring how bone morphogenetic protein (BMP) receptors and (protein) ligands interact with one another. You might say “binding affinity? why don’t we just use surface plasmon resonance (SPR) to figure this out?” but unfortunatley BMP receptors create a complex that sends a downstream gene expression signal through two type I and two type II receptors combining with a protein ligand. SPR is limited to one-to-one interactions, nullifying. Ok, no problem, why don’t we engineer some fluroescent surface marker on the proteins to indicate whether the receptors are bound or not. Besides ignoring the multiplexing you might need to do to visualize this, determining the steady-state of signal (when a type of k_d can be approximated for this complex) still only determines the rate at which a complex is created. However, extrcellular markers are only half the challenge of inferring gene expression changes - we still need to account for intracellular relays of information through SMAD proteins (and possibly others) that result in a downstream signal ofinterest.
Instead, one might be able to just simulate different realities using different parameter sets and then try to match the simualtions that match reality, then, decide which parameters most likely generated the observed data. This intutive approach is actually a very simplistic Approximate Bayesian Computation (ABC) algorithm. However, more recent methods allow for amortized inference using neural networks and some advantages when using neural networks to parameterize the likelihood, posterior, or likelihood ratio (really just a critic of which simulations and latent parameters go together). However, in truth, if very accurate posterior inferece is your goal, the ABC literature is rich and deep with methods that can fit to the particular nuance of a user’s scientific simulator: Each simulator is unique with challenges in its own way. Although, if you need to reuse or perform inference over multiple diffferent datasets, amortization makes sense, but not always.
One reason you may use a neural network to parameterize a likelihood, posterior, or likelihood ratio is when performing experimental design. In this setting, the scientific simulator not only has latent parameters $\theta$ but also a design, or decision, variable $\xi$. This decision variable changes the distribution of the likelihood in a continuous manner. Previous approaches performed this process in two-step manner: 1) approximate the inference object (define this earlier to shorten) and 2) then perform Bayesian optimziatoin wtih an acquisition function with the inference object. This is computationally expensive in learning the inference object for a given design(s) and in performing Baysian optimization. Should a simulator be very expensive to query, for example some particle physics simualtors, then this may take a long time…
This is where our recent UAI 2026 paper comes to the rescue. For normalizing flow-based SBI models (diffusion/flow matching are in waiting), we can actually jointly perform optimization of an inference object and experimental design (with the help of some tricks) over one, or multiple, designs. We show efficiency gains over alternative methods that also perform BOED but require that the simulator have differentiable designs, which may be the case where you can pre-simulate responses and use those in a grid, or, the simulator is diffrentiable w.r.t. its designs. However, the real kicker is its applicability to simulators that aren’t diffrentiable w.r.t. designs, such as a motivating example of the BMP simulator. There, we demonstrated gains over Bayesian optimization in terms of speed and quality of downstream inference.
To get this to work, we basically just leveraged contrastive learning via the InfoNCE mutual information bound with normalized density estimators (normalizing flows are normalized by design, which also makes them cumbersome to use in higher-dimensional cases). Theoretically, we should be able to take gradient steps w.r.t. both the inference object’s parameters and the experimental designs. In practice, this is difficult with a static experimental design. We overcome this issue by introducing a design distribution that provides the simple conditonal neural network with a bit of diversity in the experimental designs, indicating which direction designs may be more effective. This sprinkle of randomness helps the joint optimization process go from learning no inofrmative design to very informative ones (add figure here). We temper the design distribution over the number of optimization steps to go from wide to narrow, since a wider distirbuiton penalizes the EIG more than a narrow one. A satisfying gif of this process can be seen in Figure 2.
(figure of the ablation here as well as the gif of the design distribution optimization).
The coup de grâce of the paper is the application to the stubborn BMP simulator model. There, we demonstrate simulator efficiency over the two-step Bayesian optimization method and also better inference in the form of better median distance (the distance of all the posterior predictive samples from the observed data point - or average epsilon ball radius of error). The advantage of BO in the EIG can be explained by the biased estimate of the EIG using the InfoNCE bound that scales with (1/M TODO: look up this number more carefully) with the data. Looking at how long the optimization took for SBI-BOED also shows that the joint optimziaton is much faster for SBI-BOED.
Application to active learning in SBI
So far, this paper has really been about using inference objects that are workhorses in SBI to experimental design. One fun application in the reverse direction is for active learning of which $\theta$ samples to use to improve downstream inference. If simulations are expensive, then we want to prioritize which $\theta$ samples are used to generate joint $(y, \theta)$ samples. Here, we demonstrate an improvement on the two-moons toy simulator on the C2ST calibration metric, which is generally preferred as a measure of posterior performance (cite Leuckmann et al. paper from 2021/2022). Generally, we see improvements over simple sampling from the prior for different rounds of active learning, regardless of the $k$ and $\lambda$ parameters. One open question taht we weren’t able to answer in this paper is what value of $k$ and $\lambda$ are best for active learning. It seems that using a lambda value that decreases with the active learning round might be best, but this likely also depends on the underlying simulator used. In summary, we show how to extend the InfoNCE bound to help prioritize inferences for active learning, applying principles from BOED to benefit SBI.
Discussion
We introduced our new SBI-BOED method and shown some of the benefits on non-differentiable simualtors and also extensions to active learning. We have code to share it at (TODO: put code link here) and encourage folks to try it out. Future work will likely extend this to diffusion models and larger-scale models, where we note previous work (cite diffusion BOED paper) that could apply with modifications to the differentiable conditons.
Enjoy Reading This Article?
Here are some more articles you might like to read next: