Fast Compartment Model Calibration using Annealed and Transformed Variational Inference

Compartment models are widely used in climate science, epidemiology, and physics, among other disciplines. An important example of a compartment model is susceptible-infected-recovered (SIR) model, which can describe disease dynamics. Bayesian inference for SIR models is challenging because likelihood evaluation requires solving expensive ordinary differential equations. Although variational inference (VI) can be a fast alternative to the traditional Bayes approaches, VI has limited applicability due to boundary issues and local optima problems. To address these challenges, we propose flexible VI methods based on deep generative models that do not require parametric assumptions on the variational distribution. We embed a surjective transformation in our framework to avoid posterior truncation at the boundary. We provide theoretical conditions that guarantee the success of the algorithm. Furthermore, our temperature annealing scheme can avoid being trapped in local optima through a series of intermediate posteriors. We apply our method to variants of SIR models, illustrating that the proposed method can provide fast and accurate inference compared to its competitors.