# Gradient-based Bayesian Experimental
Design for Implicit Models using
Mutual Information Lower Bounds

Steven<EMAIL_ADDRESS>[ Michael U.
<EMAIL_ADDRESS>[ University of Edinburgh , e3

(2021)

###### Abstract

We introduce a framework for Bayesian experimental design (BED) with implicit
models, where the data-generating distribution is intractable but sampling
from it is still possible. In order to find optimal experimental designs for
such models, our approach maximises mutual information lower bounds that are
parametrised by neural networks. By training a neural network on sampled data,
we simultaneously update network parameters and designs using stochastic
gradient-ascent. The framework enables experimental design with a variety of
prominent lower bounds and can be applied to a wide range of scientific tasks,
such as parameter estimation, model discrimination and improving future
predictions. Using a set of intractable toy models, we provide a comprehensive
empirical comparison of prominent lower bounds applied to the aforementioned
tasks. We further validate our framework on a challenging system of stochastic
differential equations from epidemiology.

62K05,

62L05,

Bayesian Experimental Design,

Likelihood-Free Inference,

Mutual Information,

Implicit Models,

Parameter Estimation,

Model Discrimination,

###### keywords:

[class=MSC]

###### keywords:

††volume: TBA††issue: TBA

and

##  1 Introduction

Parametric statistical models allow us to describe and study the behaviour of
natural phenomena and processes. These models attempt to closely match the
underlying process in order to simulate realistic data and, therefore, usually
have a complex form and involve a variety of variables. Consequently, they
tend to be characterised by intractable data-generating distributions
(likelihood functions). Such models are referred to as implicit models and
have an increasingly widespread use in the natural and medical sciences.
Notable examples include high-energy physics (Agostinelli et al., 2003),
cosmology (Schafer and Freeman, 2012), epidemiology (Corander et al., 2017),
cell biology (Ross et al., 2017) and cognitive science (Palestro et al.,
2018).

Statistical inference is a key component of many scientific goals such as
estimating model parameters, comparing plausible models and predicting future
events. Unfortunately, because the likelihood function for implicit models is
intractable, we have to revert to so-called likelihood-free, or simulation-
based, inference to solve these downstream tasks (see Cranmer et al., 2020 for
a recent overview). Likelihood-free inference has gained much traction
recently, with many methods leveraging advances in machine-learning (e.g.
Gutmann and Corander, 2016; Lueckmann et al., 2017; Järvenpää et al., 2018;
Chen and Gutmann, 2019; Papamakarios et al., 2019; Thomas et al., 2020;
Järvenpää et al., 2021). Ultimately, however, the quality of the statistical
inference within a scientific downstream task depends on the data that are
available in the first place. Because gathering data is often time-consuming
and expensive, we should therefore aim to collect data that allow us to solve
our specific task in the most efficient and appropriate manner possible.

Bayesian experimental design (BED) attempts to solve this problem by
appropriately allocating resources in an experiment (see Ryan et al., 2016 for
a review). Scientific experiments generally involve controllable variables,
called experimental designs $\mathbf{d}$, that affect the data gathering
process. These might, for instance, be measurement times and locations,
initial conditions of a dynamical system, or interventions and stimuli that
are used to perturb a natural process. BED aims to find optimal experimental
designs $\mathbf{d}^{\ast}$ that allow us to solve our scientific goals in the
most efficient way. As part of this, we have to construct and optimise a
utility function $U(\mathbf{d})$ that indicates the value of an experimental
design $\mathbf{d}$ according to the specific task at hand. In order to be
fully-Bayesian, the utility function generally has to be a functional of the
posterior distribution (Ryan et al., 2016). This exacerbates the difficulty of
BED for implicit models, as the involved posterior distributions are
intractable. Proposing appropriate utility functions and devising methods to
compute them efficiently for implicit models, with the help of likelihood-free
inference, has been the focus of much present-day research in the field of BED
(e.g. Drovandi and Pettitt, 2013; Price et al., 2018; Overstall and McGree,
2020).

A popular choice of utility function with deep roots in information theory is
the mutual information (MI; Lindley, 1972) between two random variables, which
describes the uncertainty reduction in one variable when observing the other.
The MI is a standard utility function in BED for models with tractable
likelihoods (e.g. Ryan, 2003; Paninski, 2005; Overstall and Woods, 2017), but
has only recently started to gain traction in BED for implicit models (e.g.
Price et al., 2018; Kleinegesse and Gutmann, 2019; Foster et al., 2019;
Kleinegesse et al., 2020a). In the context of implicit models, the main
difficulties that arise are 1) that estimating the posterior and obtaining
samples from it via likelihood-free inference is expensive and 2) that
gradients of the MI are generally not readily available, complicating its
optimisation.

Kleinegesse and Gutmann (2020b) recently addressed the aforementioned problems
by considering gradient-ascent on mutual information lower bounds.
Importantly, they only considered one type of lower bound in their work, while
there exist other bounds with different potentially helpful properties (see
Poole et al., 2019 for an overview). Moreover, their method only considers the
specific task of parameter estimation, even though other research goals are
often of equal practical importance.

In this paper, we extend the method of Kleinegesse and Gutmann (2020b) to
other lower bounds and scientific aims, thereby devising a general framework
for Bayesian experimental design based on mutual information lower bounds.
Importantly, our framework does not require the use of variational
distributions (as in Foster et al., 2019), but fully relies on parametrisation
by neural networks. We believe that this increases practicality and
generality, allowing scientists to use this method easily in a variety of
settings. In short, our contributions are:

  1. 1.

We devise a BED framework that allows for the use of general lower bounds on
mutual information in the context of implicit models, and

  2. 2.

use it to perform BED for parameter estimation, model comparison and improving
future predictions.

Figure 1 provides a visualisation of our general framework, showcasing how
training a neural network allows us to move towards optimal experimental
designs. Our method optimises a parametrised lower bound with gradient-ascent,
while simultaneously tightening it. This avoids spending resources on
estimating the MI accurately for non-optimal designs, which can be seen in the
right plot of Figure 1, where the lower bound is only tight near the optimal
design. Furthermore, once the optimal design has been found, our approach
provides an amortised approximation of the posterior distribution, avoiding an
additional costly likelihood-free inference step. An animation corresponding
to Figure 1 can be found in the public code repository that also includes the
code to reproduce all results in this paper:
https://github.com/stevenkleinegesse/GradBED.

Figure 1: Visualisation of our proposed framework. During early times of
training (left plot), the lower bound on mutual information is loose and the
current design is far from the optimal design. We move closer to the optimal
design while tightening the lower bound during the training process (middle
plot). Finally, we find the optimal design and obtain a tight lower bound at
the end of training (right plot). Note how we only spend resources on
tightening the lower bound near the optimal design.

We provide background knowledge of Bayesian experimental design with different
scientific aims in Section 2. In Section 3 we present our BED framework for
implicit models with different scientific aims and general lower bounds. This
includes an introduction of parametrised mutual information lower bounds for
BED, an explanation of how to optimise them and some practical guidance. In
Section 4 we test our framework on a number of implicit models, followed by a
discussion in Section 5.

##  2 Bayesian Experimental Design for different Aims

In this section we provide background on Bayesian experimental design (BED)
for different scientific aims using mutual information. Generally speaking,
the aim of BED is to identify experimental designs $\mathbf{d}$ that allow us
to achieve our scientific goals most rapidly. In order to do so, we first have
to construct a utility function $U(\mathbf{d})$ that tells us the worth of
doing an experiment with design $\mathbf{d}\in\mathcal{D}$, where
$\mathcal{D}$ specifies the domain of feasible experimental designs. Optimal
experimental designs $\mathbf{d}^{\ast}$ are then found by maximising the
utility function $U(\mathbf{d})$, i.e.

$\mathbf{d}^{\ast}=\operatorname*{arg\,max}_{\mathbf{d}\in\mathcal{D}}U(\mathbf{d}).$
(2.1)

Naturally, the choice of utility function $U(\mathbf{d})$ is crucial, as it
determines the optimal designs that are found. A popular utility function in
BED is the mutual information (MI) between two random variables, which
describes the uncertainty reduction, or expected information gain, in one
variable when observing the other. It is a desirable quantitity because of its
sensitivity to non-linearity and multi-modality, which other utilities cannot
effectively deal with (Ryan et al., 2016; Kleinegesse et al., 2020a).

We here introduce the notion of a variable of interest $\bm{v}$ that we wish
to learn about by means of a scientific experiment. As such, our aim is to
compute the MI $I(\bm{v};\mathbf{y}|\mathbf{d})$ between the variable of
interest $\bm{v}$ and the observed data $\mathbf{y}$ at a particular design
$\mathbf{d}$ (henceforth abbreviated by $\mathbf{y}|\mathbf{d}$). The MI
utility function can then conveniently be adapted to different scientific aims
by changing the variable of interest $\bm{v}$ that is used in its computation
(Ryan et al., 2016). Its general form is given by

$\displaystyle U(\mathbf{d})$ $\displaystyle=I(\bm{v};\mathbf{y}|\mathbf{d})$
(2.2) $\displaystyle=\int\int
p(\bm{v},\mathbf{y}|\mathbf{d})\log\left(\frac{p(\bm{v},\mathbf{y}|\mathbf{d})}{p(\mathbf{y}|\mathbf{d})p(\bm{v})}\right)\mathrm{d}\mathbf{y}\mathrm{d}\bm{v}$
(2.3) $\displaystyle=\int\int
p(\mathbf{y}|\bm{v},\mathbf{d})p(\bm{v})\log\left(\frac{p(\bm{v}|\mathbf{y},\mathbf{d})}{p(\bm{v})}\right)\mathrm{d}\mathbf{y}\mathrm{d}\bm{v},$
(2.4)

where we make the common assumption that the prior over the variable $\bm{v}$
is unaffected by the designs $\mathbf{d}$, i.e.
$p(\bm{v}|\mathbf{d})=p(\bm{v})$. In this form, the MI can also be
reformulated as the expected Kullback-Leibler (KL) divergence (Kullback and
Leibler, 1951) between the posterior $p(\bm{v}|\mathbf{y},\mathbf{d})$ and
prior $p(\bm{v})$. This interpretation means that we are looking for
experimental designs resulting in data that maximally increase the difference
between posterior and prior (Ryan et al., 2016). In the case of implicit
models, the data-generating distribution $p(\mathbf{y}|\bm{v},\mathbf{d})$
and, hence, the posterior distribution $p(\bm{v}|\mathbf{y},\mathbf{d})$ are
intractable, severely complicating the estimation and optimisation of Equation
2.4.

##### Parameter Estimation

The aim here is to find optimal designs that allow us to estimate the model
parameters $\bm{\theta}$ most effectively. In this scenario,
$\bm{v}\rightarrow\bm{\theta}$ and the utility function is the MI
$I(\bm{\theta};\mathbf{y}|\mathbf{d})$ between the model parameters
$\bm{\theta}$ and data $\mathbf{y}|\mathbf{d}$,

$U_{\text{PE}}(\mathbf{d})=\int\int
p(\mathbf{y}|\bm{\theta},\mathbf{d})p(\bm{\theta})\log\left(\frac{p(\bm{\theta}|\mathbf{y},\mathbf{d})}{p(\bm{\theta})}\right)\mathrm{d}\mathbf{y}\mathrm{d}\bm{\theta}.$
(2.5)

Parameter estimation is perhaps the most commonly-used scientific aim in BED
for implicit models (e.g. Drovandi and Pettitt, 2013; Hainy et al., 2016;
Price et al., 2018; Kleinegesse et al., 2020a; Kleinegesse and Gutmann,
2020b).

##### Model Discrimination

In this setting we wish to find designs that allow us to optimally
discriminate between competing models. Each competing model is assigned a
discrete model indicator $m$ that determines from which competing model the
data is sampled, i.e. $\mathbf{y}\sim p(\mathbf{y}|m,\mathbf{d})$. As such, we
let $\bm{v}\rightarrow m$ and define the utility function as the MI
$I(m;\mathbf{y}|\mathbf{d})$ between the model indicator $m$ and data
$\mathbf{y}|\mathbf{d}$,

$U_{\text{MD}}(\mathbf{d})=\sum_{m}\int
p(\mathbf{y}|m,\mathbf{d})p(m)\log\left(\frac{p(m|\mathbf{y},\mathbf{d})}{p(m)}\right)\mathrm{d}\mathbf{y}.$
(2.6)

Generally, each competing model $m$ also has its own model parameters
$\bm{\theta}_{m}$ that are needed for simulating data. In this particular
case, where we are only concerned with discriminating between models, these
model parameters are marginalised out. BED for model discrimination has
recently received increased attention in the context of implicit models (e.g.
Dehideniya et al., 2018a; Hainy et al., 2018).

##### Joint MD/PE

Combining the previous two tasks, we here wish to find designs that allow us
to optimally discriminate between models $m$ and estimate a particular model’s
parameters $\bm{\theta}_{m}$. For this joint task we have
$\bm{v}\rightarrow(\bm{\theta}_{m},m)$, with the utility function being the MI
$I(\bm{\theta}_{m},m;\mathbf{y}|\mathbf{d})$ between $(\bm{\theta}_{m},m)$ and
simulated data $\mathbf{y}|\mathbf{d}$,

$\displaystyle U_{\text{MDPE}}(\mathbf{d})=\sum_{m}\int
p(\mathbf{y}|\bm{\theta}_{m},m,\mathbf{d})p(\bm{\theta}_{m},m)\log\left(\frac{p(\bm{\theta}_{m},m|\mathbf{y},\mathbf{d})}{p(\bm{\theta}_{m},m)}\right)\mathrm{d}\mathbf{y}.$
(2.7)

There have only been a few studies on joint parameter estimation and model
discrimination in the context of BED for implicit models (e.g. Dehideniya et
al., 2018b).

##### Improving Future Predictions

The premise of this setting is slightly different to the previous ones. We
assume that we already know that we can perform an experiment with designs
$\mathbf{d}_{T}$ at some time in the future, yielding (yet unknown) data
$\mathbf{y}_{T}|\mathbf{d}_{T}$. However, we do have some budget to do a few
initial experiments with designs $\mathbf{d}$ and observe initial data
$\mathbf{y}$. Our aim is then to find the experimental designs that allow us
to optimally predict the future observations $\mathbf{y}_{T}$ given the
initial observations $\mathbf{y}$, i.e.
$\bm{v}\rightarrow\mathbf{y}_{T}|\mathbf{d}_{T}$. The corresponding utility is
the MI $I(\mathbf{y}_{T}|\mathbf{d}_{T};\mathbf{y}|\mathbf{d})$ between future
observations $\mathbf{y}_{T}|\mathbf{d}_{T}$ and initial observations
$\mathbf{y}$, conditioned on initial designs $\mathbf{d}$ (Chaloner and
Verdinelli, 1995):

$\displaystyle U_{\text{FP}}(\mathbf{d})=\int\int
p(\mathbf{y}_{T}|\mathbf{y},\mathbf{d},\mathbf{d}_{T})p(\mathbf{y}|\mathbf{d})\log\left(\frac{p(\mathbf{y}_{T}|\mathbf{y},\mathbf{d},\mathbf{d}_{T})}{p(\mathbf{y}_{T}|\mathbf{d}_{T})}\right)\mathrm{d}\mathbf{y}\mathrm{d}\mathbf{y}_{T},$
(2.8)

where we have made the assumptions that
$p(\mathbf{y}_{T}|\mathbf{d},\mathbf{d}_{T})=p(\mathbf{y}_{T}|\mathbf{d}_{T})$
and $p(\mathbf{y}|\mathbf{d},\mathbf{d}_{T})=p(\mathbf{y}|\mathbf{d})$.
Importantly, the future design $\mathbf{d}_{T}$ stays constant during
optimisation because we assume knowledge of this variable. Since we are only
concerned with improving future predictions, we here marginalise out the model
parameters $\bm{\theta}$. This utility function can also be interpreted as the
expected KL divergence between the posterior predictive and prior predictive
of future observations $\mathbf{y}_{T}$. As far as we are aware, there has
only been one work on improving future predictions in the context of BED for
implicit models (Liepe et al., 2013), but there has been some older work for
explicit models where the likelihood is tractable (e.g. Martini and
Spezzaferri, 1984; Verdinelli et al., 1993).

Table 1: Variable of interests and corresponding mutual information utilities for different scientific goals. Scientific Goal | Variable of Interest $\bm{v}$ | MI Utility
---|---|---
Parameter Estimation | $\bm{\theta}$ |  $I(\bm{\theta};\mathbf{y}|\mathbf{d})$.
Model Discrimination | $m$ | $I(m;\mathbf{y}|\mathbf{d})$
Joint MD/PE | $\bm{\theta}_{m},m$ | $I(\bm{\theta}_{m},m;\mathbf{y}|\mathbf{d})$
Improving Future Predictions | $\mathbf{y}_{T}|\mathbf{d}_{T}$ | $I(\mathbf{y}_{T}|\mathbf{d}_{T};\mathbf{y}|\mathbf{d})$

We summarise the above formulations in Table 1; however, this list is by no
means exhaustive. In fact, because we only have to change the variable of
interest $\bm{v}$ when computing the MI $I(\bm{v};\mathbf{y}|\mathbf{d})$, we
can easily incorporate a multitude of scientific goals. Furthermore, the
formalism in Equation 2.4 allows us to devise a general framework for BED with
implicit models that is agnostic to a particular scientific goal.

##  3 General BED Framework with Lower Bounds on MI

In this section we explain our general BED framework for implicit models based
on mutual information lower bounds. We first provide motivation for using MI
lower bounds and rephrase the BED problem accordingly. This is followed by
examples of prominent lower bounds from literature. We identify common
structures between them and explain how to use these to derive gradients with
respect to designs, allowing us to solve the rephrased BED problem with
gradient-based optimisation. Finally, we address a few technical difficulties
and how to overcome them.

###  3.1 Rephrasing the BED Problem

Even though it is an effective and desirable metric of dependency, mutual
information is notoriously difficult to estimate. This difficulty is
exacerbated for implicit models, where we do not have access to the data-
generating distribution in Equation 2.4. Recent work devised tractable mutual
information estimators by combining variational bounds with machine-learning
(see Poole et al., 2019 for an overview). These methods generally work by
using a neural network to parametrise a lower bound on mutual information,
which is then tightened by training the neural network with gradient-ascent.
There also exist some methods that use variational approximations to
intractable distributions within lower bounds (e.g. Donsker and Varadhan,
1983; Barber and Agakov, 2003; Foster et al., 2019). However, specifying a
family of variational distributions carries the risk of mis-specification due
to choosing an overly simple variational family. Parametrising non-linear
functions such as neural networks is simpler than parametrising probability
distributions and less prone to mis-specification.111Given that the neural
network has enough capacity to capture the true distribution. This difficulty
is amplified in the context of implicit models, where we might have little
knowledge about the underlying (intractable) distributions. Upper bounds on
the mutual information generally require tractable distributions or
variational approximations as well (e.g. Poole et al., 2019; Cheng et al.,
2020).222We note that this might be because upper bounds on MI have been
consistently understudied. Thus, in this work we shall only be considering
lower bounds on mutual information that are solely parametrised by neural
networks.

Our overall goal is to maximise the MI $I(\bm{v};\mathbf{y}|\mathbf{d})$
between a variable of interest $\bm{v}$ and observed data $\mathbf{y}$ with
respect to the design $\mathbf{d}$. Following the notation of Belghazi et al.
(2018), we represent the scalar output of a neural network by
$T_{\bm{\psi}}(\bm{v},\mathbf{y})$, where $\bm{\psi}$ are the network
parameters. We denote lower bounds on MI by
$\mathcal{L}(\bm{\psi},\mathbf{d})$; these depend on $\bm{\psi}$ via the
neural network $T_{\bm{\psi}}(\bm{v},\mathbf{y})$ and on $\mathbf{d}$ via the
distributions over which expectations are taken. Our aim is then to maximise
$\mathcal{L}(\bm{\psi},\mathbf{d})$ with respect to $\bm{\psi}$ and
$\mathbf{d}$. The BED optimisation problem in Equation 2.1 becomes

$\mathbf{d}^{\ast}=\operatorname*{arg\,max}_{\mathbf{d}\in\mathcal{D}}\max_{\bm{\psi}}\mathcal{L}(\bm{\psi},\mathbf{d}).$
(3.1)

As we maximise $\mathcal{L}(\bm{\psi},\mathbf{d})$ with respect to $\bm{\psi}$
we tighten the lower bound, while maximisation with respect to $\mathbf{d}$
allows us to find the optimal design. This optimisation problem is especially
difficult for implicit models, as $\mathcal{L}(\bm{\psi},\mathbf{d})$
generally involves expectations over distributions which depend on
$\mathbf{d}$ in some complex manner. We will first discuss some prominent
lower bounds from literature in Section 3.2 and then discuss how to optimise
them in Section 3.3.

###  3.2 Prominent Lower Bounds

##### NWJ

The lower bound $\mathcal{L}_{\text{NWJ}}$ was first developed by Nguyen,
Wainwright and Jordan (NWJ; Nguyen et al., 2010) but is also known as $f$-GAN
KL (Nowozin et al., 2016) and MINE-$f$ (Belghazi et al., 2018):

$\mathcal{L}_{\text{NWJ}}(\bm{\psi},\mathbf{d})\equiv\mathbb{E}_{p(\bm{v},\mathbf{y}\mid\mathbf{d})}\left[T_{\bm{\psi}}(\bm{v},\mathbf{y})\right]-e^{-1}\mathbb{E}_{p(\bm{v})p(\mathbf{y}\mid\mathbf{d})}\left[e^{T_{\bm{\psi}}(\bm{v},\mathbf{y})}\right].$
(3.2)

The optimal critic for this lower bound, i.e. when it is tight, is
$T^{\ast}_{\bm{\psi}}(\bm{v},\mathbf{y})=1+\log{\frac{p(\bm{v}|\mathbf{y},\mathbf{d})}{p(\bm{v})}}$.
In other words, training a neural network $T_{\bm{\psi}}(\bm{v},\mathbf{y})$
with Equation 3.2 as the objective function amounts to learning an amortised
density ratio of the posterior to prior distribution. The NWJ bound has been
shown to have a low bias but high variance (Poole et al., 2019; Song and
Ermon, 2020).

##### InfoNCE

The InfoNCE lower bound $\mathcal{L}_{\text{NCE}}$ was introduced by van den
Oord et al. (2018) in the context of contrastive representation learning and
takes the form

$\displaystyle\mathcal{L}_{\text{NCE}}(\bm{\psi},\mathbf{d})$
$\displaystyle\equiv\mathbb{E}\left[\frac{1}{K}\sum_{i=1}^{K}\log{\frac{e^{T_{\bm{\psi}}(\bm{v}_{i},\mathbf{y}_{i})}}{\frac{1}{K}\sum_{j=1}^{K}\left[e^{T_{\bm{\psi}}(\bm{v}_{j},\mathbf{y}_{i})}\right]}}\right],$
(3.3)

where the expectation is over $K$ independent samples from the joint
distribution $p(\bm{v},\mathbf{y}|\mathbf{d})$. The optimal critic for this
lower bound is
$T^{\ast}_{\bm{\psi}}(\bm{v},\mathbf{y})=\log{p(\mathbf{y}|\bm{v},\mathbf{d})}+c(\mathbf{y}|\mathbf{d})$
for the form shown above, where $c(\mathbf{y}|\mathbf{d})$ is an indeterminate
function. However, if we swap the indices of $\bm{v}_{j}$ and $\mathbf{y}_{i}$
in the denominator, i.e. sum out $\mathbf{y}$, the optimal critic becomes
$T^{\ast}_{\bm{\psi}}(\bm{v},\mathbf{y})=\log{p(\bm{v}|\mathbf{y},\mathbf{d})}+c(\bm{v})$,
where $c(\bm{v})$ is again indeterminate. We use the version in Equation 3.3,
as this allows us to obtain an unnormalised posterior that we can sample from,
or normalise numerically in low dimensions (since $c(\mathbf{y}|\mathbf{d})$
is fixed for a given real-world observation). The $\mathcal{L}_{\text{NCE}}$
lower bound has a low variance but, as noted by van den Oord et al. (2018) and
verified by Poole et al. (2019), it is upper-bounded by $\log{K}$, where $K$
is the batch-size in Equation 3.3, leading to a potentially large bias.

##### JSD

We are generally concerned with maximising the mutual information and not
estimating it to a high accuracy. Hjelm et al. (2019) therefore argued to use
the Jensen-Shannon divergence (JSD) as a proxy to mutual information. They
proposed to use a lower bound on the JSD between the joint distribution
$p(\bm{v},\mathbf{y}|\mathbf{d})$ and product of marginal distributions
$p(\bm{v})p(\mathbf{y}|\mathbf{d})$, which takes the form

$\mathcal{L}_{\text{JSD}}(\bm{\psi},\mathbf{d})\equiv\mathbb{E}_{p(\bm{v},\mathbf{y}\mid\mathbf{d})}\left[-\text{sp}(-T_{\bm{\psi}}(\bm{v},\mathbf{y}))\right]-\mathbb{E}_{p(\bm{v})p(\mathbf{y}\mid\mathbf{d})}\left[\text{sp}(T_{\bm{\psi}}(\bm{v},\mathbf{y}))\right],$
(3.4)

where $\text{sp}(z)=\log(1+e^{z})$ is the softplus function and the optimal
critic of this lower bound is
$T^{\ast}_{\bm{\psi}}(\bm{v},\mathbf{y})=\log\frac{p(\bm{v}|\mathbf{y},\mathbf{d})}{p(\bm{v})}$.
The authors generally found the JSD lower bound to be more stable than other
MI lower bounds. The JSD and KL divergence are both $f$-divergences and have a
close relationship that can be derived analytically (see e.g. Hjelm et al.,
2019). The authors also showed experimentally that distributions that lead to
the highest JSD also lead to the highest MI, suggesting that maximising JSD is
an appropriate proxy for maximising MI. If one is in fact concerned with
accurate MI estimation as well, one could use the learned log density ratio in
a MC estimation of the MI, or substitute it in any of the other MI lower
bounds (as done in Poole et al., 2019; Song and Ermon, 2020). We note that the
JSD lower bound can be viewed as an expectation of the objective used in
likelihood-free inference by ratio estimation (LFIRE; Thomas et al.,
2020)333See the supplementary material for more information on this
relationship., which has been used before in the context of Bayesian
experimental design (Kleinegesse and Gutmann, 2019; Kleinegesse et al.,
2020a), and is, like InfoNCE, related to noise-contrastive estimation (NCE;
Gutmann and Hyvärinen, 2012) as discussed by Hjelm et al. (2019).

We want to emphasise again that this work aims to showcase the use of general
mutual information lower bounds, parametrised by neural networks, in Bayesian
experimental design for different scientific aims. We do not claim that the
aforementioned list of lower bounds is exhaustive and do not suggest that some
bounds are particularly superior to others. We leave this investigation to
more comprehensive studies focused exclusively on mutual information lower
bounds (such as in Poole et al., 2019).

###  3.3 Optimisation of Lower Bounds

In order to solve the rephrased BED problem in Equation 3.1, we need to be
able to optimise the lower bound $\mathcal{L}(\mathbf{d},\bm{\psi})$ with
respect to the neural network parameters $\bm{\psi}$ and the experimental
designs $\mathbf{d}$, regardless of the choice of lower bound. In the interest
of scalability (Spall, 2003), we here optimise both of these with gradient-
based methods.

Looking at MI lower bounds in literature, including the prominent ones shown
in Section 3.2, we can identify a few similar structures that allow us to
generalise gradient-based optimisation. First, lower bounds generally involve
expectations over the joint distribution and/or the product of marginal
distributions. Second, these expectations tend to contain non-linear,
differentiable functions of the neural network output. In other words, lower
bounds usually involve expectations
$\mathbb{E}_{p(\bm{v},\mathbf{y}\mid\mathbf{d})}\left[f(T_{\bm{\psi}}(\bm{v},\mathbf{y}))\right]$
and/or
$\mathbb{E}_{p(\bm{v})p(\mathbf{y}\mid\mathbf{d})}\left[g(T_{\bm{\psi}}(\bm{v},\mathbf{y}))\right]$,
where $f:\mathbb{R}\rightarrow\mathbb{R}$ and
$g:\mathbb{R}\rightarrow\mathbb{R}$ are non-linear, differentiable functions.
For instance, for the NWJ lower bound in Equation 3.2 we have $f(z)=z$ and
$g(z)=e^{z-1}$. By deriving gradients of these expectations with respect to
$\bm{\psi}$ and $\mathbf{d}$ we can then easily derive gradients for all
prominent lower bounds listed in the Section 3.2, as well as other ones that
share the same structure.

Fortunately, gradients of the aforementioned expectations with respect to the
network parameters $\bm{\psi}$ are straightforward, as the involved
distributions do not depend on them. This means that we can simply pull the
gradient operator $\nabla_{\bm{\psi}}$ inside the expectations and apply the
chain rule, yielding

$\displaystyle\nabla_{\bm{\psi}}\mathbb{E}_{p(\bm{v},\mathbf{y}\mid\mathbf{d})}\left[f(T_{\bm{\psi}}(\bm{v},\mathbf{y}))\right]$
$\displaystyle=\mathbb{E}_{p(\bm{v},\mathbf{y}\mid\mathbf{d})}\left[f^{\prime}(T)\Bigr{|}_{T=T_{\bm{\psi}}(\bm{v},\mathbf{y})}\nabla_{\bm{\psi}}T_{\bm{\psi}}(\bm{v},\mathbf{y})\right]\quad\text{and}$
(3.5)
$\displaystyle\nabla_{\bm{\psi}}\mathbb{E}_{p(\bm{v})p(\mathbf{y}\mid\mathbf{d})}\left[g(T_{\bm{\psi}}(\bm{v},\mathbf{y}))\right]$
$\displaystyle=\mathbb{E}_{p(\bm{v})p(\mathbf{y}\mid\mathbf{d})}\left[g^{\prime}(T)\Bigr{|}_{T=T_{\bm{\psi}}(\bm{v},\mathbf{y})}\nabla_{\bm{\psi}}T_{\bm{\psi}}(\bm{v},\mathbf{y})\right],$
(3.6)

where $f^{\prime}(T)=\partial f(T)/\partial T$ and $g^{\prime}(T)=\partial
g(T)/\partial T$. The first factor inside the expectations only depends on the
form of the non-linear functions, while the second factor, the gradients of
the network output with respect to its parameters, is generally computed with
automatic differentiation, i.e. back-propagation. The expectations can then be
approximated as sample averages with $N$ samples from the corresponding
distributions.

Gradients with respect to the designs $\mathbf{d}$ are more complicated, as
the intractable distributions over which expectations are taken depend on
$\mathbf{d}$. For instance, we cannot compute
$\nabla_{\mathbf{d}}\mathbb{E}_{p(\bm{v},\mathbf{y}\mid\mathbf{d})}\left[f(T_{\bm{\psi}}(\bm{v},\mathbf{y}))\right]$
in the same way that we computed the previous gradients with respect to
$\bm{\psi}$, because the joint distribution $p(\bm{v},\mathbf{y}|\mathbf{d})$
(and therefore its gradient) is intractable. Similarly, we cannot use score-
function estimators (Kleijnen and Rubinstein, 1996) to approximate these
gradients, as they require an analytic derivative of the log densities.

However, the pathwise gradient estimator lends itself to our setting with
implicit models (see Mohamed et al., 2020 for a review). Indeed, like implicit
models, this estimator assumes that sampling from the data-generation
distribution $p(\mathbf{y}|\bm{v},\mathbf{d})$ is exactly the same as sampling
from a base distribution $p(\bm{\epsilon})$ and then using that noise sample
as an input to a non-linear, deterministic function
$\mathbf{h}(\bm{\epsilon};\bm{v},\mathbf{d})$ (Mohamed et al., 2020), called
the sampling path, i.e.444For instance, sampling a Gaussian random variable
from $\mathcal{N}(y;\mu,\sigma^{2})$ is exactly the same as first sampling
noise from a standard normal $\mathcal{N}(\epsilon;0,1)$ and then computing
$y=\mu+\epsilon\sigma$.

$\mathbf{y}\sim
p(\mathbf{y}|\bm{v},\mathbf{d})\iff\mathbf{y}=\mathbf{h}(\bm{\epsilon};\bm{v},\mathbf{d}),\quad\bm{\epsilon}\sim
p(\bm{\epsilon}).$ (3.7)

This then allows us to invoke the law of the unconscious statistician (LOTUS;
e.g. Grimmett and Stirzaker, 2001) and, for instance, rephrase the expectation
over the joint $p(\bm{v},\mathbf{y}|\mathbf{d})$ in terms of the base
distribution $p(\bm{\epsilon})$ and the prior over $p(\bm{v})$,

$\mathbb{E}_{p(\bm{v},\mathbf{y}\mid\mathbf{d})}\left[f(T_{\bm{\psi}}(\bm{v},\mathbf{y}))\right]=\mathbb{E}_{p(\bm{v})p(\bm{\epsilon})}\left[f(T_{\bm{\psi}}(\bm{v},\mathbf{h}(\bm{\epsilon};\bm{v},\mathbf{d})))\right],$
(3.8)

where we have factorised the joint distribution as
$p(\bm{v},\mathbf{y}|\mathbf{d})=p(\mathbf{y}|\bm{v},\mathbf{d})p(\bm{v})$,
assuming that $p(\bm{v})$ is unaffected by $\mathbf{d}$.

For the expectation over the product of marginals we need to perform an
additional step. Equation 3.7 specifies the sampling path
$\mathbf{h}(\bm{\epsilon};\bm{v},\mathbf{d})$ of the data-generation
distribution, but we require the sampling path of the marginal data
$\mathbf{y}\sim p(\mathbf{y}|\mathbf{d})$. This means that we need to express
the marginal as an expectation of the data-generating distribution over the
prior, i.e.
$p(\mathbf{y}|\mathbf{d})=\mathbb{E}_{p(\widetilde{\bm{v}})}[p(\mathbf{y}|\widetilde{\bm{v}},\mathbf{d})]$,
where $p(\widetilde{\bm{v}})$ is exactly the same distribution as $p(\bm{v})$.
We can then specify
$\mathbf{y}=\mathbf{h}(\bm{\epsilon};\widetilde{\bm{v}},\mathbf{d})$ and
invoke LOTUS again, yielding the second rephrased expectation:

$\mathbb{E}_{p(\bm{v})p(\mathbf{y}\mid\mathbf{d})}\left[g(T_{\bm{\psi}}(\bm{v},\mathbf{y}))\right]=\mathbb{E}_{p(\bm{v})p(\widetilde{\bm{v}})p(\bm{\epsilon})}\left[g(T_{\bm{\psi}}(\bm{v},\mathbf{h}(\bm{\epsilon};\widetilde{\bm{v}},\mathbf{d})))\right].$
(3.9)

The rephrased expectations in Equations 3.8 – 3.9 are now taken over
distributions that do not directly depend on the designs $\mathbf{d}$,
allowing us to easily take gradients with respect to $\mathbf{d}$. In the
interest of space, let us define
$T=T_{\bm{\psi}}(\bm{v},\mathbf{h}(\bm{\epsilon};\bm{v},\mathbf{d}))$ and
$\widetilde{T}=T_{\bm{\psi}}(\bm{v},\mathbf{h}(\bm{\epsilon};\widetilde{\bm{v}},\mathbf{d}))$.
The required gradients are then

$\displaystyle\nabla_{\mathbf{d}}\mathbb{E}_{p(\bm{v},\mathbf{y}\mid\mathbf{d})}\left[f(T)\right]$
$\displaystyle=\mathbb{E}_{p(\bm{v})p(\bm{\epsilon})}\left[f^{\prime}(T)\nabla_{\mathbf{d}}T\right]\quad\text{and}$
(3.10)
$\displaystyle\nabla_{\mathbf{d}}\mathbb{E}_{p(\bm{v})p(\mathbf{y}\mid\mathbf{d})}\left[g(T)\right]$
$\displaystyle=\mathbb{E}_{p(\bm{v})p(\widetilde{\bm{v}})p(\bm{\epsilon})}\left[g^{\prime}(\widetilde{T})\nabla_{\mathbf{d}}\widetilde{T}\right],$
(3.11)

where the derivatives $f^{\prime}(T)=\partial f(T)/\partial T$ and
$g^{\prime}(\widetilde{T})=\partial g(\widetilde{T})/\partial\widetilde{T}$
depend on the lower bound in question. The $\nabla_{\mathbf{d}}T$ and
$\nabla_{\mathbf{d}}\widetilde{T}$ factors in the above equations are the
gradients of the network output with respect to designs, which is the most
difficult technical part of our method. We discuss different approaches and
caveats to computing these gradients in detail in Section 3.4.

We next briefly consider a special setting where the implicit model is defined
such that we can separate the data generation into a ‘known’ observational
process555With ‘known’, we mean that the model of the observational process is
known analytically in closed form. $p(\mathbf{y}|\bm{v},\mathbf{d},\bm{z})$
and a differentiable, stochastic latent process $p(\bm{z}|\mathbf{d},\bm{v})$,
where $\bm{z}$ is a latent variable. In this case we can apply a score-
function estimator to the observational process, because we assume knowledge
of the likelihood $p(\mathbf{y}|\bm{v},\mathbf{d},\bm{z})$, and a pathwise
gradient estimator to the latent process. In a similar manner to Equation 3.8,
we can, for instance, rephrase the expectation of
$f(T_{\bm{\psi}}(\bm{v},\mathbf{y}))$ over the joint distribution as

$\displaystyle\mathbb{E}_{p(\bm{v},\mathbf{y}\mid\mathbf{d})}\left[f(T_{\bm{\psi}}(\bm{v},\mathbf{y}))\right]$
$\displaystyle=\mathbb{E}_{p(\mathbf{y}|\bm{v},\mathbf{d},\bm{z})p(\bm{z}|\mathbf{d},\bm{v})p(\bm{v})}\left[f(T_{\bm{\psi}}(\bm{v},\mathbf{y}))\right]$
(3.12)
$\displaystyle=\mathbb{E}_{p(\mathbf{y}|\bm{v},\mathbf{d},\mathbf{h}(\bm{\epsilon};\bm{v},\mathbf{d}))p(\bm{\epsilon})p(\bm{v})}\left[f(T_{\bm{\psi}}(\bm{v},\mathbf{y}))\right],$
(3.13)

where in this case $\bm{\epsilon}$ is the noise random variable that defines
the sampling path of the latent stochastic variable
$\bm{z}=\mathbf{h}(\bm{\epsilon};\bm{v},\mathbf{d})$. This allows us to pull
the gradient operatior $\nabla_{\mathbf{d}}$ into the required expectations;
see the supplementary material for more details on the resulting estimator.
This rephrasing is useful because 1) we can inject extra knowledge into the
method as we know parts of the data generation process and 2) this also allows
us to deal with some implicit models that generate discrete data (in which
case the full gradient with respect to $\mathbf{d}$ is undefined). For
instance, in experiments later in the paper, we present an implicit model
where the latent process is the solution of a stochastic differential
equation, which is differentiable by nature, and the observational process is
a known discrete Poisson process.

Lastly, we apply the aforementioned methodology of computing gradients to all
of the prominent lower bounds listed in Section 3.2, as they only effectively
differ in what kind of non-linear functions $f$ and $g$ they use, yielding the
gradients shown in Table 2. A similar table for the special setting of knowing
the observational process is shown in the supplementary materials.
Furthermore, we provide an overview of the discussed framework of Bayesian
experimental design for implicit models using lower bounds on mutual
information in Algorithm 1.

Table 2: Gradients of several lower bounds with respect to experimental designs. We define $T=T_{\bm{\psi}}(\bm{v},\mathbf{h}(\bm{\epsilon};\bm{v},\mathbf{d}))$, $\widetilde{T}=T_{\bm{\psi}}(\bm{v},\mathbf{h}(\bm{\epsilon};\widetilde{\bm{v}},\mathbf{d}))$ and $T_{ij}=T_{\bm{\psi}}(\bm{v}_{j},\mathbf{h}(\bm{\epsilon}_{i};\bm{v}_{i},\mathbf{d}))$. We also use the shorthand $P=p(\bm{v})p(\bm{\epsilon})$ and $Q=p(\bm{v})p(\widetilde{\bm{v}})p(\bm{\epsilon})$. Here, $\sigma(T)$ is the logistic sigmoid function. See the supplementary material for detailed derivations and a corresponding table for the case where the observation model is analytically tractable. Lower Bound | Gradients with respect to designs
---|---
NWJ | $\mathbb{E}_{P}\left[\nabla_{\mathbf{d}}T\right]-e^{-1}\mathbb{E}_{Q}\left[e^{\widetilde{T}}\nabla_{\mathbf{d}}\widetilde{T}\right]$
InfoNCE | $\mathbb{E}_{P^{K}}\left[\frac{1}{K}\sum_{i=1}^{K}\frac{\sum_{j=1}^{K}e^{T_{ij}}(\nabla_{\mathbf{d}}T_{ii}-\nabla_{\mathbf{d}}T_{ij})}{\sum_{j=1}^{K}e^{T_{ij}}}\right]$
JSD | $\mathbb{E}_{P}\left[\sigma(-T)\nabla_{\mathbf{d}}T\right]-\mathbb{E}_{Q}\left[\sigma(\widetilde{T})\nabla_{\mathbf{d}}\widetilde{T}\right]$
Algorithm 1 Gradient-Based BED for Implicit Models using MI Lower Bounds

1:Implicit model sampling path $\mathbf{h}(\bm{\epsilon};\bm{v},\mathbf{d})$,
neural network $T_{\bm{\psi}}(\bm{v},\mathbf{y})$, MI lower bound
$\mathcal{L}(\bm{\psi},\bm{v})$, optimisers for both $\bm{\psi}$ and
$\mathbf{d}$, number of samples $N$

2:Optimal design $\mathbf{d}^{\ast}$, trained neural network
$T_{\bm{\psi}^{\ast}}(\bm{v},\mathbf{y})$, MI estimate
$\mathcal{L}(\bm{\psi}^{\ast},\mathbf{d}^{\ast})$

3:

4:Sample from the prior over the variable of interest: $\bm{v}^{(i)}\sim
p(\bm{v})$ for $i=1,\dots,N$

5:Randomly initialise the neural network parameters $\bm{\psi}$

6:Randomly initialise the experimental designs $\mathbf{d}$

7:while $\mathcal{L}(\bm{\psi},\mathbf{d})$ not converged do

8: Sample noise from the base distribution, i.e. $\bm{\epsilon}^{(i)}\sim
p(\bm{\epsilon})$ for $i=1,\dots,N$

9: Obtain joint data via the sampling path
$\mathbf{y}^{(i)}=\mathbf{h}(\bm{\epsilon}^{(i)},\bm{v}^{(i)},\mathbf{d})$ for
$i=1,\dots,N$

10: Randomly shuffle data to obtain marginal samples
$\\{\widetilde{\mathbf{y}}^{(i)}\\}_{i=1}^{N}$

11: Compute network output with joint samples:
$T_{\bm{\psi}}(\bm{v}^{(i)},\mathbf{y}^{(i)})$ for $i=1,\dots,N$

12: Compute network output with marginal samples:
$T_{\bm{\psi}}(\bm{v}^{(i)},\widetilde{\mathbf{y}}^{(i)})$ for $i=1,\dots,N$

13: Compute a sample average of the MI lower bound
$\mathcal{L}(\bm{\psi},\mathbf{d})$

14: Estimate gradients with respect to $\bm{\psi}$ and $\mathbf{d}$

15: Update $\bm{\psi}$ and $\mathbf{d}$ using two separate optimisers

###  3.4 Tackling the Sampling Path Gradients

Recall the gradients of the expectations in Equations 3.10 – 3.11,which
involve the gradient of the network output with respect to designs,
i.e.$\nabla_{\mathbf{d}}T_{\bm{\psi}}(\bm{v},\mathbf{h}(\bm{\epsilon};\bm{v},\mathbf{d}))$.
Below we discuss several approaches to computing, or approximating, this
gradient.

##### Automatic Differentiation

As mentioned previously, the gradients of lower bounds with respect to
$\bm{\psi}$ are already computed with automatic differentiation in a standard
neural network training fashion. We can similarly use automatic
differentiation to compute gradients with respect to $\mathbf{d}$ if the
implicit model is written in a differential programming framework, such as JAX
(Bradbury et al., 2018), PyTorch (Paszke et al., 2019) or TensorFlow (Abadi et
al., 2015). This allows us to attach the simulator model to the computation
graph of the framework and back-propagate from evaluations of the lower bound
to the experimental designs $\mathbf{d}$. This means that we do not have to
explicitly provide gradients, which is helpful if the implicit model is
extremely complex, and allows us to fully-utilise parallelisation with GPUs.
Furthermore, we expect the usefulness of this aspect, and the capability of
our method, to grow with the development of automatic differentiation
frameworks, which are already essential in machine learning.

##### Manual Gradient Computation

In some situations we may need to manually compute the required gradients, for
example when the implicit model is not written in a differential programming
framework. To do so, we first apply the chain rule to
$\nabla_{\mathbf{d}}T_{\bm{\psi}}(\bm{v},\mathbf{h}(\bm{\epsilon};\bm{v},\mathbf{d}))$,
yielding

$\nabla_{\mathbf{d}}T_{\bm{\psi}}(\bm{v},\mathbf{h}(\bm{\epsilon};\bm{v},\mathbf{d}))=\mathbf{J_{y}}^{\top}\nabla_{\mathbf{y}}T_{\bm{\psi}}(\bm{v},\mathbf{y})\bigr{|}_{\mathbf{y}=\mathbf{h}(\bm{\epsilon};\bm{v},\mathbf{d})}.$
(3.14)

The factor $\nabla_{\mathbf{y}}T_{\bm{\psi}}(\bm{v},\mathbf{y})$ is the
derivative of the neural network output with respect to its inputs, which can
be readily obtained via automatic differentiation in most machine-learning
frameworks. $\mathbf{J_{y}}$ is the Jacobian matrix, defined by
$(\mathbf{J_{y}})_{ij}=\partial y_{i}/\partial d_{j}$, and contains the
gradients of the sampling path with respect to the designs. We either need to
assume knowledge of the sampling path gradients or find ways to approximate
them. This includes common settings where the data generation can be written
down in one line, e.g. a linear model with a combination of noises from
different non-Gaussian sources, or we have access to the ordinary, or
stochastic, differential equations that govern data-generation.

##### Gradient-Free Optimisation

While we can facilitate gradient-based optimisation for a variety of implicit
models, there might be situations in which we can neither exactly compute nor
approximate the required Jacobian-vector product through any of the
aforementioned ways. For instance, sampling path gradients might be undefined
if the data generation process involves any discrete variables that we need to
differentiate and we cannot separate it into a known observational process and
a latent process. In such situations we can only resort to gradient-free
optimisation methods for the updates of $\mathbf{d}$. Examples of gradient-
free optimisation techniques for BED with implicit models are grid search,
random search, evolutionary algorithms (as done in e.g. Zhang et al., 2021)
and Bayesian optimisation (as done in e.g. Kleinegesse and Gutmann, 2020b). We
note that in some situations we might be able to use continuous relaxations of
discrete variables instead, which we shall leave for future work. While the
gradient-free methods allow us to deal with a larger class of implicit models,
they do not scale well with the dimensionality of experimental designs
$\mathbf{d}$ (Spall, 2003). As such, we do not consider the gradient-free
setting in this work, but refer the reader to Kleinegesse and Gutmann (2020b)
for more detailed explanations and experiments for such a setting.

##  4 Experiments

In this section we demonstrate our approach to BED for implicit models using
MI lower bounds. First, we consider a toy example that consists of a set of
models with linear, logarithmic and square-root responses. In a comprehensive
comparison, we apply all prominent lower bounds presented in Section 3.2 (NWJ,
InfoNCE and JSD) to all scientific goals presented in Section 2 (PE, MD, MD/PE
and FP). Second, we consider a setting from epidemiology where the latent
process is the solution to a stochastic differential equation and the discrete
observational process is analytically tractable. We here only apply the JSD
lower bound on the PE and MD tasks. Code to reproduce results can be found
here: https://github.com/stevenkleinegesse/GradBED.

###  4.1 Toy Models

We here assume that an observable variable $y\in\mathbb{R}$ depends on an
experimental design $d\in[-2,2]$ in either a linear ($m=1$), logarithmic
($m=2$) or square-root ($m=3$) manner, where the variable $m\in\\{1,2,3\\}$ is
the model indicator. Each model $m$ is also governed by two model parameters
$\bm{\theta}_{m}=(\theta_{m,0},\theta_{m,1})$, e.g. the offset and slope in
the case of the linear model. We assume a limited budget of $10$ measurements,
which means that we have to construct a 10-dimensional design vector
$\mathbf{d}=(d_{1},\dots,d_{10})$ with a corresponding data vector
$\mathbf{y}=(y_{1},\dots,y_{10})$. We include two separate sources of noise,
Gaussian noise $\mathcal{N}(\epsilon;0,1)$ and Gamma noise $\Gamma(\nu;2,2)$,
which leads to all toy models having likelihood functions that do not admit a
closed-form expression. The sampling paths are given by

$\mathbf{y}=\begin{cases}\theta_{1,0}\mathbf{1}+\theta_{1,1}\mathbf{d}+\bm{\epsilon}+\bm{\nu}&\text{if
}m=1,\\\
\theta_{2,0}\mathbf{1}+\theta_{2,1}\log(|\mathbf{d}|)+\bm{\epsilon}+\bm{\nu}&\text{if
}m=2,\\\
\theta_{3,0}\mathbf{1}+\theta_{3,1}\sqrt{|\mathbf{d}|}+\bm{\epsilon}+\bm{\nu}&\text{if
}m=3,\end{cases}$ (4.1)

where $\bm{\epsilon}=(\epsilon_{1},\dots,\epsilon_{10})$ and
$\bm{\nu}=(\nu_{1},\dots,\nu_{10})$ are iid samples from a Gaussian and Gamma
noise source, respectively, $\mathbf{1}$ denotes a 10-dimensional vector of
ones and $|\cdot|$ implies taking the element-wise absolute value. To ensure
numerical stability in the logarithmic model, we clipped $|\mathbf{d}|$ by
$10^{-4}$ from below, i.e. $|\mathbf{d}|\equiv\max(|\mathbf{d}|,10^{-4})$. We
generally use a Gaussian prior
$\mathcal{N}(\bm{\theta}_{m};\mathbf{0},3^{2}\mathds{1})$ over the model
parameters and a discrete uniform prior $\mathcal{U}(m)$ over the model
indicators.

Comprehensive information about neural network architectures, hyper-parameters
and implementation details can be found in the supplementary material.
Throughout this section we compare final MI estimates and posterior
distributions to reference values and distributions. Information about how
these are computed for all scientific tasks can also be found in the
supplementary materials.

####  4.1.1 Parameter Estimation for the Linear Model

Figure 2: PE results for the linear toy model. Shown are training curves for
the NWJ, InfoNCE and JSD (via NWJ) lower bounds. The top row shows lower bound
estimates as a function of training epochs, with the dotted line being the
reference MI value, and the bottom row shows the elements of the design vector
as it is being updated.

We first consider the scientific task of parameter estimation (PE) for the
linear toy model ($m=1$ in Equation B.2), where the aim is to estimate the
model parameters $\bm{\theta}_{1}=(\theta_{1,0},\theta_{1,1})$, i.e. the slope
and offset of a straight line, respectively. Figure 2 shows the training
results for the PE task, using the NWJ (left column), InfoNCE (middle column)
and JSD (right column) lower bounds. As explained in Section 3.2, we use the
density ratio learned by maximising the JSD lower bound as an input to the NWJ
lower bound, in order to get an actual lower bound on the mutual information,
which can be seen in the right column. All lower bounds converge smoothly to a
final value that is close to the reference MI value, with all lower bounds
finding the same optimal design $\mathbf{d}^{\ast}$, whose elements are
equally clustered at $d=-2$ and $d=2$. Experimental designs at the boundaries
of the design domain are useful because that is where the signal-to-noise
ratio is highest. Intuitively, this means that our data contains more signal
than noise, allowing us to better measure the effect of model parameters on
data. Interestingly, there is no noticeable difference in the convergence
behaviour of each lower bound (they all use the same neural network and design
initialisations).

Table 3: Toy model average MI estimates ($\pm$ standard error) for all scientific tasks, using the NWJ, InfoNCE and JSD (via NWJ) lower bounds on MI (higher values are better). Also shown are average reference MI estimates.  MI LB | PE | MD | MD/PE | FP
---|---|---|---|---
NWJ | 3.43 $\pm$ 0.06 | 0.726 $\pm$ 0.007 | 3.72 $\pm$ 0.06 | 1.33 $\pm$ 0.01
InfoNCE | 3.36 $\pm$ 0.06 | 0.725 $\pm$ 0.027 | 3.80 $\pm$ 0.07 | 1.33 $\pm$ 0.04
JSD | 3.48 $\pm$ 0.05 | 0.730 $\pm$ 0.007 | 3.95 $\pm$ 0.08 | 1.34 $\pm$ 0.01
Ref. | 3.55 $\pm$ 0.04 | 0.737 $\pm$ 0.007 | 3.97 $\pm$ 0.02 | 1.34 $\pm$ 0.02

The PE column in Table 3 shows the final MI estimates for each lower bound,
evaluated on several synthetic validation sets, with the JSD lower bound
having the highest estimate. This is intriguing because, while the gradients
are computed with the JSD lower bound, we evaluate it with NWJ lower bound
and, yet, its estimate is higher than that for the actual NWJ lower bound.

Figure 3: PE posterior results for the linear toy model. Shown are average
marginal posteriors for all lower bounds, including the prior, the ground
truth and the average reference posterior. Note how all lower bounds produce
similar posteriors.

As explained in Section 3.2, we can use the trained neural network to directly
compute a posterior distribution. We first generate synthetic ‘real-world’
data $\mathbf{y}^{\ast}$ that was sampled with some ground-truth
$\bm{\theta}_{1,\text{true}}=(2,3)$ at the optimal design $\mathbf{d}^{\ast}$,
i.e. $\mathbf{y}^{\ast}\sim
p(\mathbf{y}|\mathbf{d}^{\ast},\bm{\theta}_{1,\text{true}})$. The prior
distribution, the trained neural network and the observation can then be used
to compute the posterior $p(\bm{\theta}|\mathbf{y}^{\ast},\mathbf{d}^{\ast})$
according to Section 3.2. The top row in Figure 3 shows the marginal
posteriors of the offset $\theta_{1,0}$ and slope $\theta_{1,1}$ for each
lower bound, averaged over $5{,}000$ ‘real-world’ observations. Also shown are
average reference posteriors computed with the same observations.666See the
supplementary material for how we compute these. All posterior distributions
cover the ground truth well, and their MAP estimate is close to the ground
truth. However, there is no noticeable difference between any of the shown
posterior distributions. The relative performance of each lower bound is more
apparent in Table 4, where we show the average KL-Divergence between estimated
and reference posteriors. The JSD lower bound results in posterior
distributions that are closest to the reference posteriors.

Table 4: Toy model average KL-Divergences ($\pm$ standard error) between estimated and reference posteriors for all scientific tasks and all lower bounds (lower values are better).  MI LB | PE | MD | MD/PE | FP
---|---|---|---|---
NWJ | 0.079 $\pm$ 0.003 | 0.016 $\pm$ 0.001 | 0.728 $\pm$ 0.010 | 0.009 $\pm$ 0.001
InfoNCE | 0.107 $\pm$ 0.003 | 0.018 $\pm$ 0.001 | 0.053 $\pm$ 0.001 | 0.012 $\pm$ 0.001
JSD | 0.028 $\pm$ 0.001 | 0.009 $\pm$ 0.001 | 0.026 $\pm$ 0.001 | 0.005 $\pm$ 0.001

####  4.1.2 Model Discrimination

Next, we consider the task of model discrimination (MD), where the aim is to
distinguish between competing models and the variable of interest is the model
indicator $m\in\\{1,2,3\\}$. To sample data $\mathbf{y}\sim
p(\mathbf{y}|\mathbf{d},m)$, as required in Algorithm 1, we need to
marginalise over the model parameters $\bm{\theta}_{m}$. To do so, we first
obtain prior samples, use these to sample data $\mathbf{y}$ from the sampling
path in Equation B.2 and then simply discard the $\bm{\theta}_{m}$ samples.

Figure 4 shows the training results for the MD task, using the NWJ (left
column), InfoNCE (middle column) and JSD (right column) lower bounds. Similar
to the PE task, we evaluate the JSD bound by using the learned density ratio
as an input to the NWJ lower bound. All lower bounds have a similar
convergence behaviour and lead to final MI estimates that are close to the
reference MI value. Similar to the PE tasks, all lower bounds find the same
optimal design $\mathbf{d}^{\ast}$, which consists of $3$ elements at $d=-2$,
$4$ elements at $d=0$ and $4$ elements at $d=2$. The additional design cluster
around $d=0$, as compared to the PE task, is useful for model discrimination
because of the large response of the logarithmic model ($m=2$). Although we
clipped $|\mathbf{d}|$ from below by $10^{-4}$ when $m=2$, this still means
that the response of the logarithmic model is significantly larger than that
of the linear model ($m=1$) and the square-root model ($m=3$) as $d\rightarrow
0$. This allows us to determine with relative ease whether or not observed
data was generated from the logarithmic model. Similarly, the other two
clusters at the boundaries of the design domain are helpful because that is
where the data distributions of all models are most different, which is
helpful in distinguishing between them.777See the supplementary material for a
figure showing this.

Figure 4: MD results for the set of toy models. Shown are training curves for
the NWJ, InfoNCE and JSD (via NWJ) lower bounds. The top row shows lower bound
estimates as a function of training epochs, with the dotted line being the
reference MI value, and the bottom row shows the elements of the design vector
as it is being updated.

Evaluations of the lower bounds on validation sets are presented in Table 3,
showing that while all lower bounds perform well, the JSD lower bound results
in the highest MI lower bound estimate. We note that, even though all lower
bounds had the same neural network and design initialisations, the InfoNCE
lower bound in particular would sometimes find slightly different (locally)
optimal designs; we show an example of this in the supplementary materials.

Figure 5: MD posterior results for the set of toy models. Shown are average
posterior probabilities for different ground truth models (one per row) for
the NWJ, InfoNCE and JSD lower bounds.

Similar to the PE task, we can use the trained neural networks to find
posterior distributions over the model indicator $m$, which express our
posterior belief that a given model has generated the observed ‘real-world’
data $\mathbf{y}^{\ast}$. We generate three different sets of ‘real-world’
data, one for each of the three models being the underlying true model. We
also use the same model parameter ground truth
$\bm{\theta}_{m,\text{true}}=(2,3)$ for each model. Figure 5 shows average
posterior distributions for all lower bounds, including an average reference
posterior distribution. All methods result in extremely good model recovery,
meaning that they can, on average, confidently identify from which model the
data was generated. The difference in performance between the different lower
bounds becomes more pronounced when considering the average KL-Divergence
between estimated and reference posteriors, as shown in Table 4. The average
KL-Divergences show that the JSD lower bound results in posterior
distributions that are slightly closer to the reference than NWJ and InfoNCE.

####  4.1.3 Joint MD/PE

Here we consider the joint task of model discrimination and parameter
estimation (MD/PE), i.e. a combination of the previous two tasks. This means
that our variable of interest is now the set of model indicator and
corresponding model parameters, i.e. $\bm{v}=(\bm{\theta}_{m},m)$. The MD/PE
results are similar to previous results seen for the separate parameter
estimation and model discrimination tasks. As such, we only show the relevant
training curves and posterior distributions in the supplementary materials.
The training curves look similar for all lower bounds and the optimal designs
$\mathbf{d}^{\ast}$ were also approximately the same. The elements of
$\mathbf{d}^{\ast}$ were clustered around $d=-2$, $d=0$ and $d=2$, similar to
those found in the MD task. In the MD/PE column in Table 3 we present final MI
estimates for each lower bound, evaluated on several validation data sets. The
JSD lower bound performs significantly better than the NWJ and InfoNCE lower
bounds, i.e. it is closer to the reference value. Similarly, Table 4 shows
that the posteriors estimated via JSD have, on average, a lower KL-Divergence
to reference posteriors, compared to the other lower bounds. Furthermore, the
NWJ lower bound results in an extremely large average KL-Divergence, which is
in part due to a poor parameter estimation for the logarithmic model (see the
supplementary materials for a plot).

####  4.1.4 Improving Future Predictions for the Linear Model

Finally, we consider the task of improving future predictions (FP), as
discussed in Section 2, for the linear toy model. The aim here is to find
current optimal designs $\mathbf{d}^{\ast}$, and gather corresponding data
$\mathbf{y}^{\ast}$, that allow us to maximally improve our predictions of
future data $y_{T}$ gathered at a (fixed and known) future design $d_{T}$. Our
variable of interest is thus the future data $y_{T}$. As before, we construct
a 10-dimensional design vector $\mathbf{d}$ that has elements restricted to
the domain $d\in[-2,2]$, with a corresponding 10-dimensional data vector
$\mathbf{y}$. For simplicity, we assume that our future design is one-
dimensional and fixed at $d_{T}=4$, which is outside the domain of our current
design vector. This emulates a setting where, for instance, we know that we
will be able to make measurements in a specific geographical location of
interest but currently only have access to a different, limited region,
allowing us to improve our prediction at that future measurement location.
This setting naturally assumes that the implicit model extrapolates well to
future designs $d_{T}$ when they are outside the current design domain.

The training curves for the NWJ, InfoNCE and JSD lower bound again show a very
similar convergence behaviour and yield the same optimal designs. As such, we
only show the training curves in the supplementary materials. The optimal
designs $\mathbf{d}^{\ast}$ have elements that are clustered equally at the
boundaries, which is exactly the same as for the parameter estimation task.
These designs are useful for improving future predictions, as the parameter
estimation is best (see the PE section) and the absolute values are close to
that of the future design. Table 3 shows that the final MI estimates for each
lower bound are similar and nearly match the reference value, with the JSD
lower bound yielding a slightly closer estimate.

Figure 6: FP results for the linear toy model. Shown are average posterior
predictive distributions for the NWJ, InfoNCE and JSD lower bounds. Shown are
also the prior predictive and a histogram of ‘real-world’ samples at the
future design.

For the FP task, we can use the trained neural networks to obtain a posterior
predictive distribution $p(y_{T}|d_{T},\mathbf{y}^{\ast},\mathbf{d}^{\ast})$
of our variable of interest $y_{T}|d_{T}$. To generate ‘real-world’ data
$\mathbf{y}^{\ast}$ at $\mathbf{d}^{\ast}$ we use ground-truth values of
$\bm{\theta}_{\text{true}}=(2,3)$ for the offset and slope, respectively. In
Figure 6 we show the average posterior predictive distributions, averaged over
$5{,}000$ different $\mathbf{y}^{\ast}$, for all lower bounds, including the
prior predictive distribution and a histogram showing samples of the data-
generating distribution of $y_{T}$ at $d_{T}$ with $\bm{\theta}_{\text{true}}$
(which is not a predictive distribution). The posterior predictive
distributions for all lower bounds are similar to each other and all have a
mode that matches that of the histogram of ‘real-world’ data at $d_{T}$. The
differences between the lower bounds becomes more noticeable when looking at
the average KL-Divergence values shown in Table 4. As for the other scientific
tasks, the JSD lower bound results in posteriors that have the lowest average
KL-Divergence to reference posteriors.

###  4.2 SDE Epidemiology Model

Figure 7: Example simulations of the SDE-based SIR (left) and SEIR (right)
models.

We here consider the spread of a disease within a population of $N$
individuals, modelled by stochastic versions of the well-known SIR (Allen,
2008) and SEIR (Lekone and Finkenstädt, 2006) models. In the SIR model,
indicated by $m=1$, individuals start in a susceptible state $S(t)$ and can
then move to an infectious state $I(t)$ with an infection rate of $\beta$.
These infectious individuals then move to a recovered state $R(t)$ with a
recovery rate of $\gamma$, after which they can no longer be infected. The SIR
model, governed by the state changes $S(t)\rightarrow I(t)\rightarrow R(t)$,
thus has two model parameters $\bm{\theta}_{1}=(\beta,\gamma)$. In the SEIR
model, indicated by $m=2$, susceptibles first move to an additional exposed
state $E(t)$, where individuals are infected but not yet infectious.
Afterwards, they move to the infectious state $I(t)$ with a rate of $\sigma$.
The SEIR model ($m=2$), governed by $S(t)\rightarrow E(t)\rightarrow
I(t)\rightarrow R(t)$, thus has three model parameters
$\bm{\theta}_{2}=(\beta,\sigma,\gamma)$. We further make the common assumption
that the total population size $N$ stays constant.

The stochastic versions of these epidemiological processes are usually defined
by a continuous-time Markov chain (CTMC), from which we can sample via the
Gillespie algorithm (see Allen, 2017). However, this generally yields discrete
population states that have undefined gradients. In order to test our
gradient-based algorithm, we thus resort to an alternative simulation
algorithm that uses stochastic differential equations (SDEs), where gradients
can be approximated. Figure 7 shows example simulations of the SDE-based SIR
and SEIR models, generated according to the method below.

We first define population vectors $\mathbf{X}_{1}=(S_{1}(t),I_{1}(t))^{\top}$
for the SIR model and $\mathbf{X}_{2}=(S_{2}(t),E_{2}(t),I_{2}(t))^{\top}$ for
the SEIR model. We can effectively ignore the population of recovered because
the total population is fixed.888Allowing us to, e.g., compute
$R_{2}(t)=N-S_{2}(t)-E_{2}(t)-I_{2}(t)$ for the SEIR model. The system of Itô
SDEs for the above epidemiological processes is

$\mathrm{d}\mathbf{X}_{m}(t)=\mathbf{f}_{m}(\mathbf{X}_{m}(t))\mathrm{d}t+\mathbf{G}_{m}(\mathbf{X}_{m}(t))\mathrm{d}\mathbf{W}(t),$
(4.2)

where $m$ is a model indicator, $\mathbf{f}_{m}$ is called the drift vector,
$\mathbf{G}_{m}$ is called the diffusion matrix, and $\mathbf{W}(t)$ is a
vector of independent Wiener processes (also called Brownian motion). The
drift vector and diffusion matrix for both models are derived in detail in the
supplementary materials. Sample paths of the above SDEs can be simulated by
means of the Euler-Maruyama algorithm (see Allen, 2017), which is based on
finite differences and discretised Wiener processes. In our experiments we use
the torchsde package (Li et al., 2020) to do this efficiently.

Moreover, we assume that we only have access to discrete, noisy estimates of
the number of infectious individuals $I(t)$, as other population states might
be difficult to measure in reality. We thus sample a single noisy measurement
of $I_{m}(t)$ by

$y|t,m,\bm{\theta}_{m}\sim\text{Poisson}(y;\phi I_{m}(t)),$ (4.3)

where $\phi$ is assumed to be known and set to $0.95$ throughout our
experiments.999While we have not done so, our method would allow us to treat
$\phi$ as a model parameter as well. This corresponds to the setting of a
known observational model described by Equation 3.12.

The design variable in our experiments is the measurement time $t$ in Equation
4.3. We generally wish to take several measurements at once and hence define a
design vector $\mathbf{d}=(t_{1},\dots,t_{D})$ with corresponding observations
$\mathbf{y}=(y_{1},\dots,y_{D})$, where $D$ is the number of measurements. In
our experiments, we use $N=500$ with initial conditions of
$\mathbf{X}_{1}(t=0)=(498,2,0)^{\top}$ for the SIR model and
$\mathbf{X}_{2}(t=0)=(498,0,2,0)^{\top}$ for the SEIR model. We also use a
discretised design space of $t\in[0,100]$ with resolution of $\Delta
t=10^{-2}$. For the model parameter priors we use
$p(\beta)=\text{Lognorm}(0.50,0.50^{2})$,
$p(\sigma)=\text{Lognorm}(0.20,0.50^{2})$ and
$p(\gamma)=\text{Lognorm}(0.10,0.50^{2})$. Further information about neural
network architectures, hyper-parameters and experimental details are given in
the supplementary materials.

####  4.2.1 Parameter Estimation for the SIR Model

Figure 8: PE results for the SDE-based SIR model with different experimental
budgets. The left shows validation MI estimates for different number of
measurements, averaged over several validation data sets. The right plot shows
the corresponding optimal measurement times.

First, we consider the task of parameter estimation (PE) for the SIR model,
where the aim is to estimate its model parameters
$\bm{\theta}_{1}=(\beta,\gamma)$. We here use the JSD lower bound as an
objective function and test different experimental budgets of up to $10$
measurements. Figure 18 summarises the results for settings with different
budgets. The left plot shows validation estimates of the maximum MI at the
respective optimal designs (shown in the right plot), obtained through our
proposed method, for varying numbers of measurements $D$. The optimal designs
$\mathbf{d}^{\ast}$ were all in the region between $t=10$ and $t=40$, with a
slightly different spreading for different $D$. Intuitively, this region might
be useful because that is where the signal-to-noise ratio of the SIR
observations is highest.101010See the supplementary materials for a plot
showing this. As before in the PE task for the linear toy model, this means
that we are able to better measure the effect of model parameters on
observations. The maximum MI in the left of Figure 18 starts to saturate at
$D=10$ measurements, suggesting that more measurements may not improve
parameter estimation by much.

Figure 9: PE posterior distributions for the SDE-based SIR model with
different experimental budgets. The top left shows the prior distribution,
while the other plots show posteriors for an increasing number of
measurements. The posteriors are averaged over $5{,}000$ ‘real-world’
observations $\mathbf{y}^{\ast}$ generated with ground-truth $(0.60,0.15)$.

We can estimate posterior distributions of the model parameters
$(\beta,\gamma)$ by means of the trained neural network (see Section 3). To do
so, we generate ‘real-world’ observations $\mathbf{y}^{\ast}$ at
$\mathbf{d}^{\ast}$ using the ground truth
$\bm{\theta}_{\text{true}}=(0.60,0.15)$. Figure 9 shows the prior distribution
(top left) and average posterior distributions for different number of
measurements, where the average is taken over $5{,}000$ observations
$\mathbf{y}^{\ast}$. As we increase the number of measurements, the average
posterior naturally becomes more peaked and moves closer to the ground truth
parameters (indicated by a red cross). At $D=10$, the ground truth is well-
captured by the average posterior, with an average MAP estimate, taken over
all ‘real-world’ observations, of
$\widehat{\bm{\theta}}_{\text{MAP}}=(0.59,0.07)\pm(0.15,0.01)$, where the
errors indicate one standard deviation. All training curves, posterior
distributions and optimal designs are shown in the supplementary materials.

####  4.2.2 Model Discrimination

Next, we consider the task of model discrimination (MD) between the SIR and
SEIR model. We again use the JSD lower bound as a the utility function and
test different experimental budgets of up to $10$ measurements. The results
for this setting are summarised in Figure 19. The left plot shows optimal
designs for varying numbers of measurements $D$, obtained through our proposed
method. The optimal designs have elements that are clustered around smaller
measurement times between $t=15$ and $t=40$, similar to the optimal designs
for the PE task in Figure 18. These measurement times might be useful because
that is where the average signal-to-noise ratio is highest for both models and
the corresponding data distributions are significantly different.111111See the
supplementary material for plots showing this. This means that we can gather
(relatively) low-noise data that can be related to the corresponding models
more effectively. We show corresponding validation estimates of the MI as a
function of $D$ in the supplementary materials. Similar to the PE task, the MI
tends to increase as the number of measurements are increased, although the
change is much more subtle for the MD task.

Figure 10: MD results for the SDE-based epidemiology models. The left plot
shows optimal experimental designs for varying number of measurements. The
other plots show distributions of posterior densities for $D=10$ measurements,
given that the SIR model (middle plot) or the SEIR model (right plot) is the
ground truth.

For the case of $D=10$ measurements, the middle and right plots in Figure 19
show distributions of posterior densities
$p(m|\mathbf{d}^{\ast},\mathbf{y}^{\ast})$ estimated using several
$\mathbf{y}^{\ast}$, with the middle plot assuming that SIR is the ground-
truth model and the right plot assuming that SEIR is the ground-truth model.
We can recover the ground-truth model effectively in both cases, with
corresponding F1-scores of $0.996$ and $0.983$. However, the model recovery is
slightly worse when the SEIR model is the ground truth, which we discuss
further in the supplementary materials.

##  5 Conclusions

In this work we have introduced a framework for Bayesian experimental design
with implicit models. Our approach finds optimal experimental designs by
maximising general lower bounds on mutual information that are parametrised by
neural networks. We have derived gradients of prominent lower bounds with
respect to the experimental designs, allowing us to perform stochastic
gradient-ascent. In doing so, we have provided a methodology to derive
gradients of general lower bounds on mutual information, which may be applied
to other lower bounds that are introduced in the future. Furthermore, our
framework yields an amortised posterior distribution as a by-product, as long
as the utilised lower bound facilitates this, which is generally a challenging
problem for implicit models and has an entire research field dedicated to it
(see Lintusaari et al., 2017; Sisson et al., 2018; Cranmer et al., 2020 for
recent overviews).

By means of a set of intractable toy models, we have provided a thorough
experimental comparison of several prominent lower bounds for a variety of
scientific tasks, showcasing the versatility of our proposed framework. We
have also applied our method to a challenging set of implicit models from
epidemiology, where the responses are discrete observations of the solutions
to stochastic differential equations. Our approach allowed us to efficiently
obtain optimal designs, leading to suitable posterior distributions that
capture ground-truths well. Throughout our experiments we have considered
popular scientific tasks such as parameter estimation, model discrimination
and improving future predictions. Importantly, we believe that our framework
would also be able to deal with other scientific tasks that are formulated
using mutual information.

Our approach requires us to have access to gradients of the sampling path with
respect to experimental designs, or that we can compute them by means of
automatic differentiation. There may be cases where these gradients are
unknown or undefined, e.g. when data is discrete or categorical. We note that
our framework can handle those cases when the corresponding observational
process has an analytically tractable model, but currently not otherwise.
Kleinegesse and Gutmann (2020b) have shown that it is still possible to use
gradient-free optimisation techniques to find optimal designs, following the
method explained in Section 3. However, these approaches do not scale well
with design dimensions, unlike gradient-based approaches (see Spall, 2003). It
would be interesting to explore how our method can be scalably adapted to
these situations.

In related work, Foster et al. (2019) perform gradient-based Bayesian
experimental design for implicit models by using variational approximations to
likelihoods or posteriors. This allows them to use score-function estimators,
as opposed to pathwise gradient estimators. Recently, Zhang et al. (2021) have
used evolutionary strategies to approximate gradients of the SMILE lower bound
in the context of Bayesian experimental design. There are several other
interesting methods that allow us to approximate gradients during the
optimisation procedure, e.g. finite-difference or simultaneous perturbation
stochastic approximations (Spall, 2003), that have not been investigated yet.
As such, investigating avenues that mitigate or relax the need to have access
to gradients of the sampling path may increase the applicability of our method
further.

We have applied our framework to the setting of static Bayesian experimental
design, where we wish to determine a set of optimal designs prior to actually
performing the experiment. Sequential Bayesian experimental design, however,
advocates that we should update our belief of the variable of interest as we
perform the experiment (see e.g. Kleinegesse et al., 2020a). It would be
interesting to apply our framework to this setting as well. In particular,
Foster et al. (2021) have recently introduced an approach that yields a design
network which directly outputs optimal designs and, as such, allows for
amortised sequential Bayesian experimental design, greatly increasing
practicability. While their method is designed for explicit models that have
known likelihoods, we believe that it would be interesting to apply our method
of maximising parametrised lower bounds to their method as well.

Furthermore, in this work we have mostly glanced over the neural network
architecture selection. However, we believe that there is a lot still to be
explored that may help during the training procedure, e.g. residual
connections, separable critics, pooling, convolutions, and other techniques
that leverage known structures. In particular, it may be useful to directly
incorporate a dependency on the designs into the neural network, e.g. by using
the current design as an input as well. This may strengthen the extrapolation
of the neural network, improving training speed and robustness. We also
believe that the interplay between the number of dimensions of the variable of
interest $\bm{v}$ and the data variable $\mathbf{y}$ is highly important. If
one is much larger than the other, the neural network
$T_{\bm{\psi}}(\bm{v},\mathbf{y})$ may have more difficulties learning the
appropriate density ratio. It may be useful to use a separable critic, or
automatically learn summary statistics that reduce the dimensionality while
retaining all the necessary information as e.g. recently done by Chen et al.
(2021) in the field of likelihood-free inference.

##  Appendix A MI Lower bounds

###  A.1 Relationship between the JSD lower bound and LFIRE

Likelihood-free inference by ratio estimation (LFIRE, Thomas et al., 2020)
aims to approximate the density ratio $r(x,\theta)$ of the data-generating
distribution $p(x|\theta)$ and the marginal distribution $p(x)$ for implicit
models, where $p(x|\theta)$ is intractable but sampling from it is still
possible. Using a known prior distribution $p(\theta)$, this learned density
ratio can then be used to compute the posterior distribution $p(\theta|x)$.
Importantly, LFIRE only requires samples $x^{\theta}_{i}\sim p(x|\theta)$ from
the data-generating distribution and samples $x^{m}_{i}\sim p(x)$ from the
marginal distribution.

This likelihood-free inference method works by formulating a classification
problem between data sampled from $p(x|\theta)$ and $p(x)$, and then solving
this via (non-linear) logistic regression. Using their notation, the loss
function that they minimise is

$\mathcal{J}(h,\theta)=\frac{1}{n_{\theta}+n_{m}}\left\\{\sum_{i=1}^{n_{\theta}}\log\left[1+\nu
e^{-h(x^{\theta}_{i})}\right]+\sum_{i=1}^{n_{m}}\log\left[1+\frac{1}{\nu}e^{h(x^{m}_{i})}\right]\right\\},$
(A.1)

where $n_{\theta}$ and $n_{m}$ are the number of samples from $p(x|\theta)$
and $p(x)$, respectively, $\nu=n_{\theta}/n_{m}$ is used to correct unequal
class sizes, and $h$ is a non-linear, parametrised function. The authors prove
that for large $n_{\theta}$ and $n_{m}$, the function $h^{\ast}$ that
minimises $\mathcal{J}(h,\theta)$ is given by the desired log-ratio, i.e.
$h^{\ast}(x,\theta)=\log r(x,\theta)$. Importantly, this result holds for any
a fixed $\theta$, so that the result still holds if we consider the objective
$\mathcal{J}(h,\theta)$ averaged over $\theta$.

Recall that the JSD lower bound presented in the main text is given by

$\mathcal{L}_{\text{JSD}}(\bm{\psi},\mathbf{d})\equiv\mathbb{E}_{p(\bm{v},\mathbf{y}\mid\mathbf{d})}\left[-\text{sp}(-T_{\bm{\psi}}(\bm{v},\mathbf{y}))\right]-\mathbb{E}_{p(\bm{v})p(\mathbf{y}\mid\mathbf{d})}\left[\text{sp}(T_{\bm{\psi}}(\bm{v},\mathbf{y}))\right],$
(A.2)

where $\text{sp}(z)=\log(1+e^{z})$ is the softplus function and
$T_{\bm{\psi}}(\bm{v},\mathbf{y})$ is a neural network with parameters
$\bm{\psi}$. The variables $\mathbf{y},\bm{v}$ and $\mathbf{d}$ are the data,
variable of interest and experimental designs, respectively. Our BED framework
works by maximising Equation A.2 with respect to $\bm{\psi}$ and $\mathbf{d}$
to find optimal designs $\mathbf{d}$. Let us here focus on the specific case
of fixed $\mathbf{d}$, i.e. only $\bm{\psi}$ is updated. We here aim to show
the strong relationship between the JSD lower bound and the objective used in
LFIRE.

First, let us make the notational changes $\bm{v}\rightarrow\theta$,
$\mathbf{y}|\mathbf{d}\rightarrow x$ and $T_{\bm{\psi}}\rightarrow h$, in
order to match the notation of Thomas et al. (2020). Equation A.2 then becomes

$\displaystyle\mathcal{L}_{\text{JSD}}(h)$
$\displaystyle=\mathbb{E}_{p(x,\theta)}\left[-\text{sp}(-h(x,\theta)\right]-\mathbb{E}_{p(x)p(\theta)}\left[\text{sp}(h(x,\theta))\right]$
(A.3)
$\displaystyle=\mathbb{E}_{p(x,\theta)}\left[-\log\left(1+e^{-h(x,\theta)}\right)\right]-\mathbb{E}_{p(x)p(\theta)}\left[\log\left(1+e^{h(x,\theta)}\right)\right].$
(A.4)

Writing $p(x,\theta)=p(x|\theta)p(\theta)$ we can move the expectation
$\mathbb{E}_{p(\theta)}$ outside, as it is present in both terms in the above
equation. The resulting JSD lower bound can then be seen as an expectation of
the LFIRE objective in Equation A.1, with equal class sizes $n_{\theta}=n_{m}$
and large number of samples $n_{m}$, $n_{\theta}\rightarrow\infty$, i.e.

$\displaystyle\mathcal{L}_{\text{JSD}}(h)$
$\displaystyle=-\mathbb{E}_{p(\theta)}\left\\{\mathbb{E}_{p(x\mid\theta)}\left[\log\left(1+e^{-h(x,\theta)}\right)\right]+\mathbb{E}_{p(x)}\left[\log\left(1+e^{h(x,\theta)}\right)\right]\right\\}$
(A.5)
$\displaystyle=-\mathbb{E}_{p(\theta)}\left\\{\mathcal{J}(h,\theta)\right\\}.$
(A.6)

Recall that LFIRE is used to learn the density ratio $r(x,\theta)$ for a
particular value of $\theta$ (with an amortisation in $x$). As explained in
the main text, a neural network trained with the JSD lower bound can be used
to directly compute the log density ratio $\log r(x,\theta)$, but with
amortisation in $\theta$ as well as $x$. This fact and Equation A.6 imply that
maximising the JSD lower bound can be seen as minimising the LFIRE objective
in an amortised fashion. While LFIRE and the JSD lower bound have existed in
literature for quite some time, we believe that this strong relationship
between them has previously not been recorded. Moreover, LFIRE has been used
before in the context of Bayesian experimental design for implicit models
(Kleinegesse and Gutmann, 2019; Kleinegesse et al., 2020a).

###  A.2 Derivation of lower bound gradients

Using the methodology presented in Section 3.3 of the main text, we here
provide detailed derivations of the lower bound gradients shown in Table 2 of
the same section. Recall that mutual information lower bounds in literature
generally involve expectations
$\mathbb{E}_{p(\bm{v},\mathbf{y}\mid\mathbf{d})}\left[f(T_{\bm{\psi}}(\bm{v},\mathbf{y}))\right]$
over the joint and/or expectations
$\mathbb{E}_{p(\bm{v})p(\mathbf{y}\mid\mathbf{d})}\left[g(T_{\bm{\psi}}(\bm{v},\mathbf{y}))\right]$
over the marginal distribution, where $f:\mathbb{R}\rightarrow\mathbb{R}$ and
$g:\mathbb{R}\rightarrow\mathbb{R}$ are non-linear, differentiable functions.
As shown in Equations 3.10 and 3.11 in the main text, we can compute gradients
of these expectations by means of pathwise gradient estimators (i.e. the
reparametrisation trick).

In the interest of space, let us define
$T=T_{\bm{\psi}}(\bm{v},\mathbf{h}(\bm{\epsilon};\bm{v},\mathbf{d}))$,
$\widetilde{T}=T_{\bm{\psi}}(\bm{v},\mathbf{h}(\bm{\epsilon};\widetilde{\bm{v}},\mathbf{d}))$,
as well as $P=p(\bm{v})p(\bm{\epsilon})$ and
$Q=p(\bm{v})p(\widetilde{\bm{v}})p(\bm{\epsilon})$, where
$p(\widetilde{\bm{v}})$ is the same as $p(\bm{v})$. Importantly, $P$ and $Q$
do not depend on $\mathbf{d}$. Our methodology in the main text requires us to
first compute $f^{\prime}(T)=\partial f(T)/\partial T$ and
$g^{\prime}(\widetilde{T})=\partial g(\widetilde{T})/\partial\widetilde{T}$ in
order to compute gradients of lower bounds with respect to designs. Below we
show step-by-step derivations of the gradients for all lower bounds shown in
Table 2 of the main text.

##### NWJ

Defined in Equation 3.2 in the main text, the NWJ lower bound involves the
functions $f(T)=T$ and $g(\widetilde{T})=e^{\widetilde{T}-1}$, which yield the
corresponding gradients $\nabla_{T}f(T)=1$ and
$\nabla_{\widetilde{T}}g(\widetilde{T})=e^{\widetilde{T}-1}$. The desired
gradients of the NWJ lower bound with respect to designs $\mathbf{d}$ are then
computed as follows,

$\displaystyle\nabla_{\mathbf{d}}\mathcal{L}_{\text{NWJ}}(\bm{\psi},\mathbf{d})$
$\displaystyle=\nabla_{\mathbf{d}}\left\\{\mathbb{E}_{P}\left[f(T)\right]-\mathbb{E}_{Q}\left[g(\widetilde{T})\right]\right\\}$
(A.7)
$\displaystyle=\mathbb{E}_{P}\left[\nabla_{\mathbf{d}}f(T)\right]-\mathbb{E}_{Q}\left[\nabla_{\mathbf{d}}g(\widetilde{T})\right]$
(A.8)
$\displaystyle=\mathbb{E}_{P}\left[\nabla_{T}f(T)\nabla_{\mathbf{d}}T\right]-\mathbb{E}_{Q}\left[\nabla_{\widetilde{T}}g(\widetilde{T})\nabla_{\mathbf{d}}\widetilde{T}\right]$
(A.9)
$\displaystyle=\mathbb{E}_{P}\left[\nabla_{\mathbf{d}}T\right]-\mathbb{E}_{Q}\left[e^{\widetilde{T}-1}\nabla_{\mathbf{d}}\widetilde{T}\right],$
(A.10)

where $\nabla_{\mathbf{d}}T$ and $\nabla_{\mathbf{d}}\widetilde{T}$ are given
by Equation 3.12 in the main text.

##### InfoNCE

This lower bound, defined in Equation 3.3 in the main text, only involves an
expectation over the joint distribution and not the marginal distribution,
i.e. $g(T)=0$. In addition to the previous shorthands, let us also define
$T_{ij}=T_{\bm{\psi}}(\bm{v}_{j},h(\bm{\epsilon}_{i};\bm{v}_{i},\mathbf{d}))$.
In this particular case, it is simpler to directly start with the gradient
with respect to $\mathbf{d}$, yielding

$\displaystyle\nabla_{\mathbf{d}}\mathcal{L}_{\text{NCE}}(\bm{\psi},\mathbf{d})$
$\displaystyle=\nabla_{\mathbf{d}}\mathbb{E}_{P^{K}}\left[\frac{1}{K}\sum_{i=1}^{K}\log{\frac{e^{T_{ii}}}{\frac{1}{K}\sum_{j=1}^{K}e^{T_{ij}}}}\right]$
(A.11)
$\displaystyle=\mathbb{E}_{P^{K}}\left[\frac{1}{K}\sum_{i=1}^{K}\nabla_{\mathbf{d}}\log{\frac{e^{T_{ii}}}{\frac{1}{K}\sum_{j=1}^{K}e^{T_{ij}}}}\right]$
(A.12)
$\displaystyle=\mathbb{E}_{P^{K}}\left[\frac{1}{K}\sum_{i=1}^{K}\nabla_{\mathbf{d}}T_{ii}-\nabla_{\mathbf{d}}\log{\sum_{j=1}^{K}e^{T_{ij}}}+\nabla_{\mathbf{d}}\log{K}\right]$
(A.13)
$\displaystyle=\mathbb{E}_{P^{K}}\left[\frac{1}{K}\sum_{i=1}^{K}\nabla_{\mathbf{d}}T_{ii}-\nabla_{\mathbf{d}}\log{\sum_{j=1}^{K}e^{T_{ij}}}\right]$
(A.14)
$\displaystyle=\mathbb{E}_{P^{K}}\left[\frac{1}{K}\sum_{i=1}^{K}\nabla_{\mathbf{d}}T_{ii}-\frac{\sum_{j=1}^{K}e^{T_{ij}}\nabla_{\mathbf{d}}T_{ij}}{\sum_{j=1}^{K}e^{T_{ij}}}\right]$
(A.15)
$\displaystyle=\mathbb{E}_{P^{K}}\left[\frac{1}{K}\sum_{i=1}^{K}\frac{\sum_{j=1}^{K}e^{T_{ij}}(\nabla_{\mathbf{d}}T_{ii}-\nabla_{\mathbf{d}}T_{ij})}{\sum_{j=1}^{K}e^{T_{ij}}}\right],$
(A.16)

where $\nabla_{\mathbf{d}}T_{ii}$ and $\nabla_{\mathbf{d}}T_{ij}$ are again
given by Equation 3.12 in the main text. As explained in the main text, the
above formulations yields an optimal critic
$T^{\ast}_{\bm{\psi}}(\bm{v},\mathbf{y})=\log{p(\mathbf{y}|\bm{v},\mathbf{d})}+c(\mathbf{y}|\mathbf{d})$,
where $c(\mathbf{y}|\mathbf{d})$ is an indeterminate function.

##### JSD

The JSD lower bound is defined in Equation 3.4 in the main text and involves
expectations of the non-linear functions $f(T)=-\text{sp}(-T)$ and
$g(\widetilde{T})=\text{sp}(\widetilde{T})$, where
$\text{sp}(z)=\log\left(1+e^{z}\right)$ is the softplus function. The
derivative of the softplus function is simply the logistic sigmoid function,
i.e. $\nabla_{z}\text{sp}(z)=1/(1+e^{-z})\equiv\sigma(z)$, resulting in the
gradients $\nabla_{T}f(T)=\sigma(-T)$ and
$\nabla_{\widetilde{T}}g(\widetilde{T})=\sigma(\widetilde{T})$. This allows us
to compute the gradients of the JSD lower bound with respect to experimental
designs,

$\displaystyle\nabla_{\mathbf{d}}\mathcal{L}_{\text{JSD}}(\bm{\psi},\mathbf{d})$
$\displaystyle=\nabla_{\mathbf{d}}\left\\{\mathbb{E}_{P}\left[f(T)\right]-\mathbb{E}_{Q}\left[g(\widetilde{T})\right]\right\\}$
(A.17)
$\displaystyle=\mathbb{E}_{P}\left[\nabla_{\mathbf{d}}f(T)\right]-\mathbb{E}_{Q}\left[\nabla_{\mathbf{d}}g(\widetilde{T})\right]$
(A.18)
$\displaystyle=\mathbb{E}_{P}\left[\nabla_{T}f(T)\nabla_{\mathbf{d}}T\right]-\mathbb{E}_{Q}\left[\nabla_{\widetilde{T}}g(\widetilde{T})\nabla_{\mathbf{d}}\widetilde{T}\right]$
(A.19)
$\displaystyle=\mathbb{E}_{P}\left[\sigma(-T)\nabla_{\mathbf{d}}T\right]-\mathbb{E}_{Q}\left[\sigma(\widetilde{T})\nabla_{\mathbf{d}}\widetilde{T}\right],$
(A.20)

where $\nabla_{\mathbf{d}}T$ and $\nabla_{\mathbf{d}}\widetilde{T}$ are again
computed using Equation 3.12 in the main text.

###  A.3 Gradient estimator for analytically tractable observation models

We here provide the gradient estimator of mutual information lower bounds in
situations where it is possible to leverage an analytically tractable
observation model, as discussed in Section 3.4 of the main text. Recall that
we assume that we can separate the data generation into an observation process
with a model $p(\mathbf{y}|\bm{v},\mathbf{d},\bm{z})$ that is known
analytically in closed form and a differentiable, latent process
$p(\bm{z}|\bm{v},\mathbf{d})$. We can reparametrise the latent variable using
its sampling path, i.e. $\bm{z}=\bm{h}(\bm{\epsilon};\bm{v},\mathbf{d})$,
where the noise random variable $\bm{\epsilon}$ defines the sampling path.
This allows us to use score-function estimators to compute the required
gradients.

The score function estimator works by re-writing the derivative of
$p(\mathbf{y}|\bm{v},\mathbf{d},\bm{z})$ using the definition of the
logarithmic derivative, i.e.

$\nabla_{\mathbf{d}}p(\mathbf{y}|\bm{v},\mathbf{d},\bm{z})=p(\mathbf{y}|\bm{v},\mathbf{d},\bm{z})\nabla_{\mathbf{d}}\log
p(\mathbf{y}|\bm{v},\mathbf{d},\bm{z}).$ (A.21)

Following the explanations in the main text, we can then compute the gradient
of the expectation of $f(T_{\bm{\psi}}(\bm{v},\mathbf{y}))$ as follows,

$\displaystyle\nabla_{\mathbf{d}}$
$\displaystyle\mathbb{E}_{p(\bm{v},\mathbf{y}\mid\mathbf{d})}\left[f(T_{\bm{\psi}}(\bm{v},\mathbf{y}))\right]$
(A.22)
$\displaystyle=\nabla_{\mathbf{d}}\mathbb{E}_{p(\mathbf{y}|\bm{v},\mathbf{d},\bm{z})p(\bm{z}|\bm{v},\mathbf{d})p(\bm{v})}\left[f(T_{\bm{\psi}}(\bm{v},\mathbf{y}))\right]$
(A.23)
$\displaystyle=\nabla_{\mathbf{d}}\mathbb{E}_{p(\mathbf{y}|\bm{v},\mathbf{d},\mathbf{h}(\bm{\epsilon};\bm{v},\mathbf{d}))p(\bm{\epsilon})p(\bm{v})}\left[f(T_{\bm{\psi}}(\bm{v},\mathbf{y}))\right]$
(A.24)
$\displaystyle=\int\mathbb{E}_{p(\bm{\epsilon})p(\bm{v})}\left[\nabla_{\mathbf{d}}\\{p(\mathbf{y}|\bm{v},\mathbf{d},\mathbf{h}(\bm{\epsilon};\bm{v},\mathbf{d}))\\}f(T_{\bm{\psi}}(\bm{v},\mathbf{y}))\right]\mathrm{d}\mathbf{y}$
(A.25)
$\displaystyle=\mathbb{E}_{p(\mathbf{y}|\bm{v},\mathbf{d},\mathbf{h}(\bm{\epsilon};\bm{v},\mathbf{d}))p(\bm{\epsilon})p(\bm{v})}\left[\nabla_{\mathbf{d}}\\{\log
p(\mathbf{y}|\bm{v},\mathbf{d},\mathbf{h}(\bm{\epsilon};\bm{v},\mathbf{d}))\\}f(T_{\bm{\psi}}(\bm{v},\mathbf{y}))\right]$
(A.26)

The last equation above can then simply be approximated using a Monte-Carlo
sample average. Note that when computing the log derivative
$\nabla_{\mathbf{d}}\\{\log
p(\mathbf{y}|\bm{v},\mathbf{d},\mathbf{h}(\bm{\epsilon};\bm{v},\mathbf{d}))\\}$
we generally require the gradients of the sampling path
$\mathbf{h}(\bm{\epsilon};\bm{v},\mathbf{d})$ as well. Using Equation 3.9 in
the main text, we can derive a similar equation for the derivative of the
expectation of $g(T_{\bm{\psi}}(\bm{v},\mathbf{y}))$ over the marginal
distribution,

$\displaystyle\nabla_{\mathbf{d}}$
$\displaystyle\mathbb{E}_{p(\bm{v})p(\mathbf{y}\mid\mathbf{d})}\left[g(T_{\bm{\psi}}(\bm{v},\mathbf{y}))\right]$
(A.27)
$\displaystyle=\nabla_{\mathbf{d}}\mathbb{E}_{p(\mathbf{y}|\widetilde{\bm{v}},\mathbf{d})p(\widetilde{\bm{v}})p(\bm{v})}\left[g(T_{\bm{\psi}}(\bm{v},\mathbf{y}))\right]$
(A.28)
$\displaystyle=\nabla_{\mathbf{d}}\mathbb{E}_{p(\mathbf{y}|\widetilde{\bm{v}},\mathbf{d},\bm{z})p(\bm{z}|\widetilde{\bm{v}},\mathbf{d})p(\widetilde{\bm{v}})p(\bm{v})}\left[g(T_{\bm{\psi}}(\bm{v},\mathbf{y}))\right]$
(A.29)
$\displaystyle=\nabla_{\mathbf{d}}\mathbb{E}_{p(\mathbf{y}|\widetilde{\bm{v}},\mathbf{d},\mathbf{h}(\bm{\epsilon};\widetilde{\bm{v}},\mathbf{d}))p(\bm{\epsilon})p(\widetilde{\bm{v}})p(\bm{v})}\left[g(T_{\bm{\psi}}(\bm{v},\mathbf{y}))\right]$
(A.30)
$\displaystyle=\int\mathbb{E}_{p(\bm{\epsilon})p(\widetilde{\bm{v}})p(\bm{v})}\left[\nabla_{\mathbf{d}}\\{p(\mathbf{y}|\widetilde{\bm{v}},\mathbf{d},\mathbf{h}(\bm{\epsilon};\widetilde{\bm{v}},\mathbf{d}))\\}g(T_{\bm{\psi}}(\bm{v},\mathbf{y}))\right]\mathrm{d}\mathbf{y}$
(A.31)
$\displaystyle=\mathbb{E}_{p(\mathbf{y}|\widetilde{\bm{v}},\mathbf{d},\mathbf{h}(\bm{\epsilon};\widetilde{\bm{v}},\mathbf{d}))p(\bm{\epsilon})p(\widetilde{\bm{v}})p(\bm{v})}\left[\nabla_{\mathbf{d}}\\{\log
p(\mathbf{y}|\widetilde{\bm{v}},\mathbf{d},\mathbf{h}(\bm{\epsilon};\widetilde{\bm{v}},\mathbf{d}))\\}g(T_{\bm{\psi}}(\bm{v},\mathbf{y}))\right],$
(A.32)

where the last equation above can conveniently be approximated using a Monte-
Carlo sample average. Equations A.26 and A.32 do not require differentiating
$f$ or $g$ with respect to the neural network output. This allows us to
quickly write down gradients of the prominent lower bounds presented in the
main text (NWJ, InfoNCE and JSD) when the observation model is known
analytically in closed form. This yields Table 5 in a similar manner to Table
2 in the main text.

Table 5: Gradients of several lower bounds with respect to experimental designs, when the observation model is analytically tractable. We define $T=T_{\bm{\psi}}(\bm{v},\mathbf{y})$ and $T_{ij}=T_{\bm{\psi}}(\bm{v}_{j},\mathbf{y}_{i})$, as well as the shorthands $q_{\mathbf{y}}=p(\mathbf{y}|\bm{v},\mathbf{d},\mathbf{h}(\bm{\epsilon};\bm{v},\mathbf{d}))$, $\widetilde{q}_{\mathbf{y}}=p(\mathbf{y}|\widetilde{\bm{v}},\mathbf{d},\mathbf{h}(\bm{\epsilon};\widetilde{\bm{v}},\mathbf{d}))$, $P=q_{\mathbf{y}}p(\bm{\epsilon})p(\bm{v})$ and $Q=\widetilde{q}_{\mathbf{y}}p(\bm{\epsilon})p(\widetilde{\bm{v}})p(\bm{v})$. Here, $\text{sp}(T)$ is the softplus function. Lower Bound | Gradients with respect to designs
---|---
NWJ | $\mathbb{E}_{P}\left[\nabla_{\mathbf{d}}\\{\log q_{\mathbf{y}}\\}T\right]-e^{-1}\mathbb{E}_{Q}\left[\nabla_{\mathbf{d}}\\{\log\widetilde{q}_{\mathbf{y}}\\}e^{T}\right]$
InfoNCE | $\mathbb{E}_{P^{K}}\left[\nabla_{\mathbf{d}}\\{\log q_{\mathbf{y}}\\}\frac{1}{K}\sum_{i=1}^{K}\log{\frac{\exp(T_{ii})}{\frac{1}{K}\sum_{j=1}^{K}\exp(T_{ij})}}\right]$
JSD | $\mathbb{E}_{P}\left[-\nabla_{\mathbf{d}}\\{\log q_{\mathbf{y}}\\}\text{sp}(-T)\right]-\mathbb{E}_{Q}\left[\nabla_{\mathbf{d}}\\{\log\widetilde{q}_{\mathbf{y}}\\}\text{sp}(T)\right]$

##  Appendix B Toy experiments

###  B.1 Reference computations

In order to compute reference values of the mutual information (MI) and
posteriors, we rely on an approximation of the likelihood
$p(\mathbf{y}|\bm{\theta}_{m},m,\mathbf{d})$ that is based on kernel density
estimation (KDE), where $m$ is a model indicator and $\bm{\theta}_{m}$ are the
corresponding model parameters.

As described in Equation 4.1 of the main text, the sampling paths of all toy
models have additive Gaussian noise $\epsilon\sim\mathcal{N}(\epsilon;0,1)$
and additive Gamma noise $\nu\sim\Gamma(\nu;2,2)$. The overall noise
distribution $p_{\text{noise}}$ of $\epsilon+\nu$ is given by a convolution of
the individual densities and could be computed via numerical integration. We
here opt for an alternative approach, namely to compute the overall noise
distribution by a KDE of $50{,}000$ samples of $\epsilon$ and $\nu$. Once we
have fitted a KDE to noise samples and obtained an estimate
$\widehat{p}_{\text{noise}}$, we can re-arrange the sampling paths, which then
allows us to compute estimates of the likelihood. For the set of toy models we
here make the assumption that performing experiments with certain designs does
not change the data-generating distribution, thereby allowing us to write

$\widehat{p}(\mathbf{y}|\bm{\theta}_{m},m,\mathbf{d})=\prod_{j=1}^{D}\widehat{p}(y_{j}|\bm{\theta}_{m},m,d_{j}),$
(B.1)

where $y_{j}$ and $d_{j}$ are the elements of $\mathbf{y}$ and $\mathbf{d}$,
respectively. We can then focus on estimating the likelihood for each
dimension of $\mathbf{d}$ separately and then multiplying them. Using the
estimated noise distribution, the individual likelihood estimates
$\widehat{p}(y_{j}|\bm{\theta}_{m},m,d_{j})$ for the linear ($m=1$),
logarithmic ($m=2$) and square-root ($m=3$) toy models are,

$\widehat{p}(y_{j}|\bm{\theta}_{m},m,d_{j})=\begin{cases}\widehat{p}_{\text{noise}}(y_{j}-(\theta_{1,0}+\theta_{1,1}d_{j}))&\text{if
}m=1,\\\
\widehat{p}_{\text{noise}}(y_{j}-(\theta_{2,0}+\theta_{2,1}\log(|d_{j}|)))&\text{if
}m=2,\\\
\widehat{p}_{\text{noise}}(y_{j}-(\theta_{3,0}+\theta_{3,1}\sqrt{|d_{j}|}))&\text{if
}m=3.\end{cases}$ (B.2)

Below we explain how to use this likelihood approximation to compute reference
MI and posteriors for each of the scientific tasks considered in our paper.
Note that, in the main text, the equations of mutual information are generally
given in terms of the posterior to prior ratio. We here compute reference MI
values by estimating the ratio of likelihood to marginal, which is equivalent
to the posterior to prior ratio by Bayes’ rule.

##### Parameter Estimation

Here we only consider the linear model, where $m=1$, with the aim of solving
the task of estimating the model parameters $\bm{\theta}_{1}$, i.e. the slope
and offset of a straight line. The analytic MI is given by Equation 2.5 in the
main text and we here approximate it with a nested Monte-Carlo (MC) sample
average and the approximate likelihoods in Equation B.2, i.e.

$\widehat{I}(\bm{\theta};\mathbf{y}|\mathbf{d})\approx\frac{1}{N}\sum_{i=1}^{N}\log\frac{\widehat{p}(\mathbf{y}^{(i)}|\mathbf{d},\bm{\theta}_{1}^{(i)},m=1)}{\frac{1}{K}\sum_{k=1}^{K}\widehat{p}(\mathbf{y}^{(i)}|\mathbf{d},\bm{\theta}_{1}^{(k)},m=1)},$
(B.3)

where $\mathbf{y}^{(i)}\sim
p(\mathbf{y}|\mathbf{d},\bm{\theta}_{1}^{(i)},m=1)$,
$\bm{\theta}_{1}^{(i)}\sim p(\bm{\theta}_{1})$ and $\bm{\theta}_{1}^{(k)}\sim
p(\bm{\theta}_{1})$. The Gaussian prior distributions for each model are
provided in the main text. To compute the above estimate, we use $N=2{,}000$
and $K=500$. We can also use the approximate likelihood and Bayes’ rule to
compute a reference posterior,

$\widehat{p}(\bm{\theta}_{1}|\mathbf{d},\mathbf{y})=p(\bm{\theta}_{1})\frac{\widehat{p}(\mathbf{y}|\mathbf{d},\bm{\theta}_{1},m=1)}{\widehat{p}(\mathbf{y}|\mathbf{d})},$
(B.4)

where the marginal
$\widehat{p}(\mathbf{y}|\mathbf{d})=\int\widehat{p}(\mathbf{y}|\mathbf{d},\bm{\theta}_{1},m=1)p(\bm{\theta}_{1})\mathrm{d}\bm{\theta}_{1}$
can be computed with a MC sample average or via numerical integration. We here
use numerical integration by means of Simpson’s rule because of the low
dimensions of $\bm{\theta}_{1}$.

##### Model Discrimination

We can approximate the analytic MI in Equation 2.6 of the main text in a
similar manner. We use a nested MC sample average, sum over the model
indicators $m$ and approximate the relevant likelihoods using Equation B.2,
yielding

$\widehat{I}(m;\mathbf{y}|\mathbf{d})\approx\sum_{m=1}^{3}p(m)\frac{1}{N}\sum_{i=1}^{N}\log\frac{\widehat{p}(\mathbf{y}^{(i)}|\mathbf{d},m)}{\widehat{p}(\mathbf{y}^{(i)}|\mathbf{d})},$
(B.5)

where $\mathbf{y}^{(i)}\sim p(\mathbf{y}|m,\mathbf{d})$ and we use
$N=3{,}000$. We approximate the density
$\widehat{p}(\mathbf{y}^{(i)}|m,\mathbf{d})$ with a MC sample average as well,
i.e.

$\widehat{p}(\mathbf{y}^{(i)}|\mathbf{d},m)=\frac{1}{K}\sum_{k=1}^{K}\widehat{p}(\mathbf{y}^{(i)}|\bm{\theta}_{m}^{(k)},m,\mathbf{d}),\quad\bm{\theta}_{m}^{(k)}\sim
p(\bm{\theta}_{m}),$ (B.6)

where we use $K=1000$. The density $\widehat{p}(\mathbf{y}^{(i)}|\mathbf{d})$
is computed by summing over $m$,

$\widehat{p}(\mathbf{y}^{(i)}|\mathbf{d})=\sum_{m=1}^{3}p(m)\,\widehat{p}(\mathbf{y}^{(i)}|\mathbf{d},m).$
(B.7)

Using the above approximate densities and Bayes’ rule we can then compute a
reference posterior over the model indicator as well, i.e

$\widehat{p}(m|\mathbf{d},\mathbf{y})=p(m)\frac{\widehat{p}(\mathbf{y}|\mathbf{d},m)}{\widehat{p}(\mathbf{y}|\mathbf{d})}.$
(B.8)

##### Joint MD/PE

The analytic MI for the joint MD/PE task is shown in Equation 2.7 of the main
text. A nested MC sample average of that utility is

$\widehat{I}(\bm{\theta}_{m},m;\mathbf{y}|\mathbf{d})\approx\sum_{m=1}^{3}p(m)\frac{1}{N}\sum_{i=1}^{N}\log\frac{\widehat{p}(\mathbf{y}^{(i)}|\bm{\theta}_{m}^{(i)},m,\mathbf{d})}{\widehat{p}(\mathbf{y}^{(i)}|\mathbf{d})},$
(B.9)

where $\mathbf{y}^{(i)}\sim p(\mathbf{y}|\bm{\theta}_{m}^{(i)},m,\mathbf{d})$
and $\bm{\theta}_{m}^{(i)}\sim p(\bm{\theta}_{m})$. The marginal density
$\widehat{p}(\mathbf{y}^{(i)}|\mathbf{d})$ is computed by summing over $m$ and
using a MC sample average over $\bm{\theta}_{m}$, i.e.

$\widehat{p}(\mathbf{y}^{(i)}|\mathbf{d})=\sum_{i=1}^{3}p(m)\frac{1}{K}\sum_{k=1}^{K}\widehat{p}(\mathbf{y}^{(i)}|\bm{\theta}_{m}^{(k)},m,\mathbf{d}),\quad\bm{\theta}_{m}^{(k)}\sim
p(\bm{\theta}_{m}).$ (B.10)

In our computations we use $N=2{,}000$ and $K=500$. As before, we can compute
a reference joint posterior by using Bayes’ rule:

$\widehat{p}(\bm{\theta}_{m},m|\mathbf{d},\mathbf{y})=p(m)\frac{\widehat{p}(\mathbf{y}|\bm{\theta}_{m},m,\mathbf{d})}{\widehat{p}(\mathbf{y}|\mathbf{d})}.$
(B.11)

##### Improving Future Predictions

We here consider the task of improving future predictions for the linear toy
model ($m=1$). In order to approximate the analytic MI shown in Equation 2.8
of the main text, we need to compute approximations of the posterior
predictive distribution $\widehat{p}(y_{T}|\mathbf{y},\mathbf{d},d_{T})$,
which is given by

$\widehat{p}(y_{T}|\mathbf{y},\mathbf{d},d_{T})=\int\widehat{p}(y_{T}|\bm{\theta}_{1},d_{T})\widehat{p}(\bm{\theta}_{1}|\mathbf{y},\mathbf{d})\mathrm{d}\bm{\theta}_{1},$
(B.12)

where $\widehat{p}(y_{T}|\bm{\theta}_{1},d_{T})$ is computed using Equation
B.2 and $\widehat{p}(\bm{\theta}_{1}|\mathbf{y},\mathbf{d})$ is computed using
Equation B.4. We then approximate
$\widehat{p}(y_{T}|\mathbf{y},\mathbf{d},d_{T})$ by numerically integrating
over a grid of $\bm{\theta}_{1}$. We found that a $50\times 50$ grid of
$\bm{\theta}_{1}\in[-10,10]^{2}$ was sufficient. In the same manner, we
approximate the prior predictive
$\widehat{p}(y_{T}|d_{T})=\int\widehat{p}(y_{T}|\bm{\theta}_{1},d_{T})p(\bm{\theta}_{1})\mathrm{d}\bm{\theta}_{1}$
by numerical integration. Being able to compute the posterior predictive for a
certain $\\{\mathbf{d},\mathbf{y}\\}$ then allows us to approximate the
analytic MI using a combination of numerical integration and MC sample
averages, i.e.

$\widehat{I}(y_{T}|d_{T};\mathbf{y}|\mathbf{d})\approx\frac{1}{N}\sum_{i=1}^{N}\int\widehat{p}(y_{T}|\mathbf{y}^{(i)},\mathbf{d},d_{T})\log\frac{\widehat{p}(y_{T}|\mathbf{y}^{(i)},\mathbf{d},d_{T})}{\widehat{p}(y_{T}|d_{T})}\mathrm{d}y_{T},$

where $\mathbf{y}^{(i)}\sim p(\mathbf{y}|\mathbf{d})$. We used $N=800$ and
numerically integrated over $y_{T}\in[-50,50]$ using Simpson’s rule with a
grid of size $100$.

###  B.2 Data distributions

Figure 11: Toy model data distributions (top), average signal-to-noise ratio
(SNR, middle) and Jensen-Shannon (JS) divergences between different data
distributions (bottom).

In Figure 11 we summarise general information about data simulated from each
toy model presented in the main text, i.e. the linear model, the log model and
the square-root model. The top row shows the prior predictive distribution for
each model, where the solid lines represent the means and the shaded areas
indicate one standard deviation from the means. While the linear and square-
root model have a similar data distribution, the log model prior predictive is
significantly different, rapidly increasing as $d$ approaches $0$. This
signifies in which regions the data distributions are most different, under
our prior belief about the model parameters $\bm{\theta}_{m}$, i.e. at the
boundaries $d=-2,2$ and at $d=0$. The bottom plot, which shows the Jensen-
Shannon (JS) divergence between the data distributions of different toy
models, further emphasises this difference. This is in accordance with the
optimal designs found for the toy model MD and MD/PE tasks in the main text.

The middle row of Figure 11 shows the average signal-to-noise ratio (SNR) for
each of the toy models. The SNR, for a single $\bm{\theta}_{m}$, is given by

$\text{SNR}(d|\bm{\theta}_{m})=\left(\frac{\mu(d|\bm{\theta}_{m})}{\sigma(d|\bm{\theta}_{m})}\right)^{2},$
(B.13)

where $\mu(d|\bm{\theta}_{m})$ and $\sigma(d|\bm{\theta}_{m})$ are the mean
and standard deviation of the model response $y|\bm{\theta}_{m},d$,
respectively, averaged over signal noise. The average SNR is then
$\overline{\text{SNR}}(d)=\mathbb{E}_{p(\bm{\theta})}[\text{SNR}(d|\bm{\theta}_{m})]$.
While the linear and square-root model have a relatively high average SNR at
the boundaries, the log model has an extremely high average SNR at $d=0$. A
high SNR is useful for parameter estimation, because the observed data
essentially has less _relative_ noise, allowing us to better identify which
values of $\bm{\theta}_{m}$ might have generated observations. This is again
in accordance with results found in the main text.

###  B.3 Architectures and hyper-parameters

In Table 6 we show neural network (NN) architectures, including the number of
hidden layers and number of hidden units, as well as the learning rates (L.R.)
for the NN parameters $\bm{\psi}$ and experimental designs $\mathbf{d}$. We
optimise $\bm{\psi}$ and $\mathbf{d}$ with two separate Adam optimisers. with
default parameters from the PyTorch package in Python. Each epoch, we simulate
$10{,}000$ new data samples.

Table 6: Toy model hyper-parameters for each scientific goal. Scientific Goal | NN Layers | NN Units | L.R. for $\bm{\psi}$ | L.R. for $\mathbf{d}$
---|---|---|---|---
PE | 2 | 50 | $10^{-4}$ | $10^{-3}$
MD | 2 | 50 | $10^{-3}$ | $10^{-4}$
MD/PE | 2 | 70 | $10^{-4}$ | $5\times 10^{-4}$
FP | 2 | 100 | $10^{-3}$ | $10^{-2}$

###  B.4 Additional results

Here we show additional results for the toy model experiments that were
mentioned in the main text. Figure 12 shows a setting where the goal is MD and
the InfoNCE lower bound found different optimal experimental designs than the
NWJ and JSD lower bound. This is slightly surprising, as the neural networks
for each lower bound used the same architecture and the same initialisation.
Because different optimal designs were found, the final MI estimate for the
InfoNCE lower bound is lower than that of the NWJ and JSD lower bounds.
Importantly, this implies that the choice of lower bound does matter, as
different optimal designs may be found. In our work we do not, however, aim to
explore which lower bound is best, but rather show how different lower bounds
may be used in our proposed framework.

Figure 12: MD training curves when InfoNCE yields different locally optimal
design. The top row shows lower bounds as a function of training epochs, with
the dotted line being the reference MI at the final optimal design of the
respective lower bound. The bottom row shows the corresponding design elements
as the design vector is being updated. Figure 13: MD/PE results for the set of
toy models. Shown are training curves for the NWJ, InfoNCE and JSD (via NWJ)
lower bounds. The top row shows lower bound estimates as a function of
training epochs, with the dotted line being the reference MI value, and the
bottom row shows the elements of the design vector as it is being updated.
Figure 14: MD/PE posterior results of the model indicator for the set of toy
models. Shown are average posterior probabilities for different ground truths
(one per row) for the NWJ, InfoNCE and JSD lower bounds.

In Figure 13 we show the training curves of each lower bound (top row) for the
joint MD/PE task, as well as the experimental designs as they are being
updated (bottom row). All lower bounds yield the same optimal design and
converge to a final MI close to a reference MI (with the NWJ bound being
further away). Using the trained neural networks for each lower bound we can
compute posterior distributions, as described in the main text. Figure 14
shows marginal posterior distributions over the model indicator $m$ and Figure
15 shows marginal posterior distributions over model parameters
$\bm{\theta}_{m}$. Note that we show average posterior distributions over
$5{,}000$ ‘real-world’ observations
$\mathbf{y}^{\ast}|\mathbf{d}^{\ast},\bm{\theta}_{m,\text{true}}$, where
$\bm{\theta}_{m,\text{true}}=(2,3)$ for all models. Finally, for the FP task,
we show the training curves of each lower bound for the FP task (top row), as
well as corresponding design curves (bottom row), in Figure 16.

Figure 15: MD/PE posterior results of the model parameters for all toy models.
Shown are average marginal posteriors for all lower bounds, including the
prior and an average reference posterior. Figure 16: FP results for the linear
toy model. Shown are training curves for the NWJ, InfoNCE and JSD (via NWJ)
lower bounds. The top row shows lower bound estimates as a function of
training epochs, with the dotted line being the reference MI value, and the
bottom row shows the elements of the design vector as it is being updated.

##  Appendix C Epidemiology experiments

###  C.1 Discrete state continuous time Markov chain formulation

We start by defining the SIR and SEIR models in terms of continuous time
Markov chains (CTMC), where time is continuous and the state populations are
discrete.

Recall that the SIR model (e.g. Allen, 2008) is governed by the state changes
$S_{1}(t)\rightarrow I_{1}(t)\rightarrow R_{1}(t)$, where $S_{1}(t)$ are the
number of susceptible individuals, $I_{1}(t)$ are the number of infectious
individuals and $R_{1}(t)$ are the individuals that have recovered and can no
longer be infected. The dynamics of these state changes are determined by the
model parameters $\bm{\theta}_{1}=(\beta,\gamma)$, where $\beta$ is the
infection rate and $\gamma$ is the recovery rate. As done in the main text,
let us define a population vector
$\mathbf{X}_{1}(t)=(S_{1}(t),I_{1}(t))^{\top}$.121212We can ignore the
population of recovered individuals here because we assume that the total
population $N$ stays constant. The state changes $\Delta\mathbf{X}_{1}(t)$ of
the SIR model within a small time frame $\Delta t$ are then affected by two
events: _infection_ , where $\Delta\mathbf{X}_{1}(\Delta t)=(-1,+1)^{\top}$,
and _recovery_ , where $\Delta\mathbf{X}_{1}(\Delta t)=(0,-1)^{\top}$.
Denoting the $i$-th event by $(\Delta\mathbf{X}_{1}(t))^{i}$, Table 7
summarises these state changes and their corresponding transition
probabilities131313$\mathchoice{{\scriptstyle\mathcal{O}}}{{\scriptstyle\mathcal{O}}}{{\scriptscriptstyle\mathcal{O}}}{\scalebox{0.7}{$\scriptscriptstyle\mathcal{O}$}}(x)$
in Tables 7 and 8 denotes a function that goes faster to zero than $x$, i.e.
$\lim_{x\to
0}\mathchoice{{\scriptstyle\mathcal{O}}}{{\scriptstyle\mathcal{O}}}{{\scriptscriptstyle\mathcal{O}}}{\scalebox{0.7}{$\scriptscriptstyle\mathcal{O}$}}(x)/x=0$.,
which define the behaviour of the CTMC SIR model.

Table 7: CTMC SIR model state changes and their probabilities. Event | State Change | Probability
---|---|---
Infection | $(\Delta\mathbf{X}_{1}(t))^{1}=(-1,+1)^{\top}$ | $p_{1}=\beta\frac{S_{1}(t)I_{1}(t)}{N}\Delta t+\mathchoice{{\scriptstyle\mathcal{O}}}{{\scriptstyle\mathcal{O}}}{{\scriptscriptstyle\mathcal{O}}}{\scalebox{0.7}{$\scriptscriptstyle\mathcal{O}$}}(\Delta t)$
Recovery | $(\Delta\mathbf{X}_{1}(t))^{2}=(\,\,\,\,\,0,-1)^{\top}$ | $p_{2}=\gamma I_{1}(t)\Delta t+\mathchoice{{\scriptstyle\mathcal{O}}}{{\scriptstyle\mathcal{O}}}{{\scriptscriptstyle\mathcal{O}}}{\scalebox{0.7}{$\scriptscriptstyle\mathcal{O}$}}(\Delta t)$

The SEIR model (Lekone and Finkenstädt, 2006) introduces an exposed state
$E_{2}(t)$, where individuals are infected but not yet infectious, and is
governed by the state changes $S_{2}(t)\rightarrow E_{2}(t)\rightarrow
I_{2}(t)\rightarrow R_{2}(t)$. The dynamics of these changes are determined by
the model parameters $\bm{\theta}_{2}=(\beta,\gamma,\sigma)$, where $\beta$
and $\gamma$ are as before and $\sigma^{-1}$ is the average incubation period.
We again define a population vector
$\mathbf{X}_{2}(t)=(S_{2}(t),E_{2}(t),I_{2}(t))^{\top}$. The state changes
$\Delta\mathbf{X}_{2}(t)$ of the SEIR model are affected by three events:
_infection_ , _becoming infectious_ and _recovery_. These state changes,
denoted by $(\Delta\mathbf{X}_{2}(t))^{i}$, and their corresponding
probabilities, which together define the CTMC SEIR model, are summarised in
Table 8.

Table 8: CTMC SEIR model state changes and their probabilities. Event | State Change | Probability
---|---|---
Infection | $(\Delta\mathbf{X}_{2}(t))^{1}=(-1,+1,\,\,\,\,\,0)^{\top}$ | $p_{1}=\beta\frac{S_{2}(t)I_{2}(t)}{N}\Delta t+\mathchoice{{\scriptstyle\mathcal{O}}}{{\scriptstyle\mathcal{O}}}{{\scriptscriptstyle\mathcal{O}}}{\scalebox{0.7}{$\scriptscriptstyle\mathcal{O}$}}(\Delta t)$
Becoming infectious | $(\Delta\mathbf{X}_{2}(t))^{2}=(\,\,\,\,\,0,-1,+1)^{\top}$ | $p_{2}=\sigma E_{2}(t)\Delta t+\mathchoice{{\scriptstyle\mathcal{O}}}{{\scriptstyle\mathcal{O}}}{{\scriptscriptstyle\mathcal{O}}}{\scalebox{0.7}{$\scriptscriptstyle\mathcal{O}$}}(\Delta t)$
Recovery | $(\Delta\mathbf{X}_{2}(t))^{3}=(\,\,\,\,\,0,\,\,\,\,\,0,-1)^{\top}$ | $p_{3}=\gamma I_{2}(t)\Delta t+\mathchoice{{\scriptstyle\mathcal{O}}}{{\scriptstyle\mathcal{O}}}{{\scriptscriptstyle\mathcal{O}}}{\scalebox{0.7}{$\scriptscriptstyle\mathcal{O}$}}(\Delta t)$

###  C.2 Diffusion Approximations

Based on the CTMC formulation of the SIR and SEIR models in Table 7 and Table
8, we can derive continuous-state diffusion approximations of these
epidemiological processes using the method of Allen et al. (2008), which
results in systems of stochastic differential equations (SDEs).

In the main text we wrote down the system of Itô SDEs as

$\mathrm{d}\mathbf{X}_{m}(t)=\mathbf{f}_{m}(\mathbf{X}_{m}(t))\mathrm{d}t+\mathbf{G}_{m}(\mathbf{X}_{m}(t))\mathrm{d}\mathbf{W}(t),$
(C.1)

where $m$ is a model indicator, $\mathbf{f}_{m}$ the drift vector,
$\mathbf{G}_{m}$ the diffusion matrix, and $\mathbf{W}(t)$ is a vector of
independent Wiener processes. The drift vector equals the infinitesimal mean
$\lim_{\mathrm{d}t\to 0}\mathbb{E}(\mathrm{d}\mathbf{X}_{m}(t))/\mathrm{d}t$
of the diffusion, and $\mathbf{G}_{m}$ defines the infinitesimal variance
$\lim_{\mathrm{d}t\to
0}\mathbb{V}\text{ar}(\mathrm{d}\mathbf{X}_{m}(t))/\mathrm{d}t=\mathbf{G}_{m}\mathbf{G}_{m}^{\top}$.
They can be chosen such that the diffusion approximation matches the
infinitesimal mean and variance of the discrete state continuous Markov chain
models. Allen et al. (2008) provide a systematic way for doing this, which we
shall follow below.

Let us first consider the CTMC SIR model summarised in Table 7. The drift
vector is chosen to match the expected change in $\mathbf{X}_{1}(t)$ of the
CTMC formulation, i.e.

$\displaystyle\mathbf{f_{1}}(\mathbf{X}_{1}(t))$ $\displaystyle=\lim_{\Delta
t\to 0}\frac{1}{\Delta t}\mathbb{E}\left[\Delta\mathbf{X}_{1}(t)\right]$ (C.2)
$\displaystyle=\lim_{\Delta t\to 0}\frac{1}{\Delta
t}\sum_{i=1}^{2}p_{i}(\Delta\mathbf{X}_{1}(t))^{i}$ (C.3)
$\displaystyle=\beta\frac{S_{1}(t)I_{1}(t)}{N}\begin{pmatrix}-1\\\
1\end{pmatrix}+\gamma I_{1}(t)\begin{pmatrix}0\\\ -1\end{pmatrix}$ (C.4)
$\displaystyle=\begin{pmatrix}-\beta\frac{S_{1}(t)I_{1}(t)}{N}\\\\[10.00002pt]
\beta\frac{S_{1}(t)I_{1}(t)}{N}-\gamma I_{1}(t)\\\\[10.00002pt]
\end{pmatrix}.$ (C.5)

In order to derive the diffusion matrix $\mathbf{G}_{1}$, let us first define
a matrix $\bm{\Lambda}_{1}$ whose rows correspond to the state changes
$(\Delta\mathbf{X}_{1}(t))^{i}$ in Table 7, i.e.

$\bm{\Lambda}_{1}=\begin{pmatrix}-1&+1&\\\ 0&-1&\\\ \end{pmatrix}.$ (C.6)

Allen et al. (2008) show that the elements $G_{1,ij}$ of the diffusion matrix
$\mathbf{G}_{1}$ are then given by

$G_{1,ij}=\Lambda_{ji}\,p_{j}^{1/2},$ (C.7)

allowing us to derive the full diffusion matrix using Table 7 and Equation
C.7, i.e.

$\displaystyle\mathbf{G}_{1}(\mathbf{X}_{1}(t))$
$\displaystyle=\begin{pmatrix}\Lambda_{11}\,p_{1}^{1/2}&\Lambda_{21}\,p_{2}^{1/2}\\\
\Lambda_{12}\,p_{1}^{1/2}&\Lambda_{22}\,p_{2}^{1/2}\\\ \end{pmatrix}$ (C.8)
$\displaystyle=\begin{pmatrix}-\sqrt{\beta\frac{S_{1}(t)I_{1}(t)}{N}}&0\\\\[10.00002pt]
\sqrt{\beta\frac{S_{1}(t)I_{1}(t)}{N}}&-\sqrt{\gamma I_{1}(t)}\\\\[10.00002pt]
\end{pmatrix}.$ (C.9)

We can check that this diffusion matrix matches the infinitesimal variance of
the CTMC SIR model, i.e.

$\displaystyle\lim_{\Delta t\to
0}\frac{\mathbb{V}\text{ar}(\Delta\mathbf{X}_{1}(t))}{\Delta t}$
$\displaystyle=\lim_{\Delta t\to 0}\frac{1}{\Delta
t}\sum_{i=1}^{2}p_{i}(\Delta\mathbf{X}_{1}(t))^{i}(\Delta\mathbf{X}_{1}(t))^{i,\top}-\lim_{\Delta
t\to 0}\frac{1}{\Delta
t}\mathbb{E}\left[\Delta\mathbf{X}_{1}(t)\right]\mathbb{E}\left[\Delta\mathbf{X}_{1}(t)\right]^{\top}$
$\displaystyle=\lim_{\Delta t\to 0}\frac{1}{\Delta
t}\sum_{i=1}^{2}p_{i}(\Delta\mathbf{X}_{1}(t))^{i}(\Delta\mathbf{X}_{1}(t))^{i,\top}-\lim_{\Delta
t\to
0}\frac{\mathchoice{{\scriptstyle\mathcal{O}}}{{\scriptstyle\mathcal{O}}}{{\scriptscriptstyle\mathcal{O}}}{\scalebox{0.7}{$\scriptscriptstyle\mathcal{O}$}}(\Delta
t)}{\Delta t}$ (C.10)
$\displaystyle=\beta\frac{S_{1}(t)I_{1}(t)}{N}\begin{pmatrix}-1\\\
1\end{pmatrix}\begin{pmatrix}-1&1\end{pmatrix}+\gamma
I_{1}(t)\begin{pmatrix}0\\\ -1\end{pmatrix}\begin{pmatrix}0&-1\end{pmatrix}$
(C.11)
$\displaystyle=\begin{pmatrix}\beta\frac{S_{1}(t)I_{1}(t)}{N}&-\beta\frac{S_{1}(t)I_{1}(t)}{N}\\\\[10.00002pt]
-\beta\frac{S_{1}(t)I_{1}(t)}{N}&\beta\frac{S_{1}(t)I_{1}(t)}{N}+\gamma
I_{1}(t)\\\\[10.00002pt] \end{pmatrix},$ (C.12)

which exactly corresponds to $\mathbf{G}_{1}\mathbf{G}_{1}^{\top}$, as
required. We note that the diffusion matrix $\mathbf{G}_{1}$ is not uniquely
defined, as we only wish to match the variance. For instance, we could
multiply $\mathbf{G}_{1}$ by any orthogonal matrix and the product
$\mathbf{G}_{1}\mathbf{G}_{1}^{\top}$ remains unchanged.

We now follow the same procedure for the SEIR model. To do so, let us consider
the state changes and corresponding probabilities summarised in Table 8. As
done for the SIR model, given these state changes and probabilities, we can
derive the drift vector by computing the expected change in
$\mathbf{X}_{2}(t)$,

$\displaystyle\mathbf{f_{2}}(\mathbf{X}_{2}(t))$ $\displaystyle=\lim_{\Delta
t\to 0}\frac{1}{\Delta t}\mathbb{E}\left[\Delta\mathbf{X}_{2}(t)\right]$
(C.13) $\displaystyle=\lim_{\Delta t\to 0}\frac{1}{\Delta
t}\sum_{i=1}^{3}p_{i}(\Delta\mathbf{X}_{2}(t))^{i}$ (C.14)
$\displaystyle=\beta\frac{S_{2}(t)I_{2}(t)}{N}\begin{pmatrix}-1\\\ 1\\\
0\end{pmatrix}+\sigma E_{2}(t)\begin{pmatrix}0\\\ -1\\\ 1\end{pmatrix}+\gamma
I_{2}(t)\begin{pmatrix}0\\\ 0\\\ -1\end{pmatrix}$ (C.15)
$\displaystyle=\begin{pmatrix}-\beta\frac{S_{2}(t)I_{2}(t)}{N}\\\\[10.00002pt]
\beta\frac{S_{2}(t)I_{2}(t)}{N}-\sigma E_{2}(t)\\\\[10.00002pt] \sigma
E_{2}(t)-\gamma I_{2}(t)\\\\[10.00002pt] \end{pmatrix}.$ (C.16)

In order to derive the diffusion matrix $\mathbf{G}_{2}$, we define a matrix
$\bm{\Lambda}_{2}$ whose rows correspond to the state changes in Table 8, i.e.

$\bm{\Lambda}_{2}=\begin{pmatrix}-1&+1&0\\\ 0&-1&+1\\\ 0&0&-1\\\
\end{pmatrix}.$ (C.17)

Substituting $\bm{\Lambda}_{2}$ in Equation C.7 and using Table 8 then allows
us to derive $\mathbf{G}_{2}$,

$\displaystyle\mathbf{G}_{2}(\mathbf{X}_{2}(t))$
$\displaystyle=\begin{pmatrix}\Lambda_{11}\,p_{1}^{1/2}&\Lambda_{21}\,p_{2}^{1/2}&\Lambda_{31}\,p_{3}^{1/2}\\\
\Lambda_{12}\,p_{1}^{1/2}&\Lambda_{22}\,p_{2}^{1/2}&\Lambda_{32}\,p_{3}^{1/2}\\\
\Lambda_{13}\,p_{1}^{1/2}&\Lambda_{23}\,p_{2}^{1/2}&\Lambda_{33}\,p_{3}^{1/2}\\\
\end{pmatrix}$ (C.18)
$\displaystyle=\begin{pmatrix}-\sqrt{\beta\frac{S_{2}(t)I_{2}(t)}{N}}&0&0\\\
\sqrt{\beta\frac{S_{2}(t)I_{2}(t)}{N}}&-\sqrt{\sigma E_{2}(t)}&0\\\
0&\sqrt{\sigma E_{2}(t)}&-\sqrt{\gamma I_{2}(t)}\\\ \end{pmatrix}.$ (C.19)

We can again check that $G_{2}$ matches the infitesimal variance of the CTMC
SEIR model,

$\displaystyle\lim_{\Delta t\to
0}\frac{\mathbb{V}\text{ar}(\Delta\mathbf{X}_{2}(t))}{\Delta t}$
$\displaystyle=\lim_{\Delta t\to 0}\frac{1}{\Delta
t}\sum_{i=1}^{3}p_{i}(\Delta\mathbf{X}_{2}(t))^{i}(\Delta\mathbf{X}_{2}(t))^{i,\top}-\lim_{\Delta
t\to 0}\frac{1}{\Delta
t}\mathbb{E}\left[\Delta\mathbf{X}_{2}(t)\right]\mathbb{E}\left[\Delta\mathbf{X}_{2}(t)\right]^{\top}$
$\displaystyle=\lim_{\Delta t\to 0}\frac{1}{\Delta
t}\sum_{i=1}^{3}p_{i}(\Delta\mathbf{X}_{2}(t))^{i}(\Delta\mathbf{X}_{2}(t))^{i,\top}-\lim_{\Delta
t\to
0}\frac{\mathchoice{{\scriptstyle\mathcal{O}}}{{\scriptstyle\mathcal{O}}}{{\scriptscriptstyle\mathcal{O}}}{\scalebox{0.7}{$\scriptscriptstyle\mathcal{O}$}}(\Delta
t)}{\Delta t}$ (C.20) $\displaystyle\begin{aligned}
=\beta\frac{S_{2}(t)I_{2}(t)}{N}&\begin{pmatrix}-1\\\ 1\\\
0\end{pmatrix}\begin{pmatrix}-1&1&0\end{pmatrix}\\\ &+\sigma
E_{2}(t)\begin{pmatrix}0\\\ -1\\\
1\end{pmatrix}\begin{pmatrix}0&-1&1\end{pmatrix}\\\ &+\gamma
I_{2}(t)\begin{pmatrix}0\\\ 0\\\
-1\end{pmatrix}\begin{pmatrix}0&0&-1\end{pmatrix}\end{aligned}$ (C.21)
$\displaystyle=\begin{pmatrix}\beta\frac{S_{2}(t)I_{2}(t)}{N}&-\beta\frac{S_{2}(t)I_{2}(t)}{N}&0\\\\[10.00002pt]
-\beta\frac{S_{2}(t)I_{2}(t)}{N}&\beta\frac{S_{2}(t)I_{2}(t)}{N}+\sigma
E_{2}(t)&-\sigma E_{2}(t)\\\\[10.00002pt] 0&-\sigma E_{2}(t)&\sigma
E_{2}(t)+\gamma I_{2}(t)\\\ \end{pmatrix},$ (C.22)

which exactly corresponds to $\mathbf{G}_{2}\mathbf{G}_{2}^{\top}$, as
required. However, as previously said, the diffusion matrix $\mathbf{G}_{2}$
is not uniquely defined.

###  C.3 Architectures and hyper-parameters

We list the hyper-parameters for the parameter estimation (PE) and model
discrimination (MD) task of the SDE-based epidemiology experiments in Table 9
and Table 10, respectively, for all experimental budgets $D$ that were
presented in the main text. Shown are neural network (NN) architectures,
including the number of hidden layers and number of hidden units, as well as
the learning rates (L.R.) for the NN parameters $\bm{\psi}$ and experimental
designs $\mathbf{d}$. We optimise $\bm{\psi}$ and $\mathbf{d}$ with two
separate Adam optimisers with default parameters from the PyTorch package in
Python. We pre-simulate $20{,}000$ SDE solutions on a fine time grid. During
training time we can then access solutions, and their gradients, at a specific
time point by simply looking up the solutions corresponding to the nearest
point in the time grid. We found that this was generally much faster than
simulating/solving the SDEs on the fly, at the cost of higher memory usage.

Table 9: PE hyper-parameters for the SDE-based SIR model, for varying number of measurements $D$ (i.e. experimental budget). Budget $D$ | NN Layers | NN Units | L.R. for $\bm{\psi}$ | L.R. for $\mathbf{d}$
---|---|---|---|---
1 | 2 | 20 | $10^{-4}$ | $3\times 10^{-2}$
2 | 3 | 20 | $10^{-4}$ | $3\times 10^{-2}$
3 | 4 | 20 | $10^{-4}$ | $3\times 10^{-2}$
5 | 5 | 20 | $3\times 10^{-4}$ | $3\times 10^{-2}$
10 | 3 | 30 | $10^{-4}$ | $10^{-2}$
Table 10: MD hyper-parameters for the SDE-based SIR and SEIR models, for varying number of measurements $D$ (i.e. experimental budget). Budget $D$ | NN Layers | NN Units | L.R. for $\bm{\psi}$ | L.R. for $\mathbf{d}$
---|---|---|---|---
1 | 2 | 20 | $10^{-4}$ | $3\times 10^{-2}$
2 | 3 | 20 | $10^{-4}$ | $3\times 10^{-2}$
3 | 4 | 20 | $10^{-4}$ | $3\times 10^{-2}$
5 | 5 | 20 | $3\times 10^{-4}$ | $3\times 10^{-2}$
10 | 3 | 30 | $10^{-4}$ | $10^{-2}$
20 | 3 | 30 | $10^{-4}$ | $10^{-2}$

###  C.4 Data distributions

In Figure 17 we provide general information about data simulated from the SIR
and SEIR models. The top left plot shows the SIR model prior predictive
distributions of the number of susceptible individuals $S(t)$ and the number
of infectious individuals $I(t)$ as a function of time $t$. Similarly, the top
right plot shows the SEIR model prior predictive distributions of $S(t)$,
$I(t)$ and the number of exposed individuals $E(t)$. Recall, however, that we
only use the number of infectious individuals as data in our experiments. As
such, we compare $I(t)$ from the SIR and from the SEIR model more closely in
the bottom left plot. As can be seen, the number of infectious individuals for
the SEIR model peak at later times $t$. The bottom centre plot then shows the
average signal-to-noise ratio (SNR), computed by means of Equation B.13, for
the $I(t)$ response of the SIR and SEIR model. The average SNR for the SEIR
model is generally much higher than that of the SIR model, also peaking at
later measurement times. Lastly, the bottom right plot shows the Jensen-
Shannon divergence between the prior predictives of the SIR and SEIR model,
i.e. $\text{JS}(p(I_{1}(t)|t)\mid\mid p(I_{2}(t)|t))$, where $I_{1}(t)$ and
$I_{2}(t)$ are the number of infectious individuals for the SIR and SEIR
model, respectively. Interestingly, there are two peaks in the Jensen-Shannon
divergence, towards earlier measurement times and around $t=60$. The prior
predictive distributions are most similar where the means of the data
distributions cross, i.e. near $t=30$ (see the bottom left plot).

Figure 17: General summary of data simulated from the SDE-based epidemiology
models. The top row shows prior predictive distributions for the SIR model
(top left) and the SEIR model (top right). The bottom left figure shows the
number of infectious individuals for the SIR and SEIR models, while the bottom
centre plot shows the corresponding average signal-to-noise ratios (SNR). The
bottom right shows the Jensen-Shannon (JS) divergence between the prior
predictive distributions of both models.

###  C.5 Additional results

We here present and discuss several additional results for the epidemiology
experiments in the main text.

Figure 18 shows the training curves of the parameter estimation (PE) task for
the SIR model with different experimental budgets $D$. For these experiments
we maximised the JSD lower bound, which is shown in the top row as a function
of training epochs. The elements of the experimental design vector are shown
in the middle row and show a convergence towards early measurement times. The
bottom row of Figure 18 shows evaluations of the NWJ lower bound using the
density ratio learned by maximising the JSD lower bound (as discussed in the
main text). These NWJ lower bound evaluations show large fluctuations, in
stark contrast to the smooth JSD lower bound curves. These fluctuations
presumably arise because of the exponential term in the NWJ lower bound (see
Equation 3.2 in the main text). Importantly, however, we only use the NWJ
bound as a means to get an estimate of the mutual information, and we do not
use it to update the neural network parameters or experimental designs. Thus,
the training behaviour is not affected by these large fluctuations seen in the
top row. Putting a cap on the exponential term in the NWJ lower bound may help
reduce variance (as done in Song and Ermon, 2020) if that is required.

Figure 18: PE training curves for the SDE-based SIR model, with different
experimental budgets $D$. Shown are JSD lower bound values (top row), the
elements of the experimental design vector (middle row) and NWJ lower bound
evaluations (bottom row). Importantly, the JSD lower bound was used to update
neural network parameters and designs, while the NWJ lower bound was purely
used to obtain MI estimates. Figure 19: MD training curves for the SDE-based
SIR and SEIR models, with different experimental budgets $D$. Shown are JSD
lower bound values (top row), the elements of the experimental design vector
(middle row) and NWJ lower bound evaluations (bottom row). Importantly, the
JSD lower bound was used to update neural network parameters and designs,
while the NWJ lower bound was purely used to obtain MI estimates.

Similarly, Figure 19 shows the training curves of the model discrimination
(MD) task for the SDE-based SIR and SEIR models, with different experimental
budgets $D$, i.e. number of measurements. Unlike for the PE task, the NWJ
lower bound evaluations for the MD task do not have large fluctuations and are
quite smooth. Validation estimates of the mutual information for different
number of measurements are shown in Figure 20, as well as corresponding
optimal designs (that are also shown in the main text). We show corresponding
average posterior distributions for different ground truth models in Figure
21. Model recovery tends to improve as the number of measurements increases.
However, the model recovery of the SEIR model, i.e. the average posterior
probability
$\mathbb{E}[\widehat{p}(m=2|\mathbf{y}^{\ast},\mathbf{d}^{\ast},m_{\text{truth}}=2)]$,
is always worse than that of the SIR model, i.e. the average posterior
probability
$\mathbb{E}[\widehat{p}(m=1|\mathbf{y}^{\ast},\mathbf{d}^{\ast},m_{\text{truth}}=1)]$,
which are shown by the diagonal entries in Figure 21. This may be because the
optimal designs (shown in Figure 20) are all clustered towards earlier
measurement times, i.e. below $t=40$. Looking at the bottom left of Figure 17,
this is the region where the SIR model responses dominate over the SEIR model
responses.

Figure 20: MD results for the SDE-based SIR and SEIR models with different
experimental budgets. The left plot shows validation MI estimates for