Using Autodiff to Estimate Posterior Moments, Marginals and Samples

Importance sampling is a popular technique in Bayesian inference: by reweighting samples drawn from a proposal distribution we are able to obtain samples and moment estimates from a Bayesian posterior over latent variables. Recent work, however, indicates that importance sampling scales poorly --- in order to accurately approximate the true posterior, the required number of importance samples grows is exponential in the number of latent variables [Chatterjee and Diaconis, 2018]. Massively parallel importance sampling works around this issue by drawing K samples for each of the n latent variables and reasoning about all K^n combinations of latent samples. In principle, we can reason efficiently over K^n combinations of samples by exploiting conditional independencies in the generative model. Previous work only detailed how to compute an ELBO/marginal likelihood estimator by summing over all K^n combinations. However, that work did not give an approach for computing other quantities of interest, namely posterior expectations, marginals and samples, as computing these quantities is far more complex. Specifically, these computations involve iterating forward (following the generative process), then iterating backward through the generative model. These backward traversals can be very complex, and require different backward traversals for each operation of interest. Our contribution is to exploit the source term trick from physics to entirely avoid the need to hand-write backward traversals. Instead, we demonstrate how to simply and easily compute all the required quantities --- posterior expectations, marginals and samples --- by differentiating through a slightly modified marginal likelihood estimator.