Automatic structured variational inference
This method addresses the bottleneck of manual variational family selection for practitioners in probabilistic modeling, offering an automated solution with demonstrated gains.
The authors tackled the problem of choosing an appropriate variational family for stochastic variational inference in probabilistic programming by introducing automatic structured variational inference (ASVI), which automatically constructs structured families that capture complex dependencies and improve performance over methods like mean-field and inverse autoregressive flows.
Stochastic variational inference offers an attractive option as a default method for differentiable probabilistic programming. However, the performance of the variational approach depends on the choice of an appropriate variational family. Here, we introduce automatic structured variational inference (ASVI), a fully automated method for constructing structured variational families, inspired by the closed-form update in conjugate Bayesian models. These convex-update families incorporate the forward pass of the input probabilistic program and can therefore capture complex statistical dependencies. Convex-update families have the same space and time complexity as the input probabilistic program and are therefore tractable for a very large family of models including both continuous and discrete variables. We validate our automatic variational method on a wide range of low- and high-dimensional inference problems. We find that ASVI provides a clear improvement in performance when compared with other popular approaches such as the mean-field approach and inverse autoregressive flows. We provide an open source implementation of ASVI in TensorFlow Probability.