Discrete Neural Flow Samplers with Locally Equivariant Transformer

TL;DR

This paper introduces Discrete Neural Flow Samplers (DNFS), a trainable framework for efficient discrete sampling that leverages locally equivariant Transformers to improve training and sampling from complex distributions.

Contribution

The paper proposes DNFS, a novel trainable sampler that uses a control variate approach and locally equivariant Transformers for efficient discrete distribution sampling.

Findings

Effective in sampling from unnormalised distributions

Improves training efficiency with locally equivariant Transformers

Demonstrates success in energy-based models and combinatorial problems

Abstract

Sampling from unnormalised discrete distributions is a fundamental problem across various domains. While Markov chain Monte Carlo offers a principled approach, it often suffers from slow mixing and poor convergence. In this paper, we propose Discrete Neural Flow Samplers (DNFS), a trainable and efficient framework for discrete sampling. DNFS learns the rate matrix of a continuous-time Markov chain such that the resulting dynamics satisfy the Kolmogorov equation. As this objective involves the intractable partition function, we then employ control variates to reduce the variance of its Monte Carlo estimation, leading to a coordinate descent learning algorithm. To further facilitate computational efficiency, we propose locally equivaraint Transformer, a novel parameterisation of the rate matrix that significantly improves training efficiency while preserving powerful network expressiveness. Empirically, we demonstrate the efficacy of DNFS in a wide range of applications, including sampling from unnormalised distributions, training discrete energy-based models, and solving combinatorial optimisation problems.