Accelerating Stochastic Simulation with Spatiotemporal Neural Processes

Stochastic simulations such as large-scale, spatiotemporal, age-structured epidemic models are computationally expensive at fine-grained resolution. We propose Spatiotemporal Neural Processes (STNP), a neural latent variable model to mimic the spatiotemporal dynamics of stochastic simulators. To further speed up training, we use a Bayesian active learning strategy to proactively query the simulator, gather more data, and continuously improve the model. Our model can automatically infer the latent processes which describe the intrinsic uncertainty of the simulator. This also gives rise to a new acquisition function based on latent information gain. Theoretical analysis demonstrates that our approach reduces sample complexity compared with random sampling in high dimension. Empirically, we demonstrate that 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.