Pathwise Derivatives for Multivariate Distributions

TL;DR

This paper introduces efficient pathwise gradient estimators for multivariate distributions by leveraging the transport equation, reducing variance, and enabling derivatives for complex mixtures, with demonstrated improvements in variational inference tasks.

Contribution

It presents a novel method to compute low-variance pathwise derivatives for multivariate distributions, including mixtures, using solutions to the transport equation.

Findings

Gradient estimators outperform existing methods in high-dimensional variational inference.

Null solutions of the transport equation effectively reduce estimator variance.

Method applies to mixtures of multivariate Normal distributions with arbitrary means and diagonal covariances.

Abstract

We exploit the link between the transport equation and derivatives of expectations to construct efficient pathwise gradient estimators for multivariate distributions. We focus on two main threads. First, we use null solutions of the transport equation to construct adaptive control variates that can be used to construct gradient estimators with reduced variance. Second, we consider the case of multivariate mixture distributions. In particular we show how to compute pathwise derivatives for mixtures of multivariate Normal distributions with arbitrary means and diagonal covariances. We demonstrate in a variety of experiments in the context of variational inference that our gradient estimators can outperform other methods, especially in high dimensions.