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