Learning to Draw Samples with Amortized Stein Variational Gradient Descent

TL;DR

This paper introduces a neural network training algorithm that efficiently draws samples from complex distributions by leveraging Stein variational gradients, applicable to various probabilistic inference tasks.

Contribution

It presents a novel amortized Stein variational gradient descent method that trains neural networks to approximate target distributions, enabling flexible and scalable probabilistic inference.

Findings

Effective in training neural networks for sampling complex distributions

Applicable to variational autoencoders with expressive encoders

Enables adaptive hyper-parameter learning for MCMC samplers

Abstract

We propose a simple algorithm to train stochastic neural networks to draw samples from given target distributions for probabilistic inference. Our method is based on iteratively adjusting the neural network parameters so that the output changes along a Stein variational gradient direction (Liu & Wang, 2016) that maximally decreases the KL divergence with the target distribution. Our method works for any target distribution specified by their unnormalized density function, and can train any black-box architectures that are differentiable in terms of the parameters we want to adapt. We demonstrate our method with a number of applications, including variational autoencoder (VAE) with expressive encoders to model complex latent space structures, and hyper-parameter learning of MCMC samplers that allows Bayesian inference to adaptively improve itself when seeing more data.