Accelerating Stochastic Simulation with Interactive Neural Processes

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

Reviews and author responses:

 Dear Authors,

Thank you for you detailed followup.  I have already raised my score, and so while I believe there is a good conference submission in this line of work in the future, I still do not believe it is not in that place right now, and so I cannot further upgrade my score to anything engendering acceptance.  

I implore you to further tease apart the methodology and test the components individually on standard problems -- even if it is just to showcase them and help the reader understand them further.  

Good luck.  XP3S.

 Thank you for addressing my comments. I remain positive about the manuscript. Given the other reviews and the authors' rebuttal, I will maintain my rating.

 Thank you! We very much appreciate your encouraging feedback! We will expand the discussion about the limitations of our method and potential application domains. For limitations, vanilla NPs are susceptible to underfitting (Kim et al., 2019) and it could also be the case with STNP, which belongs to the NP family. For potential application domains, we'll discuss how to use our framework for stochastic
processes in finance (Lamberton & Lapeyre, 2007), chemistry (Gillespie, 2007), and meteorology (Vaughan et al., 2021). 

Please find below our answers to your comments.

> I wasn't suggesting you benchmark against Vaughan et al, but rather, apply STNP to the sort of data that Vaughan discusses, eg meteorological data.

We agree that meteorological data prediction could be an appropriate application, but to apply STNP, we need to change the model architecture to convolutional-NP for regular grid-based meteorological data. It is possible thanks to the modular design of our INP framework, which is another key advantage of our framework to be applied in various application domains.

> I push back on your claim that that you are not learning a function approximation ...

There are some misunderstandings of the NP family. As stated in (Garnelo et al., 2018), "NPs define distributions over functions... and can estimate the uncertainty in their predictions." Hence we are learning a distribution over functions, instead of the function itself. NPs share the fundamental property of Gaussian Processes (GPs) to learn to model distributions over functions with the global latent variable $z$. In STNP, we use the temporal latent process $(z_{1},...,z_{T})$ to substitute the global latent variable $z$ and inherit this property. We agree that variational inference does come with the potential issue of expressivity and poor parameterization. But this is an open question in ML that is beyond the scope of this paper. We are happy to include more discussion in the updated version.

> What happens to the predictions on outcomes if the other models learned in Figure-5-right are used? By what fraction do we over/underestimate casualties?

We will include the analysis in our paper. The general finding is that NPs fit the simulation dynamics better than GPs and the sequential neural likelihood (SNL) method for mean prediction. Also, NPs have closer confidence intervals to the truth, reflecting the simulator’s intrinsic uncertainty.  For NP models with the same architecture but different acquisition functions [Figure 5 right], the predictions look similar once they converge. But they converge at different rates.

>  Can this approach be used for estimating interventions? Does the increased accuracy under a fixed sample budget lead to more appropriate interventions?

While our approach cannot directly estimate interventions, it can evaluate "appropriate" intervention proposals. We use the learned amortized surrogate model as a situational awareness and outbreak analytics tool, which can be used by public health officials to assess the (potential) impact of past/future policies on the evolution of an epidemic. We appreciate these two raised questions as they would lead to two extensions of our method, and leave them to future work.


Reference

* Attentive Neural Processes [Kim et al., 2019].
* Introduction to stochastic calculus applied to finance [Lamberton & Lapeyre, 2007].
* Stochastic simulation of chemical kinetics [Gillespie, 2007].
* Convolutional conditional neural processes for local climate downscaling [Vaughan et al., 2021].
* Neural Processes [Garnelo et al., 2018].


If you have any additional comments, we are very eager to address them. Thank you very much!







 We would appreciate if the reviewer kindly let us know whether there were any further questions in the very limited time remaining. We are eager to do our best to address them!

 Thank you for the question. To be clear, we did not mean the knowledge of the Jacobian of the simulator.  We meant the Jacobian of the lower bound of the mutual information.  As stated in 2.1 the original paper [1],  "the computation of the Jacobian, or the Jacobian-gradient
product, is the main technical obstacle." The paper hence proposed to sample from the data distribution and then transform the samples through the "sampling path", which can be expensive.   We have updated the wording accordingly.

1. http://proceedings.mlr.press/v119/kleinegesse20a/kleinegesse20a.pdf

This paper tackles learning a flexible distribution approximator to approximate the output of a high-dimensional and computationally intensive stochastic simulator.  The contributions of this paper then reduce into three main components: 

1 - Definition of a spatiotemporal neural process that can succinctly model a more richly structured latent process.

2 - Definition of a new acquisition function and using a neural process in an active learning setting.

3 - Application to epidemiological simulators to “compile” estimation.

The method appears to work, validated on a toy-ish SEIR model, and a more sophisticated pre-existing epidemiological model.  
 ### Summary++:

This is a rangy paper that presents a lot of material.  The primary application domain the authors have in mind (expediting estimation in epidemiological simulators) is well defined and has great utility to the field.  

However, the paper itself is cluttered/jumbled, and feels very rushed.  The clarity of communication is generally well below what I would expect for publication.  There are also a lot of separate ideas in this paper that are intermingled.  In reality, I think there is probably two and a half papers worth of material (if it were to be properly and fully explored and evaluated) that is only very briefly touched on here.  As a result, it feels like a bit of a jumbled mess with a lot of content that isn’t very well explained or analysed.  It is difficult to tease apart whether the performance gains come from the STNP, from the active learning, from the acquisition function, or simply as a result of hyperparameter tuning or particular applicability to this domain.  Beyond this, I feel like the lack of analysis of the fail-cases of this method leave it vulnerable to expected failures when deployed.  

As a machine learning paper, these components are insufficiently empirically analysed and contrasted to other SotA methods.  On the more application-specific flipside, if you are suggesting this as a method for expediting analysis of epidemiological models, then the implications and downstream utility of this method are under-explored to be a viable applications paper.  


### Major Comments:

1 - As intimated above, I think there are too many disparate elements in this paper.  Crucially, these elements are untested independently.  I would like to see ablation experiments on simpler and more widely studied problems to help understand the limitations of each component (i.e. apply STNP in isolation on some meteorological data; i.e. Vaughan et al., [2021]).  The epidemiological traces (especially in SEIR) are quite simple, and so I am not convinced that the STNP has been sufficiently validated.  The ablation studies in Figure 4 are very task-specific and I have very little intuition for how well the active component is _actually_ working.  

2 - Beyond this, I think at some point in this paper the core message gets lost a little bit.  At its core, you are making a function approximator mapping from parameters to outcomes in an expensive simulator.  That is the “acceleration” discussed, and indeed the utility of the method -- being able to remove an expensive simulator from the loop.  Ultimately, it is the performance of the learned function approximator that is paramount.  The specifics of how you learn this approximator are somewhat tangential to this.  You could just learn a neural network mapping from inputs to outputs.  There is therefore, in my opinion, too little analysis of how accurate the predictions are and the limitations of this approach, especially given that the examples are quite simple.  On more complex simulators, does the NP just fail to capture complexity or become too difficult to learn?  At which point the predictions become unreliable, even spurious.  There is no evaluation or discussion of this, which I find troubling, as there is no mention that the method could silently fail and produce poor predictions (or worse, asymptotically biased predictions, unlike in, i.e., Baydin et al., [2019]).  

2.1 - The active learning element gets far too much attention.  We sort-of-know that well-tuned active learning methodologies outperform non-active methods in terms of sample complexity, especially where the “true” model is within the support of the model being studied.  

2.2 - I would also like some comparison of runtimes.  How long does the NP take to train?  How long does inference in the NP take?  How long does the optimisation in the active learning part take?  How long does running the simulator itself take?  These are important quantities to know if the main selling point is “accelerating”. 

3 - Following on from this, I find the “interactive” part of the title challenging (and further evidence that this paper is a bit jumbled, both on paper and in the mind of the authors).  There is nothing interactive about the neural process -- it is just a (novelly structured, yes) neural process.  The training method has a degree of interaction, but that is not to do with the neural process.  Just as many models can be learned _using_ interactive methods, that does not make each of those models interactive.  

4 - The theory section is just dropped in with no real justification or use.  I feel the authors have just dropped this in to make the paper feel more legitimate, while adding very little to the paper.  
 

### Minor comments:

A - You should also reference these very relevant papers: 
- _Deep probabilistic surrogate networks for universal simulator approximation_ [Munk et al., 2019].

- _Likelihood-free inference with emulator networks_ [Lueckmann et al., 2019].

- _Planning as inference in epidemiological models_ [Wood et al., 2020].

- _Simulation-Based Inference for Global Health Decisions_ [de Witt et al., 2020].

- _Etalumis: Bringing Probabilistic Programming to Scientific Simulators at Scale_ [Baydin et al., 2019].

- _Convolutional conditional neural processes for local climate downscaling_ [Vaughan et al., 2021].

- I would also like to see some (qualitative) comparison to work like _Semi-supervised Sequential Generative Models_ [Teng et al., 2020]; or more generally the wealth of VRNN work (see Chung et al., [2015] and follow-ups) for alternatives for compiled time-series modelling.  

B - Only capitalise proper nouns: i.e. “Interactive Neural Processes” -> “interactive neural processes”.

C - The text on the figures is _tiny_.

D - The last two sentences of the paper really concern me: “We plan to leverage Monte Carlo sampling and Bayesian optimization techniques to parameterize the reward function. Then we would also be able to directly optimize for the target parameters with auto-differentiation.”  These sentences make such little sense to me, that they make me lose a little bit of confidence in what is written in the rest of the document. 

E - The first two paragraphs of Section 3 are incomprehensible.
 This work is interesting, but is ultimately not ready for publication yet.  The raw quality of the manuscript is well below publication quality.  The components are presented in a bit of a jumbled way.  The methods presented seem to work, but it is impossible to identify where the performance gains come from, and the wider applicability of this method, or each of the components in isolation.  I implore the authors to tease apart the different threads presented in this work and explore them more thoroughly, completely and independently across several intermediate papers*.  

Good luck!   

(*and _then_ you can write a huge journal paper combining them all and get a big, sexy Nature paper!!)


 Dear Authors,

Thank you for your detailed response, and conducting more experiments and making changes at the request of the reviewers.

Some minor responses to your comments: 

- I wasn't suggesting you benchmark against Vaughan et al, but rather, apply STNP to the sort of data that Vaughan discusses, eg meteorological data.  

- I push back on your claim that that you are _not_ learning a function approximation.  You are learning a "function" that outputs the approximation of the true posterior.  I am generally okay with this (heck, that is what most of V.I. is), but there is too little discussion on the silent implications of this for my liking.  What if the variational family is no way near expressive enough?  What if poor parameters are learned?  This is different to, say, learning an amortized proposal that is then used to expedite inference and asymptotic unbiasedness is retained.  While I don't think this omission is terminal for the work, I think it is important to include thorough discussion of it. 

- I understand that getting a lower error in fewer time steps is often the absolute benchmark for success at ML conferences, but here it still feels a little unsatisfactory.  You are mooting that (I-)STNP allows for more expedient analysis of evolving epidemiological (or more generally, S-T scenarios), but there is very little discussion of how (I-)STNP actually allows this.  I.e. what happens to the predictions on outcomes if the other models learned in Figure-5-right are used?  By what fraction do we over/underestimate casualties?  Can this approach be used for estimating interventions?  Does the increased accuracy under a fixed sample budget lead to more appropriate interventions?  Again, these are not grounds for rejection, but I am less-than-convinced by the wider application (/epidemiological) merit of the framework.  

Furthermore, I still stand by the core of my critique: There are too many elements presented in a bit of a jumbled mess.  I do not believe the individual ideas presented in this paper have been thoroughly and independently examined enough for me to have absolute confidence in the elements and their interaction.  I do now understand that you are suggesting this as a "framework"; but then I have queries about the generality and applicability of this framework more widely, given that you have only tested it on a low-dimension pedagogical example and a highly domain specific application for which I have very little feel for.  

However, your comments have swayed me enough that I will no longer "fight" against acceptance.  I will upgrade by score to a weak reject, conditioned on the authors adding some discussion of how this approach may silently fail and other potential application domains.  

Good luck,
Reviewer XP3S
 



 Dear reviewer,

We would like to thank you once again for your thoughtful reviews and comments. 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!

 Dear reviewer,

We would like to thank you once again for your thoughtful reviews and comments. 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!

 Dear reviewer,

We would like to thank you once again for your thoughtful reviews and comments. 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!

 Dear reviewer,

We would like to thank you once again for your thoughtful reviews and comments. 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!

 We would like to thank all of the reviewers for their time and thoughtful reviews and comments. We have posted a revised manuscript to address the points raised by everyone, and posted specific responses to individual reviews. The changes made in the revised manuscript are indicated in blue.

Here is the revision summary list:
* implemented the Sequential Neural Likelihood (SNL) model proposed by Papamakarios et al. 2019 for the SEIR experiment. The result can be found in Table 5 and Figure 7 [see Appendix C.1]. The implementation detail can be found in Appendix B.4.
* revised the “Kleinegesse & Gutmann (2020)" statement and included the related work reviewers mentioned in section 2.
* added a detailed discussion of the broader impact on the application side in Appendix B.6.
* added the caption and description to Figure 2.
* added the definition of context and target points in section 3.2.
* added the definition of augmented parameters in section 3.3.
* added the dimensionality of SEIR in section 4.1.
* added the dimensionality of LEAM-US in Appendix B.3.



 Thank you very much for your feedback which has helped us improve the paper! Please find below our answers to your comments. The changes made in the revised manuscript are indicated in blue.

> the authors write Kleinegesse & Gutmann (2020)...

We have revised this statement in the manuscript to “Kleinegesse & Gutmann (2020) consider implicit models where the likelihood function is intractable, but requires computable Jacobian and are limited to low (<=10) dimensional design.”

> related to the point above, there is a large body of work in machine learning…

We appreciate the reviewer for providing the reference. We have included them in the revised manuscript [see Section 2].

> I believe the study would strengthen substantially if the method would be compared both in terms of inference quality and computational efficiency with other state-of-the-art likelihood-free inference methods …

We have implemented the Sequential Neural Likelihood (SNL) model proposed by Papamakarios et al. 2019 for the SEIR experiment. The result can be found in Table 5 and Figure 7 [see Appendix C.1]. The implementation detail can be found in Appendix B.4.

Table 5:

| 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.00 |

As you pointed out, neural density approaches (e.g. SNL) for learning the likelihood requires a relatively larger number of simulations compared to deep learning surrogate models. We find our best interactive framework setup LIG in INP is always better than SNL as our framework only requires a small number of simulations to converge to the offline performance. 

We emphasize that our framework is designed for real-world complex simulations (stochastic and high-dimensional) that are expensive computationally. Here, the SEIR experiment is a pedagogical example to show our framework’s high sample efficiency. For complex real-world simulations, the training time and inference time of deep learning surrogate model can be ignored compared with simulation time. In this case, the sample efficiency is equivalent to the computational efficiency.

> ...the dimensionality of the complex spatiotemporal LEAM-US problem…

The dimensionality of LEAM-US can be found in Appendix B.3.

 > 2. At its core, you are making a function approximator mapping... You could just learn a neural network mapping from inputs to outputs.

That’s not true. First, the model should be probabilistic as a surrogate model for stochastic simulators. NP family meets the requirement by defining the distributions over functions. Secondly, the surrogate model is required to approximate $p(x|\theta)$ globally. NP family can approximate the likelihood for any $x$ so it can generalize better. Thirdly, NP should be capable of rapid adaptation to new observations, which is required for active learning setup. Straightforward input-output mapping will suffer from overfitting, and the uncertainty of the model cannot be used to guide active learning.

> There is therefore, in my opinion, too little analysis of how accurate the predictions are and the limitations of this approach, especially given that the examples are quite simple …

Our experiment results are consistent for both simple (SEIR) and complex (LEAM-US) simulators, see Figure 3 and Figure 5.

> 2.1 The active learning element gets far too much attention…

Learning the input-output map with a deep learning surrogate model is not enough because our simulator is **stochastic** and **high-dimensional**.  It is infeasible to simulate all possible scenarios a priori. Therefore, performing active learning to “select” training data is the key. 

> 2.2 I would also like some comparison of runtimes …

For the LEAM-US simulator, a single simulation requires about 1 hour to be completed (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. 

The computational cost of INP (training  + inference + optimization) per iteration is about 0.417 GPU hours (on a GeForce RTX 2080). Therefore, the computational cost of INP is negligible compared to the cost of simulating training data. In this case, the sample efficiency is equivalent to the computational efficiency.

> 3. Following on from this, I find the “interactive” part of the title challenging …

We did not mean the neural processes model itself is interactive. In fact, the Interactive Neural Processes is the name of our designed framework, which is shown in Figure 1. We are happy to change the wording and make it clear.


 Thank you very much for your feedback which has helped us improve the paper! Please find below our answers to your comments. The changes made in the revised manuscript are indicated in blue.

> The additional condition that must be met for a strong recommendation is that either..., but I personally do not see its merits here as I will explain below ...

Our paper is innovative in both method and application. 

For method, we have performed sufficient experiments to compare our methods to SoTA machine learning benchmarks. Specifically, we utilized the SoTA Bayesian learning model Neural Processes and extended it to STNP. In Table 1, we did an ablation study to compare our STNP with SoTA baselines (NP, SNP) in offline learning setup. Under the active learning setup, we compared the neural processes family to the gaussian processes family (Table 2, Table 3, Figure 3). Moreover, we compared our proposed acquisition function LIG to existing baselines for regression (Table 2, Table 3, Table 4, Figure 3, Figure 5). According to Reviewer Epnf's comment, we have also included an experiment in the revised paper to compare our best interactive framework setup (INP+LIG) with SoTA likelihood-free inference method based on neural density estimation (Papamakarios et al. 2019) (Table 5, Figure 7).

For application, we thank the reviewer for their comment that indeed stresses the differences between our approach and the existing literature. Indeed, our goal is to be able to provide a **general surrogate model** that is not tailored/optimized to solve only for a unique “policy objective”/target outcome. The differential advantage of a model of the complexity of LEAM-US and of similar mechanistic epidemic models vis a vis (potentially black-box) statistical/machine learning models is that while the latter typically aim at predicting a specific target (e.g. number of future cases/deaths/hospitalizations) without describing the underlying state of the system.

Mechanistic models provide a complete situational awareness tool to policymakers since they inherently encode expert knowledge about the state of the system that the other models might not be able to capture due to lack of relevant data. For example, a mechanistic model calibrated to provide realistic forecasts on deaths/cases/hospitalizations can also provide insights on other unobservables such as the fraction of immune population in a certain location/age group, which might help policymakers assess the impact of past or potential impact of future policies (e.g. introducing childhood vaccination, introducing vaccine boosters, etc..). Another example is shown in Davis, J.T., Chinazzi, M., Perra, N. et al. Nature (2021) where the cryptic transmission of SARS-CoV-2 and the first COVID-19 wave has been studied, providing a mechanistic understanding of the early dispersal of infections and the temporal windows of the introduction and onset of SARS-CoV-2 local transmission in Europe and the United States. To provide those kinds of analysis, it is necessary to build an emulator that acts as a “surrogate for the entire mechanistic simulator” rather than a reduced version of it.


 Thank you very much for your feedback which has helped us improve the paper! Please find below our answers to your comments. The changes made in the revised manuscript are indicated in blue.

1. we have added the caption and description to Figure 2. $h_t$ and $H_t$ correspond to the encoder and decoder hidden state. $A$ is the adjacency matrix for the spatial graph, which is omitted in the figure for simplicity.

2. we used $z_{1:T}$ to capture the dynamics of the simulator. We implement it with RNN architecture. The details of the implementation are deferred to Appendix B.3. 

3. ${\hat{x}_1, \hat{x}_2, ..., \hat{x}_T}$ are the future states w.r.t day 0 based on the input ${\theta_1, …, \theta_T, x_0}$. 

4. We have added the definition of context and target points in section 3.2. Augmented parameters are the candidate scenarios.

5. The key difference of our model is how the latent temporal dynamics $z_t$ is modeled, as specified in the updated Figure 2 caption.   The distribution of $z_t$ is defined differently from SNP, implemented by state-level encoder and decoder, as shown in Figure 6 in Appendix B.3. The original NP/SNP does not have a spatial component. But for a fair comparison, we have adapted both models to also take in the spatial part. 

 Thank you very much for your feedback which has helped us improve the paper! Please find below our answers to your comments. The changes made in the revised manuscript are indicated in blue.

> 1. There are also a lot of separate ideas in this paper that are intermingled… 
> As a machine learning paper, these components are insufficiently empirically analysed and contrasted to other SotA methods
> I think there are too many disparate elements in this paper … understand the limitations of each component.

We respectfully disagree with the reviewer’s comment. Instead of combining multiple separate ideas together, we integrate novel ideas (active learning, neural processes, and deep generative models) to design a Bayesian active learning framework. It is able to proactively learn a deep learning surrogate model and accelerate stochastic simulation. 

Furthermore, we have performed careful ablative studies for each component.  Specifically, we utilized the SoTA Bayesian learning model Neural Processes and extended it to STNP. In Table 1, we did an ablation study to compare our STNP with SoTA baselines (NP, SNP) in offline learning setup. Under the active learning setup, we compared the neural processes family to the gaussian processes family (Table 2, Table 3, Figure 3). Moreover, we compared our proposed acquisition function LIG to existing baselines for regression (Table 2, Table 3, Table 4, Figure 3, Figure 5). According to Reviewer Epnf's comment, we have also included an experiment in the revised paper to compare our best interactive framework setup (INP+LIG) with SoTA likelihood-free inference method based on neural density estimation (Papamakarios et al. 2019) (Table 5, Figure 7).

All the listed contributions are well supported by the experimental results and theory. In fact, as pointed out by Reviewer Epnf, our paper is “technically sound, and claims are for the most part well supported by empirical evaluation”. As pointed out by Reviewer TEvq, “The experimental results show that INP works better to accelerate stochastic simulation than GP, and the proposed LIG leads to faster convergence compared with alternatives.” As pointed out by Reviewer YBHd, “the method is developed carefully with regard to implementation and calibration strategy (i.e., design of the model architecture, identification of objective function and construction of a suitable acquisition function). These choices are supported by appropriate theoretical analyses of bounds and convergence rates.”

> apply STNP in isolation on some meteorological data; i.e. Vaughan et al., [2021]

Vaughan et al., [2021] is not directly comparable to our setup. Their convolutional conditional neural processes are regular **grid-based**, whereas our STNP is **graph-based**.

> On the more application-specific flipside, ... utility of this method are under-explored to be a viable applications paper.

On the application side, we have added a detailed discussion of the broader impact in Appendix B.6.

>The epidemiological traces (especially in SEIR) are quite simple ... is actually working.

There are some misunderstandings of our experiments. We use the SEIR experiment as a pedagogical example, only to help illustrate the idea. Due to its simplicity, it abstracts a wider range of dynamics, hence is not task-specific. Also, as the SEIR simulator does not include spatial information, we used the vanilla NP model instead of STNP (see Section 4.1 Model  Implementation). 

We include this experiment to visualize the acquisition functions behavior in 2-D input space (Figure 4). It is informative as it shows our proposed LIG is able to balance exploration and exploitation. 

For the LEAM-US experiment, it is much more sophisticated. The dimension of the input space is increased to 16,264. The dimension of the output space is increased to 672 (see Appendix B.3). We found our proposed LIG works the best using our designed INP framework for both simple and complex simulators, see Figure 3 and Figure 5. To validate the performance of STNP, the comparison results in the offline setup can be seen in Table 1.  


 Thank you for addressing my comments.

I am not sure I agree with the statement: 

> We have revised this statement in the manuscript to “Kleinegesse & Gutmann (2020) consider implicit models where the likelihood function is intractable, but requires computable Jacobian and are limited to low (<=10) dimensional design.”

Maybe I am misinterpreting the statement by the authors, but Kleinegesse & Gutmann (2020) consider models for which one can only simulate: the method developed by Kleinegesse & Gutmann (2020) does not require the knowledge of the Jacobian of the simulator.

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



The paper is technically sound, and claims are for the most part well supported by empirical evaluation. However, I have a concern I would appreciate the authors to address:

-I believe the study would strengthen substantially if the method would be compared both in terms of inference quality and computational efficiency with other state-of-the-art likelihood-free inference methods, such as the ones based on neural density estimation (e.g. Papamakarios et al. 2019, Lueckmann et al. 2018). Indeed, if the simulators are not very expensive computationally, then maybe neural density approaches for learning the likelihood (Papamakarios et al. 2019; Lueckmann et al. 2018), which tend to need a large(r) number of simulations, will be competitive with the proposed approach in learning a surrogate of the simulator.



The paper is generally clear, and I would just like to point a few typos and imprecisions:

-capital letter missing in "we prove the equivalence between";

-in "The total number of dimension", it should be "dimensions". There are a few other instances of "dimension" where this should also be corrected;

-in "which is scales linearly in computation", it should be "which scales linearly in computation";

-"where it first randomly selects a set of groups and then use(s)";

-one can guess the dimensionality of the problem for the SEIR problem (if I am not mistaken, 2 parameters and 100-dimensional output). However, it would be important to explicitly state the dimensionality of the complex spatiotemporal LEAM-US problem (number of parameters of the simulator and dimensionality of the output of the simulator). Overall, the paper is technically sound and has a novel methodological contribution with convincing empirical results. This will inspire further method development. Some of the discussion of previous work needs to be revised and augmented, and it would be good to provide empirical comparisons with state-of-the-art simulation-based inference methods (in particular, ones that learn the likelihood function).

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

There needs some significant improvement in the illustration of the model:
1. Figure 2 is confusing. Some elements in the figure are not defined in the text, e.g., h_t, H_t. Basically, the graphs in all these figures don't have their matching text. The spatial component is not shown in the third column, no A. Basically, I got lost reading the modeling part and the figure 2. 
2. z_{1:T} is used to capture the dynamics, but how? Is there any dynamics relationship used to define z_{1:T}?
3. In the first paragraph of 3.1, what is "future states"? Not defined. If they refer to "{\hat{x}_1, \hat{x}_2, ..., \hat{x}_T}", there is no "future". 
4. The definition of context {theta, x} and target {theta, x} is unclear. "Augmented parameter theta" is mentioned a couple of times. Is it target parameter? What is "augmented"?
5. In the ablation study, it's claimed that the temporal latent process is the key innovation. But SNP also has the latent dynamics z. So based on the analysis, the spatial part should be more significant. Unless the temporal component is encoded in different ways, which might require a bit more extensive ablation studies to analyze. 

Minor:
1. In the second paragraph of 4.1, R is not defined. 
2. In page 9, "which is scales linearly in computation" -> "which scales linearly in computation" I think the model is novel and useful for disease forecasting. But I think the paper (especially the methodology part) is poorly written. It doesn't reach the bar for publication. 

The manuscript entitled, "Accelerating Stochastic Simulation with Interactive Neural Processes", presents a novel approach to the problem of statistical emulation for mechanistic models of epidemic disease transmission.  To this end, structured neural processes are developed to exploit and respect the temporal and spatio-temporal character of these models.  An active learning strategy is developed to train these neural processes to minimise the computational costs of generating training instances of disease simulator outputs.  The methodology developed is applied to two SEIR compartmental model examples: a minimal one with a single homogenous population and a maximal one with many age and space delimited cohorts. A strength of this paper is that the method is developed carefully with regard to implementation and calibration strategy (i.e., design of the model architecture, identification of objective function and construction of a suitable acquisition function).  These choices are supported by appropriate theoretical analyses of bounds and convergence rates.  Considering the high benchmark for papers submitted to ICLR, it is my opinion that (while highly commendable) reaching this level of quality is only a necessary but not sufficient condition to strongly recommend a paper for the conference.  The additional condition that must be met for a strong recommendation is that either the method itself is highly novel across other competing machine learning methods or the contribution to solving a problem in the applied area is highly substantive.  In this case I believe the paper is presenting an argument for itself in the latter category, but I personally do not see its merits here as I will explain below.

Mechanistic models of disease transmission can be incredibly computationally expensive to run with high overheads for both RAM and cpu time.  Even in the best case scenario that many of the parameters can be fixed leaving only one or a few free parameters to be learned against data, the calibration problem is computationally expensive.  In many cases an additional problem will be the lack of a tractable likelihood function and/or lack of control access to latent system random variables, such that the simulator needs to be treated as a black box.  Once calibrated, the analyst will typical want to generate model outputs under a range of 'tweaked' parameterisations representing scenarios such as possible interventions or future trajectories under a range of possible changes to human behaviour parameters or so on.

Machine learning methods have been used to speed up both these aspects of disease modelling in the past and broadly speaking all have followed the approach of Bayesian optimisation.  Namely, iteratively learning a mapping between the model parameter space and a univariate statistic being either the goodness-of-fit to observed data (Andrianakis et al., PLoS Comp Bio, 2015; Reiker et al., 2021 preprint: https://doi.org/10.1101/2021.01.27.21250484 ) or a utility function for policy impact (Bent et al. 2017: https://arxiv.org/abs/1712.00428 ).  Rarely has the problem been tackled from the perspective of building an emulator as a surrogate for the entire mechanistic simulator itself (though see for example Cameron et al. Nat Comms 2015 in which a functional regression model is used to predict new simulator outputs from a training library).  The reason this latter approach is less explored is because the general input to output space is vastly larger than the sub-domain required for calibration to a specific dataset or specific set of policy objectives.  If the application of the surrogate model is to address one of these problems then the Bayesian optimisation methods focussed on that single objective will inevitably out-perform a general surrogate model.

Since the method proposed here falls into the latter category of producing a general surrogate I am not convinced that the applications to which it has been explored here are sufficient to demonstrate its utility.  Application areas in which a general surrogate model might yield a particular advantage could be e.g. in transfer learning scenarios in which the calibration exercise must be performed over many different observation datasets, or perhaps as an intermediate model for a multi-resolution Bayesian optimisation (Kandasamy et al. ICML 2017). In brief, my opinion is that, while the methods are well developed, the utility of this method for solving problems in the application domain is not sufficient for a strong recommendation.

Meta-review: This paper provides some novel ideas combining neural processes, active learning, and deep sequence models toward accelerating the computationally intensive task of stochastic simulations for epidemic models. The authors incorporate aspects such as spatiotemporal dependence and age structure in the models they consider, and propose an active learning framework to leverage a neural process that can serve as a proxy for direct simulation of the stochastic dynamics. Most of the reviewers agreed that some of these ideas are novel, but the assessments together present a borderline case, and one reviewer remains convinced that the paper needs refinement and a more focused exposition to make clear the contributions and relative efficacy compared to existing state of the art. I tend to agree with some of the concerns raised by that reviewer, who responded favorably to the author response but remained not fully convinced.I also agree that important parts of the paper feel rushed, if I am to interpret "feeling rushed" as a statement on the ideas being spread thin in exposition due to the amount of ground the authors try to cover (rather than a statement on the quality of the presentation). There are many moving parts that combine good intuitions toward accelerating simulation-based methods for stochastic epidemic models. Several reviewers mentioned improving the empirical comparisons with recent existing methods, including likelihood-based methods, and though I agree with the authors too that it is not possible to exhaustively explore outcomes in stochastic process approaches, the largely empirically supported contributions can be better motivated with a more complete comparison and streamlined exposition. I encourage the authors to continue to revise and refine this manuscript to maximize the potential behind the ideas they propose here.

Decision: Reject