Advancing calibration for stochastic agent-based models in epidemiology with Stein variational inference and Gaussian process surrogates

TL;DR

Calibrating stochastic ABMs in epidemiology is computationally challenging with traditional MCMC due to high-dimensional parameter spaces and costly ABM evaluations. The study introduces Stein Variational Inference (SVI) with Gaussian Process surrogates to calibrate CityCOVID and compares results to DRAM/MCMC and ABC. It shows that the SVI-calibrated posteriors approximate the target posterior $P(theta|h0,d0)$ and yield predictive distributions for hospitalizations and deaths that are comparable in accuracy to MCMC, while offering potential scalability benefits with higher-dimensional parameter spaces. However, SVI requires careful hyperparameter tuning and monitoring of particle dynamics, and DRAM can exhibit multimodal posteriors reflecting discrete parameter structure. These results suggest SVI as a viable, scalable approach for uncertainty quantification in complex stochastic ABMs, with practical implications for public health policy.

Abstract

Accurate calibration of stochastic agent-based models (ABMs) in epidemiology is crucial to make them useful in public health policy decisions and interventions. Traditional calibration methods, e.g., Markov Chain Monte Carlo (MCMC), that yield a probability density function for the parameters being calibrated, are often computationally expensive. When applied to ABMs which are highly parametrized, the calibration process becomes computationally infeasible. This paper investigates the utility of Stein Variational Inference (SVI) as an alternative calibration technique for stochastic epidemiological ABMs approximated by Gaussian process (GP) surrogates. SVI leverages gradient information to iteratively update a set of particles in the space of parameters being calibrated, offering potential advantages in scalability and efficiency for high-dimensional ABMs. The ensemble of particles yields a joint probability density function for the parameters and serves as the calibration. We compare the performance of SVI and MCMC in calibrating CityCOVID, a stochastic epidemiological ABM, focusing on predictive accuracy and calibration effectiveness. Our results demonstrate that SVI maintains predictive accuracy and calibration effectiveness comparable to MCMC, making it a viable alternative for complex epidemiological models. We also present the practical challenges of using a gradient-based calibration such as SVI which include careful tuning of hyperparameters and monitoring of the particle dynamics.

Figures (11)

• Figure 1: Example hospitalization and death trajectories from the GP surrogate compared with the CityCOVID outputted values for a test parameter set.
• Figure 2: Histogram of median absolute relative error for the GP surrogate applied to all test parameter sets (20% of dataset).
• Figure 3: Self-consistency of SVI for increasing number of particles showing the asymptotically decreasing Cramer-von Moises criterion of posterior distributions to the distribution using 300 particles. Each posterior distribution is a combination of 10 distributions using the same number of particles but different initialization.
• Figure 4: Comparison of posterior distributions yielded by SVI for varied numbers of particles. Each posterior distribution is a combination of 10 distributions using the same number of particles but different initialization.
• Figure 5: Kernel density estimate representation of posterior distributions from SVI calibration using 200 particles and different particle initializations based on random seed.