Rao-Blackwellised Reparameterisation Gradients

TL;DR

This paper introduces the R2-G2 estimator, a Rao-Blackwellised reparameterisation gradient method, which improves gradient estimation for models with latent Gaussian variables, leading to better training performance.

Contribution

The paper proposes the R2-G2 estimator as a Rao-Blackwellised version of the reparameterisation gradient, extending its benefits to various probabilistic models.

Findings

R2-G2 outperforms standard estimators in training models with reparameterisation.

The local reparameterisation gradient for Bayesian MLPs is an instance of R2-G2.

Initial training with R2-G2 yields better model performance.

Abstract

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.