Year: 2021

Abstract:
Stochastic simulations such as large-scale, spatiotemporal, age-structured epidemic models are computationally expensive at fine-grained resolution. We propose Interactive Neural Process (INP), a Bayesian active learning framework to proactively learn a deep learning surrogate model and accelerate simulation. Our framework is based on the novel integration of neural process, deep sequence model and active learning. In particular, we develop a novel spatiotemporal neural process model to mimic the simulator dynamics. Our model automatically infers the latent process which describes the intrinsic uncertainty of the simulator. This also gives rise to a new acquisition function based on the latent information gain. We design Bayesian active learning algorithms to iteratively query the simulator, gather more data, and continuously improve the model. We perform theoretical analysis and demonstrate that our approach reduces sample complexity compared with random sampling in high dimension. Empirically, we demonstrate our framework can faithfully imitate the behavior of a complex infectious disease simulator with a small number of examples, enabling rapid simulation and scenario exploration.

heading: 1 INTRODUCTION
text: Computational modeling is now more than ever at the forefront of infectious disease research due to the COVID-19 pandemic. Stochastic simulations play a critical role in understanding and forecasting infectious disease dynamics, creating what-if scenarios, and informing public health policy making (Cramer et al., 2021). More broadly, stochastic simulation (Ripley, 2009; Asmussen & Glynn, 2007) produces forecasts about complex interactions among people, environment, space, and time given a set of parameters. They provide the numerical tools to simulate probabilistic models and stochastic processes in finance (Lamberton & Lapeyre, 2007), chemistry (Gillespie, 2007) and many scientific disciplines.
Unfortunately, stochastic simulations at fine-grained spatial and temporal resolution are extremely computationally expensive. For example, epidemic models whose aim is to simulate realistic diffusion dynamics via in-silico experiments require the exploration of a large parameter space (e.g. characteristics of a virus, policy interventions, people’s behavior). Therefore, hundreds of thousands of simulations are required to explore and calibrate the epidemic model with observed epidemic surveillance data. This process significantly hinders the adaptive capability of existing stochastic simulators, especially in “war time” emergencies, due to the lead time needed to execute new simulations and produce actionable insights that could help guide decision makers.
Learning deep surrogate models to speed up complex simulation has been explored in climate modeling and fluid dynamics for deterministic dynamics (Sanchez-Gonzalez et al., 2020; Wang et al., 2020; Holl et al., 2019; Rasp et al., 2018), but not for stochastic simulations. Furthermore, these surrogate models can only approximate specific system dynamics and fail to generalize under different parametrization. Especially for pandemic scenario planning, we desire models that can predict futuristic scenarios under different conditions. Furthermore, the majority of the surrogate models are trained passively using a simulation data set. This requires a large number of simulations to cover different parameter regimes of the simulator and ensure generalization.
We propose Interactive Neural Process (INP), a Bayesian active learning framework to speed up the stochastic simulation. Given parameters such as disease reproduction number, incubation and infectious periods, mobility dynamics, and initial state, current mechanistic simulators generate
future outbreak states with time-consuming numerical integration. Our framework accelerates the simulation by learning the input-output map between parameters and future states, hence bypassing numerical integration. We develop a novel neural process (NP) (Garnelo et al., 2018) model to mimic the mechanistic simulator with uncertainty. We extend NP to spatiotemporal setting: we introduce a time-evolving latent process for temporal dynamics and integrate graph convolution for spatial modeling.
Instead of learning passively, we design active learning algorithms to interact with the simulator and update our model in “real-time”. We derive a new acquisition function, Latent Information Gain (LIG), based on our unique model design. Our algorithm selects the parameters with the highest LIG, queries the simulator to generate new simulation data, and continuously updates our model. Compared with Gaussian processes (GP), INP is more expressive, more accurate, and scales easily to high-dimensional sequence data. We provide theoretical guarantees for the sample efficiency of this procedure over random sampling. We also demonstrate the efficacy of our method on large-scale spatiotemporal epidemic models. In summary, our contributions include:
• Interactive Neural Process: a deep Bayesian active learning framework for accelerating large-scale stochastic simulation
• A novel spatiotemporal neural process model (STNP) for high-dimensional time series data with temporal latent process and spatial graph convolution.
• New acquisition function, Latent Information Gain (LIG), based on the inferred temporal latent process to quantify uncertainty with theoretical guarantees
• Real-world application to speed up stochastic, spatiotemporal epidemic models.

heading: 2 RELATED WORK
text: Bayesian Active Learning and Experimental Design. Bayesian active learning, or experimental design is well-studied in statistics and machine learning (Chaloner & Verdinelli, 1995; Cohn et al., 1996). Gaussian Processes (GPs) are popular for posterior estimation e.g. Houlsby et al. (2011) and (Zimmer et al., 2018), but often struggle in high dimension. Deep neural networks provide scalable solutions for active learning. Deep active learning has been applied to discrete problems such as image classification (Gal et al., 2017) and sequence labeling (Siddhant & Lipton, 2018) whereas our task is continuous time series. Our problem can also be viewed as sequential experimental design where we design simulation parameters to obtain the desired outcome (imitating the simulator). Foster et al. (2021) propose deep design networks for Bayesian experiment design but they require a explicit likelihood model and conditional independence in experiments. Kleinegesse & Gutmann (2020) consider implicit models where the likelihood function is intractable, but computing the Jacobian through sampling path can be expensive and their experiments are mostly limited to low (<=10) dimensional design. In contrast, our design space is of much higher-dimension and we do not have access to an explicit likelihood model for the simulator.
Neural Processes. Neural processes (NP) (Garnelo et al., 2018) model distributions over functions and imbue neural networks with the ability of GPs to estimate uncertainty. NP has many extensions such as attentive NP (Kim et al., 2019) and functional NP (Louizos et al., 2019). However, NP implicitly assumes permutation invariance in the latent variables and can be limiting in modeling temporal dynamics. Singh et al. (2019) proposes sequential NP by incorporating a temporal transition model into NP. Still, sequential NP assumes the latent variables are independent conditioned on the hidden states. We propose spatiotemporal NP with temporal latent process and spatial graph convolution, which is well-suited for modeling the spatiotemporal dynamics of infectious disease. We apply our model to real-world large-scale Bayesian active learning. Note that even though Garnelo et al. (2018) has demonstrated NP for Bayesian optimization, it is only for toy 1-D functions.
Stochastic Simulation and Dynamics Modeling. Stochastic simulations are fundamental to many scientific fields (Ripley, 2009), especially epidemic modeling. Data-driven models of infectious diseases are increasingly used to forecast the evolution of an ongoing outbreak (Arik et al., 2020; Cramer et al., 2021; Lourenco et al., 2020). However, very few models can mimic the internal mechanism of a stochastic simulator and answer “what-if questions”. GPs are commonly used as surrogate models for expensive simulators (Meeds & Welling, 2014; Gutmann et al., 2016; Qian et al., 2020), but GPs do not scale well to high-dimensional data. Likelihood-free inference Lueckmann
et al. (2019); Papamakarios et al. (2019); Munk et al. (2019); Wood et al. (2020) aim to learn the posterior of the parameters given observed data. They use neural density estimators to learn an emulator, but require a lot of simulations. We propose a novel neural process model to learn the likelihood. For active learning, instead of relying on Monte Carlo sampling, we directly compute the information gain in the latent process. For epidemic modeling, Qian et al. (2020) only consider the simple SEIR model. In contrast, we demonstrate the use of deep sequence model as a prior distribution in Bayesian active learning for real-world, large-scale, spatiotemporal dynamics. Our framework is also compatible with other deep sequence models for time series, e.g. Deep State Space (Rangapuram et al., 2018), Neural ODE (Chen et al., 2018) and etc.

heading: 3 METHODOLOGY
text: Consider a stochastic process {X1, · · · , XT }, governed by time-varying parameters θt ∈ RK , and the initial state x0 ∈ RD. In epidemic modeling, θt can represent the effective reproduction number of the virus at a given time, the effective contact rates between individuals belonging to different age groups, the people’s degree of short- or long-range mobility, or the effects of time varying policy interventions (e.g. non-pharmaceutical interventions). The state xt ∈ RD includes both the daily prevalence and daily incidence for each compartment of the epidemic model (e.g. number of people that are infectious and number of new infected individuals at time t).
Stochastic simulation uses a mechanistic model F (θ; ξ) to simulate the process where the random variable ξ represents the randomness in the simulator. Let θ := (x0, θ1, · · · , θT ) represent the initial state and all the parameters over time. For each θ, we obtain a different set of simulation data {(x1, · · · , xT )m}Mm=1. However, realistic large-scale stochastic simulations require the exploration of a large parameter space and are extremely computationally intensive. In the following section, we describe a pipeline called Interactive Neural Process (INP) to proactively query the stochastic simulator, generate simulation data, in order to learn a fast surrogate model for rapid simulation.

heading: 3.1 INTERACTIVE NEURAL PROCESS
text: Our high-level goal is to train a deep surrogate model to mimic the stochastic simulator. As show in Figure 1, given parameters θ, we query the simulator, i.e., the mechanistic model to obtain a set of simulations {(x1, · · · , xT )m}Mm=1. We train a neural process (NP) model to learn the probabilistic map from parameters to future states. Our NP model is spatiotemporal and it captures the unique disease dynamics of the epidemic simulator. During inference, the model needs to generate predictions (x̂1, · · · , x̂T ) at the target parameters θ corresponding to different scenarios. Instead of simulating at a wide range of parameter regimes, we take a Bayesian active learning approach to proactively query the simulator and update the model incrementally. Using NP, we can infer the latent temporal process (z1, · · · , zT ) that encodes the uncertainty of the current surrogate
model. Then we propose a new reward/acquisition function, latent information gain (LIG), to select the θ? with the highest reward. We use θ? to query the simulator, and in turn generate new simulation to further improve the model. Next, we describe each of the components in detail.

heading: 3.2 SPATIOTEMPORAL NEURAL PROCESS
text: Neural process (NP) (Garnelo et al., 2018) is a type of deep generative model that represents distributions over functions. It approximates a stochastic process by combining Gaussian process with deep learning. NP introduces a global latent variable z to capture the stochasticity and learns the conditional distribution p(x1:T |θ) by optimizing the evidence lower bound (ELBO):
log p(x1:T |θ) ≥ Eq(z|x1:T ,θ) [ log p(x1:T |z, θ) ] − KL ( q(z|x1:T , θ)‖p(z) ) (1)
Here p(z) is the prior distribution for the latent variable. We use x1:T as a shorthand for (x1, · · · , xT ). The prior distribution p(z) is conditioned on a set of context points θc, xc1:T as p(z|xc1:T , θc). The set of context points is randomly sampled from the training data. Its complement is the target set.
However, the global latent variable z in NP can be limiting for non-stationary, spatiotemporal dynamics in the epidemics. We propose spatiotemporal neural process (STNP) with two extensions. First, we introduce a temporal latent process (z1, · · · , zT ) to represent the unknown dynamics. The latent process provides an expressive description of the internal mechanism of the stochastic simulator. Each latent variable zt is sampled conditioning on the past history. Second, we explicitly model the spatial dependency in xt ∈ RD. Rather than treating the dimensions in xt as independent features, we capture their correlations with a graph. For instance, the travel graph between locations can be represented as an adjacency matrix A ∈ RD×D. Given parameters {θ}, simulation data {x1:T }, and the spatial graph A as inputs, STNP models the conditional distribution p(x1:T |θ,A) by optimizing the following ELBO objective:
log p(x1:T |θ,A) ≥ Eq(z1:T |x1:T ,θ,A) log p(x1:T |z1:T , θ, A)− KL ( q(z1:T |x1:T , θ, A)‖p(z1:T ) ) (2)
where the distributions q(z1:T |x1:T , θ, A) and p(x1:T |z1:T , θ, A) are parameterized with neural networks. The prior distribution p(z1:T ) is conditioned on a set of contextual sequences p(z1:T |xc1:T , θc, A). Figure 2 visualizes the graphical models of our STNP, the vanilla NP (Garnelo et al., 2018), and Sequential NP (Singh et al., 2019).
We implement STNP following an encoder-decoder architecture. The encoder parametrizes the mean and standard deviation of the variational posterior q(z1:T |x1:T , θ, A) and the decoder approximates the predictive distribution p(x1:T |z1:T , θ, A). To incorporate the spatial graph information, we use a Diffusion Convolutional Gated Recurrent Unit (DCGRU) layer (Li et al., 2017) which integrates graph convolution in a GRU cell. We use multi-layer GRUs to obtain hidden states from the inputs. Using re-parametrization (Kingma & Welling, 2013), we sample zt from the encoder and then decode xt conditioned on zt in an auto-regressive fashion. To ensure fair comparisons, we adapt NP and SNP to graph-based settings and use the same architecture as STNP to generate the hidden states shown in Figure 6. See Appendix B.3 for architecture details.

heading: 3.3 BAYESIAN ACTIVE LEARNING
text: Algorithm 1 details a Bayesian active learning algorithm, based on Bayesian optimization (Shahriari et al., 2015; Frazier, 2018). We train a NP model to interact with the simulator and improve
Algorithm 1: Interactive Neural Process Input: Initial simulation dataset S1
1 Train the model NP(1)(S1) ; 2 for i = 1, 2, · · · do 3 Learn (z1, z2, · · · , zT ) ∼ q(i)(z1:T |x1:T , θ,Si); 4 Predict (x̂1, x̂2, · · · , x̂T ) ∼ p(i)(x1:T |z1:T , θ,Si) ; 5 Select a batch {θ(i+1)} ← arg maxθ Ep(x1:T |z1:T ,θ) [r(x̂1:T |z1:T , θ)]; 6 Simulate {x(i+1)1:t } ← Query the simulator F (θ(i+1); ξ) ; 7 Augment training set Si+1 ← Si ∪ {θ(i+1), x(i+1)1:T } ; 8 Update the model NP(i+1)(Si+1) ; 9 end
learning. Let the superscript (i) denote the i-th interaction. We start with an initial data set S1 = {θ(1), x(1)1:T } and use it to train our NP model and learn the latent process. During inference, given the augmented parameters θ that are candidate scenarios, we use the trained NP model to predict the states (x̂1, · · · , x̂T ). We evaluate the models’ predictions with an acquisition function r(x̂1:T , z1:T , θ) and select the set of parameters {θ(i+1)} with the highest reward. We query the simulator with {θ(i+1)} to augment the training data set Si+1 and update the NP model for the next iteration. The choice of the reward (acquisition) function r depends on the goal of the active learning task. For example, to find the model that best fits the data, the reward function can be the log-likelihood r = log p(x̂1:T |θ,A). To collect data and reduce model uncertainty in Bayesian experimental design, the reward function can be the mutual information. In what follows, we discuss different strategies to design the reward/acquisition function. We also propose a novel acquisition function based on information gain in the latent space tailored to our STNP model.

heading: 3.4 REWARD/ACQUISITION FUNCTIONS
text: For regression tasks, standard acquisition functions for active learning include Maximum Mean Standard Deviation (Mean STD), Maximum Entropy, Bayesian Active Learning by Disagreement (BALD), and random sampling (Gal et al., 2017). With the distributions represented by NP, we explore various acquisition functions and their approximations. The temporal latent process in STNP encodes the internal mechanism of the stochastic simulator. We introduce a new acquisition function based on our unique NP design called Latent Information Gain (LIG).
Maximum Mean STD. Mean STD (Gal & Ghahramani, 2016) is a heuristic used to estimate the model uncertainty. For each augmented parameter θ, we sample multiple z1:T and generate a set of predictions {x̂1:T }. For a length T sequence with dimension D, we compute the standard deviation σt,d for time step t and feature d. Mean STD computes σ̄ = 1TD ∑T t=1 ∑D d=1 ∑ σt,d for each parameter θ. We select the θ with the maximum σ̄. Empirically, we found that Mean STD often becomes over-conservative and tends to explore less.
Maximum Entropy. Maximum entropy computes the maximum predictive entropy as H(x̂) = −E[log p(x̂1:T )]. In general, entropy is intractable for continuous output. Our NP model implicitly assumes the predictions follow a multivariate Gaussian, which allows us to compute the differential entropy (Jaynes, 1957). We follow the same procedure as Mean STD to estimate the empirical covariance Σ ∈ RTD×TDand compute the differential entropy for each parameter as H = 12 ln |Σ|+ TD 2 (1 + ln 2π). We select the parameter θ with the maximum entropy.
BALD and EIG. BALD (Houlsby et al., 2011) quantifies the mutual information between the prediction and model posterior H(x̂1:T |θ) −H(x̂1:T |z1:T , θ), which is equivelant to the expected information gain (EIG). Computing the EIG/BALD is challenging since p(x̂1:T |z1:T , θ) cannot be found in closed form in general. The integrand is intractable and conventional MC methods are not applicable (Foster et al., 2019). One way to get around this is to employ a nested MC estimator with quadratic computational cost for sampling (Myung et al., 2013; Vincent & Rainforth, 2017). We prove the equivalence between the expected information gain for selecting the parameter θ and the
expected KL divergence in the latent processes, as illustrated by the following proposition. In this way, conventional MC method becomes applicable, which helps reduce the quadratic computational cost to linear. Proposition 1. The expected information gain (EIG) for neural process is equivalent to the KL divergence between the likelihood and posterior in the latent process, that is
EIG(x̂1:T , θ) := E[H(x̂1:T )−H(x̂1:T |z1:T , θ)] = Ep(x̂1:T |θ) [ KL ( p(z1:T |x̂1:T , θ)‖p(z1:T ) )] (3)
See proof in the appendix A.1. Inspired by this fact, we propose a novel acquisition function based on the information gain in the latent process.
Latent Information Gain. The trained NP model produces a variational posterior given the current dataset S as p(z1:T |S). For every parameters θ remained in the search space, we can predict x̂1:T with the decoder. We use x̂1:T and θ as input to the encoder to re-evaluate the posterior p(z1:T |x̂1:T , θ,S). LIG computes the distributional difference with respect to the latent process z1:T as Ep(x̂1:T |θ) [KL (p(z1:T |x̂1:T , θ,S)‖p(z1:T |S))] where KL(·‖·) denotes the KL-divergence between two distributions. We select a batch of θ with the highest latent information gain.

heading: 3.5 THEORETICAL ANALYSIS
text: We shed light onto the intuition behind choosing adaptive sample selection over random sampling via analyzing a simplifying situation. Assume that at a certain stage we have learned a feature map Ψ which maps the input θ of the neural network to the last layer. Then the output X can be modeled as X = 〈Ψ(θ), z∗〉+ , where z∗ is the true hidden variable, is the random noise. Our goal is to generate an estimate ẑ, and use it to make predictions 〈Ψ(θ), ẑ〉. A good estimate shall achieve small error in terms of ‖ẑt − z∗‖2 with high probability. In the following theorem, we prove that greedily maximizing the variance of the prediction to choose θ will lead to an error of orderO(d) less than that of random exploration in the space of θ, which is significant in high dimension. Theorem 1. For random feature map Ψ(·), greedily optimizing the KL divergence, KL (p(z|x̂, θ)‖p(z)) , or equivalently the variance of the posterior predictive distribution E [ (〈Ψ(θ), ẑ〉 − E 〈Ψ(θ), ẑ〉)2 ] in search of θ will lead to an error ‖ẑt − z∗‖2 of order O ( σd/ √ t ) with high probability.
On the other hand, random sampling of θ will lead to an error of order O ( σd2/ √ t )
with high probability.
See proofs in the Appendix A.2.

heading: 4 EXPERIMENTS
text: We start with a simple susceptible-exposed-infectious-removed (SEIR) compartmental model. We compare our method with Gaussian process for validation. We then apply our INP framework to mimic the dynamics of a real-world large-scale Local Epidemic and Mobility model (LEAM-US) simulator and report the performance using different acquisition functions. We implement GP using gpytorch (Gardner et al., 2018) with Matern kernel. All NP models are implemented using pytorch (Paszke et al., 2019).

heading: 4.1 SEIR COMPARTMENTAL MODEL
text: SEIR Simulation. To show that INP indeed provides a more flexible and faster alternative to GP, we begin with a simple stochastic, discrete, chain-binomial SEIR compartmental model as our stochastic simulator. In this model, susceptible individuals (S) become exposed (E) through interactions with infectious individuals (I) and are eventually removed (R), see Appendix B.1 for details.
We set the total populationN = S+E+I+R as 100, 000, the initial number of exposed individuals as E0 = 2, 000, and the initial number of infectious individuals as I0 = 2, 000. The latent period varies between approximately 1.5 to 4 days, i.e. latent individuals move to the infectious stage at a rate ε ∈ [0.25, 0.65] (step 0.05) that is inversely proportional to the latent period, the infectious period
µ−1 is set to be equal to 1 day, and we let the basic reproduction number R0 (which in this case coincides with the transmissibility rate β) vary between 1.1 and 4.0 (step 0.1). Here, each (β, ε) pair corresponds to a specific scenario, which determines the 2-dimensional parameters θ. We simulate the first 100 days of the epidemic for the infectious population (output dimension 100) with a total of 300 scenarios. We generate 30 samples for each scenario.
Model Implementation. We compare INP with GP for Bayesian active learning. In this toy experiment, we only predict the number of individuals in the infectious compartment. The input is (β, ε) pair and the output is the 100 days’ infection prediction. As the simulator is simple, we use the vanilla NP model with the global latent variable z. For each epoch, we randomly select 10% of the samples as context. Implementation details are deferred to Appendix B.5.
Uncertainty Quantification. Figure 3 compares the NP and GP performance on one scenario in the held-out test set. It shows the ground truth and the predicted number of infectious population for the first 50 days. We also include the confidence intervals (CI) with 5 standard deviations for ground truth and NP predictions and 1 standard deviation for GP predictions. We observe that NP fits the simulation dynamics better than GP for mean prediction. Moreover, NP has closer CIs to the truth, reflecting the simulator’s intrinsic uncertainty. GP shows larger CIs which represent the model’s own uncertainty. Note that NP is much more flexible than GP and can scale easily to high-dimensional data.
Acquisition Functions. We compare different acquisition functions in Bayesian active learning. We select 2 scenarios as the initial training dataset. Then we continue augmenting the training set using
the remaining scenarios until the validation loss converges. GP is not a latent variable model, hence LIG does not apply. We report the average test mean absolute error (MAE) over three random runs. The candidate set for training includes 270 scenarios. We set aside 15 scenarios each for validation and test. The scenarios are evenly distributed in the search space.
Figure 3 shows the testing MAE versus the percentage of data used for training for NP and GP with different acquisition functions. None of the GP methods converge after selecting 4.07% of the data for training while NP methods converge much faster. Our proposed acquisition function LIG is more sample efficient in NP. It takes only 4.07% of the data to converge and reach the NP offline performance, which uses the entire training set for training.
Exploration Exploitation Trade-off. To understand the large performance gap for LIG vs. baselines on NP, we visualize the values of test MAE and the reward from acquisition functions for each Bayesian active learning iteration, shown in Figure 4. For Mean STD and Maximum Entropy, both tend to exploit the region with large transmission rate for the first 2 iterations. Including these scenarios makes the training set unbalanced. The MAE in the region with small transmission rate become worse after 2 iterations. Meanwhile, Random is doing pure exploration. The improvement of MAE performance is not apparent after 2 iterations. Our proposed LIG is able to reach a balance by exploiting the uncertainty in the latent process and encouraging exploration. Hence, with a small number of iterations (I = 2), it has already selected “informative scenarios” in the search space.

heading: 4.2 LOCAL EPIDEMIC AND MOBILITY MODEL (LEAM-US)
text: As a real-world experiment, we use a large-scale, spatiotemporal, age-structured epidemic model, LEAM-US, as our stochastic simulator. Our goal is to use INP to mimic the internal mechanism of this simulator to accelerate disease modeling and scenario creation.
LEAM-US Simulator. The Local Epidemic and Mobility model (LEAM-US) is a stochastic, spatial, age-structured epidemic model based on a metapopulation approach which divides the US in more than 3,100 subpopulations, each one corresponding to a each US county or statistically equivalent entity. Population size and county-specific age distributions reflect Census’ annual resident population estimates for year 2019. We consider individuals divided into 10 age groups. Contact mixing patterns are age-dependent and state specific and modeled considering contact matrices that describe the interaction of individuals in different social settings (Mistry et al., 2021). LEAM-US integrates a human mobility layer, represented as a network, using both short-range (i.e., commuting) and long-range (i.e., flights) mobility data, see more details in Appendix B.2.
Dataset. We separate data in California monthly to predict the 28 days’ sequence from the 2nd to the 29th day of each month from March to December. Each θ includes the county-level parameters of LEAM-US and state level incidence and prevalence compartments. The total number of dimensions in θ is 16, 264, see details in Appendix B.2. Overall, there are 315 scenarios in the search space, corresponding to 315 different θ with total 16, 254 samples. We split 78% of the data as the candidate train set, and 11% for validation and test. We use 8.6% data for initial training. We first perform offline learning to show that the STNP model can accurately learn the LEAM-US dynamics. For Bayesian active learning, we use the candidate train set as the search space.
Model Implementation. We instantiate a STNP model to mimic an epidemic simulator that has θ at both county and state level and xt at the state level. We use county-level parameter θ together
with a county-to-county mobility graph A in California as input. We use the DCGRU layer (Li et al., 2017) to encode the mobility graph in a GRU. We use a linear layer to map the county-level output to hidden features at the state level. For both the state-level encoder and decoder, we use multi-layer GRUs. For each epoch, we randomly select 20% samples as context sequence.
STNP Predictions. Figure 5 left visualize the STNP predictions in four key compartments of a typical scenario withR0 = 3.1 from March 2nd to March 29th. The confidence interval is plotted with 2 standard deviations. We can see that both the mean and confidence interval of STNP predictions match the truth well. These results demonstrate the promise that the generative STNP model can serve as a deep surrogate model for the LEAM-US simulator.
Acquisition Functions. We compare different acquisition functions with STNP. All methods start with the same initial train set. Then we continue adding 8 scenarios to the training set until the validation loss converges. We measure the average performance over three random runs and report the test MAE. Figure 5 right shows the log scale MAE versus the percentage of samples included for training, which scales linearly in computation. It shows our proposed LIG always has the best MAE performance until the convergence. Moreover, after using 31.4% of data from the candidate train set, LIG already outperforms the offline model.
Ablation Study on NP Models. We perform ablation study on our proposed NP model, STNP. Compared with vanilla NP (Garnelo et al., 2018) and SNP (Singh et al., 2019), the key innovation of model is temporal latent process, modeled by state level encoder and decoder as in Figure 6. To ensure fair comparison, we modified NP by adding the convolutional layers with diffusion convolution (Li et al., 2018) to embed the graphs. Similarly, we modified SNP by replacing the convolutional layers with diffusion convolution. Figure 2 shows the corresponding
graphical models. Table 1 shows the testing MAE of different NP models trained in an offline fashion. Our STNP significantly improves the performance and can accurately learn the simulator dynamics.
Batch Active Learning with LIG. In the case where we add more than one scenario of data per iteration, LIG allows us to calculate the corresponding reward for a batch of scenarios due to the mean aggregation layer in the NP encoder. In our LEAM-US experiment, we add 8 scenarios per iteration (batch size 8). As an ablation study, we compared 4 different batch sizes, see Appendix Figure 8. We observe that larger batch leads to better performance. This strategy combines LIG and random search, where it first randomly selects a set of groups and then uses LIG to measure the reward for the data within each group. Random search encourages exploration whereas LIG exploits the highest reward.

heading: 5 CONCLUSION & LIMITATION
text: We present a new way to seamlessly interact with existing stochastic simulators and accelerate simulation. Our Neural Interactive Process (INP) includes a novel spatiotemporal neural process model to approximate the underlying simulation dynamics. It infers the latent process which describes the intrinsic uncertainty of the simulator. We exploit this uncertainty and propose LIG as an acquisition function in Bayesian active learning. We perform theoretical analysis and demonstrate that our approach reduces sample complexity compared with random sampling in high dimension. Using a complex epidemic simulator, we demonstrate our method provides a faster and more flexible alternative to Gaussian process for policy optimization and scenario creation.
One limitation of our current work is the optimization step. Currently, we use a brute force search mechanism to select the parameters according to the value of an acquisition function. However, this search will become expensive in high dimension. We plan to leverage Monte Carlo sampling and Bayesian optimization techniques to parameterize the reward function. Then we would also be able to directly optimize for the target parameters with auto-differentiation.

heading: REPRODUCIBILITY STATEMENT
text: The implementation code is included in the supplementary material. The readme file includes the corresponding instructions. The full proof of the Theorem 1 can be found in Appendix A.2.

heading: A THEORETICAL ANALYSIS

heading: A.1 LATENT INFORMATION GAIN
text: Proposition 1. The expected information gain (EIG) for neural process is equivalent to the KL divergence between the prior and posterior in the latent process, that is
EIG(x̂, θ) := E[H(x̂)−H(x̂|z, θ)] = Ep(x̂|θ) [ KL ( p(z|x̂, θ)‖p(z) )] (4)
Proof of Proposition 1. The information gained in the latent process z, by selecting the parameter θ and generate x̂ is the reduction in entropy from the prior to the posterior IG(θ) = H(x̂)−H(x̂|z, θ). Take the expectation of IG(x̂, θ) under the marginal distribution, we obtain from the conditional independence of z and θ that
Ep(x̂|θ) [ KL ( p(z|x̂, θ)‖p(z) )] = Ep(x̂,z|θ) [ log
p(z|x̂, θ) p(z) ] = Ep(x̂,z|θ) [ log
p(z|x̂, θ) p(z|θ) ] = Ep(x̂,z|θ) [ log p(z, x̂, θ)− log p(x̂, θ)− log p(z, θ) + log p(θ)
] = Ep(x̂,z|θ) [ log p(x̂|z, θ)− log p(x̂|θ)
] = Ep(z) [ Ep(x̂|z,θ)[log p(x̂|z, θ)]− Ep(x̂|θ)[log p(x̂|θ)]
] = Ep(z)[H(x̂|θ)−H(x̂|z, θ)] = EIG(x̂, θ).

heading: A.2 SAMPLE EFFICIENCY OF ACTIVE LEARNING
text: From the main text we know that in each round, the output random variable
X = 〈Ψ(θ), z∗〉+ . (5)
We further assume that the random noise is mean zero and σ-subGaussian.
Using this information, we treat z as an unknown parameter and define a likelihood function so that p(X|z; θ) has good coverage over the observations:
p(Xk|z; θk) ∝ exp ( − 1
2σ2 (Xk − 〈Ψ(θk), z〉)2
) .
Let the prior distribution over z be p(z|θk) = p(z) ∝ exp ( − m2σ2 ‖z‖ 2 )
. Here we use k instead of (i) in the Algorithm 1 to represent the number of iterations. We can form a posterior over z in the k-th round:
p(z|X1, θ1, . . . , Xk, θk) ∝ exp ( − m
2σ2 ‖z‖2 − 1 2σ2 k∑ s=1
(Xs − 〈Ψ(θs), z〉)2 ) .
Focusing on the random variable z ∼ p(·|X1, θ1, . . . , Xk, θk), the estimate of the hidden variable, we can express it at k-th round as:
zk = ẑk + σV −1 k ηk, (6)
where ẑk = V −1k ∑k s=1XsΨ(θs), Vk = mI + ∑k s=1 Ψ(θs)Ψ(θs)
T, and ηk is a standard normal random variable.
We can either choose action θ randomly or greedily. A random choice of θ corresponds to taking
θk ∼ N (0, I) , (7)
A greedy procedure is to choose action θk in the k-th round to optimize KL (p(z|x̂, θ)‖p(z)) = Ep(z|x̂,θ) ( log p(z|x̂,θ)p(z) ) , where we denote the estimated output variable x̂ given θ and z as x̂ = 〈Ψ(θ), z〉. This optimization procedure is equivalent to maximizing the variance of the prediction:
θk = arg max θ∈Rd
Ez∼p(·|X1,θ1,...,Xk−1,θk−1) [( 〈Ψ(θ), z〉 − Ez∼p(·|X1,θ1,...,Xk−1,θk−1) 〈Ψ(θ), z〉 )2] .
(8)
For both approaches, we assume that the features Ψ(θ) are normalized.
We compare the statistical risk of this approach with the random sampling approach.
Assume that the features are normalized, so that for all θ ∈ Rd, Ψ(θ) ∈ Sd−1. Define a matrix Ak ∈ Rd×k containing all the column vectors {Ψ(θ1), . . . ,Ψ(θk)}. We can then express the estimation error in the following lemma. Lemma 1. The estimation error ‖ẑk − z∗‖2 can be bounded as follow.
‖ẑk − z∗‖2 ≤ m ( m+ σmin ( AkA T k ))−1 · ‖z∗‖2 + min 1/ (2√m) , 1/ √σmin (AkATk )+ m√
σmin ( AkATk
)  · σ√d.
We now analyze random sampling of θ versus greedy search for θ.
If the feature map Ψ(·) = id, then from random matrix theory, we know that for θ randomly sampled from a normal distribution and normalized to ‖θ‖ = 1, σmin ( 1 kAkA T k ) will converge to(√
1/k − √ 1/d )2 for large k, which is of order Ω(1/d). This will lead to an appealing risk bound
for ‖ẑk − z∗‖2 on the order of O ( d/ √ k ) .
However, in high dimension, this feature map is often far from identity. In the proof of Theorem 1 below, we demonstrate that even when Ψ(·) is simply a linear random feature map, with i.i.d. normal entries, random exploration in θ can lead to deteriorated error bound. This setting is motivated by the analyses in wide neural networks, where the features learned from gradient descent are close to those generated from random initialization Du et al. (2019); Mei & Montanari (2019). Theorem 1 (Formal statement). Assume that the noise in equation 5 is σ-subGaussian.
For a normalized linear random feature map Ψ(·), greedily optimizing the KL divergence,

heading: KL (p(z|x̂, θ)‖p(z)) (or equivalently the variance of the posterior predictive distribution defined
text: in equation equation 8) in search of θ will lead to an error ‖ẑk − z∗‖2 = O ( σd/ √ k )
with high probability.

heading: On the other hand, random sampling of θ following equation 7 will lead to ‖ẑk − z∗‖2 =
text: O ( σd2/ √ k ) with high probability.
Proof of Theorem 1. For a linear random feature map, we can express Ψ(θ) = Ψθ, where entries in Ψ ∈ Rd×d are i.i.d. normal. The entries of Ψθ are then normalized.
• For random exploration of θ, the matrix containing the feature vectors becomes Ak = ΨΘk, where matrix Θk ∈ Rd×k collects all the k column vectors of {θ1, . . . , θk}. Then AkA T k = ΨΘkΘ T kΨ
T. From random matrix theory, we know that the condition number of Ψ is equal to d with high probability Chen & Dongarra (2005). Hence for normalized Ψ and θ, σmin ( ΨΘkΘ T kΨ T ) ≥ σ2min (Ψ)σmin ( ΘkΘ T k ) = 1d2σmin ( ΘkΘ T k ) . The inequality holds because the smallest singular value is the inverse of the norm of the inverse matrix.
We then use the fact from random matrix theory that for normalized random θ, the asymptotic distribution of the eigenvalues of 1kΘkΘ T k follow the (scaled) Marchenko–Pastur distribution, which is supported on λ ∈ [(√ 1/k − √ 1/d )2 , (√ 1/k + √ 1/d )2] , where the 1/d
scaling comes from the fact that θ is normalized Götze & Tikhomirov (2004). Hence for large k, σmin ( ΘkΘ T k ) ≥ ( 1− √ k/d )2
with high probability. This combined with the previous paragraph yields that for the random feature model,
σmin ( AkA T k ) = Ω
( 1
d2
( 1− √ k/d )2)
with high probability. Plugging this result into Lemma 1, we obtain that the error ‖ẑk − z∗‖2 for random exploration in the space of θ is of order O ( d2/ √ k ) .
• We then analyze the error associated with greedy maximization of the posterior predictive variance. We first note that the variance of the posterior predictive distribution in equation equation 8 can be expressed as follows using equation equation 6:
E [ (〈Ψ(θ), z〉 − E 〈Ψ(θ), z〉)2 ] = σ2E [(〈 Ψ(θ), V −1k−1ηk 〉)2] = σ2Ψ(θ)TV −2k−1Ψ(θ),
(9)
where the expectations are with respect to z ∼ p(·|X1, θ1, . . . , Xk−1, θk−1). We perform a singular value decomposition Ak = UkΛkWk. Then∑k
s=1 Ψ(θs)Ψ(θs) T = AkA T k = UkΛkΛ T kU T k , and that V −2 k−1 = ( mI +Ak−1A T k−1 )−2 =
Uk−1 ( mI + Λk−1Λ T k−1 )−2
UTk−1. Via this formulation, we see that maximizing Ψ(θ)TV −2k−1Ψ(θ) in equation equation 9 to choose θk is equivalent to choosing Ψ(θk) = (Uk−1) T (·,l), where l = arg mini∈{1,...,d} ( Λk−1Λ T k−1 ) (i,i)
. In words, when we use greedy method and maximize the variance of the prediction, it corresponds to taking Ψ(θk) in the direction of the smallest eigenvector of Vk−1.
Since every Ψ(θ) is normalized and we initialize uniformly: V0 = mI, the process is equivalent to scanning the orthogonal spaces of normalized vectors in Rd for bk/dc times. For large k, entries in ΛkΛTk are approximately uniform and are all larger than or equal to bk/dc. Then σmin ( AkA T k ) = Ω(k/d). Plugging into the bound of Lemma 1, we obtain that
‖ẑk − z∗‖2 = O ( σd√ k ) .
Proof of Lemma 1. We first express the estimate ẑk as follows.
ẑk = V −1 k k∑ s=1 XsΨ(θs) = V −1 k k∑ s=1 Ψ(θs)Ψ(θs) Tz∗ + V −1k k∑ s=1 sΨ(θs).
Then
‖ẑk − z∗‖2 = ∥∥∥∥∥ ( V −1k k∑ s=1 Ψ(θs)Ψ(θs) T − I ) z∗ + V −1k k∑ s=1 sΨ(θs) ∥∥∥∥∥ 2
≤ ∥∥∥∥∥ ( V −1k k∑ s=1 Ψ(θs)Ψ(θs) T − I ) z∗ ∥∥∥∥∥ 2︸ ︷︷ ︸
T1
+ ∥∥∥∥∥V −1k k∑ s=1 sΨ(θs) ∥∥∥∥∥ 2︸ ︷︷ ︸
T2
.
Define a matrix Ak ∈ Rd×k containing all the column vectors {Ψ(θ1), . . . ,Ψ(θk)} and perform a singular value decomposition Ak = UkΛkWk. Then ∑k s=1 Ψ(θs)Ψ(θs) T = AkA T k = UkΛkΛ T kU T k ,
and Vk = mI + AkATk , We further define vector ek ∈ Rs where (ek)s = s. We use this definition to simplify the two terms further.
For term T1,∥∥∥∥∥ ( V −1k k∑ s=1 Ψ(θs)Ψ(θs) T − I ) z∗ ∥∥∥∥∥ 2 = m ∥∥V −1k z∗∥∥2
≤ m ∥∥V −1k ∥∥2 · ‖z∗‖2
= m ( m+ σmin ( AkA T k ))−1 · ‖z∗‖2 . For term T2, we define a diagonal matrix Λ̄k ∈ Rk×k which satisfies ( Λ̄k ) i,i
= 1 if i ≤ d and( Λ̄k ) i,i
= 0 if i > d, when k > d. The following bound on T2 can be achieved.∥∥∥∥∥V −1k k∑ s=1 sΨ(θs) ∥∥∥∥∥ 2 = ∥∥V −1k Akek∥∥2
= ∥∥∥Uk (ΛkΛTk +mI)−1 UTk UkΛkΛ̄kWkek∥∥∥
2 ≤ ∥∥∥Uk (ΛkΛTk +mI)−1 Λk∥∥∥ 2 · ∥∥Λ̄kWkek∥∥2
= ∥∥∥(ΛkΛTk +mI)−1 Λk∥∥∥ 2 · ∥∥Λ̄kWkek∥∥2
≤ min 1/ (2√m) , 1/ √σmin (AkATk )+ m√
σmin ( AkATk
)  · ∥∥Λ̄kWkek∥∥2 .
Assuming that noise s is σ-subGaussian, then so is Wkek since Wk is a unitary matrix. Multiplied by the diagonal matrix Λ̄k which has zero, ∥∥Λ̄kWkek∥∥2 ≤ σ√d. Therefore,∥∥∥∥∥V −1k k∑ s=1 sΨ(θs) ∥∥∥∥∥ 2 ≤ min 1/ (2√m) , 1/ √σmin (AkATk )+ m√ σmin ( AkATk )  · σ√d.

heading: B EXPERIMENT DETAILS

heading: B.1 SEIR SIMULATOR
text: Our SEIR simulator is a simple stochastic, discrete, chain-binomial compartmental model. In this model, susceptible individuals (S) become exposed (E) through interactions with infectious individuals (I). Exposed individuals which are infected but not yet infectious transition to infectious compartment at a rate ε that is inversely proportional to the latent period of the disease. Lastly, infectious individuals transition to the removed compartment at a rate µwhich is inversely proportional to the infectious period. Removed individuals (R) are assumed to be no longer infectious and they are to be considered either recovered or dead. All transitions are simulated by randomly drawn from a binomial distribution.

heading: B.2 LEAM-US MODEL
text: LEAM-US integrates a human mobility layer, represented as a network, using both short-range (i.e., commuting) and long-range (i.e., flights) mobility data. Commuting flows between counties are obtained from the 2011-2015 5-Year ACS Commuting Flows survey and properly adjusted to account for differences in population totals since the creation of the dataset. Instead, long-range air traveling flows are quantified using origin-destination daily passenger flows between airport pairs as reported by the Official Aviation Guide (OAG) and IATA databases (updated in 2021) (OAG,
Aviation Worlwide Limited, 2021; IATA, International Air Transport Association, 2021). In addition, flight probabilities are age and country specific.
The model is initialized using a multi-scale modeling approach that utilizes GLEAM, the Global and Epidemic Mobility model (Balcan et al., 2009; 2010; Tizzoni et al., 2012; Zhang et al., 2017; Chinazzi et al., 2020; Davis et al., 2020), to simulate a set of 500 different initial conditions for LEAM-US starting on February 16th, 2020. The disease dynamics are modeled using a classic SEIR-like model and initial conditions are determined using the Global and Epidemic Mobility model (Balcan et al., 2009; 2010; Tizzoni et al., 2012; Zhang et al., 2017) calibrated to realistically represent the evolution of the COVID-19 pandemic (Chinazzi et al., 2020; Davis et al., 2020). Lastly, travel restrictions, mobility reductions, and government interventions are explicitly modeled to mimic the real timeline of interventions of the events that occurred during the COVID-19 pandemic.

heading: B.3 SPATIOTEMPORAL NP MODEL
text: As shown in Figure 6, our model has θ at both county and state level and xt at the state level. The total number of dimensions of θ is 16,264 (12 compartments × 2 (incidence + prevalence) + 28 days × 58 counties × 10 metadata parameters). The output dimension is 672 (12 compartments × 2 (incidence + prevalence) × 28 days) We use county-level parameter θ together with a county-to-county mobility graph A as input. We use the DCGRU layer (Li et al., 2017) to encode the graph in a GRU. We use a linear layer to map the county-level output to hidden features at the state level. For both the state-level encoder and decoder, we use multi-layer GRUs.
The input θ1:t is the county-level parameters for LEAM-US with a dimension of 10. The county level embedding uses 1 layer DCGRU with a width of 16. The internal state is at state level with a dimension of 16. The state level encoder and decoder use 3 layer GRUs with width of 128. The dimension of the latent process z1:t is 32. The dimension of output x1:t is 24, including the incidence and prevalence for 12 compartments. We trained STNP model for 500 steps with learning rate fixed at 10−3 using Adam optimizer. We perform early stopping with 50 patience for both offline learning and Bayesian active learning.

heading: B.4 SEQUENTIAL NEURAL LIKELIHOOD (SNL) MODEL IMPLEMENTATION FOR SEIR SIMULATOR.
text: The input and output of SNL are the same as INP and GP. We use the likelihood-free inference code (Durkan et al., 2020) to implement SNL model. For each iteration, we perform Monte Carlo sampling from the trained likelihood model to report MAE. SNL takes 1 hour to complete training.
B.5 IMPLEMENTATION DETAILS
For both GP and INP mimicking SEIR simulation, we ran experiments using CPU. No GPU accelerator is needed for this simple model. It takes 1.5 hours to complete training. For INP mimicking LEAM-US simulation, we ran experiments with GEFORCE RTX 2080. It takes 5 hours to complete training. For all experiments, we run with three different random seeds.

heading: B.6 BROADER IMPACT
text: The proposed approach has applications in the field of computational epidemiology as it provides a way to efficiently explore the parameter space of complex epidemiological model (such as LEAM-US) allowing to reduce the time-to-insights needed to provide actionable insights to policy makers. As an example, a single realization of the LEAM-US model requires about 1 hour to be completed (on a single thread of a AMD Ryzen Threadripper 3970x CPU, 3.7Ghz base clock) while in this work we show that using INP framework can achieve a performance comparable to the one of epidemic simulator while requiring less than one-third of the runs (as shown in Figure 5). In other words, this means that we can produce COVID-19 forecasts and projections three times faster than we would without the integration with the proposed framework. In addition, INP provides a general surrogate model that is not tailored to solve only for a specific target outcome (e.g. prediction of future deaths) but rather it retains the ability to describe the underlying state of the system as the original
mechanistic epidemic model. This feature preserves the ability of using the learned surrogate model as a situational awareness and outbreak analytics tool as it can also provide insights on unobservables (such as the fraction of immune population in a certain location/age group) which can be used by public health officials to assess the (potential) impact of past/future policies on the evolution of an epidemic.

heading: C ADDITIONAL RESULTS
text: C.1 INP VERSUS GP AND SNL
Table 2 shows the average results together with the standard deviation of INP for SEIR simulator after running experiments three times. Table 3 shows the average results together with the standard deviation of GP for SEIR simulator. Figure 7 shows the average test MAE over three random runs versus the percentage of training data for SNL, NP and GP with different acquisition functions. Table 5 shows the comparison of average MAE together with the standard deviation for the SEIR simulator. We observe that LIG is always better than SNL until convergence. Table 4 shows the average MAE together with the standard deviation of INP for the LEAM-US simulator.

heading: C.2 BATCH ACTIVE LEARNING WITH LIG.
text: Figure 8 compares 4 different setups: 8 batches (size 1), 4 batches (size 2), 2 batches (size 4), and 1 batch (size 8).

Review summaries:
This paper tackles learning a flexible distribution approximator to approximate the output of a high-dimensional and computationally intensive stochastic simulator.  The contributions of this paper then reduce into three main components:   1 - Definition of a spatiotemporal neural process that can succinctly model a more richly structured latent process.  2 - Definition of a new acquisition function and using a neural process in an active learning setting.  3 - Application to epidemiological simulators to “compile” estimation.  The method appears to work, validated on a toy-ish SEIR model, and a more sophisticated pre-existing epidemiological model.   

This study proposes a new method for learning surrogate models of stochastic simulators. The new method, Interactive Neural Process (INP), builds on Neural processes and leverages the spatiotemporal structure of the problems at hand to reduce the complexity of the inference task. On a few problems in epidemiology -- a low dimensional SEIR problem (2 parameters, 100-dimensional output), and a complex spatiotemporal LEAM-US problem --, the approach is shown to appropriately learn the surrogate in few iterations.

The authors proposed a novel active learning framework integrated with the neural process. The neural process is used to mimic the simulator dynamics which is later used for Bayesian active learning. They also proposed a new acquisition function utilizing the latent from the neural process. The experimental results show that INP works better to accelerate stochastic simulation than GP, and the proposed LIG leads to faster convergence compared with alternatives.  

The manuscript entitled, "Accelerating Stochastic Simulation with Interactive Neural Processes", presents a novel approach to the problem of statistical emulation for mechanistic models of epidemic disease transmission.  To this end, structured neural processes are developed to exploit and respect the temporal and spatio-temporal character of these models.  An active learning strategy is developed to train these neural processes to minimise the computational costs of generating training instances of disease simulator outputs.  The methodology developed is applied to two SEIR compartmental model examples: a minimal one with a single homogenous population and a maximal one with many age and space delimited cohorts.