Accelerating Stochastic Simulation with Interactive Neural Processes

1 INTRODUCTION . Computational modeling is now more than ever at the forefront of infectious disease research due to the COVID-19 pandemic . Stochastic simulations play a critical role in understanding and forecasting infectious disease dynamics , creating what-if scenarios , and informing public health policy making ( Cramer et al. , 2021 ) . More broadly , stochastic simulation ( Ripley , 2009 ; Asmussen & Glynn , 2007 ) produces forecasts about complex interactions among people , environment , space , and time given a set of parameters . They provide the numerical tools to simulate probabilistic models and stochastic processes in finance ( Lamberton & Lapeyre , 2007 ) , chemistry ( Gillespie , 2007 ) and many scientific disciplines . Unfortunately , stochastic simulations at fine-grained spatial and temporal resolution are extremely computationally expensive . For example , epidemic models whose aim is to simulate realistic diffusion dynamics via in-silico experiments require the exploration of a large parameter space ( e.g . characteristics of a virus , policy interventions , people ’ s behavior ) . Therefore , hundreds of thousands of simulations are required to explore and calibrate the epidemic model with observed epidemic surveillance data . This process significantly hinders the adaptive capability of existing stochastic simulators , especially in “ war time ” emergencies , due to the lead time needed to execute new simulations and produce actionable insights that could help guide decision makers . Learning deep surrogate models to speed up complex simulation has been explored in climate modeling and fluid dynamics for deterministic dynamics ( Sanchez-Gonzalez et al. , 2020 ; Wang et al. , 2020 ; Holl et al. , 2019 ; Rasp et al. , 2018 ) , but not for stochastic simulations . Furthermore , these surrogate models can only approximate specific system dynamics and fail to generalize under different parametrization . Especially for pandemic scenario planning , we desire models that can predict futuristic scenarios under different conditions . Furthermore , the majority of the surrogate models are trained passively using a simulation data set . This requires a large number of simulations to cover different parameter regimes of the simulator and ensure generalization . We propose Interactive Neural Process ( INP ) , a Bayesian active learning framework to speed up the stochastic simulation . Given parameters such as disease reproduction number , incubation and infectious periods , mobility dynamics , and initial state , current mechanistic simulators generate future outbreak states with time-consuming numerical integration . Our framework accelerates the simulation by learning the input-output map between parameters and future states , hence bypassing numerical integration . We develop a novel neural process ( NP ) ( Garnelo et al. , 2018 ) model to mimic the mechanistic simulator with uncertainty . We extend NP to spatiotemporal setting : we introduce a time-evolving latent process for temporal dynamics and integrate graph convolution for spatial modeling . Instead of learning passively , we design active learning algorithms to interact with the simulator and update our model in “ real-time ” . We derive a new acquisition function , Latent Information Gain ( LIG ) , based on our unique model design . Our algorithm selects the parameters with the highest LIG , queries the simulator to generate new simulation data , and continuously updates our model . Compared with Gaussian processes ( GP ) , INP is more expressive , more accurate , and scales easily to high-dimensional sequence data . We provide theoretical guarantees for the sample efficiency of this procedure over random sampling . We also demonstrate the efficacy of our method on large-scale spatiotemporal epidemic models . In summary , our contributions include : • Interactive Neural Process : a deep Bayesian active learning framework for accelerating large-scale stochastic simulation • A novel spatiotemporal neural process model ( STNP ) for high-dimensional time series data with temporal latent process and spatial graph convolution . • New acquisition function , Latent Information Gain ( LIG ) , based on the inferred temporal latent process to quantify uncertainty with theoretical guarantees • Real-world application to speed up stochastic , spatiotemporal epidemic models . 2 RELATED WORK . Bayesian Active Learning and Experimental Design . Bayesian active learning , or experimental design is well-studied in statistics and machine learning ( Chaloner & Verdinelli , 1995 ; Cohn et al. , 1996 ) . Gaussian Processes ( GPs ) are popular for posterior estimation e.g . Houlsby et al . ( 2011 ) and ( Zimmer et al. , 2018 ) , but often struggle in high dimension . Deep neural networks provide scalable solutions for active learning . Deep active learning has been applied to discrete problems such as image classification ( Gal et al. , 2017 ) and sequence labeling ( Siddhant & Lipton , 2018 ) whereas our task is continuous time series . Our problem can also be viewed as sequential experimental design where we design simulation parameters to obtain the desired outcome ( imitating the simulator ) . Foster et al . ( 2021 ) propose deep design networks for Bayesian experiment design but they require a explicit likelihood model and conditional independence in experiments . Kleinegesse & Gutmann ( 2020 ) consider implicit models where the likelihood function is intractable , but computing the Jacobian through sampling path can be expensive and their experiments are mostly limited to low ( < =10 ) dimensional design . In contrast , our design space is of much higher-dimension and we do not have access to an explicit likelihood model for the simulator . Neural Processes . Neural processes ( NP ) ( Garnelo et al. , 2018 ) model distributions over functions and imbue neural networks with the ability of GPs to estimate uncertainty . NP has many extensions such as attentive NP ( Kim et al. , 2019 ) and functional NP ( Louizos et al. , 2019 ) . However , NP implicitly assumes permutation invariance in the latent variables and can be limiting in modeling temporal dynamics . Singh et al . ( 2019 ) proposes sequential NP by incorporating a temporal transition model into NP . Still , sequential NP assumes the latent variables are independent conditioned on the hidden states . We propose spatiotemporal NP with temporal latent process and spatial graph convolution , which is well-suited for modeling the spatiotemporal dynamics of infectious disease . We apply our model to real-world large-scale Bayesian active learning . Note that even though Garnelo et al . ( 2018 ) has demonstrated NP for Bayesian optimization , it is only for toy 1-D functions . Stochastic Simulation and Dynamics Modeling . Stochastic simulations are fundamental to many scientific fields ( Ripley , 2009 ) , especially epidemic modeling . Data-driven models of infectious diseases are increasingly used to forecast the evolution of an ongoing outbreak ( Arik et al. , 2020 ; Cramer et al. , 2021 ; Lourenco et al. , 2020 ) . However , very few models can mimic the internal mechanism of a stochastic simulator and answer “ what-if questions ” . GPs are commonly used as surrogate models for expensive simulators ( Meeds & Welling , 2014 ; Gutmann et al. , 2016 ; Qian et al. , 2020 ) , but GPs do not scale well to high-dimensional data . Likelihood-free inference Lueckmann et al . ( 2019 ) ; Papamakarios et al . ( 2019 ) ; Munk et al . ( 2019 ) ; Wood et al . ( 2020 ) aim to learn the posterior of the parameters given observed data . They use neural density estimators to learn an emulator , but require a lot of simulations . We propose a novel neural process model to learn the likelihood . For active learning , instead of relying on Monte Carlo sampling , we directly compute the information gain in the latent process . For epidemic modeling , Qian et al . ( 2020 ) only consider the simple SEIR model . In contrast , we demonstrate the use of deep sequence model as a prior distribution in Bayesian active learning for real-world , large-scale , spatiotemporal dynamics . Our framework is also compatible with other deep sequence models for time series , e.g . Deep State Space ( Rangapuram et al. , 2018 ) , Neural ODE ( Chen et al. , 2018 ) and etc . 3 METHODOLOGY . Consider a stochastic process { X1 , · · · , XT } , governed by time-varying parameters θt ∈ RK , and the initial state x0 ∈ RD . In epidemic modeling , θt can represent the effective reproduction number of the virus at a given time , the effective contact rates between individuals belonging to different age groups , the people ’ s degree of short- or long-range mobility , or the effects of time varying policy interventions ( e.g . non-pharmaceutical interventions ) . The state xt ∈ RD includes both the daily prevalence and daily incidence for each compartment of the epidemic model ( e.g . number of people that are infectious and number of new infected individuals at time t ) . Stochastic simulation uses a mechanistic model F ( θ ; ξ ) to simulate the process where the random variable ξ represents the randomness in the simulator . Let θ : = ( x0 , θ1 , · · · , θT ) represent the initial state and all the parameters over time . For each θ , we obtain a different set of simulation data { ( x1 , · · · , xT ) m } Mm=1 . However , realistic large-scale stochastic simulations require the exploration of a large parameter space and are extremely computationally intensive . In the following section , we describe a pipeline called Interactive Neural Process ( INP ) to proactively query the stochastic simulator , generate simulation data , in order to learn a fast surrogate model for rapid simulation . 3.1 INTERACTIVE NEURAL PROCESS . Our high-level goal is to train a deep surrogate model to mimic the stochastic simulator . As show in Figure 1 , given parameters θ , we query the simulator , i.e. , the mechanistic model to obtain a set of simulations { ( x1 , · · · , xT ) m } Mm=1 . We train a neural process ( NP ) model to learn the probabilistic map from parameters to future states . Our NP model is spatiotemporal and it captures the unique disease dynamics of the epidemic simulator . During inference , the model needs to generate predictions ( x̂1 , · · · , x̂T ) at the target parameters θ corresponding to different scenarios . Instead of simulating at a wide range of parameter regimes , we take a Bayesian active learning approach to proactively query the simulator and update the model incrementally . Using NP , we can infer the latent temporal process ( z1 , · · · , zT ) that encodes the uncertainty of the current surrogate model . Then we propose a new reward/acquisition function , latent information gain ( LIG ) , to select the θ ? with the highest reward . We use θ ? to query the simulator , and in turn generate new simulation to further improve the model . Next , we describe each of the components in detail . 3.2 SPATIOTEMPORAL NEURAL PROCESS . Neural process ( NP ) ( Garnelo et al. , 2018 ) is a type of deep generative model that represents distributions over functions . It approximates a stochastic process by combining Gaussian process with deep learning . NP introduces a global latent variable z to capture the stochasticity and learns the conditional distribution p ( x1 : T |θ ) by optimizing the evidence lower bound ( ELBO ) : log p ( x1 : T |θ ) ≥ Eq ( z|x1 : T , θ ) [ log p ( x1 : T |z , θ ) ] − KL ( q ( z|x1 : T , θ ) ‖p ( z ) ) ( 1 ) Here p ( z ) is the prior distribution for the latent variable . We use x1 : T as a shorthand for ( x1 , · · · , xT ) . The prior distribution p ( z ) is conditioned on a set of context points θc , xc1 : T as p ( z|xc1 : T , θc ) . The set of context points is randomly sampled from the training data . Its complement is the target set . However , the global latent variable z in NP can be limiting for non-stationary , spatiotemporal dynamics in the epidemics . We propose spatiotemporal neural process ( STNP ) with two extensions . First , we introduce a temporal latent process ( z1 , · · · , zT ) to represent the unknown dynamics . The latent process provides an expressive description of the internal mechanism of the stochastic simulator . Each latent variable zt is sampled conditioning on the past history . Second , we explicitly model the spatial dependency in xt ∈ RD . Rather than treating the dimensions in xt as independent features , we capture their correlations with a graph . For instance , the travel graph between locations can be represented as an adjacency matrix A ∈ RD×D . Given parameters { θ } , simulation data { x1 : T } , and the spatial graph A as inputs , STNP models the conditional distribution p ( x1 : T |θ , A ) by optimizing the following ELBO objective : log p ( x1 : T |θ , A ) ≥ Eq ( z1 : T |x1 : T , θ , A ) log p ( x1 : T |z1 : T , θ , A ) − KL ( q ( z1 : T |x1 : T , θ , A ) ‖p ( z1 : T ) ) ( 2 ) where the distributions q ( z1 : T |x1 : T , θ , A ) and p ( x1 : T |z1 : T , θ , A ) are parameterized with neural networks . The prior distribution p ( z1 : T ) is conditioned on a set of contextual sequences p ( z1 : T |xc1 : T , θc , A ) . Figure 2 visualizes the graphical models of our STNP , the vanilla NP ( Garnelo et al. , 2018 ) , and Sequential NP ( Singh et al. , 2019 ) . We implement STNP following an encoder-decoder architecture . The encoder parametrizes the mean and standard deviation of the variational posterior q ( z1 : T |x1 : T , θ , A ) and the decoder approximates the predictive distribution p ( x1 : T |z1 : T , θ , A ) . To incorporate the spatial graph information , we use a Diffusion Convolutional Gated Recurrent Unit ( DCGRU ) layer ( Li et al. , 2017 ) which integrates graph convolution in a GRU cell . We use multi-layer GRUs to obtain hidden states from the inputs . Using re-parametrization ( Kingma & Welling , 2013 ) , we sample zt from the encoder and then decode xt conditioned on zt in an auto-regressive fashion . To ensure fair comparisons , we adapt NP and SNP to graph-based settings and use the same architecture as STNP to generate the hidden states shown in Figure 6 . See Appendix B.3 for architecture details .

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