Implicit Reparameterization Gradients

TL;DR

This paper introduces an implicit reparameterization gradient method that extends the reparameterization trick to a wider range of continuous distributions, improving efficiency and accuracy in gradient estimation.

Contribution

It presents a novel implicit differentiation approach for reparameterization gradients applicable to distributions like Gamma, Beta, Dirichlet, and von Mises, which are not suitable for the classic trick.

Findings

Faster gradient computation compared to existing methods.

More accurate gradient estimates for complex distributions.

Broader applicability of reparameterization in probabilistic models.

Abstract

By providing a simple and efficient way of computing low-variance gradients of continuous random variables, the reparameterization trick has become the technique of choice for training a variety of latent variable models. However, it is not applicable to a number of important continuous distributions. We introduce an alternative approach to computing reparameterization gradients based on implicit differentiation and demonstrate its broader applicability by applying it to Gamma, Beta, Dirichlet, and von Mises distributions, which cannot be used with the classic reparameterization trick. Our experiments show that the proposed approach is faster and more accurate than the existing gradient estimators for these distributions.