Learning Proposals for Probabilistic Programs with Inference Combinators

TL;DR

This paper introduces inference combinators, a flexible framework for constructing and training proposal distributions in probabilistic programming, enabling user-defined variational methods with neural network parameterization.

Contribution

It presents a novel grammar-based approach for composing importance samplers and integrates neural networks for parameterizing proposals, facilitating customizable variational inference.

Findings

Framework supports advanced variational methods like amortized Gibbs sampling

Proposals can be trained via variational objectives

Framework is correct by construction and adaptable to specific models

Abstract

We develop operators for construction of proposals in probabilistic programs, which we refer to as inference combinators. Inference combinators define a grammar over importance samplers that compose primitive operations such as application of a transition kernel and importance resampling. Proposals in these samplers can be parameterized using neural networks, which in turn can be trained by optimizing variational objectives. The result is a framework for user-programmable variational methods that are correct by construction and can be tailored to specific models. We demonstrate the flexibility of this framework by implementing advanced variational methods based on amortized Gibbs sampling and annealing.