U-Statistics for Importance-Weighted Variational Inference

Javier Burroni, Kenta Takatsu, Justin Domke, Daniel Sheldon

Keywords: Variational Inference, IW-VI

TL;DR: We propose to use complete and incomplete U-statistics in IW-VI to reduce the variance of the gradients and provide faster optimization.

Abstract:

We propose the use of U-statistics to estimate the (gradients of) the importance-weighted evidence lower bound (IW-ELBO), a variational objective that uses multiple samples from a proposal distribution to lower-bound the log-likelihood.
We propose a complete U-statistic estimator, which has variance that is never higher than the standard IW-ELBO estimator, and, under certain conditions, the lowest variance of any unbiased estimator.
However, it requires evaluating the objective on a large number of subsets of samples from the proposal distribution, which can be computationally expensive.
We propose to use incomplete U-statistics as practical alternatives.
We find empirically that both methods reduce estimator variance for its gradients with little computational cost, and lead to faster optimization.

Community Implementations: [ 1 code implementation](https://www.catalyzex.com/paper/arxiv:2302.13918/code)