Categorical Markov random fields sampled efficiently with discrete diffusions
Sample complexity bounds for categorical Markov random fields via Discrete Diffusions
Machine Learning
Summary
Sampling from complex systems where many things influence each other, like language patterns or protein folding, can be very tricky. The authors study a new approach using 'discrete diffusion' to make this sampling easier and more reliable. They introduce a special way to understand and learn the underlying patterns, leading to clearer guarantees about how much data you need. Their method also lets you balance speed and accuracy depending on the available computing resources. Tests show their approach works better for long sequence data than some common alternatives.
What this means in practice
- •For machine learning engineers: Train one neural sampler that adapts accuracy and cost for complex categorical data sampling when deploying models.
- •For computational biologists: Improve probabilistic models used for sampling protein folding scenarios by effectively handling local dependencies in categorical data.
Authors
Shivam Kumar, Nabarun Deb
Abstract
Many applications in statistics, economics, and physics require sampling from high-dimensional categorical distributions with local dependence structures. Examples include finite memory language models, Ising and Potts systems in statistical physics and protein folding, etc. In modern machine learning, discrete diffusions have emerged as a flexible approach for sampling such data, with strong empirical performance. Motivated by this, we develop learning methods with end-to-end sample complexity bounds for discrete diffusion with uniform noising under local dependence, which we model through low order Markov random fields (MRFs). Our main technical insight is a new \emph{pinning decomposition} of the discrete score. It shows that unlike in continuous diffusions, the score decomposes into components where the dependence on time separates multiplicatively from the dependence on the target. Building on this decomposition, we propose a \emph{weight-sharing neural score learner} and combine it with $τ$-leaping to obtain an end-to-end sampling procedure. Rather than treating score-learning error as a black-box input, as is common in existing sampling analyses, we study the score learning error from finite data and derive optimal sampling guarantees with explicit dependence on the vocabulary size, the interaction order of the MRF, and the sample size. Moreover, our strategy trains a single score network across uniform noise levels while leaving the sampling discretization to be chosen at inference-time. This allows the same trained model to trade accuracy for computational cost as inference-time budgets vary. Numerical experiments on Potts, Ising, and tree-structured models show that weight-sharing score networks outperform fully connected ones for sampling long sequences.