Discrete Neural Flow Samplers with Locally Equivariant Transformer

2025-05-23

AI Summary

Efficient sampling from unnormalized discrete distributions remains challenging. Method: This paper introduces the Differentiable Neural Flow Sampler (DNFS), the first differentiable discrete flow modeling framework grounded in the Kolmogorov forward equation, where dynamics are governed by rate matrices of continuous-time Markov chains. We propose a locally equivariant Transformer architecture for parameter-efficient and stable learning, combined with control variates for variance reduction, coordinate descent optimization, and Monte Carlo gradient estimation. Contribution/Results: Experiments demonstrate that DNFS significantly outperforms MCMC and existing deep sampling methods on unnormalized discrete sampling, training discrete energy-based models, and combinatorial optimization tasks—achieving faster convergence, higher sample quality, and strong theoretical grounding without sacrificing practical performance.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Machine Learning: Deep Neural Architectures and Foundation Models
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.

Problem

• Efficient sampling from unnormalised discrete distributions
• Overcoming slow mixing in Markov chain Monte Carlo
• Training discrete energy-based models effectively

Innovation

• Trainable framework for discrete sampling
• Control variates reduce estimation variance
• Locally equivariant Transformer enhances efficiency