Generative Models: An Interdisciplinary Perspective

Imagine we are designing a study whose purpose is to evaluate the influence of three interventions on microbiota community composition. For example, one intervention might be a diet change, and we hope to identify taxa whose abundances are systematically elevated or depressed following the intervention. The task is to guide the choice of the number of samples and the allocation of treatments across study participants. Previous studies point to substantial subject-to-subject variation in taxa composition. For this reason, a blocked design seems appropriate, where each subject is assigned each of the three interventions in random order. We argue that, by allowing estimation of a subject-level baseline, such a design will allow detection of small treatment effects even in the presence of large subject-to-subject differences. To our surprise, we face resistance to this proposal — we are told that asking each study participant to contribute three samples creates an undue burden and risks increased study dropout.

To sharpen an analysis of the trade-offs involved, a simulation can be used. Suppose that the biological specimen, sample $i$’s community composition, is drawn from

where $\mathcal{S}$ is the softmax function, $\mathcal{S}\left(z\right)=\frac{\operatorname{exp}\left(z_{d}\right)}{\sum_{d^{\prime}=1}^{D}\operatorname{exp}\left(z_{d^{\prime}}\right)}$, and $\beta_{t}$ and $\tau_{s}$ encode treatment ($t=1,2,3$) and subject effects across the $D$ species, respectively. Let $A$ designate the set of taxa that exhibit any treatment effect. Whenever $\sigma^{2}_{\tau}\gg\sigma^{2}_{\beta}$, it is clear that blocking is necessary; conversely, when $\sigma^{2}_{\tau}\ll\sigma^{2}_{\beta}$, there is no reason to create an extra burden on study participants. But what is the appropriate decision in more intermediate cases?

We study the question by simulating from the model (1) under both blocked and unblocked designs and then sample from the posteriors $p\left(\beta|x_{1:N}\right)$ using Markov Chain Monte Carlo (MCMC). We set $N=60$ as the total number of samples (60 subjects in the unblocked design, 20 for the blocked design) and $D=25$ features. We suppose that that $\sigma_{\beta}^{2}=0.5$ and $\sigma^{2}_{\tau}=2$ and that only features $d=1,\dots,10$ contain true effects.

Figure 1 shows the resulting posteriors. Both designs support recovery of the direction for each treatment effect and are shrunken towards 0, a natural consequence of the model’s hierarchical structure. However, for each taxon $d$, we note that the blocked design has narrower posteriors $p\left(\beta_{td}|x_{1:N}\right)$. Specifically, the marginal variances under the unblocked design are on average 1.3 times larger than those under the blocked design. Practically, this leads to reduced power against weak effects in the unblocked design, consider for example the difference in estimated effects for taxon 8. Perhaps even more troubling are the instances in the unblocked design where we might be misled into believing in the existence of a treatment effect; consider for instance, taxon 19 or 21. This is likely the consequence of a few subjects having large (or small) abundances for these taxa simply by chance, and then dragging their corresponding treatments up (or down). Though the unblocked design gives theoretically unbiased inferences, in finite samples we still risk spurious associations. The repeated subject-level samples in the blocked design allow these outlier taxa-subject combinations to be directly accounted for.

Overall, blocking seems natural given this hypothetical setup. However, our point is less the final choice of design than the comparison process. If experience had suggested a different ratio of $\sigma^{2}_{\beta}$ to $\sigma^{2}_{\tau}$, the same process may have led us to the opposite decision. Moreover, this simple analysis could be enriched in several directions. It is possible to model differential study dropout in the two designs. If data were available from studies with similar expected treatment effects, we could tailor the sizes of $\sigma^{2}_{\beta}$ and $\sigma^{2}_{\tau}$ to more closely match past data. Finally, an interactive version of Figure 1 would support rapid comparison across simulation parameters.

Even though the generative mechanism in the previous example is complex, the associated design space was simple — the choice was between either a blocked or an unblocked design. In more complex design spaces, it is impractical to evaluate all candidate designs. In this setting, an organized search over candidates can allow the more efficient discovery of promising designs (Müller, 2005, Huan & Marzouk, 2013). To make this search possible, we assume that utility tends to vary smoothly over the design space. Designs similar to those that are known to be ineffective can be rapidly discarded, and exploration can be concentrated on areas of higher expected utility or uncertainty.

To illustrate, consider the problem of designing an experiment to evaluate the short and long-term effects of an intervention; e.g., the effect of antibiotic treatment on a bacterial strain of interest. The intervention effect is thought to be strong but brief, see Figure 2. Sampling is costly, and in the absence of a formula for the optimal sample placement, we turn our attention to simulation.

In principle, the design space can consist of the locations of the $N=25$ sampling points directly. However, searching across $\left[-10,10\right]^{N}$ becomes difficult for even moderate $N$. Instead, we suppose that sampled timepoints are drawn from $\mathbb{P}_{w}:=\sum_{k=1}^{K}w_{k}\mathbb{I}\left({S_{k}}\right)$, a mixture of uniforms over predefined sets $S_{k}$ (here, we partition $\left[0,1\right]$ into $K=7$ equal length intervals) but with weights to be optimized. In this way, we reduce the search to the $K-1$-dimensional simplex.

Suppose that $f$ is the true effect and $\hat{f}_{x}$ is our estimate based on a specific sampling times $x$. We write the expected loss for a given set of weights by $L\left(w\right):=\mathbb{E}_{x\sim\mathbb{P}_{w}}\left[{\int_{\left[-10,10\right]}\left(\hat{f}_{x}\left(z\right)-f\left(z\right)\right)^{2}dz}\right]$. This is an ideal loss function that can be used to search over the simplex; it can be empirically approximated by sampling $B$ different designs $x^{b}$ from a proposed $\mathbb{P}_{w}$, computing an estimate $\hat{f}_{x^{b}}$ (we use a spline with six degrees of freedom), and then computing $\frac{1}{n_{z}}\sum_{b=1}^{B}\sum_{i=1}^{n_{z}}\left(\hat{f}_{x^{b}}\left(z_{i}\right)-f_{x^{b}}\left(z_{i}\right)\right)^{2}$ across a dense grid $z_{i}\in\left[-10,10\right]$.

The loss $L\left(w\right)$ can be used to compare alternative designs $w$, but evaluating it across all of $\Delta^{K-1}$ is computationally impractical. As an alternative, we apply Bayesian optimization to iteratively guide our search towards promising $w$ (Shahriari et al., 2015). This approach first places a Gaussian Process prior on the loss function $L\left(w\right)$. This encodes the belief that no particular configuration of $w$ should be favored a priori, but that the loss should vary smoothly as a function of $w$. The optimization strategy balances exploration and exploitation of competing designs, narrowing in on promising configurations without prematurely ruling out alternatives. Specifically, we first randomly sample an initial set of weights $w_{1},\dots,w_{N_{\text{init}}}$ from a logistic-normal distribution. For each weight $w_{i}$, we evaluate $L\left(w_{i}\right)$. The paired $w_{i},L\left(w_{i}\right)$ are used to compute a posterior. Given this posterior, we find $w$ with small values of $\mathbb{E}\left[{L\left(w\right)}\right]-\mathrm{Var}\left(L\left(w\right)\right)$, representing points with either low expected loss or high uncertainty. We evaluate $L\left(w\right)$ at these candidate $w$, update the posterior, and continue until convergence.

A few choices of $\mathbb{P}_{w}$ discovered through this process are given in Figure 3. These distributions place elevated weight near the intervention, matching the intuition that sampling should be densest in regions where the effect is expected to vary the most. To summarize, the overall approach is,

• •

  Specify a generative mechanism. In this example, it is the shape of the unknown treatment effect and the distribution of random variation around the true effect.

• •

  Define a design space and loss. This is necessary to formally evaluate competing designs. Here, the design space is the choice of weights in $\mathbb{P}_{w}$ and the loss is the average estimation error when using sampling times drawn from $\mathbb{P}_{w}$.

• •

  Optimize over designs. Iteratively evaluate losses on candidate designs, choosing candidates in a way that balances exploration with exploitation. In this example, we used Bayesian optimization.

Of the three steps, the final optimization is the most automatic. Specification of the generative mechanism and the loss require care. That said, the full modeling toolkit can be applied to improve the faithfulness of the specification, giving the designer considerable flexibility.

The examples above show how simulators can be built from scratch to support experimental design. However, for many interesting data types, there are already a number of available simulators. A complete inventory of available packages for simulation-based design of experiments is not possible, a list of R packages for Experimental Design can be found at https://cran.r-project.org/web/views/ExperimentalDesign.html. There are in fact hundreds of packages for doing design by simulation – we have curated a list in Supplementary Table 1. These packages make it possible to reuse code across the community, minimizing duplicated effort planning designs.

The simulators used within a community must be tailored to properties of the data it must generate. Nonetheless, for design-by-simulation, a few properties are generally useful,

1. 1.

   Calibration: We should be able to calibrate simulation outputs to match existing datasets. Rather than simulating data de novo, public datasets and pilot studies can be used to set expectations for a new experiment.

2. 2.

   Evaluation: Metrics should be available to evaluate the discrepancy between (a) real and simulated datasets and (b) analysis results derived from real and simulated data.

3. 3.

   Control: We should be able to directly manipulate both experimental and generative parameters. A simulator that produces realistic data is not useful for power calculations if it can only be run for a single sample size, for example.

4. 4.

   Transparency: The mechanisms leading to simulated data should be interpretable and refer to language already used within the application domain.

5. 5.

   Usability: A well-documented package and interactive interface can make a simulator more accessible, supporting its adoption.

To see how these principles can guide practice, consider how simulation packages guide experimental design in single-cell genomics. A variety of simulation mechanisms have been proposed to capture several properties common across single-cell datasets: overdispersed counts, a high degree of sparsity, low effective dimensionality, and the presence of batch effects (Zappia et al., 2017, Risso et al., 2017, Zhang et al., 2019a, Yu et al., 2020, Sun et al., 2021, Schmid et al., 2021, Qin et al., 2022). Most of these simulators are built from rich probabilistic or differential equations models. For example, scDesign2 models each gene marginally using a zero-inflated negative binomial distribution and then induces correlation using a copula (Li & Li, 2019, Sun et al., 2021). SymSin uses a sequential strategy to model the transformation of transcripts within a cell to observed counts in a single-cell dataset (Zhang et al., 2019a). Some simulators are designed with calibration and control built-in — for example, given a pilot dataset, they make it possible to simulate a version of the data with twice as many cells but with fewer differentially expressed genes.

The community has additionally proposed a variety of metrics for evaluating the faithfulness of these simulators. For example, gene-level distributions, cross-gene mean-variance relationships, dataset-level sparsity, and comparability of derived dimensionality reduction visualizations have all been used to measure discrepancy between real and simulated data (Cao et al., 2021, Soneson & Robinson, 2018). Both simulators and benchmarking (comparing performances in standard situations) have been encapsulated into packages for wider dissemination.

Note that these simulators were not originally designed with experimental design in mind – instead, most were written to support benchmarking. This was necessary because manual annotation of real single-cell data requires expertise and is labor-intensive. In contrast, simulators come with “ground truth” annotations. Nonetheless, after these simulators were implemented, they began to be used for design studies, and a later generation of simulators were written explicitly with design studies in mind.

When a dataset has many related, but not identically distributed, subsets, then a hierarchical structure may be applicable (Draper et al., 1993, Gelman, 2006, Teh & Jordan, 2010). In this type of module, a collection of parameters, one for each subset, is drawn from a shared prior. This implements a form of partial pooling — information is shared across subsets, but heterogeneous effects are reflected in the unique parameter estimates across subsets. Moreover, this facilitates adaptive estimation, strongly shrinking noisily estimated parameters towards a global estimate while minimally affecting parameters in subsets with clear signals. The DAG representation is provide in Figure 4b.

This structure is frequently encountered in practice. For example, when estimating differential expression across a collection of genes, the true effect for each gene can be drawn from a shared prior, which may itself be adapted according to the distribution of observed test statistics (Stephens, 2017, Wang & Stephens, 2021). When modeling multiple gene expression profiles simultaneously, a hierarchical model improves power by borrowing strength rather than fitting each profile separately (Stephens, 2013, Flutre et al., 2013, Zhao et al., 2016). Similarly, in political science, county-level effects can be partially pooled, stabilizing inference in regions with few samples (Hill, 2017). In more complex settings, hierarchical components may need to be embedded within larger generative models. For example, if the same counties were measured over time, we may believe each county’s parameter has the potential to change slightly. Viewed more generally, the reason this module is valuable is because large datasets rarely arise as a large collection of truly independent and identically distributed samples (Jordan, 2011). Instead, they tend to be an amalgamation of related, but not identical subsets — big data often arise by gluing many small datasets.

Modules based on latent variables provide a way to induce dependence structure across samples or measurements. They are applicable when it is thought that data lie within clusters or low-dimensional gradients, but the specific form of these structures is not known in advance. In the grammar of generative models, each sample is associated with a “local” latent variable and observations are drawn independently, conditionally on these latent variables (Blei, 2014, Zhang & Jordan, 2009). For example, in a generative clustering model, each sample is associated with a latent cluster indicator. For probabilistic factor models, the analogous local latent variables are low-dimensional coordinates, and in a topic model, they are mixed-memberships across a collection of topics (Nguyen & Holmes, 2017, Sankaran & Holmes, 2018). In contrast, the parameters defining the locations of clusters, shape of the gradient, or factor loadings are summarized by “global” variables.

Although different types of local and global variables induce very different types of observed data, all models with this type of local-global breakdown have the same DAG conditional independence structure, shown in Figure 4c. As in hierarchical models, this DAG module can be embedded within larger models to account for other problem-specific structure. Of particular interest is latent structure that evolves over time, which we review next.

In data with a spatial or temporal component, we often expect dependence between close by samples. For example, in a model of the audio recording of a meeting, a latent variable $z_{t}$ may learn to distinguish different speakers. Since each speaker will typically speak across an uninterrupted segment of time, we should encourage $z_{t}$ to only switch occasionally. One approach is through a latent Markov model, whose DAG structure is given in Figure 4d. In this model, the latent variables follow Markov dynamics $p\left(z_{t+1}|z_{1:t}\right)=p\left(z_{t+1}|z_{t}\right)$. Given these latent states, the observed data are conditionally independent, specified by $p\left(x_{t}|z_{t}\right)$. For simple likelihoods $p\left(x_{t}|z_{t}\right)$ and dynamics $p\left(z_{t+1}|z_{t}\right)$, this quantity can be computed through a closed-form recursion, like the Viterbi algorithm or Kalman Filter. More complex settings are discussed in Section 3.3.2. For example, assuming that $z_{i}\in\{1,\dots,K\}$, we may define a transition matrix $P\left(z_{t+1}=k^{\prime}|z_{t}=k\right)$ with elevated weights along the diagonal, encouraging a “stickiness” in states over time. If we use Gaussian likelihoods $p\left(x_{t}|z_{t}=k\right)=\mathcal{N}\left(\mu_{k},\Sigma_{k}\right)$, we arrive at a version of the Gaussian mixture model where samples tend to remain in the same component for long stretches of time. Relaxing either the Markov or discreteness assumptions leads to further enriched models. For example, a semi-Markov DAG structure supposes that that future states $z_{t}$ may be modulated by the pattern of several recent states, not simply the current $z_{t}$. If hard assignment to one of $K$ clusters is undesirable, we may adopt a mixed-membership approach, enforcing Markov dynamics for $z_{t}$ lying in the simplex.

We can reason about spatial dependence similarly. For example, imagine modeling a spatial transcriptomics dataset where each sample is a high-dimensional vector of gene expression measurements for an individual cell. We expect that, in spite of the hundreds of genes measured, a $K$-dimensional $z_{i}$ could provide a sufficient description of cell-to-cell variation. In this case, we may seek spatial consistency across neighboring $z_{i}$. This can be accomplished by using a Markov graph, inducing dependence along a spatial grid rather than temporal chain. Alternatively, several recent proposals place a Gaussian Process (GP) prior on $z_{i}$, using $K$ output dimensions and two input (spatial) dimensions (Townes & Engelhardt, 2021, Shang & Zhou, 2022, Velten et al., 2022). The resulting smoothness in latent variables over space can aid interpretability and performance.

Probabilistic programs are the stochastic analogs of usual, deterministic computer programs. Hence, they are naturally suited to Monte Carlo, and a number of packages have used this fact to support inference in freely-specified generative models. In probabilistic programming packages, both the deterministic and stochastic relationships between parameters and data are represented computationally, and a generic inference algorithm is used to approximate the distribution of all unobserved random nodes (parameters and latent variables), conditionally on the observed ones (data).

For example, Stan, Pyro, and Anglican provide interfaces for specifying generative mechanisms in this way (Carpenter et al., 2017, Bingham et al., 2019, Wood et al., 2014). In the background, they parse the conditional dependence relationships to perform various forms of inference, like Markov Chain Monte Carlo (e.g., Hybrid Monte Carlo, the No U-Turn Sampler, the Gibbs sampler), Variational Inference, Importance Sampling, and Sequential Monte Carlo. The result of any of these inference algorithms is a fitted generative model, which can be used to sample from the approximate conditional distribution of unobserved nodes. Critically, to solve a statistical problem using generative models no longer requires manual derivation of model-specific inference algorithms, which historically made the generative approach inaccessible to all but technical specialists.

Though they do not require manual derivations, skillful use of these platforms requires practice. Hands-on introductions can be found at (Clark, 2022, McElreath, 2020, Blau & MacKinlay, 2021, Pyro Contributors, 2022). Probabilistic programming languages represent distributions as abstract objects, each with their own methods for evaluating distribution-specific properties, like the log probability of a dataset. Though this additional layer of abstraction can seem foreign at first, it allows models to be constructed by manipulating distributions at a high-level. The benefits of learning to operate at this level are analogous to those gained by learning languages that free the programmer from managing low-level memory — more attention can be invested in the design and experimentation rather than the implementation phase.

Here we will give details on an example that illustrates the algorithmic machinery that operates in the background of these inference packages.

Particle filters are a widely used collection of stochastic models, and since they bridge different “generative schools,” we pause to give a somewhat detailed discussion. We focus on the latent Markov model of Figure 4d. Having observed a sequence $x_{1:T}$, we aim to characterize the unobserved process in states $z_{1:T}$ giving rise to the observed stochastic function as data.

For example, we may have observed a sequence of economic indices over time. We imagine that the indices $x_{1:T}$ can be succinctly described by latent states $z_{1:T}$, and we are especially curious about when this latent state undergoes rapid transitions. For example, these periods of transition could be interpreted as a period of rapid expansions of specific industries or deteriorations in overall economic equality (Kim & Nelson, 2017).

To this end, we need access to $p\left(z_{t}|x_{1:t}\right)$, but in general settings, there is no analytical solution. In this case, it is natural to consider a Monte Carlo approach. If we can sample $z_{t}|x_{1:t}$ through some mechanism, then the resulting empirical distribution should approximate $p\left(z_{t}|x_{1:t}\right)$ when the number of samples is large enough. These samples are called “particles,” and they play a role analogous to agents in an ABM.

The idea of importance sampling is to use a proposal $\pi\left(z_{1:t}|x_{1:t}\right)$ to come up with candidate samples $z_{1:t}^{b}$ from $p\left(z_{1:t}|x_{1:t}\right)$. For example, a common proposal is to sample according to the latent dynamics $\Pi\left(z_{1:t}|x_{1:t}\right):=\pi_{t^{\prime}=1}^{t}p\left(z_{t^{\prime}}|z_{t^{\prime}-1}\right)$. Even though it is not possible to sample from the density $p\left(z_{1:t}|x_{1:t}\right)$, we assume that it is possible to evaluate it for any given $z_{1:t}^{b}$. This makes it possible to upweight candidates with high probability under the true density. Specifically, by using weights $w\left(z_{1:t}^{b}\right)\propto\frac{p\left(z_{1:t}|x_{1:t}\right)}{\pi\left(z_{1:t}|x_{1:t}\right)}$, the weighted average $\sum_{b=1}^{B}f\left(z_{1:t}^{b}\right)w\left(z_{1:t}^{b}\right)$ estimates $\mathbb{E}_{p\left(z_{1:t}|x_{1:t}\right)}\left[{f\left(z_{1:t}\right)}\right]$.

While it is convenient to sample entire trajectories $z_{1:t}^{b}$ using a single set of weights, this approach deteriorates for large $t$. Only a small pocket of trajectories $z_{1:t}$ in the $t$-dimensional space will have large enough probability, leading to skewed weights. Moreover, since our goal was only to sample $z_{t}|x_{1:t}$, there is no need to sample entire trajectories; the marginals at each timepoint are enough. This suggests an alternative to importance sampling, called the particle filter (Doucet et al., 2001, Andrieu et al., 2004, Crisan & Doucet, 2002). First, note the recursive relationship between $p\left(z_{1:t}|x_{1:t}\right)$ across $t$,

The first term in the final expression is the likelihood of the current observation given the current latent state, the second reflects the transition dynamics, and the third is the previous iteration of the quantity of interest.

If we use a proposal that satisfies $\pi\left(z_{1:t}|x_{1:t}\right)=\tilde{\pi}\left(z_{t}|z_{1:\left(t-1\right)},x_{1:t}\right)\pi\left(z_{1:\left(t-1\right)}|x_{1:\left(t-1\right)}\right)$, then we also obtain a recursion for the weights,

The particle filter uses these intermediate weights to focus in on promising regions of the trajectory space. Specifically, suppose we have a sample of $B$ plausible trajectories $z_{1:\left(t-1\right)}^{b}$ up to time $t-1$. Using $\tilde{\pi}\left(z_{t}|z_{1:\left(t-1\right)},x_{1:t}\right)$, propose one-step extensions $z_{1:t}^{b,\ast}$. Rather than directly using weights $w\left(z_{1:t}^{b,\ast}\right)$ in an importance sampling step, we use these weights to draw,

Trajectories $z_{1:t}^{b,\ast}$ with low weight may have $n_{b}=0$. In contrast, those with high weight might have large $n_{b}$. To construct an extended set of trajectories $\left(z_{1:t}^{b}\right)_{b=1}^{B}$, we simply make $n_{b}$ copies of each $z_{1:t}^{\ast,b}$. It turns out that this is exactly a draw from $p\left(z_{t}|x_{1:t}\right)$ — no additional reweighting step is necessary, it is already accounted for in the multinomial sampling.

A useful interpretation is to consider the $\left(t-1\right)^{st}$ coordinate of the partial trajectory $z_{1:\left(t-1\right)}^{b}$ as one of $B$ particles living in a latent space. Our extension to $z_{1:t}^{b,\ast}$ propagates these particles to the next timestep; the particles also got larger or smaller depending on the weights $w\left(z_{1:t}^{b,\ast}\right)$. In the multinomial sampling step, small particles are eliminated. Large particles are split into several at the same location. The paths that these reproducing particles trace out over time defines a collection of plausible trajectories in the latent space. In this way, the incremental evolution of a collection of discrete particles supports analysis of an otherwise intractable analytical problem. Returning to the economics example, the particle filtering view is to understand the latent dynamics of the economy by simulating many alternative histories, ensuring that at each step, the simulated history is not too inconsistent with the observed data.

Approximate Bayesian Computation (ABC) was one of the earliest approaches to likelihood-free inference, and it remains widely used in practice (Tavaré et al., 1997, Pritchard et al., 1999, Sisson et al., 2018, Beaumont, 2019). It is used to answer parameter inference questions. Let $F\left(\mathbf{x}|\theta\right)$ denote the (inaccessible) probability density over trajectories $\mathbf{x}$ for simulators parameterized by $\theta$, and define a proposal $\pi\left(\theta\right)$ over the simulator’s parameter space $\Theta$. Generate an ensemble of simulated datasets, where parameters are themselves drawn from the proposal,

To compare a pair of datasets, compute a distance between sufficient statistics, $d\left(S\left(\mathbf{x}\right),S\left(\mathbf{x}^{\prime}\right)\right)$. Datasets with the same sufficient statistics are considered equivalent. Given a real dataset $\mathbf{x}$, ABC induces a posterior on $\theta$ by building a histogram from all those $\theta_{b}$ where $d\left(S\left(\mathbf{x}^{b}\right),S\left(\mathbf{x}\right)\right)<\epsilon$. The parameter $\epsilon$ governs the quality of the approximation — small $\epsilon$ leads to more faithful approximations, but results in most forward simulations being omitted from the posterior histogram’s construction.

Though simple to implement and applicable across a variety of settings, ABC can be greedy, requiring many forward simulations $B$ before arriving at useful posterior approximations. Moreover, estimates are sensitive to the choice of $S$ and $\epsilon$ (Prangle, 2018). For this reason, we turn next to alternatives that make stronger assumptions on the relationship between $\Theta$ and $x$, but which are more sample efficient.

For computationally intensive simulations, it can be prohibitive to generate more than just a few runs. To explore a wide range of configurations $\theta$ in an efficient way, it is often possible to estimate a surrogate model $\tilde{y}\left(\theta\right)$ of the relationship $y\left(\theta\right)$ between $\theta$ and summaries $S\left(\mathbf{x}\right)$ of interest. This new surrogate supports more rapid exploration of hypothetical outcomes, compared to running the entire simulation for each choice $\theta$. The main assumption is that $y$ is smooth; i.e., simulations using similar values of global parameters $\theta$ are expected to have similar summary statistics. Surrogate models are often estimated using either normalizing flow (Brehmer et al., 2020) or Gaussian Process (GP) (Gramacy, 2020, Baker et al., 2022) models. Both approaches have the advantage that they emulate the likelihood $F\left(\mathbf{x}|\theta\right)$, supporting more direct inference over $\theta$ than ABC (Cranmer et al., 2020, Dalmasso et al., 2021).

We detail the GP approach to simulation surrogates. This approach has been successfully applied to Ebola contagion modeling, fish capture-recapture population estimation, and ocean circulation models – see (Baker et al., 2022) for a comprehensive review. Suppose we have computed $B$ simulation runs $\mathbf{x}^{b}$ for parameters $\theta_{b}$. To ease notation, write $y_{b}=S\left(\mathbf{x}^{b}\right)$. For the moment, suppose that $y_{b}\in\mathbb{R}$ is one-dimensional and that we aim to approximate $y^{\ast}:=S\left(\mathbf{x}^{\ast}\right)$ at a new configuration $\theta^{\ast}$. A GP surrogate assumes a prior $\left(y_{1},\dots,y_{b},y^{\ast}\right)\sim\mathcal{N}\left(0,C\left(\theta\right)\right)$, where $C\left(\theta\right)$ is a covariance matrix depending on the simulation parameters via a pre-specified covariance function. For example, the $\left[C\left(\theta\right)\right]_{bb^{\prime}}$ could be $C\left(\theta_{b},\theta_{b^{\prime}}\right)=\gamma\operatorname{exp}\left(-\kappa\|\theta_{b}-\theta_{b^{\prime}}\|^{2}\right)$ for hyperparameters $\gamma,\kappa$. To emulate the simulator at a new parameter $\theta^{\ast}$, we compute the posterior $y^{\ast}|\left(y_{b}\right)_{b=1}^{B}$, which is available in closed form by the rules of Gaussian conditioning. In particular, the mean and variance of $S\left(\mathbf{x}^{\ast}\right)$ can be immediately derived.

In the case of multidimensional $y_{b}$, several options are available. First, we may model each coordinate of the summary with a separate GP. Alternatively, we may swap dimensions between the output and input space. For concreteness, suppose that $\left(y_{1}^{b},\dots,y_{T}^{b}\right)$ is a $T$-dimensional summary of the temporal evolution of the simulation at configuration $\theta_{b}$. We can define a covariance function $C\left(\left(\theta,t\right),\left(\theta^{\prime},t^{\prime}\right)\right)$ that correlates simulation outputs with similar parameters $\theta$ and timepoints $t$. This new covariance can be used in a GP model as before, and the conditional probability $y_{1}^{\ast},\dots,y_{T}^{\ast}|\left(y_{t}^{b}\right)_{t,b=1}^{T,B}$ can be found for any new sequence of outputs $\left(y_{1}^{\ast},\dots,y_{T}^{\ast}\right)$.

We next illustrate the use of simulation-based inference in a simple ABM from evolutionary genetics. The ABM provides an environment for reasoning about Batesian mimicry, a phenomenon where species evolve similar appearances under selection pressures. A classical example is the resemblance between viceroy and monarch butterflies. Birds are known to avoid monarch butterflies, whose diet includes plants toxic to birds. Viceroys are edible, but when they happen to have similar appearance as monarchs (black and red wings), birds avoid them.

This story can be captured by an ABM model, one that is implemented in the NetLogo package’s built-in library (Tisue & Wilensky, 2004). This ABM has three types of agents — monarchs, viceroys, and birds. Monarch and viceroy butterflies have a categorical “color” attribute with 20 possible levels. This represents how they appear to birds. Birds have vector-valued “memory” attribute, storing colors of up to three of the most recently eaten monarchs. All agents have an $\left(x,y\right)$ location attribute. At each time step, nearby birds and butterflies interact. If the butterfly’s color is contained within the bird’s memory, the butterfly survives; otherwise it is eaten. If the eaten butterfly happened to be a monarch, then the monarch’s color is added to the bird’s “memory” vector. If the vector is now longer than three, then the oldest color is removed from memory. Also at each time step, butterflies replicate with a small probability. When they replicate, they either create an exact replicate or they create a “mutant” with a color attribute drawn uniformly from one of the 20 options. The probability of creating a mutant is controlled by a global mutation rate.

The salient feature of this model is that the average colors for the two species eventually converge to one another. In this way, local rules are able to capture the emergent phenomenon of Batesian mimcry. The fact that viceroys survive longer when they have colors that overlap with monarchs creates a selection pressure that results in the species appearing indistinguishable, on average.

We next pose an inference question: Given an observed simulation trace, can we infer the mutation rate? Figure 7 shows that higher mutation rates lead to more rapid convergence, providing a basis for inference. We approach the problem using a variant of ABC based on the particle filter called Sequential Monte Carlo - Approximate Bayesian Computation (SMC-ABC) (Del Moral et al., 2012, Jabot et al., 2013), a blend of the algorithms discussed in sections 3.3.2 and 5.1.1. This requires specification of a prior on mutation rates and a summary statistic to use for the basis of similarity comparisons. We choose a uniform prior over mutation rates from 0 to 1. Our summary statistic is defined as $S\left(\mathbf{x}\right)=\left(\left(\bar{x}^{\text{viceroy}}-\bar{x}^{\text{monarch}}\right)_{1:\frac{T}{4}},\dots,\left(\bar{x}^{\text{viceroy}}-\bar{x}^{\text{monarch}}\right)_{\frac{3}{4}T:T}\right)$ giving the difference in average color for viceroy and monarch butterflies over four equally spaced time intervals.

Our “real” data summary statistic is found by simulating one run with a mutation rate of 90. We run each simulation for $T=175$ time steps and use a summary statistic tolerance of $\epsilon=5$. We continue simulating until 250 samples are accepted. The SMC-ABC posterior is shown in Figure 8. Our inference has clearly ruled out any mutation rates below 40, and though the posterior mode is near the true value, there remains high uncertainty. This is consistent with the difficulty in differentiating large mutation rates visually in Figure 7. We note that this posterior is still approximate, and lowering the tolerance and increasing the number of required posterior samples would lead to a more precise estimate, albeit at the cost of increased computation.

Our next example describes how Gaussian process surrogates can support efficient decision-making when working with a compute intensive COVID-19 simulator. The simulator we use, Covasim, is an ABM originally developed to support reasoning about hypothetical outcomes of various nonpharmaceutical interventions to COVID-19, including workplace and school closures, social distancing, testing, and contact tracing (Kerr et al., 2021). Each agent represents an individual, and their state reflects the potential progression of disease, from asymptomatic-infectious to various degrees of disease severity, and finally to recovery or death. At each time step, agents interact across a predefined set of networks, representing family, school, workplace, or random contact networks, and during these interactions, the disease can spread.

Alternative interventions inhibit these interactions in different ways – we focus here on testing and contact tracing interventions. These are parameterized by $p_{\text{test}}$, $n_{\text{test}}$, and $p_{\text{trace}}$. $p_{\text{test}}$ modulates the fraction of symptomatic individuals who receive a COVID-19 test. If their test comes back positive, the agents enter a self-imposed quarantine, during which they do not transmit the disease. The intensiveness of the contact tracing effort is controlled by two parameters, the number of days $n_{\text{test}}$ needed before contact tracing notifications are made, and the probability $p_{trace}$ that a contact of positively-tested agent is correctly traced.

Before evaluating the effects of alternative interventions, we first calibrate the simulator to a dataset of daily infections, tests, and deaths, representing the initial phase of an outbreak on a population of size 100,000. This calibration step runs trials with varying link infection and death probability, choosing a combination that best fits the observed initial outbreak. Given these background parameters, we use the simulator to evaluate hypothetical effects of different $p_{\text{test}}$ and $p_{\text{trace}}$ — these effects can be used to set targets for testing and contact tracing. Since evaluating the simulator along a fine grid of these parameters is costly, we train a GP surrogate to a coarser set of simulations, shown in Figure 9. This initial set of runs already gives the general outlines of intervention effects, and a surrogate model makes it possible to represent the analogous trajectory for any $\left(p_{\text{test}},p_{\text{trace}}\right)$ combination of interest. For example, Figure 10 provides 250 trajectories across a range of $p_{\text{test}}$ and $p_{\text{trace}}$ — this takes 2.65 minutes on a laptop with a 3.1 GHz Intel Core i5 processor and 8GB memory. In contrast, computing just one trajectory from the Covasim ABM takes 6.20 seconds, and the estimated time for all 250 is 25.8 minutes.