Automating Involutive MCMC using Probabilistic and Differentiable Programming

TL;DR

This paper introduces an automated method for implementing involutive MCMC kernels using probabilistic and differentiable programming, simplifying complex sampling tasks and reducing errors in probabilistic models.

Contribution

It presents a novel technique integrated into Gen that automates involutive MCMC kernel implementation, detects user errors, and enhances efficiency through sparsity exploitation.

Findings

Successfully implemented split-merge reversible jump move in Gaussian mixture models

Demonstrated state-dependent proposals for Gaussian process covariance functions

Automated detection of user errors in involutive MCMC kernel specifications

Abstract

Involutive MCMC is a unifying mathematical construction for MCMC kernels that generalizes many classic and state-of-the-art MCMC algorithms, from reversible jump MCMC to kernels based on deep neural networks. But as with MCMC samplers more generally, implementing involutive MCMC kernels is often tedious and error-prone, especially when sampling on complex state spaces. This paper describes a technique for automating the implementation of involutive MCMC kernels given (i) a pair of probabilistic programs defining the target distribution and an auxiliary distribution respectively and (ii) a differentiable program that transforms the execution traces of these probabilistic programs. The technique, which is implemented as part of the Gen probabilistic programming system, also automatically detects user errors in the specification of involutive MCMC kernels and exploits sparsity in the kernels for improved efficiency. The paper shows example Gen code for a split-merge reversible jump move in an infinite Gaussian mixture model and a state-dependent mixture of proposals on a combinatorial space of covariance functions for a Gaussian process.