Rethinking Variational Inference for Probabilistic Programs with Stochastic Support

Tim Reichelt, Luke Ong, Tom Rainforth

Keywords: Probabilistic Programming, Variational Inference, Stochastic Support

TL;DR: We present a new variational inference method for probabilistic programs with stochastic support that factorizes the guide as a mixture distribution over distinct program paths.

Abstract: We introduce Support Decomposition Variational Inference (SDVI), a new variational inference (VI) approach for probabilistic programs with stochastic support. Existing approaches to this problem rely on designing a single global variational guide on a variable-by-variable basis, while maintaining the stochastic control flow of the original program. SDVI instead breaks the program down into sub-programs with static support, before automatically building separate sub-guides for each. This decomposition significantly aids in the construction of suitable variational families, enabling, in turn, substantial improvements in inference performance.