LEAPS: A discrete neural sampler via locally equivariant networks

2025-02-15

AI Summary

This work addresses efficient importance sampling from discrete distributions specified only by unnormalized probabilities. We propose a novel sampling framework based on continuous-time Markov chains (CTMCs): the target distribution is modeled via a learnable transition rate matrix, and unbiased importance weights are constructed using the Radon–Nikodym derivative of path measures. Our key contribution is the introduction of a local equivariant structure—specifically, scalable, deep-learning-compatible equivariant MLP, attention, and convolutional layers—that enables compact parameterization of the rate matrix while ensuring training stability. To our knowledge, this is the first method to deeply integrate CTMC modeling with importance sampling. Empirically, it substantially reduces importance weight variance in statistical physics tasks, enabling high-accuracy, scalable discrete sampling. It outperforms existing discrete generative models in both efficiency and estimation quality.

Technology Category

• Machine Learning: Probabilistic Circuits and Graphical Models
• Search and Optimization: Sampling/Simulation-based Search
• Reasoning under Uncertainty: Relational Probabilistic Models

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: User modeling and simulation for interactive and conversational systems
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

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.

Problem

• Sampling from discrete distributions efficiently
• Reducing variance in importance weights
• Developing scalable training for rate matrices

Innovation

• Continuous-time Markov chain sampling
• Locally equivariant neural networks
• Minimized importance weight variance