LEAPS: A discrete neural sampler via locally equivariant networks

TL;DR

LEAPS introduces a novel neural sampling method using locally equivariant networks to efficiently sample from discrete distributions, with applications demonstrated in statistical physics.

Contribution

The paper develops LEAPS, a continuous-time Markov chain-based neural sampler with a new compact rate matrix representation and scalable training algorithms.

Findings

LEAPS effectively samples from complex discrete distributions.

The locally equivariant neural parameterization reduces variance in importance weights.

LEAPS outperforms traditional sampling methods in statistical physics problems.

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 "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.