Original: D EEP B AYESIAN A CTIVE L EARNING FOR A CCELERATING S TOCHASTIC S IMULATION
Revised: D EEP B AYESIAN A CTIVE L EARNING FOR A CCELERATING S TOCHASTIC S IMULATION

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Stochastic simulations such as large-scale, spatiotemporal, age-structured epidemic models are computationally expensive at fine-grained resolution.
Revised: Stochastic simulations such as large-scale, spatiotemporal, age-structured epidemic models are computationally expensive at fine-grained resolution.

Original: While deep surrogate models can speed up the simulations, doing so for stochastic simulations and with active learning approaches is an underexplored area.
Revised: While deep surrogate models can speed up the simulations, doing so for stochastic simulations and with active learning approaches is an underexplored area.

Original: We propose Interactive
Revised: We propose Interactive

Original: Neural Process (INP), a deep Bayesian active learning framework for learning deep surrogate models to accelerate stochastic simulations.
Revised: Neural Process (INP), a deep Bayesian active learning framework for learning deep surrogate models to accelerate stochastic simulations.

Original: INP consists of two components, a spatiotemporal surrogate model built upon Neural Process (NP) family and an acquisition function for active learning.
Revised: INP consists of two components, a spatiotemporal surrogate model built upon Neural Process (NP) family and an acquisition function for active learning.

Original: For surrogate modeling, we develop Spatiotemporal Neural Process (STNP) to mimic the simulator dynamics.
Revised: For surrogate modeling, we develop Spatiotemporal Neural Process (STNP) to mimic the simulator dynamics.

Original: For active learning, we propose a novel acquisition function, Latent Information
Revised: For active learning, we propose a novel acquisition function, Latent Information

Original: Gain (LIG), calculated in the latent space of NP based models.
Revised: Gain (LIG), calculated in the latent space of NP based models.

Original: We perform a theoretical analysis and demonstrate that LIG reduces sample complexity compared with random sampling in high dimensions.
Revised: We perform a theoretical analysis and demonstrate that LIG reduces sample complexity compared with random sampling in high dimensions.

Original: We also conduct empirical studies on two complex spatiotemporal simulators for reaction diffusion and infectious disease.
Revised: We also conduct empirical studies on two complex spatiotemporal simulators for reaction diffusion and infectious disease.

Original: The results demonstrate that STNP outperforms the baselines in the offline learning setting and LIG achieves the state-of-the-art for Bayesian active learning.
Revised: The results demonstrate that STNP outperforms the baselines in the offline learning setting and LIG achieves the state-of-the-art for Bayesian active learning.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Computational modeling is now more than ever at the forefront of infectious disease research due to the COVID-19 pandemic.
Revised: Computational modeling is now more than ever at the forefront of infectious disease research due to the COVID-19 pandemic.

Original: 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).
Revised: 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).

Original: More broadly, stochastic simulations (Ripley, 2009; Asmussen & Glynn, 2007) produce forecasts about complex interactions among people, environment, space, and time given a set of parameters.
Revised: More broadly, stochastic simulations (Ripley, 2009; Asmussen & Glynn, 2007) produce forecasts about complex interactions among people, environment, space, and time given a set of parameters.

Original: They provide the numerical tools to simulate stochastic processes in finance (Lamberton & Lapeyre, 2007), chemistry (Gillespie, 2007) and many other scientific disciplines.
Revised: They provide the numerical tools to simulate stochastic processes in finance (Lamberton & Lapeyre, 2007), chemistry (Gillespie, 2007) and many other scientific disciplines.

Original: Unfortunately, stochastic simulations at fine-grained spatial and temporal resolution can be extremely computationally expensive.
Revised: Unfortunately, stochastic simulations at fine-grained spatial and temporal resolution can be extremely computationally expensive.

Original: In example, epidemic models for realistic diffusion dynamics simulation via in-silico experiments require a large parameter space (e.g. characteristics of a virus, policy interventions, people’s behavior).
Revised: For example, epidemic models for realistic diffusion dynamics simulation via in-silico experiments require a large parameter space (e.g. characteristics of a virus, policy interventions, people’s behavior).

Original: Similarly, reaction-diffusion systems that play an important role in chemical reaction and bio-molecular processes also involve a large number of simulation conditions.
Revised: Similarly, reaction-diffusion systems that play an important role in chemical reaction and bio-molecular processes also involve a large number of simulation conditions.

Original: Therefore, hundreds of thousands of simulations are required to explore and calibrate the simulation model with observed experimental data.
Revised: Therefore, hundreds of thousands of simulations are required to explore and calibrate the simulation model with observed experimental data.

Original: 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.
Revised: 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.

Original: 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;
Revised: 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;

Original: Holl et al., 2019; Rasp et al., 2018; Cachay et al., 2021), but not for stochastic simulations.
Revised: Holl et al., 2019; Rasp et al., 2018; Cachay et al., 2021), but not for stochastic simulations.

Original: These surrogate models can only approximate specific system dynamics and fail to generalize under different parametrization.
Revised: These surrogate models can only approximate specific system dynamics and fail to generalize under different parametrization.

Original: Especially for pandemic scenario planning, we desire models that can predict futuristic scenarios under different conditions.
Revised: Especially for pandemic scenario planning, we desire models that can predict futuristic scenarios under different conditions.

Original: Furthermore, the majority of the surrogate models are trained passively using a simulation data set.
Revised: Furthermore, the majority of the surrogate models are trained passively using a simulation data set.

Original: This requires a large number of simulations beforehands to cover different parameter regimes of the simulator and ensure generalization.
Revised: This requires a large number of simulations beforehands to cover different parameter regimes of the simulator and ensure generalization.

Original: We propose Interactive Neural Process (INP), a deep Bayesian active learning framework to speed up stochastic simulations.
Revised: We propose Interactive Neural Process (INP), a deep Bayesian active learning framework to speed up stochastic simulations.

Original: Given parameters such as disease reproduction number, incubation and infectious periods, mechanistic simulators generate future outbreak states with time-consuming numerical integration.
Revised: Given parameters such as disease reproduction number, incubation and infectious periods, mechanistic simulators generate future outbreak states with time-consuming numerical integration.

Original: INP accelerates the simulation by guiding a surrogate model to learn the input-output map between parameters and future states, hence bypassing numerical integration.
Revised: INP accelerates the simulation by guiding a surrogate model to learn the input-output map between parameters and future states, hence bypassing numerical integration.

Original: The deep surrogate model of INP is built upon Neural Process (NP)
Revised: The deep surrogate model of INP is built upon Neural Process (NP)

Original: Garnelo et al.
Revised: Garnelo et al.

Original: (2018), which lies between Gaussian process (GP) and neural network (NN).
Revised: (2018), which lies between Gaussian process (GP) and neural network (NN).

Original: NPs can approximate stochastic processes and therefore are well-suitable for surrogate modeling of stochastic simulators.
Revised: NPs can approximate stochastic processes and therefore are well-suitable for surrogate modeling of stochastic simulators.

Original: They learn distributions over functions and can generate prediction uncertainty for Bayesian active learning.
Revised: They learn distributions over functions and can generate prediction uncertainty for Bayesian active learning.

Original: Compared with GPs, NPs are more flexible and scalable for high-dimensional data with spatiotemporal dependencies.
Revised: Compared with GPs, NPs are more flexible and scalable for high-dimensional data with spatiotemporal dependencies.

Original: We design a novel Spatiotemporal Neural Process (STNP) by introducing a timeevolving latent process for temporal dynamics and integrating spatial convolution for spatial modeling.
Revised: We design a novel Spatiotemporal Neural Process (STNP) by introducing a timeevolving latent process for temporal dynamics and integrating spatial convolution for spatial modeling.

Original: Instead of learning passively, we design active learning algorithms to interact with the simulator and update our model in “real-time”.
Revised: Instead of learning passively, we design active learning algorithms to interact with the simulator and update our model in “real-time”.

Original: We derive a new acquisition function, Latent Information Gain (LIG), based on our unique model design.
Revised: We derive a new acquisition function, Latent Information Gain (LIG), based on our unique model design.

Original: Our algorithm selects the parameters with the highest
Revised: Our algorithm selects the parameters with the highest

Original: LIG, queries the simulator to generate new simulation data, and continuously updates our model.
Revised: LIG, queries the simulator to generate new simulation data, and continuously updates our model.

Original: We provide theoretical guarantees for the sample efficiency of this procedure over random sampling.
Revised: We provide theoretical guarantees for the sample efficiency of this procedure over random sampling.

Original: We also demonstrate the efficacy of our method on large-scale spatiotemporal epidemic and reaction diffusion models.
Revised: We also demonstrate the efficacy of our method on large-scale spatiotemporal epidemic and reaction diffusion models.

Original: In summary, our contributions include:
Revised: In summary, our contributions include:

Original: • Interactive Neural Process: a deep Bayesian active learning framework for accelerating large-scale stochastic simulation.
Revised: • Interactive Neural Process: a deep Bayesian active learning framework for accelerating large-scale stochastic simulation.

Original: • A novel Spatiotemporal Neural Process model (STNP) for high-dimensional time series data that integrates temporal latent process and spatial convolution.
Revised: • A novel Spatiotemporal Neural Process model (STNP) for high-dimensional time series data that integrates temporal latent process and spatial convolution.

Original: • New acquisition function, Latent Information Gain (LIG), based on the inferred temporal latent process to quantify uncertainty with theoretical guarantees.
Revised: • New acquisition function, Latent Information Gain (LIG), based on the inferred temporal latent process to quantify uncertainty with theoretical guarantees.

Original: • Real-world application to speed up complex stochastic spatiotemporal simulations including age-structured epidemic dynamics and reaction-diffusion system.
Revised: • Real-world application to speed up complex stochastic spatiotemporal simulations including age-structured epidemic dynamics and reaction-diffusion system.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Bayesian Active Learning and Experimental Design.
Revised: Bayesian Active Learning and Experimental Design.

Original: Bayesian active learning, or experimental design is well-studied in statistics and machine learning (Chaloner & Verdinelli, 1995; Cohn et al., 1996).
Revised: Bayesian active learning, or experimental design is well-studied in statistics and machine learning (Chaloner & Verdinelli, 1995; Cohn et al., 1996).

Original: Gaussian Processes (GPs) are popular for posterior estimation e.g. Houlsby et al.
Revised: Gaussian Processes (GPs) are popular for posterior estimation e.g. Houlsby et al.

Original: (2011) and (Zimmer et al., 2018), but often struggle in high-dimension.
Revised: (2011) and (Zimmer et al., 2018), but often struggle in high-dimension.

Original: Deep neural networks provide scalable solutions for active learning.
Revised: Deep neural networks provide scalable solutions for active learning.

Original: 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.
Revised: 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.

Original: Our problem can also be viewed as sequential experimental design where we design simulation parameters to obtain the desired outcome (imitating the simulator).
Revised: Our problem can also be viewed as sequential experimental design where we design simulation parameters to obtain the desired outcome (imitating the simulator).

Original: Kleinegesse & Gutmann (2020) and Foster et al.
Revised: Foster et al.

Original: (2021) propose deep design networks for Bayesian experiment design but they require a explicit likelihood model, conditional independence in experiments, and are limited to low (1-2) dimensional design.
Revised: (2021) propose deep design networks for Bayesian experiment design but they require a explicit likelihood model and conditional independence in experiments.

Revised: 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 dimensional design.

Original: In contrast, our design space is of much higher-dimension and we do not have access to an explicit likelihood model for the simulator.
Revised: In contrast, our design space is of much higher-dimension and we do not have access to an explicit likelihood model for the simulator.

Original: Neural Processes.
Revised: Neural Processes.

Original: Neural Processes (NP) (Garnelo et al., 2018) model distributions over functions and imbue neural networks with the ability of GPs to estimate uncertainty.
Revised: Neural Processes (NP) (Garnelo et al., 2018) model distributions over functions and imbue neural networks with the ability of GPs to estimate uncertainty.

Original: NP has many extensions such as attentive NP (Kim et al., 2019) and functional NP (Louizos et al., 2019).
Revised: NP has many extensions such as attentive NP (Kim et al., 2019) and functional NP (Louizos et al., 2019).

Original: However, NP implicitly assumes permutation invariance in the latent variables and can be limiting in modeling temporal dynamics.
Revised: However, NP implicitly assumes permutation invariance in the latent variables and can be limiting in modeling temporal dynamics.

Original: Singh et al.
Revised: Singh et al.

Original: (2019) proposes sequential NP by incorporating a temporal transition model into NP.
Revised: (2019) proposes sequential NP by incorporating a temporal transition model into NP.

Original: Still, sequential NP assumes the latent variables are independent conditioned on the hidden states.
Revised: Still, sequential NP assumes the latent variables are independent conditioned on the hidden states.

Original: We propose spatiotemporal NP with temporal latent process and spatial convolution, which is well-suited for modeling the spatiotemporal dynamics of infectious disease.
Revised: We propose spatiotemporal NP with temporal latent process and spatial convolution, which is well-suited for modeling the spatiotemporal dynamics of infectious disease.

Original: We apply our model to real-world large-scale Bayesian active learning.
Revised: We apply our model to real-world large-scale Bayesian active learning.

Original: Note that even though Garnelo et al.
Revised: Note that even though Garnelo et al.

Original: (2018) has demonstrated NP for Bayesian optimization, it is only for toy 1-D functions.
Revised: (2018) has demonstrated NP for Bayesian optimization, it is only for toy 1-D functions.

Original: Stochastic Simulation and Dynamics Modeling.
Revised: Stochastic Simulation and Dynamics Modeling.

Original: 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).
Revised: Stochastic simulations are fundamental to many scientific fields (Ripley, 2009) such are increasingly used to forecast the evolution of an ongoing outbreak (Arik et al., 2020; Cramer et al., 2021; Lourenco et al., 2020).

Original: However, very few models can mimic the internal mechanism of a stochastic simulator and answer “what-if questions”.
Revised: However, very few models can mimic the internal mechanism of a stochastic simulator and answer “what-if questions”.

Revised: GPs are commonly used as surrogate models for expensive simulators (Meeds & Welling, 2014; Gutmann et al., 2016;

Revised: Järvenpää et al., 2019;

Revised: Qian et al., 2020), but GPs do not scale well to high-dimensional data.

Revised: Likelihood-free inference

Original: Recently, Qian et al.
Revised: Lueckmann et al. ; Papamakarios et al. Munk et al. ; Wood et al.

Revised: (2020) aim to learn the posterior of the parameters given the observed data.

Revised: They use neural density estimators to learn an emulator, but require a lot of simulations.

Revised: We propose a novel neural process model to learn the likelihood.

Revised: For active learning, instead of relying on Monte Carlo sampling, we directly compute the information gain in the latent process.

Revised: Qian et al.

Original: (2020) proposed to use Gaussian process (GPs) as a prior for a SEIRcompartmental model for learning lockdown policy effects, but GPs are computationally expensive and the simple SEIR model cannot capture the real-world large-scale, spatiotemporal dynamics considered in this work.
Revised: (2020) proposed to use GPs as a prior for a SEIR model for learning lockdown policy effects, but GPs are computationally expensive and the simple SEIR model cannot capture the real-world large-scale, spatiotemporal dynamics considered in this work.

Original: We demonstrate the use of deep sequence model as a prior distribution in Bayesian active learning.
Revised: We demonstrate the use of deep sequence model as a prior distribution in Bayesian active learning.

Original: Our framework is also  with other deep sequence models for time series
Revised: Our framework is also compatible with other deep sequence models for time series, e.g.

Original: e.g. Deep State Space ( et al Neural ODE (Chen et al., 2018).
Revised: Deep State Space (Rangapuram et al., 2018), Neural ODE (Chen et al., 2018).

Original:  M ETHODOLOGY
Revised:  M ETHODOLOGY

Original: Consider a stochastic process { X 1 , · · · , X T } , governed by time-varying parameters θ t ∈ R K , and the initial state x 0 ∈ R D .
Revised: Consider a stochastic process { X 1 , · · · , X T } , governed by time-varying parameters θ t ∈ R K , and the initial state x 0 ∈ R D .

Original: 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).
Revised: 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).

Original: The state x t ∈ R D 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 ).
Revised: The state x t ∈ R D 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 ).

Original: Stochastic simulation uses a mechanistic model F ( θ ; ξ ) to simulate the process where the random variable ξ represents the randomness in the simulator.
Revised: Stochastic simulation uses a mechanistic model F ( θ ; ξ ) to simulate the process where the random variable ξ represents the randomness in the simulator.

Original: Let θ :=
Revised: Let θ :=

Original: ( x 0 , θ 1 , · · · , θ T ) represent the initial state and all the parameters over time.
Revised: ( x 0 , θ 1 , · · · , θ T ) represent the initial state and all the parameters over time.

Original: For each θ , we obtain a different set of simulation data { ( x 1 , · · · , x T ) m } Mm =1 .
Revised: For each θ , we obtain a different set of simulation data { ( x 1 , · · · , x T ) m } Mm =1 .

Original: However, realistic large-scale stochastic simulations require the exploration of a large parameter space and are extremely computationally intensive.
Revised: However, realistic large-scale stochastic simulations require the exploration of a large parameter space and are extremely computationally intensive.

Original: In the following section, we describe the Interactive Neural Process (INP) framework to proactively query the stochastic simulator, generate simulation data, in order to learn a fast surrogate model for rapid simulation.
Revised: In the following section, we describe the Interactive Neural Process (INP) framework to proactively query the stochastic simulator, generate simulation data, in order to learn a fast surrogate model for rapid simulation.

Original: 3. I NTERACTIVE N EURAL P ROCESS
Revised: 3. I NTERACTIVE N EURAL P ROCESS

Original: INP is used to train a deep surrogate model to mimic the stochastic simulator.
Revised: INP is used to train a deep surrogate model to mimic the stochastic simulator.

Original: As shown in Figure 1, given parameters θ , we query the simulator, i.e., the mechanistic model to obtain a set of simulations { ( x 1 , · · · , x T ) m } Mm =1 .
Revised: As shown in Figure 1, given parameters θ , we query the simulator, i.e., the mechanistic model to obtain a set of simulations { ( x 1 , · · · , x T ) m } Mm =1 .

Original: We train a NP based model to learn the probabilistic map from parameters to future states.
Revised: We train a NP based model to learn the probabilistic map from parameters to future states.

Original: Our NP model can be spatiotemporal to capture complex dynamics such as the diseasedynamics 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 ( z 1 , · · · , z T ) that encodes the uncertainty of the current surrogate model.
Revised: Our NP model can be spatiotemporal to capture complex dynamics such as the disease

Original: Then we propose a new acquisition function, Latent Information Gain (LIG), to select the θ ⋆

Original: NP
Revised: NP

Original: SNP
Revised: SNP

Original: STNP z t x ct− 1 θ ct x ct x mt θ mt
Revised: STNP z t x ct− 1 θ ct x ct x mt θ mt

Original: h t h t+z t x mt θ
Revised: h t h t+z t x mt θ

Original: mt x mt
Revised: mt x mt

Original: θ mt h t−h t
Revised: θ mt h t−h t

Original: H t−
Revised: H t−

Original: H t x ct− 1 θ ct x ct x mt− 1 x mt− 1 x mt−z t x ct− 1 θ ct x ct with the highest reward.
Revised: H t x ct− 1 θ ct x ct x mt− 1 x mt− 1 x mt−z t x ct− 1 θ ct x ct dynamics of the epidemic simulator.

Revised: During inference, the model needs to generate predictions (ˆ x 1 , · · · , ˆ x T ) at the target parameters θ corresponding to different scenarios.

Revised: 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.

Revised: Using NP, we can infer the latent temporal process ( z 1 , · · · , z T ) that encodes the uncertainty of the current surrogate model.

Revised: Then we propose a new acquisition function, Latent Information Gain (LIG), to select the θ ⋆ with the highest reward.

Original: We use θ ⋆ to query the simulator, and in turn generate new simulation to further improve the model.
Revised: We use θ ⋆ to query the simulator, and in turn generate new simulation to further improve the model.

Original: Next, we describe each of the components in detail.
Revised: Next, we describe each of the components in detail.

Original: 3. S PATIOTEMPORAL N EURAL P ROCESS
Revised: 3. S PATIOTEMPORAL N EURAL P ROCESS

Original: Neural Process (NP) (Garnelo et al., 2018) is a type of deep generative model that represents distributions over functions.
Revised: Neural Process (NP) (Garnelo et al., 2018) is a type of deep generative model that represents distributions over functions.

Original: It introduces a global latent variable z to capture the stochasticity and learns the conditional distribution p ( x 1: T | θ ) by optimizing the evidence lower bound (ELBO):
Revised: It introduces a global latent variable z to capture the stochasticity and learns the conditional distribution p ( x 1: T | θ ) by optimizing the evidence lower bound (ELBO):

Original: [Equation]
Revised: [Equation]

Original: Here p ( z ) is the prior distribution for the latent variable.
Revised: Here p ( z ) is the prior distribution for the latent variable.

Original: We use x 1: T as a shorthand for ( x 1 , · · · , x T ) .
Revised: We use x 1: T as a shorthand for ( x 1 , · · · , x T ) .

Original: The prior distribution p ( z ) is conditioned on a set of context points θ c , x c 1: T as p ( z | x c 1: T , θ c ) .
Revised: The prior distribution p ( z ) is conditioned on a set of context points θ c , x c 1: T as p ( z | x c 1: T , θ c ) .

Original: However, the global latent variable z in NP can be limiting for non-stationary, spatiotemporal dynamics in the epidemics.
Revised: However, the global latent variable z in NP can be limiting for non-stationary, spatiotemporal dynamics in the epidemics.

Original: We propose Spatiotemporal Neural Process (STNP) with two extensions.
Revised: We propose Spatiotemporal Neural Process (STNP) with two extensions.

Original: First, we introduce a temporal latent process ( z 1 , · · · , z T ) to represent the unknown dynamics.
Revised: First, we introduce a temporal latent process ( z 1 , · · · , z T ) to represent the unknown dynamics.

Original: The latent process provides an expressive description of the internal mechanism of the stochastic simulator.
Revised: The latent process provides an expressive description of the internal mechanism of the stochastic simulator.

Original: Each latent variable z t is sampled conditioning on the past history.
Revised: Each latent variable z t is sampled conditioning on the past history.

Original: Second, we explicitly model the spatial dependency in x t ∈ R D .
Revised: Second, we explicitly model the spatial dependency in x t ∈ R D .

Original: Rather than treating the dimensions in x t as independent features, we capture their correlations with regular grids or graphs.
Revised: Rather than treating the dimensions in x t as independent features, we capture their correlations with regular grids or graphs.

Original: For instance, the travel graph between locations can be represented as an adjacency matrix A ∈ R D × D .
Revised: For instance, the travel graph between locations can be represented as an adjacency matrix A ∈ R D × D .

Original: Given parameters { θ } , simulation data { x 1: T } , and the spatial graph A as inputs, STNP models the conditional distribution p ( x 1: T | θ, A ) by optimizing the following ELBO objective:
Revised: Given parameters { θ } , simulation data { x 1: T } , and the spatial graph A as inputs, STNP models the conditional distribution p ( x 1: T | θ, A ) by optimizing the following ELBO objective:

Original: [Equation]
Revised: [Equation]

Original: where the distributions q ( z 1: T | x 1: T , θ, A ) and p ( x 1: T | z 1: T , θ, A ) are parameterized with neural networks.
Revised: where the distributions q ( z 1: T | x 1: T , θ, A ) and p ( x 1: T | z 1: T , θ, A ) are parameterized with neural networks.

Original: The prior distribution p ( z 1: T ) is conditioned on a set of contextual sequences p ( z 1: T | x c 1: T , θ c , A ) .
Revised: The prior distribution p ( z 1: T ) is conditioned on a set of contextual sequences p ( z 1: T | x c 1: T , θ c , A ) .

Original: Figure 2 visualizes the graphical models of our STNP, the original NP (Garnelo et al., 2018) model and Sequential NP (Singh et al., 2019).
Revised: Figure 2 visualizes the graphical models of our STNP, the original NP (Garnelo et al., 2018) model and Sequential NP (Singh et al., 2019).

Original: The main difference between STNP and baselines is the encoding procedure to infer the temporal latent process.
Revised: The main difference between STNP and baselines is the encoding procedure to infer the temporal latent process.

Original: Compared with STNP which directly embeds the history for z inference at the current timestamp, NP ignores the history and SNP only embeds the partial history information from the previous z .
Revised: Compared with STNP which directly embeds the history for z inference at the current timestamp, NP ignores the history and SNP only embeds the partial history information from the previous z .

Original: We implement STNP following an encoder-decoder architecture.
Revised: We implement STNP following an encoder-decoder architecture.

Original: The encoder parametrizes the mean and standard deviation of the variational posterior q ( z 1: T | x 1: T , θ, A ) and the decoder approximates the predictive distribution p ( x 1: T | z 1: T , θ, A ) .
Revised: The encoder parametrizes the mean and standard deviation of the variational posterior q ( z 1: T | x 1: T , θ, A ) and the decoder approximates the predictive distribution p ( x 1: T | z 1: T , θ, A ) .

Original: 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.
Revised: 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.

Original: We use multi-layer GRUs to obtain hidden states from the inputs.
Revised: We use multi-layer GRUs to obtain hidden states from the inputs.

Original: Using re-parametrization (Kingma & Welling, 2013), we sample z t from the encoder and then decode x t conditioned on z t in an auto-regressive fashion.
Revised: Using re-parametrization (Kingma & Welling, 2013), we sample z t from the encoder and then decode x t conditioned on z t in an auto-regressive fashion.

Revised: To ensure fair comparisons, we adapt NP and

Revised: SNP to graph-based settings and use the same architecture as STNP to generate the hidden states shown in Figure 6.

Original: Noted  if the spatial dependency is regular grid-based, then the DCGRU layer is replaced to Convolutional LSTM layer Linet al. ; Wang et al. Shi et al.
Revised: Noted that if the spatial dependency is regular grid-based, then the DCGRU layer is replaced by a Convolutional LSTM layer Lin al.

Original: ; Yao et al.
Revised: Yao et al.

Original: (2019; 2018), and there is no adjacency matrix A in Equation 2.
Revised: (2019; 2018), and there is no adjacency matrix A in Equation 2.

Original: 3. B AYESIAN A CTIVE L EARNING
Revised: 3. B AYESIAN A CTIVE L EARNING

Original: Algorithm 1 details a Bayesian active learning algorithm, based on Bayesian optimization (Shahriari et al., 2015; Frazier, 2018).
Revised: Algorithm 1 details a Bayesian active learning algorithm, based on Bayesian optimization (Shahriari et al., 2015; Frazier, 2018).

Original: We train an NP model to interact with the simulator and improve learning.
Revised: We train an NP model to interact with the simulator and improve learning.

Original: Let the superscript ( i ) denote the i -th interaction.
Revised: Let the superscript ( i ) denote the i -th interaction.

Original: We start with an initial data set S 1 = { θ (1) , x (1)1: T } and use it to train our NP model and learn the latent process.
Revised: We start with an initial data set S 1 = { θ (1) , x (1)1: T } and use it to train our NP model and learn the latent process.

Original: During inference, given the augmented parameters θ , we use the trained NP model to predict the future states (ˆ x 1 , · · · , ˆ x T ) .
Revised: During inference, given the augmented parameters θ , we use the trained NP model to predict the future states (ˆ x 1 , · · · , ˆ x T ) .

Original: We evaluate the current models’ predictions with an acquisition function r (ˆ x 1: T , z 1: T , θ ) and select the set of parameters { θ ( i +1) } with the highest reward.
Revised: We evaluate the current models’ predictions with an acquisition function r (ˆ x 1: T , z 1: T , θ ) and select the set of parameters { θ ( i +1) } with the highest reward.

Original: We query the simulator with { θ ( i +1) } to augment the training data set S i +1 and update the NP model for the next iteration.
Revised: We query the simulator with { θ ( i +1) } to augment the training data set S i +1 and update the NP model for the next iteration.

Original: The choice of the reward (acquisition) function r depends on the goal of the active learning task.
Revised: The choice of the reward (acquisition) function r depends on the goal of the active learning task.

Original: 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 ) .
Revised: 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 ) .

Original: To collect data and reduce model uncertainty in Bayesian experimental design, the reward function can be the mutual information.
Revised: To collect data and reduce model uncertainty in Bayesian experimental design, the reward function can be the mutual information.

Original: In what follows, we discuss different strategies to design the reward/acquisition function.
Revised: In what follows, we discuss different strategies to design the reward/acquisition function.

Original: We also propose a novel acquisition function based on information gain in the latent space tailored to our STNP model.
Revised: We also propose a novel acquisition function based on information gain in the latent space tailored to our STNP model.

Original: 3. R EWARD /A CQUISITION FUNCTIONS
Revised: 3. R EWARD /A CQUISITION FUNCTIONS

Original: For regression tasks, standard acquisition functions for active learning include Maximum Mean
Revised: For regression tasks, standard acquisition functions for active learning include Maximum Mean

Original: Standard Deviation (Mean STD), Maximum Entropy, Bayesian Active Learning by Disagreement (BALD) or expected information gain (EIG), and random sampling
Revised: Standard Deviation (Mean STD), Maximum Entropy, Bayesian Active Learning by Disagreement (BALD) or expected information gain (EIG), and random sampling

Original: (Gal et al., 2017).
Revised: (Gal et al., 2017).

Original: We explore various acquisition functions and their approximations in the context of NP.
Revised: We explore various acquisition functions and their approximations in the context of NP.

Original: We also introduce a new acquisition function based on our unique NP design called Latent Information Gain (LIG).
Revised: We also introduce a new acquisition function based on our unique NP design called Latent Information Gain (LIG).

Original: The details of Mean STD and Maximum Entropy are shown in the Appendix B.4.
Revised: The details of Mean STD and Maximum Entropy are shown in the Appendix B.4.

Original: BALD/Expected Information Gain (EIG).
Revised: BALD/Expected Information Gain (EIG).

Original: BALD (Houlsby et al., 2011) quantifies the mutual information between the prediction and model posterior H (ˆ x 1: T | θ ) − H (ˆ x 1: T | z 1: T , θ ) , which is equivalent to the expected information gain (EIG).
Revised: BALD (Houlsby et al., 2011) quantifies the mutual information between the prediction and model posterior H (ˆ x 1: T | θ ) − H (ˆ x 1: T | z 1: T , θ ) , which is equivalent to the expected information gain (EIG).

Original: Computing the EIG for surrogate modeling is challenging since p (ˆ x 1: T | z 1: T , θ ) cannot be found in closed form in general.
Revised: Computing the EIG for surrogate modeling is challenging since p (ˆ x 1: T | z 1: T , θ ) cannot be found in closed form in general.

Original: The integrand is intractable and conventional MC methods are not applicable (Foster et al., 2019).
Revised: The integrand is intractable and conventional MC methods are not applicable (Foster et al., 2019).

Original: 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), which is computationally infeasible.
Revised: 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), which is computationally infeasible.

Original: To reduce the computational cost, we assume p (ˆ x 1: T | z 1: T , θ ) follows multivariate Gaussian distribution.
Revised: To reduce the computational cost, we assume p (ˆ x 1: T | z 1: T , θ ) follows multivariate Gaussian distribution.

Original: Each feature of ˆ x 1: T can be parameterized with mean and standard deviation predicted from the surrogate model, assuming output features are independent with each other.
Revised: Each feature of ˆ x 1: T can be parameterized with mean and standard deviation predicted from the surrogate model, assuming output features are independent of each other.

Original: This distribution assumption can be limiting in the high-dimensional spatiotemporal domain, which makes EIG less informative.
Revised: This distribution assumption can be limiting in the high-dimensional spatiotemporal domain, which makes EIG less informative.

Original: Latent Information Gain (LIG).
Revised: Latent Information Gain (LIG).

Original: To overcome the limitations mentioned above, we propose a novel acquisition function by computing the expected information gain in the latent space rather than the observational space.
Revised: To overcome the limitations mentioned above, we propose a novel acquisition function by computing the expected information gain in the latent space rather than the observational space.

Original: To design this acquisition function, we prove the equivalence between the expected information gain in the observational space and the expected KL divergence in the latent processes w.r.t.
Revised: To design this acquisition function, we prove the equivalence between the expected information gain in the observational space and the expected KL divergence in the latent processes w.r.t.

Original: a candidate parameter θ , as illustrated by the following proposition.
Revised: a candidate parameter θ , as illustrated by the following proposition.

Original: Proposition 1.
Revised: Proposition 1.

Original: 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
Revised: 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

Original: [Equation]
Revised: [Equation]

Original: See proof in the Appendix A.1.
Revised: See proof in the Appendix A.1.

Original: Inspired by this fact, we propose a novel acquisition function computing the expected KL divergence in the latent processes and name it LIG.
Revised: Inspired by this fact, we propose a novel acquisition function computing the expected KL divergence in the latent processes and name it LIG.

Original: Specifically, the trained NP model produces a variational posterior given the current dataset S as p ( z 1: T |S ) .
Revised: Specifically, the trained NP model produces a variational posterior given the current dataset S as p ( z 1: T |S ) .

Original: For every parameter θ remained in the search space, we can predict ˆ x 1: T with the decoder.
Revised: For every parameter θ retained in the search space, we can predict ˆ x 1: T with the decoder.

Original: We use ˆ x 1: T and θ as input to the encoder to re-evaluate the posterior p ( z 1: T | ˆ x 1: T , θ, S ) .
Revised: We use ˆ x 1: T and θ as input to the encoder to re-evaluate the posterior p ( z 1: T | ˆ x 1: T , θ, S ) .

Original: LIG computes the distributional difference with respect to the latent process z 1: T as E p
Revised: LIG computes the distributional difference with respect to the latent process z 1: T as E p

Original: (ˆ x 1: T | θ ) [KL ( p ( z 1: T | ˆ x 1: T , θ, S )
Revised: (ˆ x 1: T | θ ) [KL ( p ( z 1: T | ˆ x 1: T , θ, S )

Original: ∥ p ( z 1: T |S ))] , where KL( ·∥· ) denotes the KL-divergence between two distributions.
Revised: ∥ p ( z 1: T |S ))] , where KL( ·∥· ) denotes the KL-divergence between two distributions.

Original: In this way, conventional MC method becomes applicable, which helps reduce the quadratic computational cost to linear.
Revised: In this way, conventional MC method becomes applicable, which helps reduce the quadratic computational cost to linear.

Original: At the same time, although z 1: T are assumed to be multivariate Gaussian and are parameterized with mean and standard deviation, they are only in the latent space not the observational space.
Revised: At the same time, although z 1: T are assumed to be multivariate Gaussian and are parameterized with mean and standard deviation, they are only in the latent space not the observational space.

Original: Moreover, LIG is also more computationally efficient and accurate for batch active learning.
Revised: Moreover, LIG is also more computationally efficient and accurate for batch active learning.

Original: Due to the context aggregation mechanism of NP, we can directly calculate LIG with respect to a batch of θ in the candidate set.
Revised: Due to the context aggregation mechanism of NP, we can directly calculate LIG with respect to a batch of θ in the candidate set.

Original: This is not available for baseline acquisition functions.
Revised: This is not available for baseline acquisition functions.

Original: They all require calculating the scores one by one for all θ in the candidate set and select a batch of θ based on their scores.
Revised: They all require calculating the scores one by one for all θ in the candidate set and select a batch of θ based on their scores.

Original: Such approach is both slow and inaccurate asacquiring points that are informative individually are not necessarily informative jointly (Kirsch et al., 2019).
Revised: Such approach is both slow and inaccurate as points that are informative individually are not necessarily informative jointly (Kirsch et al., 2019).

Original: 3. T HEORETICAL A NALYSIS
Revised: 3. T HEORETICAL A NALYSIS

Original: We shed light onto the intuition behind choosing adaptive sample selection over random sampling via analyzing a simplifying situation.
Revised: We shed light onto the intuition behind choosing adaptive sample selection over random sampling via analyzing a simplified situation.

Original: Assume that at a certain stage we have learned a feature map Ψ which maps the input θ of the neural network to the last layer.
Revised: Assume that at a certain stage we have learned a feature map Ψ which maps the input θ of the neural network to the last layer.

Original: Then the output X can be modeled as
Revised: Then the output X can be modeled as

Original: X =
Revised: X =

Original: ⟨ Ψ( θ ) , z ∗ ⟩ + ϵ , where z ∗ is the true hidden variable, ϵ is the random noise.
Revised: ⟨ Ψ( θ ) , z ∗ ⟩ + ϵ , where z ∗ is the true hidden variable, ϵ is the random noise.

Original: Our goal is to generate an estimate ˆ z , and use it to make predictions ⟨ Ψ( θ ) , ˆ z ⟩ .
Revised: Our goal is to generate an estimate ˆ z , and use it to make predictions ⟨ Ψ( θ ) , ˆ z ⟩ .

Original: A good estimate shall achieve small error in terms of ∥ ˆ z t − z ∗ ∥ 2 with high probability.
Revised: A good estimate shall achieve small error in terms of ∥ ˆ z t − z ∗ ∥ 2 with high probability.

Original: In the following theorem, we prove that greedily maximizing the variance of the prediction to choose θ will lead to an error of order O
Revised: In the following theorem, we prove that greedily maximizing the variance of the prediction to choose θ will lead to an error of order O

Original: ( d ) less than that of random exploration in the space of θ , which is significant in high dimension.
Revised: ( d ) less than that of random exploration in the space of θ , which is significant in high dimension.

Original: Theorem 1.
Revised: Theorem 1.

Original: For random feature map Ψ( · ) , greedily optimizing the KL divergence,
Revised: For random feature map Ψ( · ) , greedily optimizing the KL divergence,

Original: KL ( p ( z | ˆ x, θ ) ∥ p ( z )) , or equivalently the variance of the posterior predictive distribution
Revised: KL ( p ( z | ˆ x, θ ) ∥ p ( z )) , or equivalently the variance of the posterior predictive distribution

Original: E (cid:104) ( ⟨ Ψ( θ ) , ˆ z ⟩ − E ⟨ Ψ( θ ) , ˆ z ⟩ ) 2 (cid:105) in search of θ will lead to an error ∥
Revised: E (cid:104) ( ⟨ Ψ( θ ) , ˆ z ⟩ − E ⟨ Ψ( θ ) , ˆ z ⟩ ) 2 (cid:105) in search of θ will lead to an error ∥

Original: ˆ z t − z ∗ ∥ 2 of order O (cid:0) σd/ √ t (cid:1) with high probability.
Revised: ˆ z t − z ∗ ∥ 2 of order O (cid:0) σd/ √ t (cid:1) with high probability.

Original: On the other hand, random sampling of θ will lead to an error of order
Revised: On the other hand, random sampling of θ will lead to an error of order

Original: O (cid:0) σd 2 / √ t (cid:1) with high probability.
Revised: O (cid:0) σd 2 / √ t (cid:1) with high probability.

Original: See proofs in the Appendix A.2.
Revised: See proofs in the Appendix A.2.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We evaluate our proposed STNP for its surrogate modeling performance in the offline learning setting and LIG acquisition function for active learning performance.
Revised: We evaluate our proposed STNP for its surrogate modeling performance in the offline learning setting and LIG acquisition function for active learning performance.

Original: We aim to verify that (a)
Revised: We aim to verify that (a)

Original: LIG outperforms other acquisition functions in the NP and GP model setting for deep Bayesian active learning on non-spatiotemporal surrogate modeling, (b) STNP outperforms other existing
Revised: LIG outperforms other acquisition functions in the NP and GP model setting for deep Bayesian active learning on non-spatiotemporal surrogate modeling, (b) STNP outperforms other existing

Original: NP baselines for spatiotemporal surrogate modeling in the offline learning setting, and (c) LIG outperforms other acquisition functions in the STNP model setting for deep Bayesian active learning on spatiotemporal surrogate modeling.
Revised: NP baselines for spatiotemporal surrogate modeling in the offline learning setting, and (c) LIG outperforms other acquisition functions in the STNP model setting for deep Bayesian active learning on spatiotemporal surrogate modeling.

Original: 4. E XPERIMENTAL S ETUP
Revised: 4. E XPERIMENTAL S ETUP

Original: We experiment with the following three stochastic simulators.
Revised: We experiment with the following three stochastic simulators.

Original: SEIR Compartmental Model.
Revised: SEIR Compartmental Model.

Original: To highlight the difference between NP and GP, we begin with a simple stochastic, discrete, chain-binomial SEIR compartmental model as our stochastic simulator.
Revised: To highlight the difference between NP and GP, we begin with a simple stochastic, discrete, chain-binomial SEIR compartmental model as our stochastic simulator.

Original: In this model, susceptible individuals ( S ) become exposed ( E ) through interactions with infectious individuals ( I ) and are eventually removed ( R ), details are deferred to the Appendix B.1.
Revised: In this model, susceptible individuals ( S ) become exposed ( E ) through interactions with infectious individuals ( I ) and are eventually removed ( R ), details are deferred to the Appendix B.1.

Original: We set the total population N = S + E + I + R as 100 , 000 , the initial number of exposed individuals as E 0 = 2 , 000 , and the initial number of infectious individuals as I 

Revised: We set the total population N = S + E + I + R as 100 , 000 , the initial number of exposed individuals as E 0 = 2 , 000 , and the initial number of infectious individuals as I

Original: = 2 , 000 .
Revised: = 2 , 000 .

Original: We assume latent individuals move to the infectious stage at a rate ε ∈ [0 .
Revised: We assume latent individuals move to the infectious stage at a rate ε ∈ [0 .

Original: 65] (step 0.05), the infectious period µ − 1 is set to be equal to 1 day, and we let the basic reproduction number R 0 (which in this case coincides with the transmissibility rate β ) vary between 1 . 1 and 4 . 

Revised: 65] (step 0.05), the infectious period µ − 1 is set to be equal to 1 day, and we let the basic reproduction number R 0 (which in this case coincides with the transmissibility rate β ) vary between 1 . 1 and 4 .

Original: (step 0 . 1 ).
Revised: (step 0 . 1 ).

Original: Here, each ( β, ε ) pair corresponds to a specific scenario, which determines the parameters θ .
Revised: Here, each ( β, ε ) pair corresponds to a specific scenario, which determines the parameters θ .

Original: We simulate the first 100 days of the epidemic with a total of 300 scenarios and generate 30 samples for each scenario.
Revised: We simulate the first 100 days of the epidemic with a total of 300 scenarios and generate 30 samples for each scenario.

Original: We predict the number of individuals in the infectious compartment.
Revised: We predict the number of individuals in the infectious compartment.

Original: The input is ( β, ε ) pair and the output is the 100 days’ infection prediction.
Revised: The input is ( β, ε ) pair and the output is the 100 days’ infection prediction.

Original: As the simulator is not spatiotemporal, we use the vanilla
Revised: As the simulator is not spatiotemporal, we use the vanilla

Original: NP model with the global latent variable z .
Revised: NP model with the global latent variable z .

Original: For each epoch, we randomly select 10% of the samples as context.
Revised: For each epoch, we randomly select 10% of the samples as context.

Original: Implementation details are deferred to Appendix B.5.
Revised: Implementation details are deferred to Appendix B.5.

Original: Reaction Diffusion Model.
Revised: Reaction Diffusion Model.

Original: The reaction-diffusion (RD) system (Turing, 1990) is a spatiotemporal model that simulates how two chemicals might react to each other as they diffuse through a medium together.
Revised: The reaction-diffusion (RD) system (Turing, 1990) is a spatiotemporal model that simulates how two chemicals might react to each other as they diffuse through a medium together.

Original: The simulation is based on initial pattern, feed rate ( θ 0 ), removal rate ( θ 1 ) and reaction between two substances.
Revised: The simulation is based on initial pattern, feed rate ( θ 0 ), removal rate ( θ 1 ) and reaction between two substances.

Original: We use an RD simulator to generate sequences from 0 to 500 timestamps, sampled every 100 timestamps, resulting into 5 timestamps for each simulated sequence.
Revised: We use an RD simulator to generate sequences from 0 to 500 timestamps, sampled every 100 timestamps, resulting into 5 timestamps for each simulated sequence.

Original: Every timestamp is a 3D tensor (2 × 32 × 32) with dimension 0 corresponds to the two substances in the reaction and dimension 1 , 2 are the image representation of the reaction diffusion processes.
Revised: Every timestamp is a 3D tensor (2 × 32 × 32) with dimension 0 corresponds to the two substances in the reaction and dimension 1 , 2 are the image representation of the reaction diffusion processes.

Original: Each sequence is simulated with a unique feed rate θ 0 ∈ [0 . 029 , 0 . 045] and kill rate θ 1 ∈ [0 . 055 , 0 .
Revised: Each sequence is simulated with a unique feed rate θ 0 ∈ [0 . 029 , 0 . 045] and kill rate θ 1 ∈ [0 . 055 , 0 .

Original: 062] combination.
Revised: 062] combination.

Original: There are 200 uniformly sampled scenarios, corresponding to ( θ 0 , θ 1 ) combinations.
Revised: There are 200 uniformly sampled scenarios, corresponding to ( θ 0 , θ 1 ) combinations.

Original: We implement STNP to mimic the reaction diffusion simulator with feed rate ( θ 0 ) and kill rate ( θ 1 ) as input.

Original: The initial state of the reaction is fixed. We use multiple convolutional layers with a linear layer to encode the spatial data into latent space. We use an LSTM layer to encode the latent spatial data with θ 0 , θ 1 to map the input-output pairs to hidden features z 1:5 . With ( θ 0 , θ 1 ) , and z 1:5 sampled from the posterior distribution, we use an LSTM layer and deconvolutional layers to simulate reaction diffusion sequence.
Revised: Implementation details are deferred to Appendix B.5.

Original: For each epoch, we randomly select 20% samples as context sequence.

Original: Local Epidemic and Mobility model.
Revised: Local Epidemic and Mobility model.

Original: 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.
Revised: 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.

Original: Population size and county-specific age distributions reflect Census’ annual resident population estimates for year 2019.
Revised: Population size and county-specific age distributions reflect Census’ annual resident population estimates for year 2019.

Original: We consider individuals divided into 10 age groups.
Revised: We consider individuals divided into 10 age groups.

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Each θ includes the county-level parameters of LEAM-US and state level incidence and prevalence compartments.
Revised: Each θ includes the county-level parameters of LEAM-US and state level incidence and prevalence compartments.

Original: The total number of dimension in θ is 16 , 912 , see details in Appendix B.2.
Revised: The total number of dimensions in θ is 16 , 912 , see details in Appendix B.2.

Original: Overall, there are 315 scenarios in the search space, corresponding to 315 different θ with total 16 , 254 samples.
Revised: Overall, there are 315 scenarios in the search space, corresponding to 315 different θ with total 16 , 254 samples.

Original: We split 78% of the data as the candidate set, and 11% for validation and test.
Revised: We split 78% of the data as the candidate set, and 11% for validation and test.

Original: For active learning, we use the candidate set as the search space.
Revised: For active learning, we use the candidate set as the search space.

Original: We instantiate an STNP model to mimic an epidemic simulator that has θ at both county and state level and x t at the state level.
Revised: We instantiate an STNP model to mimic an epidemic simulator that has θ at both county and state level and x t at the state level.

Original: We use county-level parameter θ together with a county-to-county mobility graph A in California as input.
Revised: We use county-level parameter θ together with a county-to-county mobility graph A in California as input.

Original: We use the DCGRU layer (Li et al., 2017) to encode the mobility graph in a GRU.
Revised: We use the DCGRU layer (Li et al., 2017) to encode the mobility graph in a GRU.

Original: We use a linear layer to map the county-level output to hidden features at the state level.
Revised: We use a linear layer to map the county-level output to hidden features at the state level.

Original: For both the state-level encoder and decoder, we use multi-layer GRUs.
Revised: For both the state-level encoder and decoder, we use multi-layer GRUs.

Original: For each epoch, we randomly select 20% samples as context sequence.
Revised: For each epoch, we randomly select 20% samples as context sequence.

Original: 4. O FFLINE L EARNING P ERFORMANCE
Revised: 4. O FFLINE L EARNING P ERFORMANCE

Original: We compared our proposed STNP with vanilla NP (Garnelo et al., 2018) and SNP (Singh et al., 2019).
Revised: We compared our proposed STNP with vanilla NP (Garnelo et al., 2018) and SNP (Singh et al., 2019).

Original: The key innovation of STNP is the introduced temporal latent process.
Revised: The key innovation of STNP is the introduced temporal latent process.

Original: To ensure fair comparison, we modified NP for the RD model by adding convolutional layers for data encoding and deconvolutional layers for sequence generation.
Revised: To ensure fair comparison, we modified NP for the RD model by adding convolutional layers for data encoding [Figure]

Original: For the LEAM-US model, we modified NP by addingthe convolutional layers with diffusion convolution (Li et al., 2018) to embed the graphs.
Revised: adding

Original: Similarly, we modified SNP by replacing the convolutional layers with diffusion convolution.
Revised: Similarly, we modified SNP by replacing the convolutional layers with diffusion convolution.

Original: Tableshows the testing MAE of different NP models trained in an offline fashion.
Revised: Tableshows the testing MAE of different NP models trained in an offline fashion.

Original: Our STNP significantly improves the performance and can accurately learn the simulator dynamics for both experiments.
Revised: Our STNP significantly improves the performance and can accurately learn the simulator dynamics for both experiments.

Original: [Figure]

Original: [Table]
Revised: [Table]

Original:  Figure 3 left compares the NP and GP performance on one scenario in the held-out test set.
Revised: STNP 2.84 ± 0. 6.3 ± 0. Figure 3 left compares the NP and GP performance on one scenario in the held-out test set.

Original: It shows the ground truth and the predicted number of infectious population for the first 50 days.
Revised: It shows the ground truth and the predicted number of infectious population for the first 50 days.

Original: We also include the confidence intervals (CI) with 5 standard deviations for ground truth and NP predictions and 1 standard deviation for GP predictions.
Revised: We also include the confidence intervals (CI) with 5 standard deviations for ground truth and NP predictions and 1 standard deviation for GP predictions.

Original: We observe that NP fits the simulation dynamics better than GP for mean prediction.
Revised: We observe that NP fits the simulation dynamics better than GP for mean prediction.

Original: Moreover, NP has closer CIs to the truth, reflecting the simulator’s intrinsic uncertainty.
Revised: Moreover, NP has closer CIs to the truth, reflecting the simulator’s intrinsic uncertainty.

Original: GP shows larger CIs which represent the model’s own uncertainty.
Revised: GP shows larger CIs which represent the model’s own uncertainty.

Original: Note that NP is much more flexible than GP and can scale easily to high-dimensional data.
Revised: Note that NP is much more flexible than GP and can scale easily to high-dimensional data.

Original: Figure 3 middle indicates STNP can accurately predict various patterns corresponding to different ( θ 0 , θ 1 ) .
Revised: Figure 3 middle indicates STNP can accurately predict various patterns corresponding to different ( θ 0 , θ 1 ) .

Original: This confirms that our STNP is able to capture the high-dimensional spatiotemporal dependencies in RD simulations.
Revised: This confirms that our STNP is able to capture the high-dimensional spatiotemporal dependencies in RD simulations.

Original: Figure 3 right visualize the STNP predictions in four key compartments of a typical scenario with R 0{+
+} = 3 .{+
+}
Revised: Figure 3 right shows the STNP predictions in four key compartments of a typical scenario with R {+
+} = {+
+}

Original: 1 from March 2nd to March 29th.
Revised: .from March 2nd to March 29th.

Original: The confidence interval is plotted with 2 standard deviations.
Revised: The confidence interval is plotted with 2 standard deviations.

Original: We can see that both the mean and confidence interval of STNP predictions match the truth well.
Revised: We can see that both the mean and confidence interval of STNP predictions match the truth well.

Original: These two results demonstrate the promise that the generative STNP model can serve as a deep surrogate model for RD and LEAM-US simulator.
Revised: These two results demonstrate the promise that the generative STNP model can serve as a deep surrogate model for RD and LEAM-US simulator.

Original: 4. A CTIVE L EARNING
Revised: 4. A CTIVE L EARNING

Original: Implementation Details.
Revised: Implementation Details.

Original: We compare 6 different acquisition functions with NP for SEIR model and
Revised: We compare 6 different acquisition functions with NP for SEIR model and

Original: STNP for RD and LEAM-US model.
Revised: STNP for RD and LEAM-US model.

Original: For SEIR, the initial training dataset has 2 scenarios and we continue adding 1 scenario per iteration to the training set until the test loss converges to the offline modeling performance.
Revised: For SEIR, the initial training dataset has 2 scenarios and we continue adding 1 scenario per iteration to the training set until the test loss converges to the offline modeling performance.

Original: We also include GP with 3 different acquisition functions.
Revised: We also include GP with 3 different acquisition functions.

Original: For the RD model, all acquisition functions start with the same 5 scenarios randomly picked from the training dataset.
Revised: For the RD model, all acquisition functions start with the same 5 scenarios randomly picked from the training dataset.

Original: Then we continue adding 5 scenarios per iteration to the training set until the test loss converges.
Revised: Then we continue adding 5 scenarios per iteration to the training set until the test loss converges.

Original: Similarly, the LEAM-US model begins with 27 training data and we continue adding 8 scenarios per iteration to the training set until the validation loss converges.

Original: We measure the average performance over three random runs and report the MAE for the test set.

Original: Active Learning Performance.

Original: Figure 4 shows the testing MAE versus the percentage of samples included for training.

Original: The percentage of data is linearly proportional to the overall running time.

Original: This

Original: [Figure]
Revised: [Figure]

Original: LIGEIGRandomMaxEntropyMeanSTDOﬄine
Revised: LIGEIGRandomMaxEntropyMeanSTDOﬄine

Original: [Figure]
Revised: [Figure]

Original: [Figure]
Revised: [Figure]

Original: I n v e r s e o f t h e L a t e n t P e r i o d , ε
Revised: I n v e r s e o f t h e L a t e n t P e r i o d , ε

Original: [Figure]
Revised: [Figure]

Original: [Figure]

Original: u i s i t

Original: [Figure]
Revised: [Figure]

Original: n
Revised: c

Original: [Figure]
Revised: [Figure]

Original: c t i

Original: [Figure]

Original: n r
Revised: r

Original: [Figure]
Revised: [Figure]

Original: a r d

Original: I n f e c t i o u s P o p u l a t
Revised: I n f e c t i o u s P o p u l a t

Original: i o n M
Revised: i o n M

Revised: A E

Revised: Active Learning Performance.

Revised: Figure 4 shows the testing MAE versus the percentage of samples included for training.

Revised: The percentage of data is linearly proportional to the overall running time.

Original: A E figure shows our proposed LIG always has the best MAE performance until the convergence for all three experiments.
Revised: This figure shows our proposed LIG always has the best MAE performance until the convergence for all three experiments.

Original: Specifically, as shown in figure 4 left, we compare different acquisition functions on both NP and GP for SEIR model.
Revised: Specifically, as shown in Figure 4 left, we compare different acquisition functions on both NP and GP for SEIR model.

Original: It shows none of the GP methods converge after selecting 4 .
Revised: It shows none of the GP methods converge after selecting 4 .

Original: 07% of the data for training while NP methods converge much faster.
Revised: 07% of the data for training while NP methods converge much faster.

Original: Our proposed acquisition function LIG is the most sample efficient in acquisition functions used for NP.
Revised: Our proposed acquisition function LIG is the most sample efficient in acquisition functions used for NP.

Original: It takes only 4 .
Revised: It takes only 4 .

Original: 07% of the data to converge and reach the NP offline performance, which uses the entire training set for training.
Revised: 07% of the data to converge and reach the NP offline performance, which uses the entire training set for training.

Original: Moreover, there is an enormous gap between LIG and EIG with respect to the active learning performance.
Revised: Moreover, there is an enormous gap between LIG and EIG with respect to the active learning performance.

Original: This validates our theory that the uncertainty of the deep surrogate model is better measured on the latent space instead of the predictions.
Revised: This validates our theory that the uncertainty of the deep surrogate model is better measured on the latent space instead of the predictions.

Original: Similarly in figure 4 middle
Revised: Similarly in Figure 4 middle

Original: and right, we compared LIG with other acquisition functions on STNP for RD and LEAM-US model.
Revised: and right, we compared LIG with other acquisition functions on STNP for RD and LEAM-US model.

Original: It shows
Revised: It shows

Original: LIG converges to the offline performance using only 21 .
Revised: LIG converges to the offline performance using only 21 .

Original: 87% of data for RD experiment and 31 .
Revised: 87% of data for RD experiment and 31 .

Original: 4% of data for LEAM-US experiment.
Revised: 4% of data for LEAM-US experiment.

Original: Therefore, it is consistent among all three experiments that our proposed LIG always has the best MAE performance until convergence.
Revised: Therefore, it is consistent among all three experiments that our proposed LIG always has the best MAE performance until convergence.

Original: Notice that for figure 4 right, it shows the log scale MAE versus the percentage of samples included for training.
Revised: Notice that for Figureright, it shows the log scale MAE versus the percentage of samples included for training.

Original: Exploration Exploitation Trade-off.
Revised: Exploration Exploitation Trade-off.

Original: To understand the large performance gap for LIG vs. baselines, we visualize the values of test MAE and the acquisition function score for each Bayesian active learning iteration for SEIR model, shown in Figure 5.
Revised: To understand the large performance gap for LIG vs. baselines, we visualize the values of test MAE and the acquisition function score for each Bayesian active learning iteration for SEIR model, shown in Figure 5.

Original: For EIG, Mean STD, and Maximum Entropy, they all tend to exploit the region with large transmission rate for the first 2 iterations.
Revised: For EIG, Mean STD, and Maximum Entropy, they all tend to exploit the region with large transmission rate for the first 2 iterations.

Original: Including these scenarios makes the training set unbalanced.
Revised: Including these scenarios makes the training set unbalanced.

Original: The MAE in the region with small transmission rate become worse after 2 iterations.
Revised: The MAE in the region with small transmission rate become worse after 2 iterations.

Original: Meanwhile, Random is doing pure exploration.
Revised: Meanwhile, Random is doing pure exploration.

Original: The improvement of
Revised: The improvement of

Original: MAE performance is not apparent after 2 iterations.
Revised: MAE performance is not apparent after 2 iterations.

Original: Our proposed LIG is able to reach a balance by exploiting the uncertainty in the latent process and encouraging exploration.
Revised: Our proposed LIG is able to reach a balance by exploiting the uncertainty in the latent process and encouraging exploration.

Original: Hence, with a small number of iterations ( I = 2 ), it has already selected “informative scenarios” in the search space.
Revised: Hence, with a small number of iterations ( I = 2 ), it has already selected “informative scenarios” in the search space.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: We propose a unified framework Interactive Neural Processes (INP) for deep Bayesian active learning, that can seamlessly interact with existing stochastic simulators and accelerate simulation.
Revised: We propose a unified framework Interactive Neural Processes (INP) for deep Bayesian active learning, that can seamlessly interact with existing stochastic simulators and accelerate simulation.

Original: Specifically, we design STNP to approximate the underlying simulation dynamics.
Revised: Specifically, we design STNP to approximate the underlying simulation dynamics.

Original: It infers the latent process which describes the intrinsic uncertainty of the simulator.
Revised: It infers the latent process which describes the intrinsic uncertainty of the simulator.

Original: We exploit this uncertainty and propose LIG as a powerful acquisition function in deep Bayesian active learning.
Revised: We exploit this uncertainty and propose LIG as a powerful acquisition function in deep Bayesian active learning.

Original: We perform a theoretical analysis and demonstrate that our approach reduces sample complexity compared with random sampling in high dimension.
Revised: We perform a theoretical analysis and demonstrate that our approach reduces sample complexity compared with random sampling in high dimension.

Original: We also did extensive empirical evaluations on several complex real-world spatiotemporal simulators to demonstrate the superior performance of our proposed STNP and LIG.
Revised: We also did extensive empirical evaluations on several complex real-world spatiotemporal simulators to demonstrate the superior performance of our proposed STNP and LIG.

Original: For the future work, we plan to leverage Bayesian optimization techniques to directly optimize for the target parameters with auto-differentiation.
Revised: For the future work, we plan to leverage Bayesian optimization techniques to directly optimize for the target parameters with auto-differentiation.

Original: R EPRODUCIBILITY S TATEMENT
Revised: R EPRODUCIBILITY S TATEMENT

Original: The implementation code is included in the supplementary material.
Revised: The implementation code is included in the supplementary material.

Original: The readme file includes the corresponding instructions.
Revised: The readme file includes the corresponding instructions.

Original: The full proof of the Theorem 1 can be found in Appendix A.2.
Revised: The full proof of the Theorem 1 can be found in Appendix A.2.

Original: [Figure]

Original: A T HEORETICAL A NALYSIS A. L ATENT I NFORMATION G AIN
Revised: A T HEORETICAL A NALYSIS A. L ATENT I NFORMATION G AIN

Original: Proposition 1.
Revised: Proposition 1.

Original: 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
Revised: 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

Original: [Equation]
Revised: [Equation]

Original: Proof of Proposition 1.
Revised: Proof of Proposition 1.

Original: 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( θ ) =
Revised: 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( θ ) =

Original: H (ˆ x ) − H (ˆ x | z, θ ) .
Revised: H (ˆ x ) − H (ˆ x | z, θ ) .

Original: Take the expectation of IG(ˆ x, θ ) under the marginal distribution
Revised: Take the expectation of IG(ˆ x, θ ) under the marginal distribution

Original: [Equation]
Revised: [Equation]

Original: we obtain from the conditional independence of z and θ that
Revised: we obtain from the conditional independence of z and θ that

Original: [Equation]
Revised: [Equation]

Original: A. S AMPLE E FFICIENCY OF A CTIVE L EARNING
Revised: A. S AMPLE E FFICIENCY OF A CTIVE L EARNING

Original: From the main text we know that in each round, the output random variable
Revised: From the main text we know that in each round, the output random variable

Original: [Equation]
Revised: [Equation]

Original: We further assume that the random noise ϵ is mean zero and σ -subGaussian.
Revised: We further assume that the random noise ϵ is mean zero and σ -subGaussian.

Original: 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:
Revised: 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:

Original: [Equation]
Revised: [Equation]

Original: Let the prior distribution over z be p ( z | θ k ) 
Revised: Let the prior distribution over z be p ( z | θ k ) = p

Original: =
Revised: [Equation]

Original: p ( z ) ∝ exp (cid:16) − m 2 σ 2 ∥ z ∥ 2 (cid:17) .
Revised: z ) ∝ exp (cid:16) − m 2 σ 2 ∥ z ∥ 2 (cid:17) .

Original: Here we use k instead of ( i ) in the Algorithm 1 to represent the number of iterations.
Revised: Here we use k instead of ( i ) in the Algorithm 1 to represent the number of iterations.

Original: We can form a posterior over z in the k -th round:
Revised: We can form a posterior over z in the k -th round:

Original: [Equation]
Revised: [Equation]

Original: Focusing on the random variable z ∼ p ( ·| X 1 , θ 1 , . .
Revised: Focusing on the random variable z ∼ p ( ·| X 1 , θ 1 , . .

Original: , X k , θ k ) , the estimate of the hidden variable, we can express it at k -th round as:
Revised: , X k , θ k ) , the estimate of the hidden variable, we can express it at k -th round as:

Original: [Equation]
Revised: [Equation]

Original: where ˆ z k =
Revised: where ˆ z k =

Original: V − 1 k (cid:80) ks =
Revised: V − 1 k (cid:80) ks =

Original: 1 X s Ψ( θ s ) ,
Revised: 1 X s Ψ( θ s ) ,

Original: V k
Revised: V k

Original: = m I + (cid:80) ks =1 Ψ
Revised: = m I + (cid:80) ks =1 Ψ

Original: ( θ s )Ψ( θ s ) T , and η k is a standard normal random variable.
Revised: ( θ s )Ψ( θ s ) T , and η k is a standard normal random variable.

Original: We can either choose action θ randomly or greedily.
Revised: We can either choose action θ randomly or greedily.

Original: A random choice of θ corresponds to taking
Revised: A random choice of θ corresponds to taking

Original: [Equation]
Revised: [Equation]

Original: A greedy procedure is to choose action θ k in the k -th round to optimize KL ( p ( z | ˆ x, θ ) ∥ p ( z )) =
Revised: A greedy procedure is to choose action θ k in the k -th round to optimize KL ( p ( z | ˆ x, θ ) ∥ p ( z )) =

Original: E p ( z | ˆ x,θ ) (cid:16) log p ( z | ˆ x,θ ) p ( z ) (cid:17) , where we denote the estimated output variable ˆ x given θ and z as ˆ x = ⟨ Ψ( θ ) , z ⟩ .
Revised: E p ( z | ˆ x,θ ) (cid:16) log p ( z | ˆ x,θ ) p ( z ) (cid:17) , where we denote the estimated output variable ˆ x given θ and z as ˆ x = ⟨ Ψ( θ ) , z ⟩ .

Original: This optimization procedure is equivalent to maximizing the variance of the prediction:
Revised: This optimization procedure is equivalent to maximizing the variance of the prediction:

Original: [Equation]
Revised: [Equation]

Original: For both approaches, we assume that the features Ψ( θ ) are normalized.
Revised: For both approaches, we assume that the features Ψ( θ ) are normalized.

Original: We compare the statistical risk of this approach with the random sampling approach.
Revised: We compare the statistical risk of this approach with the random sampling approach.

Original: Assume that the features are normalized, so that for all θ ∈ R d , Ψ( θ ) ∈ S d − 1 .
Revised: Assume that the features are normalized, so that for all θ ∈ R d , Ψ( θ ) ∈ S d − 1 .

Original: Define a matrix
Revised: Define a matrix

Original: A k ∈ R d × k containing all the column vectors { Ψ( θ 1 ) , . . .
Revised: A k ∈ R d × k containing all the column vectors { Ψ( θ 1 ) , . . .

Original: , Ψ( θ k ) } .
Revised: , Ψ( θ k ) } .

Original: We can then express the estimation error in the following lemma.
Revised: We can then express the estimation error in the following lemma.

Original: Lemma 1.
Revised: Lemma 1.

Original: The estimation error ∥
Revised: The estimation error ∥

Original: ˆ z k − z ∗ ∥ 2 can be bounded as follow.
Revised: ˆ z k − z ∗ ∥ 2 can be bounded as follow.

Original: [Equation]
Revised: [Equation]

Original: We now analyze random sampling of θ versus greedy search for θ .
Revised: We now analyze random sampling of θ versus greedy search for θ .

Original: 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 (cid:0) 

Revised: 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 (cid:0)

Original: k
Revised: k

Original: A k A T k (cid:1) will converge to (cid:16)(cid:112) 1 /k − (cid:112) 1 /d (cid:17) 2 for large k , which is of order Ω(1 /d ) .
Revised: A k A T k (cid:1) will converge to (cid:16)(cid:112) 1 /k − (cid:112) 1 /d (cid:17) 2 for large k , which is of order Ω(1 /d ) .

Original: This will lead to an appealing risk bound for ∥
Revised: This will lead to an appealing risk bound for ∥

Original: ˆ z k − z ∗ ∥ 2 on the order of O (cid:16) d/ √ k (cid:17) .
Revised: ˆ z k − z ∗ ∥ 2 on the order of O (cid:16) d/ √ k (cid:17) .

Original: However, in high dimension, this feature map is often far from identity.
Revised: However, in high dimension, this feature map is often far from identity.

Original: In the proof of Theorembelow, we demonstrate that even when Ψ( · ) is simply a linear random feature map, with i.i.d.
Revised: In the proof of Theorembelow, we demonstrate that even when Ψ( · ) is simply a linear random feature map, with i.i.d.

Original: normal entries, random exploration in θ can lead to deteriorated error bound.
Revised: normal entries, random exploration in θ can lead to deteriorated error bound.

Original: 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.
Revised: 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.

Original: ; Mei & Montanari (2019).
Revised: ; Mei & Montanari (2019).

Original: Theorem 1 (Formal statement) .
Revised: Theorem 1 (Formal statement) .

Original: Assume that the noise ϵ in equation 5 is σ -subGaussian.
Revised: Assume that the noise ϵ in equation 5 is σ -subGaussian.

Original: For a normalized linear random feature map Ψ( · ) , greedily optimizing the KL divergence,
Revised: For a normalized linear random feature map Ψ( · ) , greedily optimizing the KL divergence,

Original: KL ( p ( z | ˆ x, θ ) ∥ p ( z )) (or equivalently the variance of the posterior predictive distribution defined in equation equation 8) in search of θ will lead to an error ∥
Revised: KL ( p ( z | ˆ x, θ ) ∥ p ( z )) (or equivalently the variance of the posterior predictive distribution defined in equation equation 8) in search of θ will lead to an error ∥

Original: ˆ z k − z ∗ ∥ 2 = O (cid:16) σd/ √ k (cid:17) with high probability.
Revised: ˆ z k − z ∗ ∥ 2 = O (cid:16) σd/ √ k (cid:17) with high probability.

Original: On the other hand, random sampling of θ following equation 7 will lead to ∥ ˆ z k − z ∗ ∥ 2 =
Revised: On the other hand, random sampling of θ following equation 7 will lead to ∥ ˆ z k − z ∗ ∥ 2 =

Original: O (cid:16) σd 2 / √ k (cid:17) with high probability.
Revised: O (cid:16) σd 2 / √ k (cid:17) with high probability.

Original: Proof of Theorem 1.
Revised: Proof of Theorem 1.

Original: For a linear random feature map, we can express Ψ( θ ) =
Revised: For a linear random feature map, we can express Ψ( θ ) =

Original: Ψ θ , where entries in
Revised: Ψ θ , where entries in

Original: Ψ ∈ R d
Revised: Ψ ∈ R d

Original: × d
Revised: × d

Original: are i.i.d.
Revised: are i.i.d.

Original: normal.
Revised: normal.

Original: The entries of Ψ θ are then normalized.
Revised: The entries of Ψ θ are then normalized.

Original: • For random exploration of θ , the matrix containing the feature vectors becomes A k = ΨΘ k , where matrix Θ k ∈ R d × k collects all the k column vectors of { θ 1 , . . .
Revised: • For random exploration of θ , the matrix containing the feature vectors becomes A k = ΨΘ k , where matrix Θ k ∈ R d × k collects all the k column vectors of { θ 1 , . . .

Original: , θ k } .
Revised: , θ k } .

Original: Then
Revised: Then

Original: A k A T k =
Revised: A k A T k =

Original: ΨΘ k Θ T k Ψ T .
Revised: ΨΘ k Θ T k Ψ T .

Original: From random matrix theory, we know that the condition number of Ψ is equal to d with high probability Chen & Dongarra (2005).
Revised: From random matrix theory, we know that the condition number of Ψ is equal to d with high probability Chen & Dongarra (2005).

Original: Hence for normalized Ψ and θ , σ min (cid:0) ΨΘ k Θ T k Ψ T (cid:1) ≥ σ 2min (Ψ) σ min (cid:0) Θ k Θ T k (cid:1) = 1 d 2 σ min (cid:0) Θ k Θ T k (cid:1) .
Revised: Hence for normalized Ψ and θ , σ min (cid:0) ΨΘ k Θ T k Ψ T (cid:1) ≥ σ 2min (Ψ) σ min (cid:0) Θ k Θ T k (cid:1) = 1 d 2 σ min (cid:0) Θ k Θ T k (cid:1) .

Original: The inequality holds because the smallest singular value is the inverse of the norm of the inverse matrix.
Revised: The inequality holds because the smallest singular value is the inverse of the norm of the inverse matrix.

Original: We then use the
Revised: We then use the

Original: fact
Revised: fact

Original: from random matrix theory
Revised: from random matrix theory

Original: that for
Revised: that for

Original: normalized random θ ,
Revised: normalized random θ ,

Original: the
Revised: the

Original: asymptotic
Revised: asymptotic

Original: distribution of the eigenvalues of 1 k Θ k Θ T k
Revised: distribution of the eigenvalues of 1 k Θ k Θ T k

Original: follow
Revised: follow

Original: the (scaled)
Revised: the (scaled)

Original: Marchenko–Pastur
Revised: Marchenko–Pastur

Original: distribution, which is supported on
Revised: distribution, which is supported on

Original: λ ∈ (cid:20)(cid:16)(cid:112) 1 /k − (cid:112) 1 /d (cid:17) 2 , (cid:16)(cid:112) 1 /k + (cid:112) 1 /d (cid:17) 2 (cid:21) , where the 1 /d
Revised: λ ∈ (cid:20)(cid:16)(cid:112) 1 /k − (cid:112) 1 /d (cid:17) 2 , (cid:16)(cid:112) 1 /k + (cid:112) 1 /d (cid:17) 2 (cid:21) , where the 1 /d

Original: scaling comes
Revised: scaling comes

Original: from
Revised: from

Original: the fact that θ
Revised: the fact that θ

Original: is normalized Götze & Tikhomirov (
Revised: is normalized Götze & Tikhomirov (

Original: Hence for large
Revised: Hence for large

Original: k σ min (cid:0) Θ k Θ T k (cid:1) ≥ (cid:16) 1 − (cid:112) k/d (cid:17) 2 with high
Revised: k σ min (cid:0) Θ k Θ T k (cid:1) ≥ (cid:16) 1 − (cid:112) k/d (cid:17) 2 with high

Original: probability
Revised: probability

Original: This
Revised: This

Original: combined with the
Revised: combined with the

Original: previous
Revised: previous

Original: paragraph yields that for the random feature model,
Revised: paragraph yields that for the random feature model,

Original: [Equation]
Revised: [Equation]

Original: with high probability.
Revised: with high probability.

Original: Plugging this result into Lemma 1, we obtain that the error ∥
Revised: Plugging this result into Lemma 1, we obtain that the error ∥

Original: ˆ z k − z ∗ ∥for random exploration in the space of θ is of order O (cid:16) d 2 / √ k (cid:17) .
Revised: ˆ z k − z ∗ ∥for random exploration in the space of θ is of order O (cid:16) d 2 / √ k (cid:17) .

Original: • We then analyze the error associated with greedy maximization of the posterior predictive variance.
Revised: • We then analyze the error associated with greedy maximization of the posterior predictive variance.

Original: We first note that the variance of the posterior predictive distribution in equation equation 8 can be expressed as follows using equation equation 6:
Revised: We first note that the variance of the posterior predictive distribution in equation equation 8 can be expressed as follows using equation equation 6:

Original: [Equation]
Revised: [Equation]

Original: where the expectations are with respect to z ∼
Revised: where the expectations are with respect to z ∼

Original: p ( ·| X
Revised: p ( ·| X

Original: [Equation]
Revised: [Equation]

Original: , X k − 1 , θ k − 1 ) .
Revised: , X k − 1 , θ k − 1 ) .

Original: We perform a singular value decomposition A k =
Revised: We perform a singular value decomposition A k =

Original: U k Λ k W k .
Revised: U k Λ k W k .

Original: Then (cid:80) ks =1 Ψ
Revised: Then (cid:80) ks =1 Ψ

Original: ( θ s )Ψ( θ s ) T =
Revised: ( θ s )Ψ( θ s ) T =

Original: A k A T k
Revised: A k A T k

Original: = U
Revised: = U

Original: k Λ k Λ T
Revised: k Λ k Λ T

Original: k U T k , and that V − 2 k − 1 = (cid:0) m I
Revised: k U T k , and that V − 2 k − 1 = (cid:0) m I

Original: + A k − 

Revised: + A k −

Original: A T k − 1 (cid:1) − 2 =
Revised: A T k − 1 (cid:1) − 2 =

Original: U k − 1 (cid:0) m I + Λ k − 1 Λ T k − 1 (cid:1) − 2 U T k − 1 .
Revised: U k − 1 (cid:0) m I + Λ k − 1 Λ T k − 1 (cid:1) − 2 U T k − 1 .

Original: Via this formulation, we see that maximizing Ψ( θ ) T V − 2 k − 1 Ψ( θ ) in equation equation 9 to choose θ k is equivalent to choosing Ψ( θ k ) =
Revised: Via this formulation, we see that maximizing Ψ( θ ) T V − 2 k − 1 Ψ( θ ) in equation equation 9 to choose θ k is equivalent to choosing Ψ( θ k ) =

Original: ( U k − 1 ) T( · ,l ) , where l = arg min i ∈{ 1 ,...,d } (cid:0) Λ k − 1 Λ T k − 1 (cid:1)
Revised: ( U k − 1 ) T( · ,l ) , where l = arg min i ∈{ 1 ,...,d } (cid:0) Λ k − 1 Λ T k − 1 (cid:1)

Original: ( i,i ) .
Revised: ( i,i ) .

Original: 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 V k − 1 .
Revised: 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 V k − 1 .

Original: Since every Ψ( θ ) is normalized and we initialize uniformly: V 

Revised: Since every Ψ( θ ) is normalized and we initialize uniformly: V

Original: = m I , the process is equivalent to scanning the orthogonal spaces of normalized vectors in R d for ⌊ k/d ⌋ times.
Revised: = m I , the process is equivalent to scanning the orthogonal spaces of normalized vectors in R d for ⌊ k/d ⌋ times.

Original: For large k , entries in Λ k Λ T k are approximately uniform and are all larger than or equal to ⌊ k/d ⌋ .
Revised: For large k , entries in Λ k Λ T k are approximately uniform and are all larger than or equal to ⌊ k/d ⌋ .

Original: Then σ min (cid:0) A k A T k (cid:1) =
Revised: Then σ min (cid:0) A k A T k (cid:1) =

Original: Ω( k/d ) .
Revised: Ω( k/d ) .

Original: Plugging into the bound of Lemma 1, we obtain that
Revised: Plugging into the bound of Lemma 1, we obtain that

Original: [Equation]
Revised: [Equation]

Original: Proof of Lemma 1.
Revised: Proof of Lemma 1.

Original: We first express the estimate ˆ z k as follows.
Revised: We first express the estimate ˆ z k as follows.

Original: [Equation]
Revised: [Equation]

Original: z ∗
Revised: z ∗

Original: [Equation]
Revised: [Equation]

Original: k

Original: [Equation]

Original: =
Revised: =

Original: [Equation]
Revised: [Equation]

Original: ( θ
Revised: ( θ

Original: [Equation]
Revised: [Equation]

Original: Then
Revised: Then

Original: [Equation]
Revised: [Equation]

Original: Define a matrix
Revised: Define a matrix

Original: A k ∈ R d × k containing all the column vectors { Ψ( θ 1 ) , . .
Revised: A k ∈ R d × k containing all the column vectors { Ψ( θ 1 ) , . .

Original: , Ψ( θ k ) } and perform a singular value decomposition
Revised: , Ψ( θ k ) } and perform a singular value decomposition

Original: A k =
Revised: A k =

Original: U k
Revised: U k

Original: Λ k W k .
Revised: Λ k W k .

Original: Then (cid:80) ks =1 Ψ
Revised: Then (cid:80) ks =1 Ψ

Original: ( θ s )Ψ( θ s ) T =
Revised: ( θ s )Ψ( θ s ) T =

Original: A k A T k
Revised: A k A T k

Original: = U
Revised: = U

Original: k Λ k Λ T
Revised: k Λ k Λ T

Original: k U T k , and V k
Revised: k U T k , and V k

Original: = m I + A k A T k , We further define vector e k ∈ R s where ( e k ) s =
Revised: = m I + A k A T k , We further define vector e k ∈ R s where ( e k ) s =

Original: ϵ s .
Revised: ϵ s .

Original: We use this definition to simplify the two terms further.
Revised: We use this definition to simplify the two terms further.

Original: For term T 1 ,
Revised: For term T 1 ,

Original: [Equation]
Revised: [Equation]

Original: For term T 2 , we define a diagonal matrix ¯Λ k ∈ R k × k which satisfies (cid:0) ¯Λ k (cid:1)
Revised: For term T 2 , we define a diagonal matrix ¯Λ k ∈ R k × k which satisfies (cid:0) ¯Λ k (cid:1)

Original: i,i = 1 if i ≤ d and (cid:0) ¯Λ k (cid:1)
Revised: i,i = 1 if i ≤ d and (cid:0) ¯Λ k (cid:1)

Original: i,i = 

Revised: i,i =

Original: if i > d , when k > d .
Revised: if i > d , when k > d .

Original: The following bound on T 2 can be achieved.
Revised: The following bound on T 2 can be achieved.

Original: [Equation]
Revised: [Equation]

Original: Assuming that noise ϵ s is σ -subGaussian
Revised: Assuming that noise ϵ s is σ -subGaussian

Original: [Equation]
Revised: [Equation]

Original: then so is W k e k since W k is a unitary matrix.
Revised: then so is W k e k since W k is a unitary matrix.

Original: Multiplied by the diagonal matrix ¯Λ k which has zero, (cid:13)(cid:13)
Revised: Multiplied by the diagonal matrix ¯Λ k which has zero, (cid:13)(cid:13)

Original: ¯Λ k W k e k (cid:13)(cid:13)
Revised: ¯Λ k W k e k (cid:13)(cid:13)

Original: 2 ≤ σ √ d .
Revised: 2 ≤ σ √ d .

Original: Therefore,
Revised: Therefore,

Original: [Equation]
Revised: [Equation]

Original: A k A T k (cid:1) 
Revised: A k A T k (cid:1) 

Original:    · σ √ d.
Revised:    · σ √ d.

Original: B E XPERIMENT D ETAILS B. SEIR M ODEL
Revised: B E XPERIMENT D ETAILS B. SEIR M ODEL

Original: Our SEIR simulator is a simple stochastic, discrete, chain-binomial compartmental model.
Revised: Our SEIR simulator is a simple stochastic, discrete, chain-binomial compartmental model.

Original: In this model, susceptible individuals ( S ) become exposed ( E ) through interactions with infectious individuals ( I ).
Revised: In this model, susceptible individuals ( S ) become exposed ( E ) through interactions with infectious individuals ( I ).

Original: 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.
Revised: 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.

Original: Lastly, infectious individuals transition to the removed compartment at a rate µ which is inversely proportional to the infectious period.
Revised: Lastly, infectious individuals transition to the removed compartment at a rate µ which is inversely proportional to the infectious period.

Original: Removed individuals ( R ) are assumed to be no longer infectious and they are to be considered either recovered or dead.
Revised: Removed individuals ( R ) are assumed to be no longer infectious and they are to be considered either recovered or dead.

Original: All transitions are simulated by randomly drawn from a binomial distribution.
Revised: All transitions are simulated by randomly drawn from a binomial distribution.

Original: B. LEAM-US M ODEL
Revised: B. LEAM-US M ODEL

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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,
Revised: 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,

Original: Aviation Worlwide Limited, 2021; IATA, International Air Transport Association, 2021).
Revised: Aviation Worlwide Limited, 2021; IATA, International Air Transport Association, 2021).

Original: In addition, flight probabilities are age and country specific.
Revised: In addition, flight probabilities are age and country specific.

Original: The model is initialized using a multi-scale modeling approach that utilizes GLEAM, the Global and
Revised: The model is initialized using a multi-scale modeling approach that utilizes GLEAM, the Global and

Original: Epidemic Mobility model (Balcan et al., 2009; 2010; Tizzoni et al., 2012;
Revised: Epidemic Mobility model (Balcan et al., 2009; 2010; Tizzoni et al., 2012;

Original: 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.
Revised: 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.

Original: 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;
Revised: 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;

Original: Zhang et al., 2017) calibrated to realistically represent the evolution of the COVID-19 pandemic (Chinazzi et al., 2020; Davis et al., 2020).
Revised: Zhang et al., 2017) calibrated to realistically represent the evolution of the COVID-19 pandemic (Chinazzi et al., 2020; Davis et al., 2020).

Original: 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.
Revised: 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.

Original: B. 
Revised: B.

Original: [Figure]
Revised: [Figure]

Original: As shown in Figure 6, our model has θ at both county and state level and x t at the state level.
Revised: As shown in Figure 6, our model has θ at both county and state level and x t at the state level.

Original: We use county-level parameter θ together with a county-to-county mobility graph A as input.
Revised: We use county-level parameter θ together with a county-to-county mobility graph A as input.

Original: We use the DCGRU layer (Li et al., 2017) to encode the graph in a GRU.
Revised: We use the DCGRU layer (Li et al., 2017) to encode the graph in a GRU.

Original: We use a linear layer to map the county-level output to hidden features at the state level.
Revised: We use a linear layer to map the county-level output to hidden features at the state level.

Original: For both the state-level encoder and decoder, we use multi-layer GRUs.
Revised: For both the state-level encoder and decoder, we use multi-layer GRUs.

Original: The input θ 1: t is the county-level parameters for LEAM-US with a dimension of 10.
Revised: The input θ 1: t is the county-level parameters for LEAM-US with a dimension of 10.

Original: The county level embedding uses 1 layer DCGRU with a width of 16.
Revised: The county level embedding uses 1 layer DCGRU with a width of 16.

Original: The internal state is at state level with dimension of 16.
Revised: The internal state is at state level with dimension of 16.

Original: The state level encoder and decoder use 3 layer GRUs with width of 128.
Revised: The state level encoder and decoder use 3 layer GRUs with width of 128.

Original: The dimension of the latent process z 1: t is 32.
Revised: The dimension of the latent process z 1: t is 32.

Original: The dimension of output x 1: t is 24, including the incidence and prevalence for 12 compartments.
Revised: The dimension of output x 1: t is 24, including the incidence and prevalence for 12 compartments.

Original: We trained STNP model for 500 steps with learning rate fixed at 10 − 3 using
Revised: We trained STNP model for 500 steps with learning rate fixed at 10 − 3 using

Original: Adam optimizer.
Revised: Adam optimizer.

Original: We perform early stopping with 50 patience for both offline learning and Bayesian active learning.
Revised: We perform early stopping with 50 patience for both offline learning and Bayesian active learning.

Original: B. A CQUISITION F UNCTION
Revised: B. A CQUISITION F UNCTION

Original: Maximum Mean STD.
Revised: Maximum Mean STD.

Original: Mean STD (Gal & Ghahramani, 2016) is a heuristic used to estimate the model uncertainty.
Revised: Mean STD (Gal & Ghahramani, 2016) is a heuristic used to estimate the model uncertainty.

Original: For each augmented parameter θ , we sample multiple z 1: T and generate a set of predictions { ˆ x 1: T } .
Revised: For each augmented parameter θ , we sample multiple z 1: T and generate a set of predictions { ˆ x 1: T } .

Original: For a length T sequence with dimension D , we compute the standard deviation σ t,d for time step t and feature d .
Revised: For a length T sequence with dimension D , we compute the standard deviation σ t,d for time step t and feature d .

Original: Mean STD computes ¯ σ = 1 TD (cid:80) Tt =1 (cid:80) Dd =1 (cid:80) σ t,d for each parameter θ .
Revised: Mean STD computes ¯ σ = 1 TD (cid:80) Tt =1 (cid:80) Dd =1 (cid:80) σ t,d for each parameter θ .

Original: We select the θ with the maximum ¯ σ .
Revised: We select the θ with the maximum ¯ σ .

Original: Empirically, we found that Mean STD often becomes over-conservative and tends to explore less.[Table] [Table] [Table]
Revised: Empirically, we found that Mean STD often becomes over-conservative and tends to explore less.

Original: [Table]

Original: Maximum Entropy.
Revised: Maximum Entropy.

Original: Maximum entropy computes the maximum predictive entropy as H (ˆ x ) = − E [log p (ˆ x 1: T )] .
Revised: Maximum entropy computes the maximum predictive entropy as H (ˆ x ) = − E [log p (ˆ x 1: T )] .

Original: In general, entropy is intractable for continuous output.
Revised: In general, entropy is intractable for continuous output.

Original: Our NP model implicitly assumes the predictions follow a multivariate Gaussian, which allows us to compute the differential entropy (Jaynes, 1957).
Revised: Our NP model implicitly assumes the predictions follow a multivariate Gaussian, which allows us to compute the differential entropy (Jaynes, 1957).

Original: We follow the same procedure as Mean STD to estimate the empirical covariance Σ ∈ R TD × TD and compute the differential entropy for each parameter as H = 12 ln | Σ | +
Revised: We follow the same procedure as Mean STD to estimate the empirical covariance Σ ∈ R TD × TD and compute the differential entropy for each parameter as H = 12 ln | Σ | +

Original: TD 2 (1 + ln 2 π ) .
Revised: TD 2 (1 + ln 2 π ) .

Original: We select the parameter θ with the maximum entropy.
Revised: We select the parameter θ with the maximum entropy.

Original: B. I MPLEMENTATION D ETAILS 
Revised: B. I MPLEMENTATION D ETAILS FOR SEIR AND R EACTION DIFFUSION SIMULATOR .

Original: For both GP and INP model mimicking SEIR simulation, we ran experiments using CPU.
Revised: For both GP and INP model mimicking SEIR simulation, we ran experiments using CPU.

Original: No GPU accelerator is needed for this simple model.
Revised: No GPU accelerator is needed for this simple model.

Original: It takes 5 hours to converge.
Revised: It takes 5 hours to converge.

Original: For INP model mimicking
Revised: For INP model mimicking

Original: LEAM-US simulation, we ran experiments with GEFORCE RTX 2080.
Revised: LEAM-US simulation, we ran experiments with GEFORCE RTX 2080.

Original: It takes one day for the training to converge.
Revised: It takes one day for the training to converge.

Original: For all experiments, we run with three different random seeds.
Revised: For all experiments, we run with three different random seeds.

Revised: We implement STNP to mimic the reaction diffusion simulator with feed rate ( θ 0 ) and kill rate ( θ 1 ) as input.

Revised: The initial state of the reaction is fixed.

Revised: We use multiple convolutional layers with a linear layer to encode the spatial data into latent space.

Original: C A DDITIONAL R ESULTS C. INP M ODEL AND GP M ODEL
Revised: We use an LSTM layer to encode the latent spatial data with θ 0 , θ 1 to map the input-output pairs to hidden features z 1:5 . With ( θ 0 , θ 1 ) , and z 1:5 sampled from the posterior distribution, we use an LSTM layer and deconvolutional layers to simulate reaction diffusion sequence. For each epoch, we randomly select 20% samples as context sequence. B. I MPLEMENTATION D ETAILS FOR MAF, SNL, AND DMFAL.

Revised: We use the likelihood-free inference code (Durkan et al., 2020) to implement Masked Autoregressive

Revised: Flow (MAF) model and Sequential Neural Likelihood (SNL) framework.

Revised: We use the VAE-based deep surrogate model code for Deep Multi-fidelity Active Learning (DMFAL) (Li et al., 2020) to implement DMFAL.

Revised: For DMFAL, we set the number of fidelity levels to 1 to meet our task setting.

Revised: Note that we only use DMFAL model for offline learning test as their proposed acquisition function is EIG applied to multiple fidelity levels, which is equivalent to EIG once we reduce the number of fidelity levels to 1 .

Revised: We follow the same hyperparameter setting used in DMFAL (Li et al., 2020) for the standard Heat simulation task.

Revised: For other experiments, we tune the hyperparameters of baseline models including the learning rate and the hidden state size to optimize their performance.

Revised: C A DDITIONAL R ESULTS C. O FFLINE LEARNING PERFORMANCE COMPARISON FOR MAF, SNL, DMFAL ON R EACTION -D IFFUSION AND H EAT SIMULATOR .

Revised: We include MAF and DMFAL as two baselines for offline learning on Reaction-Diffusion and Heat simulation tasks.

Revised: Table 2 shows that our proposed STNP performs much better than MAF, RNN,

Revised: DMFAL, and other NP baselines for both RD and Heat tasks.

Revised: There is no need to evaluate active learning baselines (e.g. SNL) based on underperforming surrogate models (e.g. MAF) since their poor performance is guaranteed.

Revised: We didn’t test DMFAL and MAF on LEAM-US simulation task as they do not support graph-based spatiotemporal surrogate modeling.

Revised: C. A CTIVE LEARNING PERFORMANCE COMPARISON ON SEIR AND H EAT SIMULATOR .

Revised: We test the SNL baseline for active learning on SEIR simulation task.

Revised: Table 3 shows that the performance of our proposed LIG is much better than SNL at each iteration.

Revised: Table 4 shows the active learning performance comparison results on Heat simulation task.

Revised: It shows that LIG always outperforms the baselines at each iteration.

Revised: C. INP M ODEL AND GP M ODEL

Original: Table 2 shows the average results together with the standard deviation of INP model for SEIR simulator after running experiments three times. 
Revised: Table 5 shows the average results together with the standard deviation of INP model for SEIR simulator after running experiments three times. Table 6 shows the average results together with the [Table] [Table]

Revised: [Table]

Original: Table 3 shows the average results together with the standard deviation of GP model for SEIR simulator.
Revised: [Table] [Table] [Table] standard deviation of GP model for SEIR simulator.

Original: Table 5 shows the average results together with the standard deviation of INP model for LEAM-US simulator.
Revised: Table 8 shows the average results together with the standard deviation of INP model for LEAM-US simulator.

Original: C. B ATCH A CTIVE L EARNING WITH LIG.
Revised: C. B ATCH A CTIVE L EARNING WITH LIG.

Original: Figure 7 compares 4 different setups: 8 batches (size 1), 4 batches (size 2), 2 batches (size 4), andbatch (size 8).
Revised: Figure 7 compares 4 different setups: 8 batches [Table]

Original: [Figure]
Revised: [Figure]