Title: Deep Bayesian Active Learning for Accelerating Stochastic Simulation
Authors: Dongxia Wu
Ruijia Niu
Matteo Chinazzi
Alessandro Vespignani
Yian Ma
Rose Yu
Venue: ICLR 2023

Abstract: Stochastic simulations such as large-scale, spatiotemporal, age-structured epidemic models are computationally expensive at fine-grained resolution. While deep surrogate models can speed up the simulations, doing so for stochastic simulations and with active learning approaches is an underexplored area. We propose Interactive Neural Process (INP), a deep Bayesian active learning framework for learning deep surrogate models to accelerate stochastic simulations. INP consists of two components, a spatiotemporal surrogate model built upon Neural Process (NP) family and an acquisition function for active learning. For surrogate modeling, we develop Spatiotemporal Neural Process (STNP) to mimic the simulator dynamics. For active learning, we propose a novel acquisition function, Latent Information Gain (LIG), calculated in the latent space of NP based models. We perform a theoretical analysis and demonstrate that LIG reduces sample complexity compared with random sampling in high dimensions. We also conduct empirical studies on two complex spatiotemporal simulators for reaction diffusion and infectious disease. 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.


Review discussion:
Pros:
The paper uses a known and well-founded variant of the Neural Process as the surrogate model; they provide a delta of the model for better capturing spatiotemporal dependencies. The authors correctly resort to active learning to allow for the model to improve on-the-fly. They also correctly recognize the limitations of the venerable Expected Information Gain criterion used as acquisition function in Bayesian active learning. The suggested solution of computing the same criterion in the latent space (instead of the observation space) is reasonable.
Cons: The provided experimental comparisons are limited; I would like to see comparison to more deep learning alternatives. The experiments to now explain in depth all the trade-offs between complexity and accuracy, as well as in terms of how long active learning takes to reach performance. The experimental scenarios are limited as are the alternative surrogates they compare to.
The writing is clear; some minor typos here and there, but good structure overall. The research is of certainly good quality. The novelly is mediocre, but carries some value. The experiments are reproducible.
NO.
Thank you for your comments. We appreciated you thinking our paper is “a good paper” and our research “is of certainly good quality”. We also have some clarification questions. It is appreciated if you could answer them so that we can provide more details to improve the paper. C: comments; R: response; Q: question.
*C1: The novelty is mediocre.*
R1: Reviewer k2B6 and gfsx both think  “the paper is technically sound and has a novel methodological contribution with convincing empirical results”, “the paper is reasonably well written, timely, and novel”, "it is clear how this work differs from these previous contributions".
We would like to emphasize the novelty of our proposed unified framework INP with respect to the model architecture, acquisition function, and application. For model architecture, we are first to propose STNP (Spatiotemporal Neural Process), which enables spatiotemporal surrogate modeling for graph-based stochastic simulators. Its novel design integrates the time evolving latent process for temporal dynamics and graph convolution for spatial dependency.
Furthermore, we are the first to conduct a systematic study on three widely used real-world stochastic simulation types: temporal simulation (e.g. SEIR),  spatiotemporal simulation with regular grids (e.g. Reaction Diffusion), and spatiotemporal simulation with irregular grids (e.g. Local Epidemic and Mobility simulator).
For the acquisition function, we proposed a novel acquisition function Latent Information Gain (LIG). To our best knowledge, this is the first acquisition function calculated in the latent space using deep latent variable modeling methods. LIG is both computationally efficient and accurate compared with the baseline acquisition functions calculated in the output space. We provide theoretical guarantees together with a wide range of empirical experiments to verify the state-of-the-art performance.
For application, we successfully applied our framework to accelerate the challenging real-world stochastic simulator (Local Epidemic and Mobility simulator) for COVID-19 modeling. Both the inputs and outputs of the simulator are high-dimensional with complex spatiotemporal dependencies, while previous works for deep surrogate modeling and Bayesian active learning focus on either low-dimensional global inputs with high-dimensional outputs or high-dimensional inputs (e.g., image) with low-dimensional outputs (e.g., classification label).
__Q1: Could you be more specific as to why the novelty of our paper is "mediocre"?__
*C2: I would like to see comparison to more deep learning alternatives. The experimental scenarios are limited as are the alternative surrogates they compare to.*
R2: For experimental scenarios, we would like to highlight that the three simulators included in our experiments are highly representative of real-world stochastic simulations, including temporal simulation, spatiotemporal simulation with regular grids, and spatiotemporal simulation with irregular grids. It shows our proposed model has the potential to be applied to an abroad range of simulations within these three categories. For surrogate models, our goal is to design deep Bayesian model that is scalable and supports acquisition function computation (e.g., latent information gain calculated in the latent space).
Therefore, we focus on latent variable modeling methods, i.e., neural process family, which are scalable and have low-dimensional latent processes for acquisition function calculation. We have extended NP to STNP and compared it with all existing methods in the NP family that can be applied to spatiotemporal modeling.
That being said, we are more than willing to include more simulation scenarios and deep learning baselines outside the NP family. As reviewer k2B6 and gfsx recommended, we are running the deep learning baselines including the RNN model and the normalizing flow based model from sequential neural likelihood framework. Note this framework is for the likelihood-free inference problem (do posterior inference in the parameter space given the observation), which is the “inverse” problem of our surrogate modeling problem. That’s why we did not include it in our original paper. We will also test out methods on another standard task, heat equation, to compare our proposed model with the existing baselines.
__Q2: We would like to know if there are any other experiments and/or baselines (e.g. “experimental scenarios” and/or “alternative surrogates”) in your mind that we could cover to make our paper stronger.__
*C3: The experiments to now explain in depth all the trade-offs between complexity and accuracy*
__Q3: We assume you mean “The experiments do not explain …”. Could you be more precise about what kind of complexity and accuracy analysis you want us to include?__
*C4: how long active learning takes to reach performance*
R4: Thanks for the question. The implementation details are in Appendix B.5 “For both GP and INP model mimicking SEIR simulation, we ran experiments using CPU. No GPU accelerator is needed for this simple model. It takes 5 hours to converge. For INP model mimicking LEAM-US simulation, we ran experiments with GEFORCE RTX 2080. It takes one day for the training to converge.” For INP model mimicking RD simulation, we ran experiments with GEFORCE RTX 2080. It takes 12 hours to converge. As we retrained the surrogate model based on the learned weights from the previous iteration, the model is fast to converge at each iteration and the overall training time is similar to the offline training time.
Additionally, in section 4.3, we mentioned “The percentage of data is linearly proportional to the overall running time.” The reason is that the computation time of training surrogate model and calculating acquisition functions is negligible compared to the cost of simulating training data in real-world applications with computationally complex simulators.
For example,  the Local Epidemic and Mobility simulator requires about 1 hour to complete a single simulation (on a single thread of an AMD Ryzen Threadripper 3970x CPU, 3.7Ghz base clock). Note it requires around 50 such runs to generate samples for one scenario. In our experiment, we add 8 scenarios per iteration, so the stochastic simulator would take 400 CPU hours.
On the other hand, the average computational cost of INP (training + calculating acquisition functions) per iteration is about 0.417 GPU hours (on a GeForce RTX 2080). In this case, we can assume the sample efficiency is equivalent to the computational efficiency.
Dear reviewer,
We have followed your suggestions to include more alternative surrogates and experimental scenarios.
* For surrogate models, we added three types of deep learning models, including normalizing flow-based model Masked Autoregressive Flow [1], RNN-based model (RNN with variational dropout), VAE-based model DMFAL [2] to compare with our proposed STNP.
* For active learning approaches, we added Sequential Neural Likelihood (SNL) [3] to compare with our proposed LIG.
* For experimental scenarios, we added one standard task for solving Heat simulation used in [2] and used it for both offline learning and active learning test.
The new experimental results are listed in “Summary and Revision”.
[1] Papamakarios, G., Pavlakou, T., & Murray, I. (2017). Masked autoregressive flow for density estimation.
[2] Li, S., Kirby, R. M., & Zhe, S. (2020). Deep Multi-Fidelity Active Learning of High-dimensional Outputs.
[3] Papamakarios, G., Sterratt, D., & Murray, I. (2019). Sequential neural likelihood: Fast likelihood-free inference with autoregressive flows.
Please let us know if our revised manuscript and replies have addressed your concerns. If you have any additional comments, we would be more than happy to follow up.
Thank you very much!
# Strengths
The basic premise of the paper is reasonably clear, it is reasonably well written and referenced, and is timely.  The model itself seems well-grounded and sensible.  Evaluating the acquisition function in the latent space is a nice idea.  The experiments use realistic models (if not real data).  The theoretical contributions are not huge, but are reasonably well-suited to the topic, and provide some nice additional material to compliment the empirical results.
# Weaknesses
I have three main questions/comments on the work.
## (1)  What is being acquired?
Something that didn’t come across particularly clearly is how the acquisition function is actually used. In S3. 4, the acquisition function is defined in terms of $x$, but in F5, the acquisition function is plotted as a function of $\theta$. There are also many $\hat{x}$’s defined across different values of $\theta$. Are you able to clarify for me how the reduction in entropy between the posterior and prior (eq (3)) is a sensible way to select a theta, and how this reduction is actually computed? The reduction in entropy is given a particular $\theta$, and so we are looking to reduce the uncertainty over the latent state values for the observed data?
Does this not “reward” parameter values for which the prior has enormous variance, even if they have low probability? If the data are generated from a wildly different parameter set, then I would expect either a huge entropy (the model is uncertain) or a tiny entropy (the model is in some way “degenerate”), which is going to skew the LIG value.
It is not clear to me that minimising the uncertainty in a latent space is the right way to go.  I invite the authors to restate the objectives of the LIG, how it works, and the weaknesses of it (in light of my comments above).
## (2)  Deployment Pipeline
The introduction talks about flexible approximators and fast inference.  Are the authors able to re-state how the envisage this pipeline being used?
When reading the introduction, i thought you were trying to learn a universal approximator to the mapping from parameters to outcomes, which can then be polled in ``war time’’ quickly.
Later on though, when talking about active learning, it is like you are trying to learn a function approximator _quickly_ _during_ war time.
If you are learning a fixed function approximator for all of parameters and state space, then you are hugely limited by the expressivity of the model.  If you are trying to fast model learning at deployment time, then I query why the recent work on Synthetic Neural Likelihoods was not compared to, as this to me seems to be the gold standard.
## (3)  Empirical Evaluation
To my eye, the empirical validation is not especially strong.  It performs better on some non-standard tasks, but I still don’t have a great handle on how well it is actually doing – are the experiments just tuned to the STNP method?  Are the baselines poorly tuned?  I don’t currently know.  I would encourage the authors to find an example to apply the STNP to from the literature, and compare to existing results.
I would even be satisfied by seeing something of an unreasonable baseline;  e.g. train an RNN (using VI) to predict the mean and variances of the trajectories,  train an MLP to predict the final compartment counts in the SEIR as a function of the input parameters – just something to give the reader a bearing on how an “alternative” and “simple” method would perform.
The active learning component is nice, but again, I would like to see it compared on a standard task – evaluating the efficacy of the acquisition function is difficult, and I can’t really tell what is happening in F5.  Is there a 1-D model (in time as well) that we can apply this to where the latent space is easily visualizable, we can determine how well the STNP is able to reproduce the latent space etc.
I also think a readily interpretable/toy example would actually increase the reach of this work, demonstrating the utility of the method in a general example, as opposed to highly specific domains.  I do not expect this to be added during the rebuttal period, but would expect to see something like this for any camera ready version.
# Minor Weaknesses / Typographical Comments.
(a)  The paper is a little wordy in places.  I think a lot of the material could be simplified and shortened, and actually dramatically improve the quality of the paper.
Although this criticism is throughout the paper, some sections I thought were particularly wordy/can be scythed in length:  BALD paragraph in S3.4;  the whole of S3.2;  The first two paragraphs of the introduction;  Most of S4.1 can be dropped to the supplement – it adds very little to the general reader.
(b)  Algorithm 1 is really not as useful as it could be.  I would like to see this algorithm blown up, with inline comments, in its own section etc.  It also seems like the labels in Alg 1 are not correct (particularly “learn”).
(c)  \citep and \citet are used incorrectly in places (i.e. Garnelo et al, pg2).
(d)  You’ve vspaced section 3.3.  Please try to avoid doing this.  (especially cf. (a)!)
(e)  There are generally quite a few typos and grammatical mistakes.  Any camera ready would require extensive proofreading beyond the submitted version.  “In example” (pg 1), “remained” (pg 6), “a simplifying situation” (pg 7) for example.
(f)  In the proof of proposition 1, are you able to re-iterate how the $\theta$ appears in the denominator of the first line (I may well just be having a mental blank here).
(g)  “figure” should be capitalised, i.e. “Figure 1”.
The work appears novel, and builds on recent methods in the field.  The method itself is reasonably well presented, and I could understand a lot of the material, despite not being an NP expert.  Code is released.
NO.
Thank you for your valuable and insightful comments. We group our responses into 3 categories: __1. Methodology Clarification__, __2. Pipeline Clarification__, __3. Experiment Clarification and New Results__. The changes made in the revised manuscript are indicated in blue.
Here are our responses. C: comments; R: response.
__1. Methodology Clarification__
*C1: What is being acquired? how the acquisition function is actually used. In S3.4, the acquisition function is defined in terms of $x$, but in F5, the acquisition function is plotted as a function of $\theta$. There are also many $\hat{x}$’s defined across different values of $\theta$.*
R1:  We are acquiring more training data because we cannot afford to simulate all scenarios beforehand. All of the acquisition functions are used to select the parameters $\theta$ from the candidate set. They are functions of $\theta$ and $\hat{x}$, where $\hat{x}$ are estimated given $\theta$. We use the decoder of the trained NP model to approximate $p(\hat{x}|\theta)$ and generate samples of $\hat{x}$.
Then we use the selected parameters $\theta^*$ to query the simulator and get the corresponding true $x$ (see figure 1).  The second and fourth rows in Figure 5 show the acquisition function score in the 2D $\theta$ mesh. Yellow dots are existing $\theta$ parameters selected from the previous iteration. Red stars are the newly selected $\theta$ parameters at the current iteration.
*C2: Are you able to clarify for me how the reduction in entropy between the posterior and prior (eq (3)) is a sensible way to select a theta, and how this reduction is actually computed?*
R2: $\mathrm{EIG}(\hat{x}, \theta) :=\mathbb{E} [H(\hat{x}) - H(\hat{x}| z, \theta)]$ is the definition of expected information gain, which is calculated in the output space. It is the state-of-the-art acquisition function used for pool-based active learning [1]. Our contribution is to prove $\mathbb{E} [H(\hat{x}) - H(\hat{x}| z, \theta)] =
\mathbb{E}_{p(\hat{x}|\theta)} [ \mathrm{KL}\big( p( z| \hat{x}, \theta) \| p(z)  \big) ]$. Intuitively, It means “The expected information gain (EIG) for Neural Process is equivalent to the KL divergence between the prior and posterior in the latent process”. We name this acquisition function Latent Information Gain (LIG).
To compute LIG, we assume the latent processes $z$ are multivariate Gaussian. The prior of the latent process is the variational posterior given the current training dataset $\mathcal{S}$. The posterior of the latent process is the variational posterior given the current training dataset $\mathcal{S}$ plus candidate $(\hat{x},\theta)$ pair. Both of them are computed through a forward pass of the NP model encoder. Intuitively, LIG measures the statistical distance between the prior and posterior distribution of the latent processes. Please refer to section 3.4 Latent Information Gain (LIG) for more details.
[1] Gal, Y., Islam, R., & Ghahramani, Z. (2017). Deep bayesian active learning with image data.
*C3: The reduction in entropy is given a particular $\theta$, and so we are looking to reduce the uncertainty over the latent state values for the observed data? It is not clear to me that minimizing the uncertainty in a latent space is the right way to go.*
R3: We didn’t reduce the uncertainty over the latent state values. Our LIG acquisition function is the expected KL divergence between the prior and posterior in the latent process $\mathbb{E}_{p(\hat{x}|\theta)} [ \mathrm{KL}\big( p( z| \hat{x}, \theta) \| p(z)  \big)]$. We are looking for the $(\hat{x},\theta)$ pairs that maximize the statistical distance between the prior and posterior distribution of the latent processes.
*C4: Does this not “reward” parameter values for which the prior has enormous variance, even if they have low probability?*
R4: The prior of the latent process is not for a single parameter value. It is the variational posterior given the current training dataset $\mathcal{S}$.
*C5: In the proof of proposition 1, are you able to re-iterate how the  $\theta$ appears in the denominator of the first line?*
R5: $\theta$ here is the candidate which has not been added to the training set.
Therefore, $\theta$ and $z$ are independent. We have $p(z|\theta) = p(z)$. We rewrite in this way to introduce the term $p(z,\theta)$ and $p(\theta)$.
__2. Pipeline Clarification__
*C6: The introduction talks about flexible approximators and fast inference. Are the authors able to re-state how the envisage this pipeline being used?*
R6: Yes, we want to learn a universal surrogate model to map from parameters to outcomes, which can then be polled in “war time” quickly.
At the same time, training such model needs to be data-efficient. We realize it using active learning. For real-world applications, the deployment pipeline will be: 1. Use active learning to train a surrogate model. 2. Use the surrogate model instead of the original simulator to help domain experts quickly explore the parameter space and create hypothetical scenarios.
*C7: If you are learning a fixed function approximator for all of parameters and state space, then you are hugely limited by the expressivity of the model. If you are trying to fast model learning at deployment time, then I query why the recent work on Synthetic Neural Likelihoods was not compared to, as this to me seems to be the gold standard.*
R7: We predefine the parameter range in our problem setup. This is a common practice for simulation and the range is usually decided by domain experts.
Also, we assume the data is uniformly distributed as our goal is to ask what-if questions and all possible scenarios in the predefined parameter range are equally important. That being said, we also test synthetic neural likelihood baseline sequential neural likelihood (SNL) [1] on SEIR simulation task. __Table 3__ in the appendix shows that the averaged performance of our proposed LIG is much better than SNL at each iteration.
__Table 3: Active learning performance comparison using MAE on SEIR simulator (LIG vs SNL)__
| Percentage of samples | LIG           | SNL          |
|-----------------------|---------------|--------------|
| 1.11%                 | 365.87±142.87 | 707.61±44.42 |
| 1.85%                 | 236.9±50.6    | 669.03±73.19 |
| 2.96%                 | 119.26±14.22  | 668.67±72.42 |
| 4.07%                 | 96.73±17.07   | 685.28±53    |
__3. Experiment Clarification and New Results__
*C8: Are the experiments just tuned to the STNP method? Are the baselines poorly tuned? I don’t currently know. I would encourage the authors to find an example to apply the STNP to from the literature, and compare to existing results.*
R8: The comparison among NP, SNP, STNP is a fair comparison, we adapted NP and SNP baselines to graph-based settings and use the same architecture as STNP to generate the hidden states shown in Figure 6. I would even be satisfied by seeing something of an unreasonable baseline; e.g. train an RNN (using VI) to predict the mean and variances of the trajectories.
We have added Masked Autoregressive Flow (MAF) [1,2], RNN (with variational dropout), and Deep Multi-fidelity Active Learning (DMFAL) [3] as three baselines for offline learning on Reaction-Diffusion and Heat simulation tasks. Heat simulation is the standard task included in DMFAL paper. We set the number of fidelity levels to $1$ to meet our task setting. We follow the same hyperparameter setting used in DMFAL [3] for the standard Heat simulation task. For other experiments, we tune the hyperparameters of baseline models including the learning rate and the hidden state size to optimize their performance. __Table 2__ in appendix shows that our proposed STNP always has the best performance compared with MAF, RNN, DMFAL, and other NP baselines for both RD and Heat tasks.
__Table 2: Offline learning performance comparison using MAE on Reaction-Diffusion and Heat simulator__
| Model              | NP               | SNP            | MAF             | RNN             | DMFAL          | STNP             |
|--------------------|------------------|----------------|-----------------|-----------------|----------------|------------------|
| Reaction-Diffusion | 3.37 ± 0.18      | 3.11 ± 0.07    | 7.22 ± 0.66     | 3.44 ± 0.07     | 4.1 ± 0.02     | 2.84 ± 0.17      |
| Heat               | 1.05e-2 ± 1.1e-3 | 9.62e-3 ± 1e-3 | 2.1e-2 ± 4.8e-3 | 3.2e-2 ± 1.7e-3 | 1.35e-2 ± 3e-4 | 7.36e-3 ± 6.7e-4 |
[1] Papamakarios, G., Sterratt, D., & Murray, I. (2019). Sequential neural likelihood: Fast likelihood-free inference with autoregressive flows.
[2] Papamakarios, G., Pavlakou, T., & Murray, I. (2017). Masked autoregressive flow for density estimation.
[3] Li, S., Kirby, R. M., & Zhe, S. (2020). Deep Multi-Fidelity Active Learning of High-dimensional Outputs.
*C9: The active learning component is nice, but again, I would like to see it compared on a standard task ...*
R9: For active learning, we included sequential neural likelihood baseline [1] for SEIR simulation task. The result is shown on __Table 3__. Its performance is much worse compared to our proposed LIG. We also added the standard Heat simulation task. Similar result is shown on __Table 4__ in appendix where LIG always has the best performance. Note that there is no need to evaluate active learning baselines SNL and DMFAL on this task since their surrogate models’ offline modeling performance is much worse than NP based offline modeling performance.
__Table 4: Active learning performance comparison using MAE on Heat simulator__
| Percentage of samples | LIG              | EIG              | Random           | MeanSTD          | MaxEntropy       |
|-----------------------|------------------|------------------|------------------|------------------|------------------|
| 5.21%                 | 1.55e-2 ± 1.9e-3 | 1.64e-2 ± 1.9e-3 | 1.74e-2 ± 2.2e-3 | 1.77e-2 ± 2e-3   | 1.83e-2 ± 1.1e-3 |
| 7.81%                 | 1.33e-2 ± 1.7e-3 | 1.56e-2 ± 1.1e-3 | 1.49e-2 ± 3.3e-3 | 1.60e-2 ± 3.7e-3 | 1.58e-2 ± 4e-4   |
| 10.42%                | 1.02e-2 ± 3.4e-3 | 1.38e-2 ± 1.6e-3 | 1.23e-2 ± 1.5e-3 | 1.30e-2 ± 2.4e-3 | 1.61e-2 ± 8e-4   |
| 13.02%                | 9.3e-3 ± 3.1e-3  | 1.05e-2 ± 2.2e-3 | 1.14e-2 ± 1.3e-3 | 1.27e-2 ± 2.5e-3 | 1.46e-2 ± 6e-4   |
| 15.62%                | 6.7e-3 ± 5e-4    | 1.08e-2 ± 1.8e-3 | 1.05e-2 ± 9e-4   | 1.17e-2 ± 1.2e-3 | 1.43e-2 ± 2e-4   |
In Fig 5, the second and fourth rows show the acquisition function score in the 2D $\theta$ mesh.  It shows the selected parameter $\theta$ marked in red star has the highest acquisition function score. To our best knowledge, there is no existing surrogate and simulator pairs whose latent space can match each other. We think it is a good direction to study in the future.
[1] Papamakarios, G., Sterratt, D., & Murray, I. (2019). Sequential neural likelihood: Fast likelihood-free inference with autoregressive flows.
*C10: I also think a readily interpretable/toy example would actually increase the reach of this work ...*
R10: Thanks for the suggestion. We use the commonly used SEIR as a toy example and Figure 5 to illustrate the whole active learning process. Please let us know if you have any other questions about Figure 5.
Other than that, we will think about adding other interpretable examples for the camera-ready version.
Please let us know if our revised manuscript and replies have addressed your concerns. If you have any additional comments, we would be more than happy to follow up.
Thank you very much!
Hi Authors,
Thank you for your response and the additional experiments.  I think these improve the strength of the empirical validation.
While I am more positive about the paper, I do not think I feel strongly enough to upgrade my score.
I think the empirical evaluation of the STNP is lacking, maybe requiring a whole raft of new experiments comparing it to more function approximation methods on a wider range of tasks (synthetic and real).  I almost wish you separate the paper in to two:  one posing and really evaluating the performance of the STNP, and one evaluating the active learning elements, really digging in to the pros/cons, performance limits, and application domains of each method, instead of providing experiments that sort-of speak to both.  This would also allow you to spend more time comparing to other methods on standard benchmarks.  I think my gripes about presentation also stem from this.  The paper still feels cluttered and jumbled, and I'd like it thinned out.
I think there is a (or maybe two!) good conference papers in this line of work, I just feel it is not currently in that place.  Good luck.
gfsx
Dear reviewer,
We appreciate your time to read our rebuttal and updated results.
Could you provide more concrete suggestions on the ``empirical evaluation of STNP''?  We have compared with almost **all** the related functional approximations methods that we can think of, including NP, SNP, MAF, RNN, DMFAL.  If there is a specific related method that is critical for you to upgrade score, pls provide the pointers.
Similarly to the tasks, we have already included 3 different tasks (synthetic and real with different spatial geometry and scales).
We are eager to address any concerns you might have regarding our experiments.
But we would not even know where to start if you don't give us more concrete suggestions.
Regarding separating the paper into two. Bayesian deep active learning is our main focus, and STNP is only an element within the framework.  These two are tightly integrated into each other and can not be separated. If you can suggest a Bayesian deep active learning framework without the surrogate model, pls give us a pointer.
We sincerely hope to receive a fair review based on scientific evidence. We look forward to seeing your response.
The paper is technically sound, and claims are well supported by empirical evaluation.
Some previous work is cited and discussed, and it is clear how this work differs from these previous contributions.
However, I believe it would be important to revise the section on related work.
In particular:
-the authors write "Kleinegesse & Gutmann (2020) [...] require a explicit likelihood model, conditional independence in experiments, and are limited to low (1-2) dimensional design". I would urge the authors to clarify this statement.
In particular, Kleinegesse & Gutmann (2020) deal with simulators for which no likelihood is available (also called "implicit models"), which is contrary to what is stated in this manuscript;
-related to the point above, there is a large body of work in machine learning on surrogate models for stochastic simulators (e.g. Meeds and Welling 2014 "GPS-ABC: Gaussian Process Surrogate Approximate Bayesian Computation"; Gutmann and Corander 2016 "Bayesian Optimization for Likelihood-Free Inference of Simulator-Based Statistical Models"; Lueckmann et al. 2018 "Likelihood-free inference with emulator networks"; Järvenpää et al. 2019 "Efficient acquisition rules for model-based approximate bayesian computation"; Papamakarios et al. 2019 "Sequential Neural Likelihood: Fast Likelihood-free Inference with Autoregressive Flows"), some of which use active learning schemes to adaptively choose model parameters for acquiring new simulations (e.g. Meeds and Welling 2014; Gutmann and Corander 2016; Lueckmann et al. 2018; Järvenpää et al. 2019). It would be important to appropriately discuss this work in the manuscript.
To my knowledge, the proposed method is novel.
Furthermore, the paper is technically sound.
The paper is generally clear, and I would just like to point a few typos and imprecisions:
-"Stochastic simulations are fundamental to many scientific fields (Ripley, 2009), especially epidemic modeling". As it is phrased, it seems to suggest that stochastic stimulations are more important in epidemiology than in other scientific fields (which is arguably not true), whereas the authors probably wanted to write that stochastic simulations are important in epidemiology;
-"Noted if the spatial dependency", instead "Note that if the spatial dependency";
-"is replaced to Convolutional LSTM layer", instead "is replaced by a Convolutional LSTM layer";
-"are independent with each other", instead "are independent of each other";
-probably better to drop the word "acquiring" in "as acquiring points that are informative individually are not necessarily informative jointly";
-in "The total number of dimension", it should be "dimensions". There are a few other instances of "dimension" where this should also be corrected;
-"Figure 3 right visualize the STNP predictions", instead "Figure 3 right shows the STNP predictions";
-Figure 5, yellow dots and red stars need to be larger for easier visualisation.
NO.
Thank you for your valuable and insightful comments. Here are our responses. C: comments; R: response. The changes made in the revised manuscript are indicated in blue.
C1: I believe it would be important to revise the section on related work.
In particular…
R1: We appreciated your suggestions to include more likelihood-free inference works on related work. We have modified the claim about “Kleinegesse & Gutmann (2020)” and included the provided reference in the revised manuscript (see section 2). We also added more experiments to compare our proposed STNP and LIG with likelihood-free inference baselines in both offline learning and active learning setup. For offline learning, we show STNP outperforms Masked Autoregressive Flow (MAF) [1] on both the Reaction-Diffusion task and the added standard task for Heat simulation. MAF is the state-of-the-art surrogate model used for Sequential Neural Likelihood (SNL) [2]. For active learning, we show LIG outperforms SNL on the SEIR task. The corresponding table results are listed in “Summary and Revision”.
[1] Papamakarios, G., Pavlakou, T., & Murray, I. (2017). Masked autoregressive flow for density estimation.
[2] Papamakarios, G., Sterratt, D., & Murray, I. (2019). Sequential neural likelihood: Fast likelihood-free inference with autoregressive flows.
Please let us know if our revised manuscript and replies have addressed your concerns. If you have any additional comments, we would be more than happy to follow up.
Thank you very much!
I would like to thank the authors for addressing my reviews to satisfaction.
However, taking into account all the reviews, I will maintain my score of 6.
Strength: the paper is well motivated and well written.
Weakness:
1. Novelty. I appreciate the effort to propose the spatio-temporal extension to the original NP as well as the latent extension to EIG, but I am not convinced that this is sufficiently novel for ICLR standard.
2. The experimental part is also limited. I would expect more ablation study so that we can fully understand the sensitivity and robustness of the proposed approach.
3. Some implementation details are not provided it seems.  How does the active learning get actually implemented?
When a new set of data is added, does the INP get retrained from scratch, or fine-tuning from the existing weights?
Overall this paper is clear to read, well written. The novelty is, as previously explained, IMO, not sufficient. The code is included, so I assume that the reproducibility should be high.
NO.
Thank you for your comments. We appreciated you thinking our paper is “well motivated and well written”. We also have some clarification questions. It is appreciated if you could answer them so that we can provide more details to improve the paper. C: comments; R: response; Q: question.
*C1: The novelty is not sufficient.*
R1: Reviewer k2B6 and gfsx both think  “the paper is technically sound and has a novel methodological contribution with convincing empirical results”, “the paper is reasonably well written, timely, and novel”, "it is clear how this work differs from these previous contributions".
We would like to emphasize the novelty of our proposed unified framework INP with respect to the model architecture, acquisition function, and application. For model architecture, we are first to propose STNP (Spatiotemporal Neural Process), which enables spatiotemporal surrogate modeling for graph-based stochastic simulators. Its novel design integrates the time evolving latent process for temporal dynamics and graph convolution for spatial dependency.
Furthermore, we are the first to conduct a systematic study on three widely used real-world stochastic simulation types: temporal simulation (e.g. SEIR),  spatiotemporal simulation with regular grids (e.g. Reaction Diffusion), and spatiotemporal simulation with irregular grids (e.g. Local Epidemic and Mobility simulator).
For the acquisition function, we proposed a novel acquisition function Latent Information Gain (LIG). To our best knowledge, this is the first acquisition function calculated in the latent space using deep latent variable modeling methods. LIG is both computationally efficient and accurate compared with the baseline acquisition functions calculated in the output space. We provide theoretical guarantees together with a wide range of empirical experiments to verify the state-of-the-art performance.
For application, we successfully applied our framework to accelerate the challenging real-world stochastic simulator (Local Epidemic and Mobility simulator) for COVID-19 modeling. Both the inputs and outputs of the simulator are high-dimensional with complex spatiotemporal dependencies, while previous works for deep surrogate modeling and Bayesian active learning focus on either low-dimensional global inputs with high-dimensional outputs or high-dimensional inputs (e.g., image) with low-dimensional outputs (e.g., classification label).
__Q1: Could you be more specific as to why the novelty of our paper is ‘not sufficient?__
*C2: the sensitivity and robustness of the proposed approach.*
R2: Our experiments already show that LIG performance is robust to random seeds, the choice of the model architecture (MLP based encoder and decoder for SEIR modeling, convolutional RNN based encoder and decoder for Reaction-Diffusion modeling, graph convolutional RNN based encoder and decoder for LEAM modeling), and the problem dimensionality (SEIR with small input and output dimensionality,  Reaction-Diffusion with small input and large output dimensionality, LEAM with large input and output dimensionality). Results are shown in Table 2,4,5 in appendix and Figure 4.
__Q2: Could you provide more details on what additional ablation study is needed?__
*C3: How does the active learning get actually implemented?*
R3: We have included the active learning implementation details in Section 4.3 Active Learning Implementation Details. The active learning procedure and framework are included in Figure 1 and Algorithm 1. The acquisition function score calculation is included in Section 3.4 and Appendix B.4. Note that we assume the output and the latent processes follow multivariate Gaussian distributions for baseline acquisition function and LIG computation, respectively. If there are still things that are missing in our paper, please let us know and we will include them in the updated version.
*C4: When a new set of data is added, does the INP get retrained from scratch, or fine-tuning from the existing weights?*
R4: We fine-tune the model from the existing weights. The corresponding trained model performance at the final iteration is slightly better and the converge rate is faster compared with retraining from scratch. We will include this detail in our updated paper.
Dear reviewer,
We have added more experiments to show the robustness of our proposed STNP and LIG and revised our manuscript. The new experimental results are listed in “Summary and Revision”.
Please let us know if our revised manuscript and replies have addressed your concerns. If you have any additional comments, we would be more than happy to follow up.
Thank you very much!
Thank you to all our reviewers. In response to your questions and suggestions, we have made several changes which we feel strengthen our paper. We have uploaded a revised draft with changes marked in blue. Our changes include:
* Added a standard heat simulation task used in [1] for both offline learning and active learning.
* Added three baselines (Masked Autoregressive Flow [2], RNN with variational dropout, DMFAL [1]) to compare with our proposed STNP for offline learning.
* Added Sequential Neural Likelihood baseline [3] to compare with our proposed LIG in active learning setup.
* Modified the related work section to include the likelihood-free inference works.
[1] Li, S., Kirby, R. M., & Zhe, S. (2020). Deep Multi-Fidelity Active Learning of High-dimensional Outputs.
[2] Papamakarios, G., Pavlakou, T., & Murray, I. (2017). Masked autoregressive flow for density estimation.
[3] Papamakarios, G., Sterratt, D., & Murray, I. (2019). Sequential neural likelihood: Fast likelihood-free inference with autoregressive flows.
We summarize the new experimental results below in three tables. The results show that our proposed STNP and LIG are robust and have the best performance among all of the tasks. For details about model implementation and results, please refer to Appendix B.6, C.1, C.2 in our revised manuscript.
__Table 2: Offline learning performance comparison using MAE on Reaction-Diffusion and Heat simulator__
| Model              | NP               | SNP            | MAF             | RNN             | DMFAL          | STNP             |
|--------------------|------------------|----------------|-----------------|-----------------|----------------|------------------|
| Reaction-Diffusion | 3.37 ± 0.18      | 3.11 ± 0.07    | 7.22 ± 0.66     | 3.44 ± 0.07     | 4.1 ± 0.02     | 2.84 ± 0.17      |
| Heat               | 1.05e-2 ± 1.1e-3 | 9.62e-3 ± 1e-3 | 2.1e-2 ± 4.8e-3 | 3.2e-2 ± 1.7e-3 | 1.35e-2 ± 3e-4 | 7.36e-3 ± 6.7e-4 |
__Table 3: Active learning performance comparison using MAE on SEIR simulator (LIG vs SNL)__
| Percentage of samples | LIG           | SNL          |
|-----------------------|---------------|--------------|
| 1.11%                 | 365.87±142.87 | 707.61±44.42 |
| 1.85%                 | 236.9±50.6    | 669.03±73.19 |
| 2.96%                 | 119.26±14.22  | 668.67±72.42 |
| 4.07%                 | 96.73±17.07   | 685.28±53    |
__Table 4: Active learning performance comparison using MAE on Heat simulator__
| Percentage of samples | LIG              | EIG              | Random           | MeanSTD          | MaxEntropy       |
|-----------------------|------------------|------------------|------------------|------------------|------------------|
| 5.21%                 | 1.55e-2 ± 1.9e-3 | 1.64e-2 ± 1.9e-3 | 1.74e-2 ± 2.2e-3 | 1.77e-2 ± 2e-3   | 1.83e-2 ± 1.1e-3 |
| 7.81%                 | 1.33e-2 ± 1.7e-3 | 1.56e-2 ± 1.1e-3 | 1.49e-2 ± 3.3e-3 | 1.60e-2 ± 3.7e-3 | 1.58e-2 ± 4e-4   |
| 10.42%                | 1.02e-2 ± 3.4e-3 | 1.38e-2 ± 1.6e-3 | 1.23e-2 ± 1.5e-3 | 1.30e-2 ± 2.4e-3 | 1.61e-2 ± 8e-4   |
| 13.02%                | 9.3e-3 ± 3.1e-3  | 1.05e-2 ± 2.2e-3 | 1.14e-2 ± 1.3e-3 | 1.27e-2 ± 2.5e-3 | 1.46e-2 ± 6e-4   |
| 15.62%                | 6.7e-3 ± 5e-4    | 1.08e-2 ± 1.8e-3 | 1.05e-2 ± 9e-4   | 1.17e-2 ± 1.2e-3 | 1.43e-2 ± 2e-4   |

Meta-review:
role: PC
data: metareview:_summary,_strengths_and_weaknesses: Summary:

The authors present a full-stack solution to accelerating learning, estimation and inference in complex simulators. They present a novel neural process architecture with spatial and temporal dependencies. They also present a new Bayesian active learning acquisition function that is particularly elegant in the context of STNP. Some theoretical results are also presented. The method is benchmarked with three examples (two epidemiological models and one reaction model), and their model appears to perform well.

Strengths:

- Overall this paper is clear to read, well written. 

- The paper uses a known and well-founded variant of the Neural Process as the surrogate model; they provide a delta of the model for better capturing spatiotemporal dependencies. 

Weaknesses:

- I appreciate the effort to propose the spatio-temporal extension to the original NP as well as the latent extension to EIG, but I am not convinced that this is sufficiently novel for ICLR standard.

- The experimental part is also limited. I would expect more ablation study so that we can fully understand the sensitivity and robustness of the proposed approach.

Some implementation details are not provided it seems. How does the active learning get actually implemented? When a new set of data is added, does the INP get retrained from scratch, or fine-tuning from the existing weights?

- It is not clear to me that minimising the uncertainty in a latent space is the right way to go. I invite the authors to restate the objectives of the LIG, how it works, and the weaknesses of it (in light of my comments above).

- the recent work on Synthetic Neural Likelihoods was not compared to

-  the empirical validation is not especially strong.

- The provided experimental comparisons are limited; I would like to see comparison to more deep learning alternatives

Recommendation:

This is a paper with mixed reviews. Two reviewers vote slightly towards acceptance while 2 vote strongly for rejection. I, therefore, recommend to reject the paper. I encourage the authors to use the feedback provided to improve the paper and resubmit to another venue.

justification_for_why_not_higher_score: Two reviewers vote strongly towards rejection.
justification_for_why_not_lower_score: N/A

Decision: Reject