Rao-Blackwellised Reparameterisation Gradients
This work addresses gradient estimation issues for researchers in probabilistic machine learning, offering an incremental improvement by extending Rao-Blackwellisation benefits to a broader suite of models.
The paper tackles the problem of high-variance gradients in variational inference for latent Gaussian models by proposing the R2-G2 estimator, which applies Rao-Blackwellisation to reparameterisation gradients, resulting in consistently better performance in models with multiple reparameterisation steps.
Latent Gaussian variables have been popularised in probabilistic machine learning. In turn, gradient estimators are the machinery that facilitates gradient-based optimisation for models with latent Gaussian variables. The reparameterisation trick is often used as the default estimator as it is simple to implement and yields low-variance gradients for variational inference. In this work, we propose the R2-G2 estimator as the Rao-Blackwellisation of the reparameterisation gradient estimator. Interestingly, we show that the local reparameterisation gradient estimator for Bayesian MLPs is an instance of the R2-G2 estimator and Rao-Blackwellisation. This lets us extend benefits of Rao-Blackwellised gradients to a suite of probabilistic models. We show that initial training with R2-G2 consistently yields better performance in models with multiple applications of the reparameterisation trick.