Discrete Neural Flow Samplers with Locally Equivariant Transformer
This addresses a fundamental challenge in machine learning and statistics for researchers and practitioners dealing with discrete data, though it appears incremental as it builds on existing neural flow and Transformer concepts.
The paper tackles the problem of sampling from unnormalised discrete distributions, which is slow with traditional methods, by proposing Discrete Neural Flow Samplers (DNFS) that learn efficient continuous-time Markov chains, achieving improved performance in applications like energy-based models and combinatorial optimisation.
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.