LEAPS: A discrete neural sampler via locally equivariant networks

Peter Holderrieth, Michael Samuel Albergo, Tommi Jaakkola

TL;DR: A discrete neural sampler for unnormalized densities via locally equivariant networks.

Abstract: We propose *LEAPS*, an algorithm to sample from discrete distributions known up to normalization by learning a rate matrix of a continuous-time Markov chain (CTMC). LEAPS can be seen as a continuous-time formulation of annealed importance sampling and sequential Monte Carlo methods, extended so that the variance of the importance weights is offset by the inclusion of the CTMC. To derive these importance weights, we introduce a set of Radon-Nikodym derivatives of CTMCs over their path measures. Because the computation of these weights is intractable with standard neural network parameterizations of rate matrices, we devise a new compact representation for rate matrices via what we call \textit{locally equivariant} functions. To parameterize them, we introduce a family of locally equivariant multilayer perceptrons, attention layers, and convolutional networks, and provide an approach to make deep networks that preserve the local equivariance. This property allows us to propose a scalable training algorithm for the rate matrix such that the variance of the importance weights associated to the CTMC are minimal. We demonstrate the efficacy of LEAPS on problems in statistical physics. We provide code in https://github.com/malbergo/leaps/.

Lay Summary: We provide a numerically scalable method for learning to sample from a probability distribution over a discrete random variable when given access only to the unnormalized energy function of the distribution.

Link To Code: https://github.com/malbergo/leaps/

Primary Area: Probabilistic Methods

Keywords: CTMC, Bayesian inference, MCMC, sampling

Submission Number: 15045