Universal Marginaliser for Deep Amortised Inference for Probabilistic Programs

TL;DR

The paper introduces the Universal Marginaliser, a neural network-based amortised inference method for probabilistic programming, enabling efficient approximation of conditional marginals across diverse models with theoretical guarantees.

Contribution

It presents a novel neural network approach that can approximate any conditional marginal in probabilistic programs, improving inference efficiency and scalability.

Findings

Successfully trained a single neural network for multiple probabilistic models.

Achieved significant reductions in inference time compared to traditional methods.

Demonstrated effectiveness across various models in Pyro.

Abstract

Probabilistic programming languages (PPLs) are powerful modelling tools which allow to formalise our knowledge about the world and reason about its inherent uncertainty. Inference methods used in PPL can be computationally costly due to significant time burden and/or storage requirements; or they can lack theoretical guarantees of convergence and accuracy when applied to large scale graphical models. To this end, we present the Universal Marginaliser (UM), a novel method for amortised inference, in PPL. We show how combining samples drawn from the original probabilistic program prior with an appropriate augmentation method allows us to train one neural network to approximate any of the corresponding conditional marginal distributions, with any separation into latent and observed variables, and thus amortise the cost of inference. Finally, we benchmark the method on multiple probabilistic programs, in Pyro, with different model structure.